diff --git a/src/agents/pi-embedded-runner/compact.hooks.test.ts b/src/agents/pi-embedded-runner/compact.hooks.test.ts index a35060173ff..4e1f1cd3799 100644 --- a/src/agents/pi-embedded-runner/compact.hooks.test.ts +++ b/src/agents/pi-embedded-runner/compact.hooks.test.ts @@ -93,6 +93,8 @@ vi.mock("@mariozechner/pi-ai/oauth", () => ({ vi.mock("@mariozechner/pi-coding-agent", () => { return { + AuthStorage: class AuthStorage {}, + ModelRegistry: class ModelRegistry {}, createAgentSession: vi.fn(async () => { const session = { sessionId: "session-1", @@ -324,6 +326,57 @@ import { getApiProvider, unregisterApiProviders } from "@mariozechner/pi-ai"; import { getCustomApiRegistrySourceId } from "../custom-api-registry.js"; import { compactEmbeddedPiSessionDirect, compactEmbeddedPiSession } from "./compact.js"; +const TEST_SESSION_ID = "session-1"; +const TEST_SESSION_KEY = "agent:main:session-1"; +const TEST_SESSION_FILE = "/tmp/session.jsonl"; +const TEST_WORKSPACE_DIR = "/tmp"; +const TEST_CUSTOM_INSTRUCTIONS = "focus on decisions"; + +function mockResolvedModel() { + resolveModelMock.mockReset(); + resolveModelMock.mockReturnValue({ + model: { provider: "openai", api: "responses", id: "fake", input: [] }, + error: null, + authStorage: { setRuntimeApiKey: vi.fn() }, + modelRegistry: {}, + }); +} + +function compactionConfig(mode: "await" | "off" | "async") { + return { + agents: { + defaults: { + compaction: { + postIndexSync: mode, + }, + }, + }, + } as never; +} + +function directCompactionArgs(overrides: Record = {}) { + return { + sessionId: TEST_SESSION_ID, + sessionKey: TEST_SESSION_KEY, + sessionFile: TEST_SESSION_FILE, + workspaceDir: TEST_WORKSPACE_DIR, + customInstructions: TEST_CUSTOM_INSTRUCTIONS, + ...overrides, + }; +} + +function wrappedCompactionArgs(overrides: Record = {}) { + return { + sessionId: TEST_SESSION_ID, + sessionKey: TEST_SESSION_KEY, + sessionFile: TEST_SESSION_FILE, + workspaceDir: TEST_WORKSPACE_DIR, + customInstructions: TEST_CUSTOM_INSTRUCTIONS, + enqueue: (task: () => unknown) => task(), + ...overrides, + }; +} + const sessionHook = (action: string) => triggerInternalHook.mock.calls.find( (call) => call[0]?.type === "session" && call[0]?.action === action, @@ -336,13 +389,7 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { hookRunner.hasHooks.mockReset(); hookRunner.runBeforeCompaction.mockReset(); hookRunner.runAfterCompaction.mockReset(); - resolveModelMock.mockReset(); - resolveModelMock.mockReturnValue({ - model: { provider: "openai", api: "responses", id: "fake", input: [] }, - error: null, - authStorage: { setRuntimeApiKey: vi.fn() }, - modelRegistry: {}, - }); + mockResolvedModel(); sessionCompactImpl.mockReset(); sessionCompactImpl.mockResolvedValue({ summary: "summary", @@ -376,14 +423,12 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { unregisterApiProviders(getCustomApiRegistrySourceId("ollama")); }); - async function runDirectCompaction(customInstructions = "focus on decisions") { - return await compactEmbeddedPiSessionDirect({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions, - }); + async function runDirectCompaction(customInstructions = TEST_CUSTOM_INSTRUCTIONS) { + return await compactEmbeddedPiSessionDirect( + directCompactionArgs({ + customInstructions, + }), + ); } it("bootstraps runtime plugins with the resolved workspace", async () => { @@ -594,26 +639,15 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { }, }); - const result = await compactEmbeddedPiSessionDirect({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - config: { - agents: { - defaults: { - compaction: { - postIndexSync: "await", - }, - }, - }, - } as never, - }); + const result = await compactEmbeddedPiSessionDirect( + directCompactionArgs({ + config: compactionConfig("await"), + }), + ); expect(result.ok).toBe(true); expect(resolveSessionAgentIdMock).toHaveBeenCalledWith({ - sessionKey: "agent:main:session-1", + sessionKey: TEST_SESSION_KEY, config: expect.any(Object), }); expect(getMemorySearchManagerMock).not.toHaveBeenCalled(); @@ -629,22 +663,11 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { getMemorySearchManagerMock.mockResolvedValue({ manager: { sync } }); let settled = false; - const resultPromise = compactEmbeddedPiSessionDirect({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - config: { - agents: { - defaults: { - compaction: { - postIndexSync: "await", - }, - }, - }, - } as never, - }); + const resultPromise = compactEmbeddedPiSessionDirect( + directCompactionArgs({ + config: compactionConfig("await"), + }), + ); void resultPromise.then(() => { settled = true; @@ -652,7 +675,7 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { await vi.waitFor(() => { expect(sync).toHaveBeenCalledWith({ reason: "post-compaction", - sessionFiles: ["/tmp/session.jsonl"], + sessionFiles: [TEST_SESSION_FILE], }); }); expect(settled).toBe(false); @@ -666,22 +689,11 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { const sync = vi.fn(async () => {}); getMemorySearchManagerMock.mockResolvedValue({ manager: { sync } }); - const result = await compactEmbeddedPiSessionDirect({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - config: { - agents: { - defaults: { - compaction: { - postIndexSync: "off", - }, - }, - }, - } as never, - }); + const result = await compactEmbeddedPiSessionDirect( + directCompactionArgs({ + config: compactionConfig("off"), + }), + ); expect(result.ok).toBe(true); expect(resolveSessionAgentIdMock).not.toHaveBeenCalled(); @@ -698,22 +710,11 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { getMemorySearchManagerMock.mockImplementation(() => managerGate); let settled = false; - const resultPromise = compactEmbeddedPiSessionDirect({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - config: { - agents: { - defaults: { - compaction: { - postIndexSync: "async", - }, - }, - }, - } as never, - }); + const resultPromise = compactEmbeddedPiSessionDirect( + directCompactionArgs({ + config: compactionConfig("async"), + }), + ); await vi.waitFor(() => { expect(getMemorySearchManagerMock).toHaveBeenCalledTimes(1); @@ -730,7 +731,7 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { await vi.waitFor(() => { expect(sync).toHaveBeenCalledWith({ reason: "post-compaction", - sessionFiles: ["/tmp/session.jsonl"], + sessionFiles: [TEST_SESSION_FILE], }); }); const result = await resultPromise; @@ -790,35 +791,25 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => { reason: undefined, result: { summary: "engine-summary", tokensAfter: 50 }, }); - resolveModelMock.mockReset(); - resolveModelMock.mockReturnValue({ - model: { provider: "openai", api: "responses", id: "fake", input: [] }, - error: null, - authStorage: { setRuntimeApiKey: vi.fn() }, - modelRegistry: {}, - }); + mockResolvedModel(); }); it("fires before_compaction with sentinel -1 and after_compaction on success", async () => { hookRunner.hasHooks.mockReturnValue(true); - const result = await compactEmbeddedPiSession({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - messageChannel: "telegram", - customInstructions: "focus on decisions", - enqueue: (task) => task(), - }); + const result = await compactEmbeddedPiSession( + wrappedCompactionArgs({ + messageChannel: "telegram", + }), + ); expect(result.ok).toBe(true); expect(result.compacted).toBe(true); expect(hookRunner.runBeforeCompaction).toHaveBeenCalledWith( - { messageCount: -1, sessionFile: "/tmp/session.jsonl" }, + { messageCount: -1, sessionFile: TEST_SESSION_FILE }, expect.objectContaining({ - sessionKey: "agent:main:session-1", + sessionKey: TEST_SESSION_KEY, messageProvider: "telegram", }), ); @@ -827,10 +818,10 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => { messageCount: -1, compactedCount: -1, tokenCount: 50, - sessionFile: "/tmp/session.jsonl", + sessionFile: TEST_SESSION_FILE, }, expect.objectContaining({ - sessionKey: "agent:main:session-1", + sessionKey: TEST_SESSION_KEY, messageProvider: "telegram", }), ); @@ -843,30 +834,19 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => { getMemorySearchManagerMock.mockResolvedValue({ manager: { sync } }); try { - const result = await compactEmbeddedPiSession({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: " /tmp/session.jsonl ", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - enqueue: (task) => task(), - config: { - agents: { - defaults: { - compaction: { - postIndexSync: "await", - }, - }, - }, - } as never, - }); + const result = await compactEmbeddedPiSession( + wrappedCompactionArgs({ + sessionFile: ` ${TEST_SESSION_FILE} `, + config: compactionConfig("await"), + }), + ); expect(result.ok).toBe(true); expect(listener).toHaveBeenCalledTimes(1); - expect(listener).toHaveBeenCalledWith({ sessionFile: "/tmp/session.jsonl" }); + expect(listener).toHaveBeenCalledWith({ sessionFile: TEST_SESSION_FILE }); expect(sync).toHaveBeenCalledWith({ reason: "post-compaction", - sessionFiles: ["/tmp/session.jsonl"], + sessionFiles: [TEST_SESSION_FILE], }); } finally { cleanup(); @@ -884,14 +864,7 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => { result: undefined, }); - const result = await compactEmbeddedPiSession({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - enqueue: (task) => task(), - }); + const result = await compactEmbeddedPiSession(wrappedCompactionArgs()); expect(result.ok).toBe(false); expect(hookRunner.runBeforeCompaction).toHaveBeenCalled(); @@ -910,23 +883,11 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => { }); try { - const result = await compactEmbeddedPiSession({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - enqueue: (task) => task(), - config: { - agents: { - defaults: { - compaction: { - postIndexSync: "await", - }, - }, - }, - } as never, - }); + const result = await compactEmbeddedPiSession( + wrappedCompactionArgs({ + config: compactionConfig("await"), + }), + ); expect(result.ok).toBe(true); expect(listener).not.toHaveBeenCalled(); @@ -940,14 +901,7 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => { hookRunner.hasHooks.mockReturnValue(true); hookRunner.runBeforeCompaction.mockRejectedValue(new Error("hook boom")); - const result = await compactEmbeddedPiSession({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - enqueue: (task) => task(), - }); + const result = await compactEmbeddedPiSession(wrappedCompactionArgs()); expect(result.ok).toBe(true); expect(result.compacted).toBe(true);