diff --git a/src/agents/pi-extensions/compaction-safeguard.test.ts b/src/agents/pi-extensions/compaction-safeguard.test.ts index 3864abc7e5d..929268b8b2a 100644 --- a/src/agents/pi-extensions/compaction-safeguard.test.ts +++ b/src/agents/pi-extensions/compaction-safeguard.test.ts @@ -131,13 +131,22 @@ const createCompactionEvent = (params: { messageText: string; tokensBefore: numb const createCompactionContext = (params: { sessionManager: ExtensionContext["sessionManager"]; - getApiKeyMock: ReturnType; + getApiKeyAndHeadersMock?: ReturnType; + getApiKeyMock?: ReturnType; }) => ({ model: undefined, sessionManager: params.sessionManager, modelRegistry: { - getApiKey: params.getApiKeyMock, + getApiKeyAndHeaders: + params.getApiKeyAndHeadersMock ?? + vi.fn(async (model) => { + const legacyGetApiKey = params.getApiKeyMock as + | undefined + | ((model: NonNullable) => Promise); + const apiKey = await legacyGetApiKey?.(model); + return apiKey !== undefined ? { ok: true, apiKey } : { ok: false, error: "missing auth" }; + }), }, }) as unknown as Partial; @@ -147,10 +156,16 @@ async function runCompactionScenario(params: { apiKey: string | null; }) { const compactionHandler = createCompactionHandler(); - const getApiKeyMock = vi.fn().mockResolvedValue(params.apiKey ?? undefined); + const getApiKeyAndHeadersMock = vi + .fn() + .mockResolvedValue( + params.apiKey !== null + ? { ok: true, apiKey: params.apiKey } + : { ok: false, error: "missing auth" }, + ); const mockContext = createCompactionContext({ sessionManager: params.sessionManager, - getApiKeyMock, + getApiKeyAndHeadersMock, }); const result = (await compactionHandler(params.event, mockContext)) as { cancel?: boolean; @@ -160,7 +175,7 @@ async function runCompactionScenario(params: { tokensBefore: number; }; }; - return { result, getApiKeyMock }; + return { result, getApiKeyAndHeadersMock }; } function expectCompactionResult(result: { @@ -1634,7 +1649,7 @@ describe("compaction-safeguard extension model fallback", () => { messageText: "test message", tokensBefore: 1000, }); - const { result, getApiKeyMock } = await runCompactionScenario({ + const { result, getApiKeyAndHeadersMock } = await runCompactionScenario({ sessionManager, event: mockEvent, apiKey: null, @@ -1645,7 +1660,7 @@ describe("compaction-safeguard extension model fallback", () => { // KEY ASSERTION: Prove the fallback path was exercised // The handler should have resolved request auth with runtime.model // (via ctx.model ?? runtime?.model). - expect(getApiKeyMock).toHaveBeenCalledWith(model); + expect(getApiKeyAndHeadersMock).toHaveBeenCalledWith(model); // Verify runtime.model is still available (for completeness) const retrieved = getCompactionSafeguardRuntime(sessionManager); @@ -1661,7 +1676,7 @@ describe("compaction-safeguard extension model fallback", () => { messageText: "test", tokensBefore: 500, }); - const { result, getApiKeyMock } = await runCompactionScenario({ + const { result, getApiKeyAndHeadersMock } = await runCompactionScenario({ sessionManager, event: mockEvent, apiKey: null, @@ -1670,7 +1685,7 @@ describe("compaction-safeguard extension model fallback", () => { expect(result).toEqual({ cancel: true }); // Verify early return: request auth should NOT have been resolved when both models are missing. - expect(getApiKeyMock).not.toHaveBeenCalled(); + expect(getApiKeyAndHeadersMock).not.toHaveBeenCalled(); }); }); @@ -1691,7 +1706,7 @@ describe("compaction-safeguard double-compaction guard", () => { customInstructions: "", signal: new AbortController().signal, }; - const { result, getApiKeyMock } = await runCompactionScenario({ + const { result, getApiKeyAndHeadersMock } = await runCompactionScenario({ sessionManager, event: mockEvent, apiKey: "sk-test", // pragma: allowlist secret @@ -1705,7 +1720,7 @@ describe("compaction-safeguard double-compaction guard", () => { expect(compaction.summary).toContain("## Open TODOs"); expect(compaction.firstKeptEntryId).toBe("entry-1"); expect(compaction.tokensBefore).toBe(1500); - expect(getApiKeyMock).not.toHaveBeenCalled(); + expect(getApiKeyAndHeadersMock).not.toHaveBeenCalled(); }); it("returns compaction result with structured fallback summary sections", async () => { @@ -1816,13 +1831,13 @@ describe("compaction-safeguard double-compaction guard", () => { messageText: "real message", tokensBefore: 1500, }); - const { result, getApiKeyMock } = await runCompactionScenario({ + const { result, getApiKeyAndHeadersMock } = await runCompactionScenario({ sessionManager, event: mockEvent, apiKey: null, }); expect(result).toEqual({ cancel: true }); - expect(getApiKeyMock).toHaveBeenCalled(); + expect(getApiKeyAndHeadersMock).toHaveBeenCalled(); }); it("treats tool results as real conversation only when linked to a meaningful user ask", async () => { diff --git a/src/agents/pi-extensions/compaction-safeguard.ts b/src/agents/pi-extensions/compaction-safeguard.ts index 9d71d4422d3..597fe90604a 100644 --- a/src/agents/pi-extensions/compaction-safeguard.ts +++ b/src/agents/pi-extensions/compaction-safeguard.ts @@ -70,10 +70,23 @@ type ToolFailure = { meta?: string; }; -type ModelRegistryWithLegacyAuthLookup = { - getApiKey?: (model: NonNullable) => Promise; +type ModelRegistryWithRequestAuthLookup = { + getApiKeyAndHeaders?: ( + model: NonNullable, + ) => Promise; }; +type ResolvedRequestAuth = + | { + ok: true; + apiKey?: string; + headers?: Record; + } + | { + ok: false; + error: string; + }; + function clampNonNegativeInt(value: unknown, fallback: number): number { const normalized = typeof value === "number" && Number.isFinite(value) ? value : fallback; return Math.max(0, Math.floor(normalized)); @@ -618,13 +631,13 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { return { cancel: true }; } - let apiKey: string | undefined; + let requestAuth: ResolvedRequestAuth; try { - const modelRegistry = ctx.modelRegistry as ModelRegistryWithLegacyAuthLookup; - if (typeof modelRegistry.getApiKey !== "function") { + const modelRegistry = ctx.modelRegistry as ModelRegistryWithRequestAuthLookup; + if (typeof modelRegistry.getApiKeyAndHeaders !== "function") { throw new Error("model registry auth lookup unavailable"); } - apiKey = await modelRegistry.getApiKey(model); + requestAuth = await modelRegistry.getApiKeyAndHeaders(model); } catch (err) { const error = err instanceof Error ? err.message : String(err); log.warn( @@ -636,7 +649,18 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { ); return { cancel: true }; } - const headers = model.headers; + if (!requestAuth.ok) { + log.warn( + `Compaction safeguard: request credential resolution failed for ${model.provider}/${model.id}: ${requestAuth.error}`, + ); + setCompactionSafeguardCancelReason( + ctx.sessionManager, + `Compaction safeguard could not resolve request credentials for ${model.provider}/${model.id}: ${requestAuth.error}`, + ); + return { cancel: true }; + } + const apiKey = requestAuth.apiKey; + const headers = requestAuth.headers; if (!apiKey && !headers) { log.warn( "Compaction safeguard: no request credentials available; cancelling compaction to preserve history.",