diff --git a/src/agents/pi-embedded-runner/compact.hooks.harness.ts b/src/agents/pi-embedded-runner/compact.hooks.harness.ts index fa796a1a59d..e065b0105b3 100644 --- a/src/agents/pi-embedded-runner/compact.hooks.harness.ts +++ b/src/agents/pi-embedded-runner/compact.hooks.harness.ts @@ -6,6 +6,11 @@ type MockResolvedModel = { authStorage: { setRuntimeApiKey: Mock<(provider?: string, apiKey?: string) => void> }; modelRegistry: Record; }; +type MockMemorySearchManager = { + manager: { + sync: (params?: unknown) => Promise; + }; +}; export const contextEngineCompactMock = vi.fn(async () => ({ ok: true as boolean, @@ -45,9 +50,11 @@ export const triggerInternalHook: Mock<(event?: unknown) => void> = vi.fn(); export const sanitizeSessionHistoryMock = vi.fn( async (params: { messages: unknown[] }) => params.messages, ); -export const getMemorySearchManagerMock = vi.fn(async () => ({ +export const getMemorySearchManagerMock: Mock< + (params?: unknown) => Promise +> = vi.fn(async () => ({ manager: { - sync: vi.fn(async () => {}), + sync: vi.fn(async (_params?: unknown) => {}), }, })); export const resolveMemorySearchConfigMock = vi.fn(() => ({ diff --git a/src/agents/pi-embedded-runner/compact.hooks.test.ts b/src/agents/pi-embedded-runner/compact.hooks.test.ts index 8989176330a..c5806609c0d 100644 --- a/src/agents/pi-embedded-runner/compact.hooks.test.ts +++ b/src/agents/pi-embedded-runner/compact.hooks.test.ts @@ -34,6 +34,23 @@ type SessionHookEvent = { sessionKey?: string; context?: Record; }; +type PostCompactionSyncParams = { + reason: string; + sessionFiles: string[]; +}; +type PostCompactionSync = (params?: unknown) => Promise; +type Deferred = { + promise: Promise; + resolve: (value: T) => void; +}; + +function createDeferred(): Deferred { + let resolve!: (value: T) => void; + const promise = new Promise((promiseResolve) => { + resolve = promiseResolve; + }); + return { promise, resolve }; +} function mockResolvedModel() { resolveModelMock.mockReset(); @@ -388,11 +405,12 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { }); it("awaits post-compaction memory sync in await mode when postCompactionForce is true", async () => { - let releaseSync: (() => void) | undefined; - const syncGate = new Promise((resolve) => { - releaseSync = resolve; + const syncStarted = createDeferred(); + const syncRelease = createDeferred(); + const sync = vi.fn(async (params) => { + syncStarted.resolve(params as PostCompactionSyncParams); + await syncRelease.promise; }); - const sync = vi.fn(() => syncGate); getMemorySearchManagerMock.mockResolvedValue({ manager: { sync } }); let settled = false; @@ -405,14 +423,12 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { void resultPromise.then(() => { settled = true; }); - await vi.waitFor(() => { - expect(sync).toHaveBeenCalledWith({ - reason: "post-compaction", - sessionFiles: [TEST_SESSION_FILE], - }); + await expect(syncStarted.promise).resolves.toEqual({ + reason: "post-compaction", + sessionFiles: [TEST_SESSION_FILE], }); expect(settled).toBe(false); - releaseSync?.(); + syncRelease.resolve(undefined); const result = await resultPromise; expect(result.ok).toBe(true); expect(settled).toBe(true); @@ -435,12 +451,17 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { }); it("fires post-compaction memory sync without awaiting it in async mode", async () => { - const sync = vi.fn(async () => {}); - let resolveManager: ((value: { manager: { sync: typeof sync } }) => void) | undefined; - const managerGate = new Promise<{ manager: { sync: typeof sync } }>((resolve) => { - resolveManager = resolve; + const sync = vi.fn(async () => {}); + const managerRequested = createDeferred(); + const managerGate = createDeferred<{ manager: { sync: PostCompactionSync } }>(); + const syncStarted = createDeferred(); + sync.mockImplementation(async (params) => { + syncStarted.resolve(params as PostCompactionSyncParams); + }); + getMemorySearchManagerMock.mockImplementation(async () => { + managerRequested.resolve(undefined); + return await managerGate.promise; }); - getMemorySearchManagerMock.mockImplementation(() => managerGate); let settled = false; const resultPromise = compactEmbeddedPiSessionDirect( @@ -449,26 +470,19 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { }), ); - await vi.waitFor(() => { - expect(getMemorySearchManagerMock).toHaveBeenCalledTimes(1); - }); + await managerRequested.promise; void resultPromise.then(() => { settled = true; }); - await vi.waitFor(() => { - expect(settled).toBe(true); - }); + await resultPromise; + expect(getMemorySearchManagerMock).toHaveBeenCalledTimes(1); + expect(settled).toBe(true); expect(sync).not.toHaveBeenCalled(); - resolveManager?.({ manager: { sync } }); - await managerGate; - await vi.waitFor(() => { - expect(sync).toHaveBeenCalledWith({ - reason: "post-compaction", - sessionFiles: [TEST_SESSION_FILE], - }); + managerGate.resolve({ manager: { sync } }); + await expect(syncStarted.promise).resolves.toEqual({ + reason: "post-compaction", + sessionFiles: [TEST_SESSION_FILE], }); - const result = await resultPromise; - expect(result.ok).toBe(true); }); it("registers the Ollama api provider before compaction", async () => {