openclaw-openbrain-context/tests/engine.test.ts

import { describe, expect, it, vi } from "vitest";

import { OpenBrainConfigError } from "../src/errors.js";
import { OpenBrainContextEngine } from "../src/engine.js";
import type { AgentMessage } from "../src/openclaw-types.js";
import type {
  OpenBrainClientLike,
  OpenBrainThought,
  OpenBrainThoughtInput,
} from "../src/openbrain-client.js";

function makeThought(
  id: string,
  content: string,
  sessionId: string,
  role: string,
  createdAt: string,
): OpenBrainThought {
  return {
    id,
    content,
    source: `openclaw:${sessionId}`,
    metadata: {
      sessionId,
      role,
      type: "message",
    },
    createdAt,
    updatedAt: undefined,
    score: undefined,
  };
}

function makeMockClient(): OpenBrainClientLike {
  return {
    createThought: vi.fn(async (input: OpenBrainThoughtInput) => ({
      id: `thought-${Math.random().toString(36).slice(2)}`,
      content: input.content,
      source: input.source,
      metadata: input.metadata,
      createdAt: new Date().toISOString(),
      updatedAt: undefined,
      score: undefined,
    })),
    search: vi.fn(async () => []),
    listRecent: vi.fn(async () => []),
    updateThought: vi.fn(async (id, payload) => ({
      id,
      content: payload.content ?? "",
      source: "openclaw:session",
      metadata: payload.metadata,
      createdAt: new Date().toISOString(),
      updatedAt: undefined,
      score: undefined,
    })),
    deleteThought: vi.fn(async () => undefined),
    deleteThoughts: vi.fn(async () => undefined),
  };
}

const sessionId = "session-main";

const userMessage: AgentMessage = {
  role: "user",
  content: "What did we decide yesterday?",
  timestamp: Date.now(),
};

describe("OpenBrainContextEngine", () => {
  it("throws OpenBrainConfigError at bootstrap when baseUrl/apiKey are missing", async () => {
    const engine = new OpenBrainContextEngine({});

    await expect(
      engine.bootstrap({
        sessionId,
        sessionFile: "/tmp/session.json",
      }),
    ).rejects.toBeInstanceOf(OpenBrainConfigError);
  });

  it("ingests messages with session metadata", async () => {
    const mockClient = makeMockClient();
    const engine = new OpenBrainContextEngine(
      {
        baseUrl: "https://brain.example.com",
        apiKey: "secret",
        source: "openclaw",
      },
      {
        createClient: () => mockClient,
      },
    );

    await engine.bootstrap({ sessionId, sessionFile: "/tmp/session.json" });
    const result = await engine.ingest({ sessionId, message: userMessage });

    expect(result.ingested).toBe(true);
    expect(mockClient.createThought).toHaveBeenCalledWith(
      expect.objectContaining({
        source: "openclaw:session-main",
        metadata: expect.objectContaining({
          sessionId,
          role: "user",
          type: "message",
          turn: 0,
        }),
      }),
    );
  });

  it("ingests batches and returns ingested count", async () => {
    const mockClient = makeMockClient();
    const engine = new OpenBrainContextEngine(
      {
        baseUrl: "https://brain.example.com",
        apiKey: "secret",
      },
      { createClient: () => mockClient },
    );

    await engine.bootstrap({ sessionId, sessionFile: "/tmp/session.json" });
    const result = await engine.ingestBatch({
      sessionId,
      messages: [
        { role: "user", content: "one", timestamp: 1 },
        { role: "assistant", content: "two", timestamp: 2 },
      ],
    });

    expect(result.ingestedCount).toBe(2);
    expect(mockClient.createThought).toHaveBeenCalledTimes(2);
  });

  it("assembles context from recent + semantic search, deduped and budget-aware", async () => {
    const mockClient = makeMockClient();
    vi.mocked(mockClient.listRecent).mockResolvedValue([
      makeThought("t1", "recent user context", sessionId, "user", "2026-03-06T12:00:00.000Z"),
      makeThought(
        "t2",
        "recent assistant context",
        sessionId,
        "assistant",
        "2026-03-06T12:01:00.000Z",
      ),
    ]);
    vi.mocked(mockClient.search).mockResolvedValue([
      makeThought(
        "t2",
        "recent assistant context",
        sessionId,
        "assistant",
        "2026-03-06T12:01:00.000Z",
      ),
      makeThought("t3", "semantic match", sessionId, "assistant", "2026-03-06T12:02:00.000Z"),
    ]);

    const engine = new OpenBrainContextEngine(
      {
        baseUrl: "https://brain.example.com",
        apiKey: "secret",
        recentMessages: 10,
        semanticSearchLimit: 10,
      },
      { createClient: () => mockClient },
    );

    await engine.bootstrap({ sessionId, sessionFile: "/tmp/session.json" });

    const result = await engine.assemble({
      sessionId,
      messages: [
        {
          role: "user",
          content: "Find the semantic context",
          timestamp: Date.now(),
        },
      ],
      tokenBudget: 40,
    });

    expect(mockClient.search).toHaveBeenCalledWith(
      expect.objectContaining({
        query: "Find the semantic context",
        limit: 10,
      }),
    );
    expect(result.estimatedTokens).toBeLessThanOrEqual(40);
    expect(result.messages.map((message) => String(message.content))).toEqual([
      "recent user context",
      "recent assistant context",
      "semantic match",
    ]);
  });

  it("compact archives a summary thought", async () => {
    const mockClient = makeMockClient();
    vi.mocked(mockClient.listRecent).mockResolvedValue([
      makeThought("t1", "first message", sessionId, "user", "2026-03-06T12:00:00.000Z"),
      makeThought("t2", "second message", sessionId, "assistant", "2026-03-06T12:01:00.000Z"),
    ]);

    const engine = new OpenBrainContextEngine(
      {
        baseUrl: "https://brain.example.com",
        apiKey: "secret",
      },
      { createClient: () => mockClient },
    );

    await engine.bootstrap({ sessionId, sessionFile: "/tmp/session.json" });

    const result = await engine.compact({
      sessionId,
      sessionFile: "/tmp/session.json",
      tokenBudget: 128,
    });

    expect(result.ok).toBe(true);
    expect(result.compacted).toBe(true);
    expect(mockClient.createThought).toHaveBeenCalledWith(
      expect.objectContaining({
        source: "openclaw:session-main",
        metadata: expect.objectContaining({
          sessionId,
          type: "summary",
        }),
      }),
    );
  });

  it("prepares subagent spawn and rollback deletes seeded context", async () => {
    const mockClient = makeMockClient();
    vi.mocked(mockClient.listRecent).mockResolvedValue([
      makeThought("t1", "parent context", sessionId, "assistant", "2026-03-06T12:00:00.000Z"),
    ]);
    vi.mocked(mockClient.createThought).mockResolvedValue({
      id: "seed-thought",
      content: "seed",
      source: "openclaw:child",
      metadata: undefined,
      createdAt: "2026-03-06T12:01:00.000Z",
      updatedAt: undefined,
      score: undefined,
    });

    const engine = new OpenBrainContextEngine(
      {
        baseUrl: "https://brain.example.com",
        apiKey: "secret",
      },
      { createClient: () => mockClient },
    );

    await engine.bootstrap({ sessionId, sessionFile: "/tmp/session.json" });

    const prep = await engine.prepareSubagentSpawn({
      parentSessionKey: sessionId,
      childSessionKey: "child-session",
    });

    expect(prep).toBeDefined();
    expect(mockClient.createThought).toHaveBeenCalledWith(
      expect.objectContaining({
        source: "openclaw:child-session",
      }),
    );

    await prep?.rollback();
    expect(mockClient.deleteThought).toHaveBeenCalledWith("seed-thought");
  });

  it("stores child outcome back into parent on subagent end", async () => {
    const mockClient = makeMockClient();
    vi.mocked(mockClient.listRecent)
      .mockResolvedValueOnce([
        makeThought("p1", "parent context", sessionId, "assistant", "2026-03-06T12:00:00.000Z"),
      ])
      .mockResolvedValueOnce([
        makeThought("c1", "child result detail", "child-session", "assistant", "2026-03-06T12:05:00.000Z"),
      ]);

    const engine = new OpenBrainContextEngine(
      {
        baseUrl: "https://brain.example.com",
        apiKey: "secret",
      },
      { createClient: () => mockClient },
    );

    await engine.bootstrap({ sessionId, sessionFile: "/tmp/session.json" });
    await engine.prepareSubagentSpawn({
      parentSessionKey: sessionId,
      childSessionKey: "child-session",
    });

    await engine.onSubagentEnded({
      childSessionKey: "child-session",
      reason: "completed",
    });

    expect(mockClient.createThought).toHaveBeenLastCalledWith(
      expect.objectContaining({
        source: "openclaw:session-main",
        metadata: expect.objectContaining({
          type: "subagent-result",
          sessionId,
        }),
      }),
    );
  });
});