Fix compaction safeguard request auth lookup

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<typeof vi.fn>;
+  getApiKeyAndHeadersMock?: ReturnType<typeof vi.fn>;
+  getApiKeyMock?: ReturnType<typeof vi.fn>;
 }) =>
   ({
     model: undefined,
     sessionManager: params.sessionManager,
     modelRegistry: {
-      getApiKey: params.getApiKeyMock,
+      getApiKeyAndHeaders:
+        params.getApiKeyAndHeadersMock ??
+        vi.fn(async (model) => {
+          const legacyGetApiKey = params.getApiKeyMock as
+            | undefined
+            | ((model: NonNullable<ExtensionContext["model"]>) => Promise<string | undefined>);
+          const apiKey = await legacyGetApiKey?.(model);
+          return apiKey !== undefined ? { ok: true, apiKey } : { ok: false, error: "missing auth" };
+        }),
     },
   }) as unknown as Partial<ExtensionContext>;
 
@@ -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 () => {

src/agents/pi-extensions/compaction-safeguard.ts

@@ -70,10 +70,23 @@ type ToolFailure = {
   meta?: string;
 };
 
-type ModelRegistryWithLegacyAuthLookup = {
-  getApiKey?: (model: NonNullable<ExtensionContext["model"]>) => Promise<string | undefined>;
+type ModelRegistryWithRequestAuthLookup = {
+  getApiKeyAndHeaders?: (
+    model: NonNullable<ExtensionContext["model"]>,
+  ) => Promise<ResolvedRequestAuth>;
 };
 
+type ResolvedRequestAuth =
+  | {
+      ok: true;
+      apiKey?: string;
+      headers?: Record<string, string>;
+    }
+  | {
+      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.",