From 7a786d88aa3a683d22f8abd5184ec317ca62f5eb Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Fri, 27 Mar 2026 00:19:00 +0100 Subject: [PATCH] fix(coding-agent): resolve models.json auth per request closes #1835 --- AGENTS.md | 1 + packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/models.md | 6 + .../examples/extensions/custom-compaction.ts | 21 +- .../examples/extensions/handoff.ts | 7 +- .../coding-agent/examples/extensions/qna.ts | 7 +- .../examples/extensions/summarize.ts | 19 +- .../coding-agent/src/core/agent-session.ts | 87 ++++---- .../coding-agent/src/core/auth-storage.ts | 8 +- .../core/compaction/branch-summarization.ts | 6 +- .../src/core/compaction/compaction.ts | 13 +- .../coding-agent/src/core/model-registry.ts | 207 +++++++++++------- .../coding-agent/src/core/model-resolver.ts | 8 +- .../src/core/resolve-config-value.ts | 56 ++++- packages/coding-agent/src/core/sdk.ts | 40 ++-- .../compaction-extensions-example.test.ts | 2 +- .../test/compaction-extensions.test.ts | 2 +- .../coding-agent/test/model-registry.test.ts | 133 +++++++---- 18 files changed, 396 insertions(+), 228 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 4f197118..38a32682 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,6 +17,7 @@ read README.md, then ask which module(s) to work on. Based on the answer, read t - **NEVER use inline imports** - no `await import("./foo.js")`, no `import("pkg").Type` in type positions, no dynamic imports for types. Always use standard top-level imports. - NEVER remove or downgrade code to fix type errors from outdated dependencies; upgrade the dependency instead - Always ask before removing functionality or code that appears to be intentional +- Do not preserve backward compatibility unless the user explicitly asks for it - Never hardcode key checks with, eg. `matchesKey(keyData, "ctrl+x")`. All keybindings must be configurable. Add default to matching object (`DEFAULT_EDITOR_KEYBINDINGS` or `DEFAULT_APP_KEYBINDINGS`) ## Commands diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index f07b3e99..0ada31a2 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -8,6 +8,7 @@ ### Fixed +- Fixed `models.json` shell-command auth and headers to resolve at request time instead of being cached into long-lived model state. pi now leaves TTL, caching, and recovery policy to user-provided wrapper commands because arbitrary shell commands need provider-specific strategies ([#1835](https://github.com/badlogic/pi-mono/issues/1835)) - Added missing `ajv` direct dependency; previously relied on transitive install via `@mariozechner/pi-ai` which broke standalone installs ([#2252](https://github.com/badlogic/pi-mono/issues/2252)) - Fixed `/export` HTML backgrounds to honor `theme.export.pageBg`, `cardBg`, and `infoBg` instead of always deriving them from `userMessageBg` ([#2565](https://github.com/badlogic/pi-mono/issues/2565)) - Fixed interactive bash execution collapsed previews to recompute visual line wrapping at render time, so previews respect the current terminal width after resizes and split-pane width changes ([#2569](https://github.com/badlogic/pi-mono/issues/2569)) diff --git a/packages/coding-agent/docs/models.md b/packages/coding-agent/docs/models.md index ab56794e..4f99f04d 100644 --- a/packages/coding-agent/docs/models.md +++ b/packages/coding-agent/docs/models.md @@ -131,6 +131,12 @@ The `apiKey` and `headers` fields support three formats: "apiKey": "sk-..." ``` +For `models.json`, shell commands are resolved at request time. pi intentionally does not apply built-in TTL, stale reuse, or recovery logic for arbitrary commands. Different commands need different caching and failure strategies, and pi cannot infer the right one. + +If your command is slow, expensive, rate-limited, or should keep using a previous value on transient failures, wrap it in your own script or command that implements the caching or TTL behavior you want. + +`/model` availability checks use configured auth presence and do not execute shell commands. + ### Custom Headers ```json diff --git a/packages/coding-agent/examples/extensions/custom-compaction.ts b/packages/coding-agent/examples/extensions/custom-compaction.ts index 1a1760bc..f374839f 100644 --- a/packages/coding-agent/examples/extensions/custom-compaction.ts +++ b/packages/coding-agent/examples/extensions/custom-compaction.ts @@ -31,9 +31,13 @@ export default function (pi: ExtensionAPI) { return; } - // Resolve API key for the summarization model - const apiKey = await ctx.modelRegistry.getApiKey(model); - if (!apiKey) { + // Resolve request auth for the summarization model + const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model); + if (!auth.ok) { + ctx.ui.notify(`Compaction auth failed: ${auth.error}`, "warning"); + return; + } + if (!auth.apiKey) { ctx.ui.notify(`No API key for ${model.provider}, using default compaction`, "warning"); return; } @@ -83,7 +87,16 @@ ${conversationText} try { // Pass signal to honor abort requests (e.g., user cancels compaction) - const response = await complete(model, { messages: summaryMessages }, { apiKey, maxTokens: 8192, signal }); + const response = await complete( + model, + { messages: summaryMessages }, + { + apiKey: auth.apiKey, + headers: auth.headers, + maxTokens: 8192, + signal, + }, + ); const summary = response.content .filter((c): c is { type: "text"; text: string } => c.type === "text") diff --git a/packages/coding-agent/examples/extensions/handoff.ts b/packages/coding-agent/examples/extensions/handoff.ts index feecf8ed..3f115a92 100644 --- a/packages/coding-agent/examples/extensions/handoff.ts +++ b/packages/coding-agent/examples/extensions/handoff.ts @@ -80,7 +80,10 @@ export default function (pi: ExtensionAPI) { loader.onAbort = () => done(null); const doGenerate = async () => { - const apiKey = await ctx.modelRegistry.getApiKey(ctx.model!); + const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model!); + if (!auth.ok || !auth.apiKey) { + throw new Error(auth.ok ? `No API key for ${ctx.model!.provider}` : auth.error); + } const userMessage: Message = { role: "user", @@ -96,7 +99,7 @@ export default function (pi: ExtensionAPI) { const response = await complete( ctx.model!, { systemPrompt: SYSTEM_PROMPT, messages: [userMessage] }, - { apiKey, signal: loader.signal }, + { apiKey: auth.apiKey, headers: auth.headers, signal: loader.signal }, ); if (response.stopReason === "aborted") { diff --git a/packages/coding-agent/examples/extensions/qna.ts b/packages/coding-agent/examples/extensions/qna.ts index fc80c41f..04f14604 100644 --- a/packages/coding-agent/examples/extensions/qna.ts +++ b/packages/coding-agent/examples/extensions/qna.ts @@ -77,7 +77,10 @@ export default function (pi: ExtensionAPI) { // Do the work const doExtract = async () => { - const apiKey = await ctx.modelRegistry.getApiKey(ctx.model!); + const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model!); + if (!auth.ok || !auth.apiKey) { + throw new Error(auth.ok ? `No API key for ${ctx.model!.provider}` : auth.error); + } const userMessage: UserMessage = { role: "user", content: [{ type: "text", text: lastAssistantText! }], @@ -87,7 +90,7 @@ export default function (pi: ExtensionAPI) { const response = await complete( ctx.model!, { systemPrompt: SYSTEM_PROMPT, messages: [userMessage] }, - { apiKey, signal: loader.signal }, + { apiKey: auth.apiKey, headers: auth.headers, signal: loader.signal }, ); if (response.stopReason === "aborted") { diff --git a/packages/coding-agent/examples/extensions/summarize.ts b/packages/coding-agent/examples/extensions/summarize.ts index 46768cb3..14c90b71 100644 --- a/packages/coding-agent/examples/extensions/summarize.ts +++ b/packages/coding-agent/examples/extensions/summarize.ts @@ -165,12 +165,15 @@ export default function (pi: ExtensionAPI) { ctx.ui.notify("Model openai/gpt-5.2 not found", "warning"); } - const apiKey = model ? await ctx.modelRegistry.getApiKey(model) : undefined; - if (!apiKey && ctx.hasUI) { + const auth = model ? await ctx.modelRegistry.getApiKeyAndHeaders(model) : undefined; + if (auth && !auth.ok && ctx.hasUI) { + ctx.ui.notify(auth.error, "warning"); + } + if (auth?.ok && !auth.apiKey && ctx.hasUI) { ctx.ui.notify("No API key for openai/gpt-5.2", "warning"); } - if (!model || !apiKey) { + if (!model || !auth?.ok || !auth.apiKey) { return; } @@ -182,7 +185,15 @@ export default function (pi: ExtensionAPI) { }, ]; - const response = await complete(model, { messages: summaryMessages }, { apiKey, reasoningEffort: "high" }); + const response = await complete( + model, + { messages: summaryMessages }, + { + apiKey: auth.apiKey, + headers: auth.headers, + reasoningEffort: "high", + }, + ); const summary = response.content .filter((c): c is { type: "text"; text: string } => c.type === "text") diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 09271db4..157da913 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -317,6 +317,32 @@ export class AgentSession { return this._modelRegistry; } + private async _getRequiredRequestAuth(model: Model): Promise<{ + apiKey: string; + headers?: Record; + }> { + const result = await this._modelRegistry.getApiKeyAndHeaders(model); + if (!result.ok) { + throw new Error(result.error); + } + if (result.apiKey) { + return { apiKey: result.apiKey, headers: result.headers }; + } + + const isOAuth = this._modelRegistry.isUsingOAuth(model); + if (isOAuth) { + throw new Error( + `Authentication failed for "${model.provider}". ` + + `Credentials may have expired or network is unavailable. ` + + `Run '/login ${model.provider}' to re-authenticate.`, + ); + } + throw new Error( + `No API key found for ${model.provider}.\n\n` + + `Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}`, + ); + } + /** * Install tool hooks once on the Agent instance. * @@ -946,9 +972,7 @@ export class AgentSession { ); } - // Validate API key - const apiKey = await this._modelRegistry.getApiKey(this.model); - if (!apiKey) { + if (!this._modelRegistry.hasConfiguredAuth(this.model)) { const isOAuth = this._modelRegistry.isUsingOAuth(this.model); if (isOAuth) { throw new Error( @@ -1384,12 +1408,11 @@ export class AgentSession { /** * Set model directly. - * Validates API key, saves to session and settings. - * @throws Error if no API key available for the model + * Validates that auth is configured, saves to session and settings. + * @throws Error if no auth is configured for the model */ async setModel(model: Model): Promise { - const apiKey = await this._modelRegistry.getApiKey(model); - if (!apiKey) { + if (!this._modelRegistry.hasConfiguredAuth(model)) { throw new Error(`No API key for ${model.provider}/${model.id}`); } @@ -1418,30 +1441,12 @@ export class AgentSession { return this._cycleAvailableModel(direction); } - private async _getScopedModelsWithApiKey(): Promise; thinkingLevel?: ThinkingLevel }>> { - const apiKeysByProvider = new Map(); - const result: Array<{ model: Model; thinkingLevel?: ThinkingLevel }> = []; - - for (const scoped of this._scopedModels) { - const provider = scoped.model.provider; - let apiKey: string | undefined; - if (apiKeysByProvider.has(provider)) { - apiKey = apiKeysByProvider.get(provider); - } else { - apiKey = await this._modelRegistry.getApiKeyForProvider(provider); - apiKeysByProvider.set(provider, apiKey); - } - - if (apiKey) { - result.push(scoped); - } - } - - return result; + private _getScopedModelsWithAuth(): Array<{ model: Model; thinkingLevel?: ThinkingLevel }> { + return this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model)); } private async _cycleScopedModel(direction: "forward" | "backward"): Promise { - const scopedModels = await this._getScopedModelsWithApiKey(); + const scopedModels = this._getScopedModelsWithAuth(); if (scopedModels.length <= 1) return undefined; const currentModel = this.model; @@ -1481,11 +1486,6 @@ export class AgentSession { const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len; const nextModel = availableModels[nextIndex]; - const apiKey = await this._modelRegistry.getApiKey(nextModel); - if (!apiKey) { - throw new Error(`No API key for ${nextModel.provider}/${nextModel.id}`); - } - const thinkingLevel = this._getThinkingLevelForModelSwitch(); this.agent.setModel(nextModel); this.sessionManager.appendModelChange(nextModel.provider, nextModel.id); @@ -1633,10 +1633,7 @@ export class AgentSession { throw new Error("No model selected"); } - const apiKey = await this._modelRegistry.getApiKey(this.model); - if (!apiKey) { - throw new Error(`No API key for ${this.model.provider}`); - } + const { apiKey, headers } = await this._getRequiredRequestAuth(this.model); const pathEntries = this.sessionManager.getBranch(); const settings = this.settingsManager.getCompactionSettings(); @@ -1690,6 +1687,7 @@ export class AgentSession { preparation, this.model, apiKey, + headers, customInstructions, this._compactionAbortController.signal, ); @@ -1853,11 +1851,12 @@ export class AgentSession { return; } - const apiKey = await this._modelRegistry.getApiKey(this.model); - if (!apiKey) { + const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model); + if (!authResult.ok || !authResult.apiKey) { this._emit({ type: "auto_compaction_end", result: undefined, aborted: false, willRetry: false }); return; } + const { apiKey, headers } = authResult; const pathEntries = this.sessionManager.getBranch(); @@ -1907,6 +1906,7 @@ export class AgentSession { preparation, this.model, apiKey, + headers, undefined, this._autoCompactionAbortController.signal, ); @@ -2155,8 +2155,7 @@ export class AgentSession { refreshTools: () => this._refreshToolRegistry(), getCommands, setModel: async (model) => { - const key = await this.modelRegistry.getApiKey(model); - if (!key) return false; + if (!this.modelRegistry.hasConfiguredAuth(model)) return false; await this.setModel(model); return true; }, @@ -2856,14 +2855,12 @@ export class AgentSession { let summaryDetails: unknown; if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { const model = this.model!; - const apiKey = await this._modelRegistry.getApiKey(model); - if (!apiKey) { - throw new Error(`No API key for ${model.provider}`); - } + const { apiKey, headers } = await this._getRequiredRequestAuth(model); const branchSummarySettings = this.settingsManager.getBranchSummarySettings(); const result = await generateBranchSummary(entriesToSummarize, { model, apiKey, + headers, signal: this._branchSummaryAbortController.signal, customInstructions, replaceInstructions, diff --git a/packages/coding-agent/src/core/auth-storage.ts b/packages/coding-agent/src/core/auth-storage.ts index d8b89969..dc435924 100644 --- a/packages/coding-agent/src/core/auth-storage.ts +++ b/packages/coding-agent/src/core/auth-storage.ts @@ -421,7 +421,7 @@ export class AuthStorage { * 4. Environment variable * 5. Fallback resolver (models.json custom providers) */ - async getApiKey(providerId: string): Promise { + async getApiKey(providerId: string, options?: { includeFallback?: boolean }): Promise { // Runtime override takes highest priority const runtimeKey = this.runtimeOverrides.get(providerId); if (runtimeKey) { @@ -477,7 +477,11 @@ export class AuthStorage { if (envKey) return envKey; // Fall back to custom resolver (e.g., models.json custom providers) - return this.fallbackResolver?.(providerId) ?? undefined; + if (options?.includeFallback !== false) { + return this.fallbackResolver?.(providerId) ?? undefined; + } + + return undefined; } /** diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index d52b2d92..9152f1ef 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -67,6 +67,8 @@ export interface GenerateBranchSummaryOptions { model: Model; /** API key for the model */ apiKey: string; + /** Request headers for the model */ + headers?: Record; /** Abort signal for cancellation */ signal: AbortSignal; /** Optional custom instructions for summarization */ @@ -282,7 +284,7 @@ export async function generateBranchSummary( entries: SessionEntry[], options: GenerateBranchSummaryOptions, ): Promise { - const { model, apiKey, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options; + const { model, apiKey, headers, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options; // Token budget = context window minus reserved space for prompt + response const contextWindow = model.contextWindow || 128000; @@ -322,7 +324,7 @@ export async function generateBranchSummary( const response = await completeSimple( model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, - { apiKey, signal, maxTokens: 2048 }, + { apiKey, headers, signal, maxTokens: 2048 }, ); // Check if aborted or errored diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 213ab6bd..235e564f 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -525,6 +525,7 @@ export async function generateSummary( model: Model, reserveTokens: number, apiKey: string, + headers?: Record, signal?: AbortSignal, customInstructions?: string, previousSummary?: string, @@ -558,8 +559,8 @@ export async function generateSummary( ]; const completionOptions = model.reasoning - ? { maxTokens, signal, apiKey, reasoning: "high" as const } - : { maxTokens, signal, apiKey }; + ? { maxTokens, signal, apiKey, headers, reasoning: "high" as const } + : { maxTokens, signal, apiKey, headers }; const response = await completeSimple( model, @@ -713,6 +714,7 @@ export async function compact( preparation: CompactionPreparation, model: Model, apiKey: string, + headers?: Record, customInstructions?: string, signal?: AbortSignal, ): Promise { @@ -739,12 +741,13 @@ export async function compact( model, settings.reserveTokens, apiKey, + headers, signal, customInstructions, previousSummary, ) : Promise.resolve("No prior history."), - generateTurnPrefixSummary(turnPrefixMessages, model, settings.reserveTokens, apiKey, signal), + generateTurnPrefixSummary(turnPrefixMessages, model, settings.reserveTokens, apiKey, headers, signal), ]); // Merge into single summary summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; @@ -755,6 +758,7 @@ export async function compact( model, settings.reserveTokens, apiKey, + headers, signal, customInstructions, previousSummary, @@ -785,6 +789,7 @@ async function generateTurnPrefixSummary( model: Model, reserveTokens: number, apiKey: string, + headers?: Record, signal?: AbortSignal, ): Promise { const maxTokens = Math.floor(0.5 * reserveTokens); // Smaller budget for turn prefix @@ -802,7 +807,7 @@ async function generateTurnPrefixSummary( const response = await completeSimple( model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, - { maxTokens, signal, apiKey }, + { maxTokens, signal, apiKey, headers }, ); if (response.stopReason === "error") { diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index 514fe832..c23139ad 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -24,7 +24,12 @@ import { existsSync, readFileSync } from "fs"; import { join } from "path"; import { getAgentDir } from "../config.js"; import type { AuthStorage } from "./auth-storage.js"; -import { clearConfigValueCache, resolveConfigValue, resolveHeaders } from "./resolve-config-value.js"; +import { + clearConfigValueCache, + resolveConfigValueOrThrow, + resolveConfigValueUncached, + resolveHeadersOrThrow, +} from "./resolve-config-value.js"; const Ajv = (AjvModule as any).default || AjvModule; const ajv = new Ajv(); @@ -143,14 +148,29 @@ ajv.addSchema(ModelsConfigSchema, "ModelsConfig"); type ModelsConfig = Static; -/** Provider override config (baseUrl, headers, apiKey, compat) without custom models */ +/** Provider override config (baseUrl, compat) without request auth/headers */ interface ProviderOverride { baseUrl?: string; - headers?: Record; - apiKey?: string; compat?: Model["compat"]; } +interface ProviderRequestConfig { + apiKey?: string; + headers?: Record; + authHeader?: boolean; +} + +export type ResolvedRequestAuth = + | { + ok: true; + apiKey?: string; + headers?: Record; + } + | { + ok: false; + error: string; + }; + /** Result of loading custom models from models.json */ interface CustomModelsResult { models: Model[]; @@ -220,12 +240,6 @@ function applyModelOverride(model: Model, override: ModelOverride): Model[] = []; - private customProviderApiKeys: Map = new Map(); + private providerRequestConfigs: Map = new Map(); + private modelRequestHeaders: Map> = new Map(); private registeredProviders: Map = new Map(); private loadError: string | undefined = undefined; @@ -248,16 +263,6 @@ export class ModelRegistry { readonly authStorage: AuthStorage, private modelsJsonPath: string | undefined = join(getAgentDir(), "models.json"), ) { - // Set up fallback resolver for custom provider API keys - this.authStorage.setFallbackResolver((provider) => { - const keyConfig = this.customProviderApiKeys.get(provider); - if (keyConfig) { - return resolveConfigValue(keyConfig); - } - return undefined; - }); - - // Load models this.loadModels(); } @@ -265,7 +270,8 @@ export class ModelRegistry { * Reload models from disk (built-in + custom from models.json). */ refresh(): void { - this.customProviderApiKeys.clear(); + this.providerRequestConfigs.clear(); + this.modelRequestHeaders.clear(); this.loadError = undefined; // Ensure dynamic API/OAuth registrations are rebuilt from current provider state. @@ -329,11 +335,9 @@ export class ModelRegistry { // Apply provider-level baseUrl/headers/compat override if (providerOverride) { - const resolvedHeaders = resolveHeaders(providerOverride.headers); model = { ...model, baseUrl: providerOverride.baseUrl ?? model.baseUrl, - headers: resolvedHeaders ? { ...model.headers, ...resolvedHeaders } : model.headers, compat: mergeCompat(model.compat, providerOverride.compat), }; } @@ -388,23 +392,20 @@ export class ModelRegistry { const modelOverrides = new Map>(); for (const [providerName, providerConfig] of Object.entries(config.providers)) { - // Apply provider-level baseUrl/headers/apiKey/compat override to built-in models when configured. - if (providerConfig.baseUrl || providerConfig.headers || providerConfig.apiKey || providerConfig.compat) { + if (providerConfig.baseUrl || providerConfig.compat) { overrides.set(providerName, { baseUrl: providerConfig.baseUrl, - headers: providerConfig.headers, - apiKey: providerConfig.apiKey, compat: providerConfig.compat, }); } - // Store API key for fallback resolver. - if (providerConfig.apiKey) { - this.customProviderApiKeys.set(providerName, providerConfig.apiKey); - } + this.storeProviderRequestConfig(providerName, providerConfig); if (providerConfig.modelOverrides) { modelOverrides.set(providerName, new Map(Object.entries(providerConfig.modelOverrides))); + for (const [modelId, modelOverride] of Object.entries(providerConfig.modelOverrides)) { + this.storeModelHeaders(providerName, modelId, modelOverride.headers); + } } } @@ -469,29 +470,12 @@ export class ModelRegistry { const modelDefs = providerConfig.models ?? []; if (modelDefs.length === 0) continue; // Override-only, no custom models - // Store API key config for fallback resolver - if (providerConfig.apiKey) { - this.customProviderApiKeys.set(providerName, providerConfig.apiKey); - } - for (const modelDef of modelDefs) { const api = modelDef.api || providerConfig.api; if (!api) continue; - // Merge headers: provider headers are base, model headers override - // Resolve env vars and shell commands in header values - const providerHeaders = resolveHeaders(providerConfig.headers); - const modelHeaders = resolveHeaders(modelDef.headers); const compat = mergeCompat(providerConfig.compat, modelDef.compat); - let headers = providerHeaders || modelHeaders ? { ...providerHeaders, ...modelHeaders } : undefined; - - // If authHeader is true, add Authorization header with resolved API key - if (providerConfig.authHeader && providerConfig.apiKey) { - const resolvedKey = resolveConfigValue(providerConfig.apiKey); - if (resolvedKey) { - headers = { ...headers, Authorization: `Bearer ${resolvedKey}` }; - } - } + this.storeModelHeaders(providerName, modelDef.id, modelDef.headers); // Provider baseUrl is required when custom models are defined. // Individual models can override it with modelDef.baseUrl. @@ -507,7 +491,7 @@ export class ModelRegistry { cost: modelDef.cost ?? defaultCost, contextWindow: modelDef.contextWindow ?? 128000, maxTokens: modelDef.maxTokens ?? 16384, - headers, + headers: undefined, compat, } as Model); } @@ -529,7 +513,7 @@ export class ModelRegistry { * This is a fast check that doesn't refresh OAuth tokens. */ getAvailable(): Model[] { - return this.models.filter((m) => this.authStorage.hasAuth(m.provider)); + return this.models.filter((m) => this.hasConfiguredAuth(m)); } /** @@ -542,15 +526,100 @@ export class ModelRegistry { /** * Get API key for a model. */ - async getApiKey(model: Model): Promise { - return this.authStorage.getApiKey(model.provider); + hasConfiguredAuth(model: Model): boolean { + return ( + this.authStorage.hasAuth(model.provider) || + this.providerRequestConfigs.get(model.provider)?.apiKey !== undefined + ); + } + + private getModelRequestKey(provider: string, modelId: string): string { + return `${provider}:${modelId}`; + } + + private storeProviderRequestConfig( + providerName: string, + config: { + apiKey?: string; + headers?: Record; + authHeader?: boolean; + }, + ): void { + if (!config.apiKey && !config.headers && !config.authHeader) { + return; + } + + this.providerRequestConfigs.set(providerName, { + apiKey: config.apiKey, + headers: config.headers, + authHeader: config.authHeader, + }); + } + + private storeModelHeaders(providerName: string, modelId: string, headers?: Record): void { + const key = this.getModelRequestKey(providerName, modelId); + if (!headers || Object.keys(headers).length === 0) { + this.modelRequestHeaders.delete(key); + return; + } + this.modelRequestHeaders.set(key, headers); + } + + /** + * Get API key and request headers for a model. + */ + async getApiKeyAndHeaders(model: Model): Promise { + try { + const providerConfig = this.providerRequestConfigs.get(model.provider); + const apiKeyFromAuthStorage = await this.authStorage.getApiKey(model.provider, { includeFallback: false }); + const apiKey = + apiKeyFromAuthStorage ?? + (providerConfig?.apiKey + ? resolveConfigValueOrThrow(providerConfig.apiKey, `API key for provider "${model.provider}"`) + : undefined); + + const providerHeaders = resolveHeadersOrThrow(providerConfig?.headers, `provider "${model.provider}"`); + const modelHeaders = resolveHeadersOrThrow( + this.modelRequestHeaders.get(this.getModelRequestKey(model.provider, model.id)), + `model "${model.provider}/${model.id}"`, + ); + + let headers = + model.headers || providerHeaders || modelHeaders + ? { ...model.headers, ...providerHeaders, ...modelHeaders } + : undefined; + + if (providerConfig?.authHeader) { + if (!apiKey) { + return { ok: false, error: `No API key found for "${model.provider}"` }; + } + headers = { ...headers, Authorization: `Bearer ${apiKey}` }; + } + + return { + ok: true, + apiKey, + headers: headers && Object.keys(headers).length > 0 ? headers : undefined, + }; + } catch (error) { + return { + ok: false, + error: error instanceof Error ? error.message : String(error), + }; + } } /** * Get API key for a provider. */ async getApiKeyForProvider(provider: string): Promise { - return this.authStorage.getApiKey(provider); + const apiKey = await this.authStorage.getApiKey(provider, { includeFallback: false }); + if (apiKey !== undefined) { + return apiKey; + } + + const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey; + return providerApiKey ? resolveConfigValueUncached(providerApiKey) : undefined; } /** @@ -586,7 +655,6 @@ export class ModelRegistry { unregisterProvider(providerName: string): void { if (!this.registeredProviders.has(providerName)) return; this.registeredProviders.delete(providerName); - this.customProviderApiKeys.delete(providerName); this.refresh(); } @@ -637,10 +705,7 @@ export class ModelRegistry { ); } - // Store API key for auth resolution - if (config.apiKey) { - this.customProviderApiKeys.set(providerName, config.apiKey); - } + this.storeProviderRequestConfig(providerName, config); if (config.models && config.models.length > 0) { // Full replacement: remove existing models for this provider @@ -649,19 +714,7 @@ export class ModelRegistry { // Parse and add new models for (const modelDef of config.models) { const api = modelDef.api || config.api; - - // Merge headers - const providerHeaders = resolveHeaders(config.headers); - const modelHeaders = resolveHeaders(modelDef.headers); - let headers = providerHeaders || modelHeaders ? { ...providerHeaders, ...modelHeaders } : undefined; - - // If authHeader is true, add Authorization header - if (config.authHeader && config.apiKey) { - const resolvedKey = resolveConfigValue(config.apiKey); - if (resolvedKey) { - headers = { ...headers, Authorization: `Bearer ${resolvedKey}` }; - } - } + this.storeModelHeaders(providerName, modelDef.id, modelDef.headers); this.models.push({ id: modelDef.id, @@ -674,7 +727,7 @@ export class ModelRegistry { cost: modelDef.cost, contextWindow: modelDef.contextWindow, maxTokens: modelDef.maxTokens, - headers, + headers: undefined, compat: modelDef.compat, } as Model); } @@ -686,15 +739,13 @@ export class ModelRegistry { this.models = config.oauth.modifyModels(this.models, cred); } } - } else if (config.baseUrl) { - // Override-only: update baseUrl/headers for existing models - const resolvedHeaders = resolveHeaders(config.headers); + } else if (config.baseUrl || config.headers) { + // Override-only: update baseUrl for existing models. Request headers are resolved per request. this.models = this.models.map((m) => { if (m.provider !== providerName) return m; return { ...m, baseUrl: config.baseUrl ?? m.baseUrl, - headers: resolvedHeaders ? { ...m.headers, ...resolvedHeaders } : m.headers, }; }); } diff --git a/packages/coding-agent/src/core/model-resolver.ts b/packages/coding-agent/src/core/model-resolver.ts index 89938a8d..1825274a 100644 --- a/packages/coding-agent/src/core/model-resolver.ts +++ b/packages/coding-agent/src/core/model-resolver.ts @@ -565,10 +565,10 @@ export async function restoreModelFromSession( ): Promise<{ model: Model | undefined; fallbackMessage: string | undefined }> { const restoredModel = modelRegistry.find(savedProvider, savedModelId); - // Check if restored model exists and has a valid API key - const hasApiKey = restoredModel ? !!(await modelRegistry.getApiKey(restoredModel)) : false; + // Check if restored model exists and still has auth configured + const hasConfiguredAuth = restoredModel ? modelRegistry.hasConfiguredAuth(restoredModel) : false; - if (restoredModel && hasApiKey) { + if (restoredModel && hasConfiguredAuth) { if (shouldPrintMessages) { console.log(chalk.dim(`Restored model: ${savedProvider}/${savedModelId}`)); } @@ -576,7 +576,7 @@ export async function restoreModelFromSession( } // Model not found or no API key - fall back - const reason = !restoredModel ? "model no longer exists" : "no API key available"; + const reason = !restoredModel ? "model no longer exists" : "no auth configured"; if (shouldPrintMessages) { console.error(chalk.yellow(`Warning: Could not restore model ${savedProvider}/${savedModelId} (${reason}).`)); diff --git a/packages/coding-agent/src/core/resolve-config-value.ts b/packages/coding-agent/src/core/resolve-config-value.ts index 2f89accd..646042e1 100644 --- a/packages/coding-agent/src/core/resolve-config-value.ts +++ b/packages/coding-agent/src/core/resolve-config-value.ts @@ -65,24 +65,50 @@ function executeWithDefaultShell(command: string): string | undefined { } } +function executeCommandUncached(commandConfig: string): string | undefined { + const command = commandConfig.slice(1); + return process.platform === "win32" + ? (() => { + const configuredResult = executeWithConfiguredShell(command); + return configuredResult.executed ? configuredResult.value : executeWithDefaultShell(command); + })() + : executeWithDefaultShell(command); +} + function executeCommand(commandConfig: string): string | undefined { if (commandResultCache.has(commandConfig)) { return commandResultCache.get(commandConfig); } - const command = commandConfig.slice(1); - const result = - process.platform === "win32" - ? (() => { - const configuredResult = executeWithConfiguredShell(command); - return configuredResult.executed ? configuredResult.value : executeWithDefaultShell(command); - })() - : executeWithDefaultShell(command); - + const result = executeCommandUncached(commandConfig); commandResultCache.set(commandConfig, result); return result; } +/** + * Resolve all header values using the same resolution logic as API keys. + */ +export function resolveConfigValueUncached(config: string): string | undefined { + if (config.startsWith("!")) { + return executeCommandUncached(config); + } + const envValue = process.env[config]; + return envValue || config; +} + +export function resolveConfigValueOrThrow(config: string, description: string): string { + const resolvedValue = resolveConfigValueUncached(config); + if (resolvedValue !== undefined) { + return resolvedValue; + } + + if (config.startsWith("!")) { + throw new Error(`Failed to resolve ${description} from shell command: ${config.slice(1)}`); + } + + throw new Error(`Failed to resolve ${description}`); +} + /** * Resolve all header values using the same resolution logic as API keys. */ @@ -98,6 +124,18 @@ export function resolveHeaders(headers: Record | undefined): Rec return Object.keys(resolved).length > 0 ? resolved : undefined; } +export function resolveHeadersOrThrow( + headers: Record | undefined, + description: string, +): Record | undefined { + if (!headers) return undefined; + const resolved: Record = {}; + for (const [key, value] of Object.entries(headers)) { + resolved[key] = resolveConfigValueOrThrow(value, `${description} header "${key}"`); + } + return Object.keys(resolved).length > 0 ? resolved : undefined; +} + /** Clear the config value command cache. Exported for testing. */ export function clearConfigValueCache(): void { commandResultCache.clear(); diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 0d83c15e..d9ee45c6 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -1,6 +1,6 @@ import { join } from "node:path"; import { Agent, type AgentMessage, type ThinkingLevel } from "@mariozechner/pi-agent-core"; -import type { Message, Model } from "@mariozechner/pi-ai"; +import { type Message, type Model, streamSimple } from "@mariozechner/pi-ai"; import { getAgentDir, getDocsPath } from "../config.js"; import { AgentSession } from "./agent-session.js"; import { AuthStorage } from "./auth-storage.js"; @@ -194,7 +194,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // If session has data, try to restore model from it if (!model && hasExistingSession && existingSession.model) { const restoredModel = modelRegistry.find(existingSession.model.provider, existingSession.model.modelId); - if (restoredModel && (await modelRegistry.getApiKey(restoredModel))) { + if (restoredModel && modelRegistry.hasConfiguredAuth(restoredModel)) { model = restoredModel; } if (!model) { @@ -293,6 +293,17 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} tools: [], }, convertToLlm: convertToLlmWithBlockImages, + streamFn: async (model, context, options) => { + const auth = await modelRegistry.getApiKeyAndHeaders(model); + if (!auth.ok) { + throw new Error(auth.error); + } + return streamSimple(model, context, { + ...options, + apiKey: auth.apiKey, + headers: auth.headers || options?.headers ? { ...auth.headers, ...options?.headers } : undefined, + }); + }, onPayload: async (payload, _model) => { const runner = extensionRunnerRef.current; if (!runner?.hasHandlers("before_provider_request")) { @@ -311,31 +322,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} transport: settingsManager.getTransport(), thinkingBudgets: settingsManager.getThinkingBudgets(), maxRetryDelayMs: settingsManager.getRetrySettings().maxDelayMs, - getApiKey: async (provider) => { - // Use the provider argument from the in-flight request; - // agent.state.model may already be switched mid-turn. - const resolvedProvider = provider || agent.state.model?.provider; - if (!resolvedProvider) { - throw new Error("No model selected"); - } - const key = await modelRegistry.getApiKeyForProvider(resolvedProvider); - if (!key) { - const model = agent.state.model; - const isOAuth = model && modelRegistry.isUsingOAuth(model); - if (isOAuth) { - throw new Error( - `Authentication failed for "${resolvedProvider}". ` + - `Credentials may have expired or network is unavailable. ` + - `Run '/login ${resolvedProvider}' to re-authenticate.`, - ); - } - throw new Error( - `No API key found for "${resolvedProvider}". ` + - `Set an API key environment variable or run '/login ${resolvedProvider}'.`, - ); - } - return key; - }, }); // Restore messages if session has existing data diff --git a/packages/coding-agent/test/compaction-extensions-example.test.ts b/packages/coding-agent/test/compaction-extensions-example.test.ts index 88159994..4f5e3359 100644 --- a/packages/coding-agent/test/compaction-extensions-example.test.ts +++ b/packages/coding-agent/test/compaction-extensions-example.test.ts @@ -23,7 +23,7 @@ describe("Documentation example", () => { expect(typeof isSplitTurn).toBe("boolean"); expect(typeof tokensBefore).toBe("number"); expect(typeof sessionManager.getEntries).toBe("function"); - expect(typeof modelRegistry.getApiKey).toBe("function"); + expect(typeof modelRegistry.getApiKeyAndHeaders).toBe("function"); expect(typeof firstKeptEntryId).toBe("string"); expect(Array.isArray(branchEntries)).toBe(true); diff --git a/packages/coding-agent/test/compaction-extensions.test.ts b/packages/coding-agent/test/compaction-extensions.test.ts index b110150e..23ca726b 100644 --- a/packages/coding-agent/test/compaction-extensions.test.ts +++ b/packages/coding-agent/test/compaction-extensions.test.ts @@ -380,7 +380,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { // sessionManager, modelRegistry, and model are now on ctx, not event // Verify they're accessible via session expect(typeof session.sessionManager.getEntries).toBe("function"); - expect(typeof session.modelRegistry.getApiKey).toBe("function"); + expect(typeof session.modelRegistry.getApiKeyAndHeaders).toBe("function"); const entries = session.sessionManager.getEntries(); expect(Array.isArray(entries)).toBe(true); diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index ae54e49d..40d5fe20 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -116,7 +116,7 @@ describe("ModelRegistry", () => { } }); - test("overriding headers merges with model headers", () => { + test("overriding headers resolves at request time", async () => { writeRawModelsJson({ anthropic: overrideConfig("https://my-proxy.example.com/v1", { "X-Custom-Header": "custom-value", @@ -127,7 +127,11 @@ describe("ModelRegistry", () => { const anthropicModels = getModelsForProvider(registry, "anthropic"); for (const model of anthropicModels) { - expect(model.headers?.["X-Custom-Header"]).toBe("custom-value"); + const auth = await registry.getApiKeyAndHeaders(model); + expect(auth.ok).toBe(true); + if (auth.ok) { + expect(auth.headers?.["X-Custom-Header"]).toBe("custom-value"); + } } }); @@ -627,7 +631,7 @@ describe("ModelRegistry", () => { expect(sonnet?.cost.output).toBeGreaterThan(0); }); - test("model override can add headers", () => { + test("model override can add headers at request time", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -641,8 +645,13 @@ describe("ModelRegistry", () => { const registry = new ModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); + expect(sonnet).toBeDefined(); - expect(sonnet?.headers?.["X-Custom-Model-Header"]).toBe("value"); + const auth = await registry.getApiKeyAndHeaders(sonnet!); + expect(auth.ok).toBe(true); + if (auth.ok) { + expect(auth.headers?.["X-Custom-Model-Header"]).toBe("value"); + } }); test("refresh() picks up model override changes", () => { @@ -954,9 +963,8 @@ describe("ModelRegistry", () => { expect(apiKey).toBe("hello-world"); }); - describe("caching", () => { - test("command is only executed once per process", async () => { - // Use a command that writes to a file to count invocations + describe("request-time resolution", () => { + test("command is executed on every provider lookup", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0"); @@ -967,18 +975,15 @@ describe("ModelRegistry", () => { }); const registry = new ModelRegistry(authStorage, modelsJsonPath); - - // Call multiple times await registry.getApiKeyForProvider("custom-provider"); await registry.getApiKeyForProvider("custom-provider"); await registry.getApiKeyForProvider("custom-provider"); - // Command should have only run once const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); - expect(count).toBe(1); + expect(count).toBe(3); }); - test("cache persists across registry instances", async () => { + test("commands are re-executed across registry instances", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0"); @@ -988,41 +993,17 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(command), }); - // Create multiple registry instances const registry1 = new ModelRegistry(authStorage, modelsJsonPath); await registry1.getApiKeyForProvider("custom-provider"); const registry2 = new ModelRegistry(authStorage, modelsJsonPath); await registry2.getApiKeyForProvider("custom-provider"); - // Command should still have only run once - const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); - expect(count).toBe(1); - }); - - test("clearApiKeyCache allows command to run again", async () => { - const counterFile = join(tempDir, "counter"); - writeFileSync(counterFile, "0"); - - const counterPath = toShPath(counterFile); - const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo "key-value"'`; - writeRawModelsJson({ - "custom-provider": providerWithApiKey(command), - }); - - const registry = new ModelRegistry(authStorage, modelsJsonPath); - await registry.getApiKeyForProvider("custom-provider"); - - // Clear cache and call again - clearApiKeyCache(); - await registry.getApiKeyForProvider("custom-provider"); - - // Command should have run twice const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(2); }); - test("different commands are cached separately", async () => { + test("different commands resolve independently", async () => { writeRawModelsJson({ "provider-a": providerWithApiKey("!echo key-a"), "provider-b": providerWithApiKey("!echo key-b"), @@ -1037,7 +1018,7 @@ describe("ModelRegistry", () => { expect(keyB).toBe("key-b"); }); - test("failed commands are cached (not retried)", async () => { + test("failed commands are retried", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0"); @@ -1048,17 +1029,14 @@ describe("ModelRegistry", () => { }); const registry = new ModelRegistry(authStorage, modelsJsonPath); - - // Call multiple times - all should return undefined const key1 = await registry.getApiKeyForProvider("custom-provider"); const key2 = await registry.getApiKeyForProvider("custom-provider"); expect(key1).toBeUndefined(); expect(key2).toBeUndefined(); - // Command should have only run once despite failures const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); - expect(count).toBe(1); + expect(count).toBe(2); }); test("environment variables are not cached (changes are picked up)", async () => { @@ -1077,7 +1055,6 @@ describe("ModelRegistry", () => { const key1 = await registry.getApiKeyForProvider("custom-provider"); expect(key1).toBe("first-value"); - // Change env var process.env[envVarName] = "second-value"; const key2 = await registry.getApiKeyForProvider("custom-provider"); @@ -1090,6 +1067,76 @@ describe("ModelRegistry", () => { } } }); + + test("getAvailable does not execute command-backed apiKey resolution", async () => { + const counterFile = join(tempDir, "counter"); + writeFileSync(counterFile, "0"); + + const counterPath = toShPath(counterFile); + const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo "key-value"'`; + writeRawModelsJson({ + "custom-provider": providerWithApiKey(command), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const available = registry.getAvailable(); + + expect(available.some((m) => m.provider === "custom-provider")).toBe(true); + const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); + expect(count).toBe(0); + }); + + test("getApiKeyAndHeaders resolves authHeader on every request", async () => { + const tokenFile = join(tempDir, "token"); + writeFileSync(tokenFile, "token-1"); + const tokenPath = toShPath(tokenFile); + + writeRawModelsJson({ + "custom-provider": { + ...providerWithApiKey(`!sh -c 'cat "${tokenPath}"'`), + authHeader: true, + }, + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const model = registry.find("custom-provider", "test-model"); + expect(model).toBeDefined(); + + const auth1 = await registry.getApiKeyAndHeaders(model!); + expect(auth1).toEqual({ + ok: true, + apiKey: "token-1", + headers: { Authorization: "Bearer token-1" }, + }); + + writeFileSync(tokenFile, "token-2"); + + const auth2 = await registry.getApiKeyAndHeaders(model!); + expect(auth2).toEqual({ + ok: true, + apiKey: "token-2", + headers: { Authorization: "Bearer token-2" }, + }); + }); + + test("getApiKeyAndHeaders returns an error for failed authHeader resolution", async () => { + writeRawModelsJson({ + "custom-provider": { + ...providerWithApiKey("!exit 1"), + authHeader: true, + }, + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const model = registry.find("custom-provider", "test-model"); + expect(model).toBeDefined(); + + const auth = await registry.getApiKeyAndHeaders(model!); + expect(auth.ok).toBe(false); + if (!auth.ok) { + expect(auth.error).toContain('Failed to resolve API key for provider "custom-provider"'); + } + }); }); }); });