feat: implement OpenBrain ContextEngine plugin

tests/engine.test.ts

@@ -0,0 +1,311 @@
+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,
+        }),
+      }),
+    );
+  });
+});