From b80bba8e12b1150ac0a065ff19d6dab85b6e551e Mon Sep 17 00:00:00 2001 From: Joey Krug Date: Sat, 14 Mar 2026 18:37:49 -0400 Subject: [PATCH] test: cover timeout compaction retry cap --- .../run.overflow-compaction.harness.ts | 28 ++++- .../run.timeout-triggered-compaction.test.ts | 117 ++++++++++++++---- src/agents/pi-embedded-runner/run.ts | 2 +- 3 files changed, 116 insertions(+), 31 deletions(-) diff --git a/src/agents/pi-embedded-runner/run.overflow-compaction.harness.ts b/src/agents/pi-embedded-runner/run.overflow-compaction.harness.ts index 9e7853ef7d5..b1664434d67 100644 --- a/src/agents/pi-embedded-runner/run.overflow-compaction.harness.ts +++ b/src/agents/pi-embedded-runner/run.overflow-compaction.harness.ts @@ -136,6 +136,15 @@ export const mockedIsLikelyContextOverflowError = vi.fn((msg?: string) => { export const mockedPickFallbackThinkingLevel = vi.fn<(params?: unknown) => ThinkLevel | null>( () => null, ); +export const mockedGetApiKeyForModel = vi.fn( + async ({ profileId }: { profileId?: string } = {}) => ({ + apiKey: "test-key", + profileId: profileId ?? "test-profile", + source: "test", + mode: "api-key" as const, + }), +); +export const mockedResolveAuthProfileOrder = vi.fn(() => [] as string[]); export const overflowBaseRunParams = { sessionId: "test-session", @@ -223,6 +232,17 @@ export function resetRunOverflowCompactionHarnessMocks(): void { }); mockedPickFallbackThinkingLevel.mockReset(); mockedPickFallbackThinkingLevel.mockReturnValue(null); + mockedGetApiKeyForModel.mockReset(); + mockedGetApiKeyForModel.mockImplementation( + async ({ profileId }: { profileId?: string } = {}) => ({ + apiKey: "test-key", + profileId: profileId ?? "test-profile", + source: "test", + mode: "api-key", + }), + ); + mockedResolveAuthProfileOrder.mockReset(); + mockedResolveAuthProfileOrder.mockReturnValue([]); } export async function loadRunOverflowCompactionHarness(): Promise<{ @@ -322,12 +342,8 @@ export async function loadRunOverflowCompactionHarness(): Promise<{ vi.doMock("../model-auth.js", () => ({ applyLocalNoAuthHeaderOverride: vi.fn((model: unknown) => model), ensureAuthProfileStore: vi.fn(() => ({})), - getApiKeyForModel: vi.fn(async () => ({ - apiKey: "test-key", - profileId: "test-profile", - source: "test", - })), - resolveAuthProfileOrder: vi.fn(() => []), + getApiKeyForModel: mockedGetApiKeyForModel, + resolveAuthProfileOrder: mockedResolveAuthProfileOrder, })); vi.doMock("../models-config.js", () => ({ diff --git a/src/agents/pi-embedded-runner/run.timeout-triggered-compaction.test.ts b/src/agents/pi-embedded-runner/run.timeout-triggered-compaction.test.ts index bccae68bfeb..647d449ceca 100644 --- a/src/agents/pi-embedded-runner/run.timeout-triggered-compaction.test.ts +++ b/src/agents/pi-embedded-runner/run.timeout-triggered-compaction.test.ts @@ -9,7 +9,9 @@ import { mockedCompactDirect, mockedContextEngine, mockedRunEmbeddedAttempt, + mockedGetApiKeyForModel, mockedPickFallbackThinkingLevel, + mockedResolveAuthProfileOrder, resetRunOverflowCompactionHarnessMocks, mockedSessionLikelyHasOversizedToolResults, mockedTruncateOversizedToolResultsInSession, @@ -18,6 +20,16 @@ import { let runEmbeddedPiAgent: typeof import("./run.js").runEmbeddedPiAgent; +const useTwoAuthProfiles = () => { + mockedResolveAuthProfileOrder.mockReturnValue(["profile-a", "profile-b"]); + mockedGetApiKeyForModel.mockImplementation(async ({ profileId } = {}) => ({ + apiKey: `test-key-${profileId ?? "profile-a"}`, + profileId: profileId ?? "profile-a", + source: "test", + mode: "api-key", + })); +}; + describe("timeout-triggered compaction", () => { beforeAll(async () => { ({ runEmbeddedPiAgent } = await loadRunOverflowCompactionHarness()); @@ -57,6 +69,13 @@ describe("timeout-triggered compaction", () => { }); mockedPickFallbackThinkingLevel.mockReturnValue(undefined); mockedGlobalHookRunner.hasHooks.mockImplementation(() => false); + mockedGetApiKeyForModel.mockImplementation(async ({ profileId }) => ({ + apiKey: "test-key", + profileId: profileId ?? "test-profile", + source: "test", + mode: "api-key", + })); + mockedResolveAuthProfileOrder.mockReturnValue([]); }); it("attempts compaction when LLM times out with high prompt token usage (>65%)", async () => { @@ -346,45 +365,95 @@ describe("timeout-triggered compaction", () => { ); }); - it("increments attempt counter even when compaction returns compacted:false", async () => { - // First timeout: high prompt usage, compaction fails (compacted:false) - mockedRunEmbeddedAttempt.mockResolvedValueOnce( - makeAttemptResult({ - timedOut: true, - lastAssistant: { - usage: { input: 150000 }, - } as never, - }), - ); + it("counts compacted:false timeout compactions against the retry cap across profile rotation", async () => { + useTwoAuthProfiles(); + mockedRunEmbeddedAttempt + .mockResolvedValueOnce( + makeAttemptResult({ + timedOut: true, + aborted: true, + lastAssistant: { + usage: { input: 150000 }, + } as never, + }), + ) + .mockResolvedValueOnce( + makeAttemptResult({ + timedOut: true, + aborted: true, + lastAssistant: { + usage: { input: 150000 }, + } as never, + }), + ); mockedCompactDirect.mockResolvedValueOnce({ ok: false, compacted: false, reason: "nothing to compact", }); - // The failed compaction falls through to timeout error; the runner - // returns with an error payload (no retry because compacted was false). + const result = await runEmbeddedPiAgent(overflowBaseRunParams); expect(mockedCompactDirect).toHaveBeenCalledTimes(1); - expect(result.payloads?.[0]?.isError).toBe(true); - }); - - it("increments attempt counter when compact() throws, blocking subsequent attempts", async () => { - // First timeout: high prompt usage, compact() throws - mockedRunEmbeddedAttempt.mockResolvedValueOnce( - makeAttemptResult({ - timedOut: true, - lastAssistant: { - usage: { input: 150000 }, - } as never, + expect(mockedCompactDirect).toHaveBeenCalledWith( + expect.objectContaining({ + runtimeContext: expect.objectContaining({ + authProfileId: "profile-a", + attempt: 1, + maxAttempts: 1, + }), }), ); + expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2); + expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ authProfileId: "profile-a" }), + ); + expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ authProfileId: "profile-b" }), + ); + expect(result.payloads?.[0]?.isError).toBe(true); + expect(result.payloads?.[0]?.text).toContain("timed out"); + }); + + it("counts thrown timeout compactions against the retry cap across profile rotation", async () => { + useTwoAuthProfiles(); + mockedRunEmbeddedAttempt + .mockResolvedValueOnce( + makeAttemptResult({ + timedOut: true, + aborted: true, + lastAssistant: { + usage: { input: 150000 }, + } as never, + }), + ) + .mockResolvedValueOnce( + makeAttemptResult({ + timedOut: true, + aborted: true, + lastAssistant: { + usage: { input: 150000 }, + } as never, + }), + ); mockedCompactDirect.mockRejectedValueOnce(new Error("engine crashed")); - // Falls through to timeout error on first attempt + const result = await runEmbeddedPiAgent(overflowBaseRunParams); expect(mockedCompactDirect).toHaveBeenCalledTimes(1); + expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2); + expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ authProfileId: "profile-a" }), + ); + expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ authProfileId: "profile-b" }), + ); expect(result.payloads?.[0]?.isError).toBe(true); + expect(result.payloads?.[0]?.text).toContain("timed out"); }); it("uses prompt/input tokens for ratio, not total tokens", async () => { diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 1a39f383ff8..5c4c1cd1665 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -1108,7 +1108,7 @@ export async function runEmbeddedPiAgent( : 0; if (timeoutCompactionAttempts >= MAX_TIMEOUT_COMPACTION_ATTEMPTS) { log.warn( - `[timeout-compaction] already compacted ${timeoutCompactionAttempts} time(s) for timeouts; falling through to failover rotation`, + `[timeout-compaction] already attempted timeout compaction ${timeoutCompactionAttempts} time(s); falling through to failover rotation`, ); } else if (tokenUsedRatio > 0.65) { const timeoutDiagId = createCompactionDiagId();