From 39fae14c72dc570cf739c3feb97a2cf7e3dcf150 Mon Sep 17 00:00:00 2001
From: Ayaan Zaidi <hi@obviy.us>
Date: Fri, 27 Mar 2026 10:51:16 +0530
Subject: [PATCH] fix(agents): adapt compaction request auth

---
 src/agents/compaction.retry.test.ts           | 21 +++++--
 src/agents/compaction.ts                      | 17 +-----
 .../pi-extensions/compaction-safeguard.ts     | 56 +++++--------------
 3 files changed, 33 insertions(+), 61 deletions(-)

diff --git a/src/agents/compaction.retry.test.ts b/src/agents/compaction.retry.test.ts
index 481b6087b96..3ff3a79039f 100644
--- a/src/agents/compaction.retry.test.ts
+++ b/src/agents/compaction.retry.test.ts
@@ -50,13 +50,26 @@ describe("compaction retry integration", () => {
     } satisfies AssistantMessage,
   ];
 
-  const testModel = {
+  const testModel: NonNullable<ExtensionContext["model"]> = {
+    id: "claude-3-opus",
+    name: "Claude 3 Opus",
+    api: "anthropic-messages",
     provider: "anthropic",
-    model: "claude-3-opus",
-  } as unknown as NonNullable<ExtensionContext["model"]>;
+    baseUrl: "https://api.anthropic.com",
+    reasoning: false,
+    input: ["text"],
+    cost: {
+      input: 0,
+      output: 0,
+      cacheRead: 0,
+      cacheWrite: 0,
+    },
+    contextWindow: 200_000,
+    maxTokens: 8_192,
+  };
 
   const invokeGenerateSummary = (signal = new AbortController().signal) =>
-    mockGenerateSummary(testMessages, testModel, 1000, "test-api-key", signal, undefined);
+    mockGenerateSummary(testMessages, testModel, 1000, "test-api-key", undefined, signal);
 
   const runSummaryRetry = (options: Parameters<typeof retryAsync>[1]) =>
     retryAsync(() => invokeGenerateSummary(), options);
diff --git a/src/agents/compaction.ts b/src/agents/compaction.ts
index af2f522a220..0a9f93f81d8 100644
--- a/src/agents/compaction.ts
+++ b/src/agents/compaction.ts
@@ -69,19 +69,6 @@ export function buildCompactionSummarizationInstructions(
   return `${identifierPreservation}\n\nAdditional focus:\n${custom}`;
 }
 
-type GenerateSummaryCompat = (
-  currentMessages: AgentMessage[],
-  model: NonNullable<ExtensionContext["model"]>,
-  reserveTokens: number,
-  apiKey: string,
-  headers?: Record<string, string>,
-  signal?: AbortSignal,
-  customInstructions?: string,
-  previousSummary?: string,
-) => Promise<string>;
-
-const generateSummaryCompat = generateSummary as unknown as GenerateSummaryCompat;
-
 export function estimateMessagesTokens(messages: AgentMessage[]): number {
   // SECURITY: toolResult.details can contain untrusted/verbose payloads; never include in LLM-facing compaction.
   const safe = stripToolResultDetails(messages);
@@ -265,12 +252,12 @@ async function summarizeChunks(params: {
   for (const chunk of chunks) {
     summary = await retryAsync(
       () =>
-        generateSummaryCompat(
+        generateSummary(
           chunk,
           model,
           params.reserveTokens,
           params.apiKey,
-          params.headers,
+          model.headers,
           params.signal,
           effectiveInstructions,
           summary,
diff --git a/src/agents/pi-extensions/compaction-safeguard.ts b/src/agents/pi-extensions/compaction-safeguard.ts
index 3d0d4eccca8..61373c3de7f 100644
--- a/src/agents/pi-extensions/compaction-safeguard.ts
+++ b/src/agents/pi-extensions/compaction-safeguard.ts
@@ -1,7 +1,6 @@
 import fs from "node:fs";
 import path from "node:path";
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
-import type { Api, Model } from "@mariozechner/pi-ai";
 import type { ExtensionAPI, FileOperations } from "@mariozechner/pi-coding-agent";
 import { extractSections } from "../../auto-reply/reply/post-compaction-context.js";
 import { openBoundaryFile } from "../../infra/boundary-file-read.js";
@@ -64,40 +63,6 @@ const compactionSafeguardDeps = {
   summarizeInStages,
 };
 
-type ModelApiKeyResolver = {
-  getApiKeyAndHeaders?: (model: Model<Api>) => Promise<{
-    ok: boolean;
-    apiKey?: string;
-    headers?: Record<string, string>;
-  }>;
-  getApiKey?: (model: Model<Api>) => Promise<string | undefined>;
-  getApiKeyForProvider?: (provider: string) => Promise<string | undefined>;
-};
-
-async function resolveModelAuth(
-  modelRegistry: unknown,
-  model: Model<Api>,
-): Promise<{ apiKey?: string; headers?: Record<string, string> }> {
-  const registry = modelRegistry as ModelApiKeyResolver;
-  if (typeof registry.getApiKeyAndHeaders === "function") {
-    const resolved = await registry.getApiKeyAndHeaders(model);
-    if (resolved?.ok) {
-      return {
-        apiKey: resolved.apiKey,
-        headers: resolved.headers,
-      };
-    }
-    return {};
-  }
-  if (typeof registry.getApiKey === "function") {
-    return { apiKey: await registry.getApiKey(model) };
-  }
-  if (typeof registry.getApiKeyForProvider === "function") {
-    return { apiKey: await registry.getApiKeyForProvider(model.provider) };
-  }
-  return {};
-}
-
 type ToolFailure = {
   toolCallId: string;
   toolName: string;
@@ -649,13 +614,20 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void {
       return { cancel: true };
     }
 
-    const fallbackHeaders =
-      model.headers && typeof model.headers === "object" && !Array.isArray(model.headers)
-        ? model.headers
-        : undefined;
-    const resolvedAuth = await resolveModelAuth(ctx.modelRegistry, model);
-    const apiKey = resolvedAuth.apiKey ?? "";
-    const headers = resolvedAuth.headers ?? fallbackHeaders;
+    const requestAuth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+    if (!requestAuth.ok) {
+      log.warn(
+        `Compaction safeguard: request auth resolution failed for ${model.provider}/${model.id}: ${requestAuth.error}`,
+      );
+      setCompactionSafeguardCancelReason(
+        ctx.sessionManager,
+        `Compaction safeguard could not resolve request auth for ${model.provider}/${model.id}: ${requestAuth.error}`,
+      );
+      return { cancel: true };
+    }
+
+    const apiKey = requestAuth.apiKey ?? "";
+    const headers = requestAuth.headers ?? model.headers;
     if (!apiKey && !headers) {
       log.warn(
         "Compaction safeguard: no request auth available; cancelling compaction to preserve history.",