feat: add provider-scoped environment overrides (#5807)

This commit is contained in:
Armin Ronacher
2026-06-16 17:19:08 +02:00
committed by GitHub
parent 3039f3e17d
commit 7f29e7a369
33 changed files with 511 additions and 215 deletions

View File

@@ -357,6 +357,7 @@ export class AgentSession {
private async _getRequiredRequestAuth(model: Model<any>): Promise<{
apiKey: string;
headers?: Record<string, string>;
env?: Record<string, string>;
}> {
const result = await this._modelRegistry.getApiKeyAndHeaders(model);
if (!result.ok) {
@@ -366,7 +367,7 @@ export class AgentSession {
throw new Error(result.error);
}
if (result.apiKey) {
return { apiKey: result.apiKey, headers: result.headers };
return { apiKey: result.apiKey, headers: result.headers, env: result.env };
}
const isOAuth = this._modelRegistry.isUsingOAuth(model);
@@ -383,13 +384,14 @@ export class AgentSession {
private async _getCompactionRequestAuth(model: Model<any>): Promise<{
apiKey?: string;
headers?: Record<string, string>;
env?: Record<string, string>;
}> {
if (this.agent.streamFn === streamSimple) {
return this._getRequiredRequestAuth(model);
}
const result = await this._modelRegistry.getApiKeyAndHeaders(model);
return result.ok ? { apiKey: result.apiKey, headers: result.headers } : {};
return result.ok ? { apiKey: result.apiKey, headers: result.headers, env: result.env } : {};
}
/**
@@ -1649,7 +1651,7 @@ export class AgentSession {
throw new Error(formatNoModelSelectedMessage());
}
const { apiKey, headers } = await this._getCompactionRequestAuth(this.model);
const { apiKey, headers, env } = await this._getCompactionRequestAuth(this.model);
const pathEntries = this.sessionManager.getBranch();
const settings = this.settingsManager.getCompactionSettings();
@@ -1708,6 +1710,7 @@ export class AgentSession {
this._compactionAbortController.signal,
this.thinkingLevel,
this.agent.streamFn,
env,
);
summary = result.summary;
firstKeptEntryId = result.firstKeptEntryId;
@@ -1898,6 +1901,7 @@ export class AgentSession {
let apiKey: string | undefined;
let headers: Record<string, string> | undefined;
let env: Record<string, string> | undefined;
if (this.agent.streamFn === streamSimple) {
const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model);
if (!authResult.ok || !authResult.apiKey) {
@@ -1912,8 +1916,9 @@ export class AgentSession {
}
apiKey = authResult.apiKey;
headers = authResult.headers;
env = authResult.env;
} else {
({ apiKey, headers } = await this._getCompactionRequestAuth(this.model));
({ apiKey, headers, env } = await this._getCompactionRequestAuth(this.model));
}
const pathEntries = this.sessionManager.getBranch();
@@ -1981,6 +1986,7 @@ export class AgentSession {
this._autoCompactionAbortController.signal,
this.thinkingLevel,
this.agent.streamFn,
env,
);
summary = compactResult.summary;
firstKeptEntryId = compactResult.firstKeptEntryId;
@@ -2784,12 +2790,13 @@ export class AgentSession {
let summaryDetails: unknown;
if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) {
const model = this.model!;
const { apiKey, headers } = await this._getRequiredRequestAuth(model);
const { apiKey, headers, env } = await this._getRequiredRequestAuth(model);
const branchSummarySettings = this.settingsManager.getBranchSummarySettings();
const result = await generateBranchSummary(entriesToSummarize, {
model,
apiKey,
headers,
env,
signal: this._branchSummaryAbortController.signal,
customInstructions,
replaceInstructions,

View File

@@ -24,6 +24,7 @@ import { resolveConfigValue } from "./resolve-config-value.ts";
export type ApiKeyCredential = {
type: "api_key";
key: string;
env?: Record<string, string>;
};
export type OAuthCredential = {
@@ -303,6 +304,14 @@ export class AuthStorage {
return this.data[provider] ?? undefined;
}
/**
* Get provider-scoped environment values for an API key credential.
*/
getProviderEnv(provider: string): Record<string, string> | undefined {
const cred = this.data[provider];
return cred?.type === "api_key" && cred.env ? { ...cred.env } : undefined;
}
/**
* Set credential for a provider.
*/
@@ -471,7 +480,7 @@ export class AuthStorage {
const cred = this.data[providerId];
if (cred?.type === "api_key") {
return resolveConfigValue(cred.key);
return resolveConfigValue(cred.key, cred.env);
}
if (cred?.type === "oauth") {

View File

@@ -69,6 +69,8 @@ export interface GenerateBranchSummaryOptions {
apiKey: string;
/** Request headers for the model */
headers?: Record<string, string>;
/** Provider-scoped environment values for the model */
env?: Record<string, string>;
/** Abort signal for cancellation */
signal: AbortSignal;
/** Optional custom instructions for summarization */
@@ -290,6 +292,7 @@ export async function generateBranchSummary(
model,
apiKey,
headers,
env,
signal,
customInstructions,
replaceInstructions,
@@ -335,7 +338,7 @@ export async function generateBranchSummary(
// request behavior (timeouts, retries, attribution headers) stays consistent
// without running through agent state/events.
const context = { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages };
const requestOptions: SimpleStreamOptions = { apiKey, headers, signal, maxTokens: 2048 };
const requestOptions: SimpleStreamOptions = { apiKey, headers, env, signal, maxTokens: 2048 };
const response = streamFn
? await (await streamFn(model, context, requestOptions)).result()
: await completeSimple(model, context, requestOptions);

View File

@@ -528,10 +528,11 @@ function createSummarizationOptions(
maxTokens: number,
apiKey: string | undefined,
headers: Record<string, string> | undefined,
env: Record<string, string> | undefined,
signal: AbortSignal | undefined,
thinkingLevel: ThinkingLevel | undefined,
): SimpleStreamOptions {
const options: SimpleStreamOptions = { maxTokens, signal, apiKey, headers };
const options: SimpleStreamOptions = { maxTokens, signal, apiKey, headers, env };
if (model.reasoning && thinkingLevel && thinkingLevel !== "off") {
options.reasoning = thinkingLevel;
}
@@ -566,6 +567,7 @@ export async function generateSummary(
previousSummary?: string,
thinkingLevel?: ThinkingLevel,
streamFn?: StreamFn,
env?: Record<string, string>,
): Promise<string> {
const maxTokens = Math.min(
Math.floor(0.8 * reserveTokens),
@@ -598,7 +600,7 @@ export async function generateSummary(
},
];
const completionOptions = createSummarizationOptions(model, maxTokens, apiKey, headers, signal, thinkingLevel);
const completionOptions = createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel);
const response = await completeSummarization(
model,
@@ -753,6 +755,7 @@ export async function compact(
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
streamFn?: StreamFn,
env?: Record<string, string>,
): Promise<CompactionResult> {
const {
firstKeptEntryId,
@@ -783,6 +786,7 @@ export async function compact(
previousSummary,
thinkingLevel,
streamFn,
env,
)
: Promise.resolve("No prior history."),
generateTurnPrefixSummary(
@@ -791,6 +795,7 @@ export async function compact(
settings.reserveTokens,
apiKey,
headers,
env,
signal,
thinkingLevel,
streamFn,
@@ -811,6 +816,7 @@ export async function compact(
previousSummary,
thinkingLevel,
streamFn,
env,
);
}
@@ -839,6 +845,7 @@ async function generateTurnPrefixSummary(
reserveTokens: number,
apiKey: string | undefined,
headers?: Record<string, string>,
env?: Record<string, string>,
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
streamFn?: StreamFn,
@@ -861,7 +868,7 @@ async function generateTurnPrefixSummary(
const response = await completeSummarization(
model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
createSummarizationOptions(model, maxTokens, apiKey, headers, signal, thinkingLevel),
createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel),
streamFn,
);

View File

@@ -240,6 +240,7 @@ export type ResolvedRequestAuth =
ok: true;
apiKey?: string;
headers?: Record<string, string>;
env?: Record<string, string>;
}
| {
ok: false;
@@ -684,17 +685,27 @@ export class ModelRegistry {
async getApiKeyAndHeaders(model: Model<Api>): Promise<ResolvedRequestAuth> {
try {
const providerConfig = this.providerRequestConfigs.get(model.provider);
const providerEnv = this.authStorage.getProviderEnv(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}"`)
? resolveConfigValueOrThrow(
providerConfig.apiKey,
`API key for provider "${model.provider}"`,
providerEnv,
)
: undefined);
const providerHeaders = resolveHeadersOrThrow(providerConfig?.headers, `provider "${model.provider}"`);
const providerHeaders = resolveHeadersOrThrow(
providerConfig?.headers,
`provider "${model.provider}"`,
providerEnv,
);
const modelHeaders = resolveHeadersOrThrow(
this.modelRequestHeaders.get(this.getModelRequestKey(model.provider, model.id)),
`model "${model.provider}/${model.id}"`,
providerEnv,
);
let headers =
@@ -713,6 +724,7 @@ export class ModelRegistry {
ok: true,
apiKey,
headers: headers && Object.keys(headers).length > 0 ? headers : undefined,
env: providerEnv && Object.keys(providerEnv).length > 0 ? providerEnv : undefined,
};
} catch (error) {
return {
@@ -777,7 +789,9 @@ export class ModelRegistry {
}
const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey;
return providerApiKey ? resolveConfigValueUncached(providerApiKey) : undefined;
return providerApiKey
? resolveConfigValueUncached(providerApiKey, this.authStorage.getProviderEnv(provider))
: undefined;
}
/**

View File

@@ -85,8 +85,8 @@ function parseConfigValueReference(config: string): ConfigValueReference {
return { type: "template", parts: parseConfigValueTemplate(config) };
}
function resolveEnvConfigValue(name: string): string | undefined {
return process.env[name] || undefined;
function resolveEnvConfigValue(name: string, env?: Record<string, string>): string | undefined {
return env?.[name] || process.env[name] || undefined;
}
function getTemplateEnvVarNames(parts: TemplatePart[]): string[] {
@@ -98,14 +98,14 @@ function getTemplateEnvVarNames(parts: TemplatePart[]): string[] {
return names;
}
function resolveTemplate(parts: TemplatePart[]): string | undefined {
function resolveTemplate(parts: TemplatePart[], env?: Record<string, string>): string | undefined {
let resolved = "";
for (const part of parts) {
if (part.type === "literal") {
resolved += part.value;
continue;
}
const envValue = resolveEnvConfigValue(part.name);
const envValue = resolveEnvConfigValue(part.name, env);
if (envValue === undefined) return undefined;
resolved += envValue;
}
@@ -123,16 +123,16 @@ export function getConfigValueEnvVarNames(config: string): string[] {
return reference.type === "template" ? getTemplateEnvVarNames(reference.parts) : [];
}
export function getMissingConfigValueEnvVarNames(config: string): string[] {
return getConfigValueEnvVarNames(config).filter((name) => resolveEnvConfigValue(name) === undefined);
export function getMissingConfigValueEnvVarNames(config: string, env?: Record<string, string>): string[] {
return getConfigValueEnvVarNames(config).filter((name) => resolveEnvConfigValue(name, env) === undefined);
}
export function isCommandConfigValue(config: string): boolean {
return parseConfigValueReference(config).type === "command";
}
export function isConfigValueConfigured(config: string): boolean {
return getMissingConfigValueEnvVarNames(config).length === 0;
export function isConfigValueConfigured(config: string, env?: Record<string, string>): boolean {
return getMissingConfigValueEnvVarNames(config, env).length === 0;
}
/**
@@ -142,12 +142,12 @@ export function isConfigValueConfigured(config: string): boolean {
* - In non-command values, "$$" escapes a literal "$" and "$!" escapes a literal "!"
* - Otherwise treats the value as a literal
*/
export function resolveConfigValue(config: string): string | undefined {
export function resolveConfigValue(config: string, env?: Record<string, string>): string | undefined {
const reference = parseConfigValueReference(config);
if (reference.type === "command") {
return executeCommand(reference.config);
}
return resolveTemplate(reference.parts);
return resolveTemplate(reference.parts, env);
}
function executeWithConfiguredShell(command: string): { executed: boolean; value: string | undefined } {
@@ -216,16 +216,16 @@ function executeCommand(commandConfig: string): string | undefined {
/**
* Resolve all header values using the same resolution logic as API keys.
*/
export function resolveConfigValueUncached(config: string): string | undefined {
export function resolveConfigValueUncached(config: string, env?: Record<string, string>): string | undefined {
const reference = parseConfigValueReference(config);
if (reference.type === "command") {
return executeCommandUncached(reference.config);
}
return resolveTemplate(reference.parts);
return resolveTemplate(reference.parts, env);
}
export function resolveConfigValueOrThrow(config: string, description: string): string {
const resolvedValue = resolveConfigValueUncached(config);
export function resolveConfigValueOrThrow(config: string, description: string, env?: Record<string, string>): string {
const resolvedValue = resolveConfigValueUncached(config, env);
if (resolvedValue !== undefined) {
return resolvedValue;
}
@@ -236,7 +236,7 @@ export function resolveConfigValueOrThrow(config: string, description: string):
}
if (reference.type === "template") {
const missingEnvVars = getMissingConfigValueEnvVarNames(config);
const missingEnvVars = getMissingConfigValueEnvVarNames(config, env);
if (missingEnvVars.length === 1) {
throw new Error(`Failed to resolve ${description} from environment variable: ${missingEnvVars[0]}`);
}
@@ -251,11 +251,14 @@ export function resolveConfigValueOrThrow(config: string, description: string):
/**
* Resolve all header values using the same resolution logic as API keys.
*/
export function resolveHeaders(headers: Record<string, string> | undefined): Record<string, string> | undefined {
export function resolveHeaders(
headers: Record<string, string> | undefined,
env?: Record<string, string>,
): Record<string, string> | undefined {
if (!headers) return undefined;
const resolved: Record<string, string> = {};
for (const [key, value] of Object.entries(headers)) {
const resolvedValue = resolveConfigValue(value);
const resolvedValue = resolveConfigValue(value, env);
if (resolvedValue) {
resolved[key] = resolvedValue;
}
@@ -266,11 +269,12 @@ export function resolveHeaders(headers: Record<string, string> | undefined): Rec
export function resolveHeadersOrThrow(
headers: Record<string, string> | undefined,
description: string,
env?: Record<string, 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}"`);
resolved[key] = resolveConfigValueOrThrow(value, `${description} header "${key}"`, env);
}
return Object.keys(resolved).length > 0 ? resolved : undefined;
}

View File

@@ -303,6 +303,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
if (!auth.ok) {
throw new Error(auth.error);
}
const env = auth.env || options?.env ? { ...(auth.env ?? {}), ...(options?.env ?? {}) } : undefined;
const providerRetrySettings = settingsManager.getProviderRetrySettings();
const httpIdleTimeoutMs = settingsManager.getHttpIdleTimeoutMs();
// SDKs treat timeout=0 as 0ms (immediate timeout), not "no timeout".
@@ -314,6 +315,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
return streamSimple(model, context, {
...options,
apiKey: auth.apiKey,
env,
timeoutMs,
websocketConnectTimeoutMs,
maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries,