fix(agents): adapt compaction request auth

This commit is contained in:
Ayaan Zaidi
2026-03-27 10:51:16 +05:30
parent 9368f834c0
commit 39fae14c72
3 changed files with 33 additions and 61 deletions

View File

@@ -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);

View File

@@ -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,

View File

@@ -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.",