shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(coding-agent): resolve models.json auth per request closes #1835

Browse Source
This commit is contained in:
Mario Zechner
2026-03-27 00:19:00 +01:00
parent fb10d9aef9
commit 7a786d88aa
18 changed files with 396 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
AGENTS.md
View File

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

1
packages/coding-agent/CHANGELOG.md
View File

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

6
packages/coding-agent/docs/models.md
View File

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

21
packages/coding-agent/examples/extensions/custom-compaction.ts
View File

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

7
packages/coding-agent/examples/extensions/handoff.ts
View File

@@ -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") {

7
packages/coding-agent/examples/extensions/qna.ts
View File

@@ -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") {

19
packages/coding-agent/examples/extensions/summarize.ts
View File

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

87
packages/coding-agent/src/core/agent-session.ts
View File

@@ -317,6 +317,32 @@ export class AgentSession {
return this._modelRegistry;
}
private async _getRequiredRequestAuth(model: Model<any>): Promise<{
apiKey: string;
headers?: Record<string, string>;
}> {
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<any>): Promise<void> {
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<Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>> {
const apiKeysByProvider = new Map<string, string | undefined>();
const result: Array<{ model: Model<any>; 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<any>; thinkingLevel?: ThinkingLevel }> {
return this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model));
}
private async _cycleScopedModel(direction: "forward" | "backward"): Promise<ModelCycleResult | undefined> {
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,

8
packages/coding-agent/src/core/auth-storage.ts
View File

@@ -421,7 +421,7 @@ export class AuthStorage {
* 4. Environment variable
* 5. Fallback resolver (models.json custom providers)
*/
async getApiKey(providerId: string): Promise<string | undefined> {
async getApiKey(providerId: string, options?: { includeFallback?: boolean }): Promise<string | undefined> {
// 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;
}
/**

6
packages/coding-agent/src/core/compaction/branch-summarization.ts
View File

@@ -67,6 +67,8 @@ export interface GenerateBranchSummaryOptions {
model: Model<any>;
/** API key for the model */
apiKey: string;
/** Request headers for the model */
headers?: Record<string, string>;
/** Abort signal for cancellation */
signal: AbortSignal;
/** Optional custom instructions for summarization */
@@ -282,7 +284,7 @@ export async function generateBranchSummary(
entries: SessionEntry[],
options: GenerateBranchSummaryOptions,
): Promise<BranchSummaryResult> {
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

13
packages/coding-agent/src/core/compaction/compaction.ts
View File

@@ -525,6 +525,7 @@ export async function generateSummary(
model: Model<any>,
reserveTokens: number,
apiKey: string,
headers?: Record<string, string>,
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<any>,
apiKey: string,
headers?: Record<string, string>,
customInstructions?: string,
signal?: AbortSignal,
): Promise<CompactionResult> {
@@ -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<any>,
reserveTokens: number,
apiKey: string,
headers?: Record<string, string>,
signal?: AbortSignal,
): Promise<string> {
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") {

207
packages/coding-agent/src/core/model-registry.ts
View File

@@ -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<typeof ModelsConfigSchema>;
/** 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<string, string>;
apiKey?: string;
compat?: Model<Api>["compat"];
}
interface ProviderRequestConfig {
apiKey?: string;
headers?: Record<string, string>;
authHeader?: boolean;
}
export type ResolvedRequestAuth =
| {
ok: true;
apiKey?: string;
headers?: Record<string, string>;
}
| {
ok: false;
error: string;
};
/** Result of loading custom models from models.json */
interface CustomModelsResult {
models: Model<Api>[];
@@ -220,12 +240,6 @@ function applyModelOverride(model: Model<Api>, override: ModelOverride): Model<A
};
}
// Merge headers
if (override.headers) {
const resolvedHeaders = resolveHeaders(override.headers);
result.headers = resolvedHeaders ? { ...model.headers, ...resolvedHeaders } : model.headers;
}
// Deep merge compat
result.compat = mergeCompat(model.compat, override.compat);
@@ -240,7 +254,8 @@ export const clearApiKeyCache = clearConfigValueCache;
*/
export class ModelRegistry {
private models: Model<Api>[] = [];
private customProviderApiKeys: Map<string, string> = new Map();
private providerRequestConfigs: Map<string, ProviderRequestConfig> = new Map();
private modelRequestHeaders: Map<string, Record<string, string>> = new Map();
private registeredProviders: Map<string, ProviderConfigInput> = 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<string, Map<string, ModelOverride>>();
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<Api>);
}
@@ -529,7 +513,7 @@ export class ModelRegistry {
* This is a fast check that doesn't refresh OAuth tokens.
*/
getAvailable(): Model<Api>[] {
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<Api>): Promise<string | undefined> {
return this.authStorage.getApiKey(model.provider);
hasConfiguredAuth(model: Model<Api>): 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<string, string>;
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<string, string>): 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<Api>): Promise<ResolvedRequestAuth> {
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<string | undefined> {
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<Api>);
}
@@ -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,
};
});
}

8
packages/coding-agent/src/core/model-resolver.ts
View File

@@ -565,10 +565,10 @@ export async function restoreModelFromSession(
): Promise<{ model: Model<Api> | 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}).`));

56
packages/coding-agent/src/core/resolve-config-value.ts
View File

@@ -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<string, string> | undefined): Rec
return Object.keys(resolved).length > 0 ? resolved : undefined;
}
export function resolveHeadersOrThrow(
headers: Record<string, string> | undefined,
description: string,
): Record<string, string> | undefined {
if (!headers) return undefined;
const resolved: Record<string, string> = {};
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();

40
packages/coding-agent/src/core/sdk.ts
View File

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

2
packages/coding-agent/test/compaction-extensions-example.test.ts
View File

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

2
packages/coding-agent/test/compaction-extensions.test.ts
View File

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

133
packages/coding-agent/test/model-registry.test.ts
View File

@@ -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"');
}
});
});
});
});