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("ingests batches in parallel chunks of five", async () => {
    const mockClient = makeMockClient();
    let inFlight = 0;
    let maxInFlight = 0;
    let createdCount = 0;

    vi.mocked(mockClient.createThought).mockImplementation(async (input: OpenBrainThoughtInput) => {
      inFlight += 1;
      maxInFlight = Math.max(maxInFlight, inFlight);
      await new Promise((resolve) => {
        setTimeout(resolve, 20);
      });
      inFlight -= 1;
      createdCount += 1;
      return {
        id: `thought-${createdCount}`,
        content: input.content,
        source: input.source,
        metadata: input.metadata,
        createdAt: new Date().toISOString(),
        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 result = await engine.ingestBatch({
      sessionId,
      messages: Array.from({ length: 10 }, (_, index) => ({
        role: index % 2 === 0 ? "user" : "assistant",
        content: `message-${index + 1}`,
        timestamp: index + 1,
      })),
    });

    expect(result.ingestedCount).toBe(10);
    expect(maxInFlight).toBe(5);
    expect(mockClient.createThought).toHaveBeenCalledTimes(10);
  });

  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 and deletes summarized inputs", async () => {
    const mockClient = makeMockClient();
    vi.mocked(mockClient.listRecent).mockResolvedValue(
      Array.from({ length: 12 }, (_, index) => {
        return makeThought(
          `t${index + 1}`,
          `message ${index + 1}`,
          sessionId,
          index % 2 === 0 ? "user" : "assistant",
          `2026-03-06T12:${String(index).padStart(2, "0")}: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",
        }),
      }),
    );
    const deletedIds = vi
      .mocked(mockClient.deleteThought)
      .mock.calls.map(([id]) => id)
      .sort((left, right) => left.localeCompare(right));
    expect(deletedIds).toEqual([
      "t10",
      "t11",
      "t12",
      "t3",
      "t4",
      "t5",
      "t6",
      "t7",
      "t8",
      "t9",
    ]);
  });

  it("stops trimming once the newest message exceeds budget", async () => {
    const mockClient = makeMockClient();
    const oversizedNewest = "z".repeat(400);
    vi.mocked(mockClient.listRecent).mockResolvedValue([
      makeThought("t1", "small older message", sessionId, "assistant", "2026-03-06T12:00:00.000Z"),
      makeThought("t2", oversizedNewest, 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.assemble({
      sessionId,
      messages: [
        {
          role: "user",
          content: "query",
          timestamp: Date.now(),
        },
      ],
      tokenBudget: 12,
    });

    expect(result.messages.map((message) => String(message.content))).toEqual([oversizedNewest]);
  });

  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,
        }),
      }),
    );
  });
});