From dcc336f2d03101d15edc14be4cd1611c093b74bb Mon Sep 17 00:00:00 2001 From: yushen Date: Tue, 3 Feb 2026 16:56:04 +0800 Subject: [PATCH] feat(agent): integrate auth profile rotation into runner and resolver Extend credentials.ts with llm.order and listProfileIdsForProvider. Add resolveApiKeyForProfile and resolveApiKeyForProvider to resolver. Modify runner to support dynamic API key swapping and automatic rotation on auth/rate_limit/billing errors with retry. Co-Authored-By: Claude Opus 4.5 --- src/agent/credentials.ts | 26 +++++ src/agent/providers/index.ts | 2 + src/agent/providers/resolver.ts | 71 ++++++++++++ src/agent/runner.ts | 184 +++++++++++++++++++++++++++++--- src/agent/types.ts | 2 + 5 files changed, 273 insertions(+), 12 deletions(-) diff --git a/src/agent/credentials.ts b/src/agent/credentials.ts index 6e1e3dc5..5f4c7555 100644 --- a/src/agent/credentials.ts +++ b/src/agent/credentials.ts @@ -21,6 +21,8 @@ export type CredentialsConfig = { llm?: { provider?: string | undefined; providers?: Record | undefined; + /** Explicit profile ordering per provider (e.g. { anthropic: ["anthropic", "anthropic:backup"] }) */ + order?: Record | undefined; } | undefined; tools?: Record | undefined; }; @@ -185,6 +187,30 @@ export class CredentialManager { return name in process.env; } + /** + * Get explicit profile order for a provider from credentials.json5 `llm.order`. + * Returns undefined if no explicit order is configured. + */ + getLlmOrder(provider: string): string[] | undefined { + this.loadCore(); + return this.coreConfig?.llm?.order?.[provider]; + } + + /** + * List all profile IDs from `llm.providers` that belong to a given provider. + * A profile matches if its key equals the provider exactly or starts with "provider:". + */ + listProfileIdsForProvider(provider: string): string[] { + this.loadCore(); + const providers = this.coreConfig?.llm?.providers; + if (!providers) return []; + + const prefix = `${provider}:`; + return Object.keys(providers).filter( + (key) => key === provider || key.startsWith(prefix), + ); + } + getResolvedEnvSnapshot(): Record { return { ...this.getResolvedSkillsEnv() }; } diff --git a/src/agent/providers/index.ts b/src/agent/providers/index.ts index 25108916..9cce3f56 100644 --- a/src/agent/providers/index.ts +++ b/src/agent/providers/index.ts @@ -28,6 +28,8 @@ export { type ProviderConfig, resolveProviderConfig, resolveApiKey, + resolveApiKeyForProfile, + resolveApiKeyForProvider, resolveBaseUrl, resolveModelId, resolveModel, diff --git a/src/agent/providers/resolver.ts b/src/agent/providers/resolver.ts index 7ec8dd14..7a18ef88 100644 --- a/src/agent/providers/resolver.ts +++ b/src/agent/providers/resolver.ts @@ -18,6 +18,12 @@ import { isOAuthProvider, } from "./registry.js"; import type { AgentOptions } from "../types.js"; +import { + loadAuthProfileStore, + resolveAuthProfileOrder, + isProfileInCooldown, +} from "../auth-profiles/index.js"; +import type { ResolvedProfileAuth } from "../auth-profiles/index.js"; // ============================================================ // Types @@ -128,6 +134,71 @@ export function resolveModelId(provider: string, explicitModel?: string): string return credentialManager.getLlmProviderConfig(provider)?.model ?? getDefaultModel(provider); } +// ============================================================ +// Profile-aware API Key Resolution +// ============================================================ + +/** + * Resolve API key for a specific auth profile ID. + * Profile IDs follow the convention: "provider" or "provider:label". + */ +export function resolveApiKeyForProfile(profileId: string): string | undefined { + const config = credentialManager.getLlmProviderConfig(profileId); + return config?.apiKey; +} + +/** + * Resolve API key by iterating auth profiles for a provider. + * Returns the first available (non-cooldown) profile with a valid key. + * Falls back to the legacy single-key resolution if no profiles are configured. + */ +export function resolveApiKeyForProvider( + provider: string, + explicitKey?: string, +): ResolvedProfileAuth | undefined { + if (explicitKey) { + return { apiKey: explicitKey, profileId: provider, provider }; + } + + // Try OAuth providers first + const providerConfig = resolveProviderConfig(provider); + if (providerConfig?.apiKey || providerConfig?.accessToken) { + const key = providerConfig.apiKey ?? providerConfig.accessToken; + if (key) return { apiKey: key, profileId: provider, provider }; + } + + // Try auth profiles (multi-key rotation) + const store = loadAuthProfileStore(); + const candidates = resolveAuthProfileOrder(provider, store); + + if (candidates.length > 0) { + for (const profileId of candidates) { + const stats = store.usageStats?.[profileId]; + if (stats && isProfileInCooldown(stats)) continue; + + const apiKey = resolveApiKeyForProfile(profileId); + if (apiKey) { + return { apiKey, profileId, provider }; + } + } + // All in cooldown — return the first one (will be retried when cooldown expires) + for (const profileId of candidates) { + const apiKey = resolveApiKeyForProfile(profileId); + if (apiKey) { + return { apiKey, profileId, provider }; + } + } + } + + // Fall back to single-key credentials.json5 + const fallbackKey = credentialManager.getLlmProviderConfig(provider)?.apiKey; + if (fallbackKey) { + return { apiKey: fallbackKey, profileId: provider, provider }; + } + + return undefined; +} + // ============================================================ // Model Resolution // ============================================================ diff --git a/src/agent/runner.ts b/src/agent/runner.ts index 90f5416a..695a0467 100644 --- a/src/agent/runner.ts +++ b/src/agent/runner.ts @@ -3,7 +3,13 @@ import { v7 as uuidv7 } from "uuid"; import type { AgentOptions, AgentRunResult } from "./types.js"; import { createAgentOutput } from "./cli/output.js"; import { resolveModel, resolveTools } from "./tools.js"; -import { resolveApiKey, resolveBaseUrl, resolveModelId } from "./providers/index.js"; +import { + resolveApiKey, + resolveApiKeyForProfile, + resolveApiKeyForProvider, + resolveBaseUrl, + resolveModelId, +} from "./providers/index.js"; import { SessionManager } from "./session/session-manager.js"; import { ProfileManager } from "./profile/index.js"; import { SkillManager } from "./skills/index.js"; @@ -14,6 +20,42 @@ import { type ContextWindowGuardResult, } from "./context-window/index.js"; import { mergeToolsConfig, type ToolsConfig } from "./tools/policy.js"; +import { + loadAuthProfileStore, + resolveAuthProfileOrder, + isProfileInCooldown, + markAuthProfileFailure, + markAuthProfileUsed, + markAuthProfileGood, +} from "./auth-profiles/index.js"; +import type { AuthProfileFailureReason } from "./auth-profiles/index.js"; + +// ============================================================ +// Error classification for auth profile rotation +// ============================================================ + +function classifyError(error: unknown): AuthProfileFailureReason { + const msg = error instanceof Error ? error.message.toLowerCase() : String(error).toLowerCase(); + + if (msg.includes("401") || msg.includes("403") || msg.includes("unauthorized") || msg.includes("invalid api key") || msg.includes("authentication")) { + return "auth"; + } + if (msg.includes("429") || msg.includes("rate limit") || msg.includes("rate_limit") || msg.includes("too many requests")) { + return "rate_limit"; + } + if (msg.includes("billing") || msg.includes("quota") || msg.includes("insufficient") || msg.includes("payment")) { + return "billing"; + } + if (msg.includes("timeout") || msg.includes("timed out") || msg.includes("econnreset") || msg.includes("etimedout")) { + return "timeout"; + } + return "unknown"; +} + +/** Check if an error is potentially retryable via profile rotation */ +function isRotatableError(reason: AuthProfileFailureReason): boolean { + return reason === "auth" || reason === "rate_limit" || reason === "billing"; +} export class Agent { private readonly agent: PiAgentCore; @@ -23,24 +65,67 @@ export class Agent { private readonly skillManager?: SkillManager; private readonly contextWindowGuard: ContextWindowGuardResult; private readonly debug: boolean; + private readonly stderr: NodeJS.WritableStream; + + // Auth profile rotation state + private readonly resolvedProvider: string; + private currentApiKey: string | undefined; + private currentProfileId: string | undefined; + private profileCandidates: string[]; + private profileIndex: number; + private readonly pinnedProfile: boolean; /** Current session ID */ readonly sessionId: string; constructor(options: AgentOptions = {}) { const stdout = options.logger?.stdout ?? process.stdout; - const stderr = options.logger?.stderr ?? process.stderr; - this.output = createAgentOutput({ stdout, stderr }); + this.stderr = options.logger?.stderr ?? process.stderr; + this.output = createAgentOutput({ stdout, stderr: this.stderr }); this.debug = options.debug ?? false; // Resolve provider and model from options > env vars > defaults - const resolvedProvider = options.provider ?? credentialManager.getLlmProvider() ?? "kimi-coding"; - const resolvedModel = resolveModelId(resolvedProvider, options.model); - const apiKey = resolveApiKey(resolvedProvider, options.apiKey); + this.resolvedProvider = options.provider ?? credentialManager.getLlmProvider() ?? "kimi-coding"; + const resolvedModel = resolveModelId(this.resolvedProvider, options.model); + + // === Auth profile resolution === + this.pinnedProfile = !!(options.authProfileId || options.apiKey); + + if (options.apiKey) { + // Explicit API key — no rotation + this.currentApiKey = options.apiKey; + this.currentProfileId = this.resolvedProvider; + this.profileCandidates = []; + this.profileIndex = 0; + } else if (options.authProfileId) { + // Pinned profile — no rotation + this.currentApiKey = resolveApiKeyForProfile(options.authProfileId) + ?? resolveApiKey(this.resolvedProvider); + this.currentProfileId = options.authProfileId; + this.profileCandidates = []; + this.profileIndex = 0; + } else { + // Profile-aware resolution with rotation support + const resolved = resolveApiKeyForProvider(this.resolvedProvider); + if (resolved) { + this.currentApiKey = resolved.apiKey; + this.currentProfileId = resolved.profileId; + } else { + this.currentApiKey = undefined; + this.currentProfileId = undefined; + } + + // Load full candidate list for rotation + const store = loadAuthProfileStore(); + this.profileCandidates = resolveAuthProfileOrder(this.resolvedProvider, store); + this.profileIndex = this.currentProfileId + ? Math.max(0, this.profileCandidates.indexOf(this.currentProfileId)) + : 0; + } this.agent = new PiAgentCore( - apiKey - ? { getApiKey: (_provider: string) => apiKey } + this.currentApiKey + ? { getApiKey: (_provider: string) => this.currentApiKey! } : {}, ); @@ -86,7 +171,7 @@ export class Agent { return tempSession.getMeta(); })(); - const effectiveProvider = resolvedModel ? resolvedProvider : (options.provider ?? storedMeta?.provider); + const effectiveProvider = resolvedModel ? this.resolvedProvider : (options.provider ?? storedMeta?.provider); const effectiveModel = resolvedModel ?? options.model ?? storedMeta?.model; let model = resolveModel({ ...options, provider: effectiveProvider, model: effectiveModel }); @@ -112,7 +197,7 @@ export class Agent { // 警告：context window 较小 if (this.contextWindowGuard.shouldWarn) { - stderr.write( + this.stderr.write( `[Context Window Guard] WARNING: Low context window: ${this.contextWindowGuard.tokens} tokens (source: ${this.contextWindowGuard.source})\n`, ); } @@ -130,7 +215,7 @@ export class Agent { // 获取 API Key（用于 summary 模式） const summaryApiKey = compactionMode === "summary" - ? resolveApiKey(resolvedProvider, options.apiKey) + ? resolveApiKey(this.resolvedProvider, options.apiKey) : undefined; // 创建 SessionManager（带 context window 配置） @@ -208,6 +293,10 @@ export class Agent { this.output.handleEvent(event); this.handleSessionEvent(event); }); + + if (this.debug && this.currentProfileId) { + console.error(`[debug] Auth profile: ${this.currentProfileId} (pinned=${this.pinnedProfile}, candidates=${this.profileCandidates.length})`); + } } /** Subscribe to raw AgentEvent from the underlying engine */ @@ -217,10 +306,81 @@ export class Agent { async run(prompt: string): Promise { this.output.state.lastAssistantText = ""; - await this.agent.prompt(prompt); + + try { + await this.agent.prompt(prompt); + } catch (error) { + // Attempt auth profile rotation on retryable errors + if (!this.pinnedProfile && this.profileCandidates.length > 1 && this.currentProfileId) { + const reason = classifyError(error); + if (isRotatableError(reason)) { + markAuthProfileFailure(this.currentProfileId, reason); + + if (this.debug) { + this.stderr.write( + `[auth-profile] Profile "${this.currentProfileId}" failed (${reason}), attempting rotation...\n`, + ); + } + + if (this.advanceAuthProfile()) { + if (this.debug) { + this.stderr.write( + `[auth-profile] Rotated to profile "${this.currentProfileId}"\n`, + ); + } + // Retry with new profile + this.output.state.lastAssistantText = ""; + await this.agent.prompt(prompt); + } else { + throw error; // No more profiles to try + } + } else { + throw error; // Non-rotatable error + } + } else { + throw error; // Pinned profile or single profile + } + } + + // Mark success + if (this.currentProfileId) { + markAuthProfileUsed(this.currentProfileId); + markAuthProfileGood(this.resolvedProvider, this.currentProfileId); + } + return { text: this.output.state.lastAssistantText, error: this.agent.state.error }; } + /** + * Advance to the next non-cooldown auth profile. + * Returns true if a new profile was activated, false if exhausted. + */ + private advanceAuthProfile(): boolean { + const store = loadAuthProfileStore(); + const startIndex = this.profileIndex; + + for (let i = 1; i < this.profileCandidates.length; i++) { + const nextIndex = (startIndex + i) % this.profileCandidates.length; + const candidateId = this.profileCandidates[nextIndex] as string | undefined; + if (!candidateId) continue; + + // Skip profiles in cooldown + const stats = store.usageStats?.[candidateId]; + if (stats && isProfileInCooldown(stats)) continue; + + // Try to resolve API key + const apiKey = resolveApiKeyForProfile(candidateId); + if (!apiKey) continue; + + this.currentApiKey = apiKey; + this.currentProfileId = candidateId; + this.profileIndex = nextIndex; + return true; + } + + return false; + } + private handleSessionEvent(event: AgentEvent) { if (event.type === "message_end") { const message = event.message as AgentMessage; diff --git a/src/agent/types.ts b/src/agent/types.ts index 75e53ad1..c7e37658 100644 --- a/src/agent/types.ts +++ b/src/agent/types.ts @@ -21,6 +21,8 @@ export type AgentOptions = { model?: string | undefined; /** Custom API key (overrides environment variable) */ apiKey?: string | undefined; + /** Pin a specific auth profile ID (e.g. "anthropic:backup"). Disables rotation. */ + authProfileId?: string | undefined; /** Custom base URL for the provider endpoint */ baseUrl?: string | undefined; /** System prompt, if profileId is set will auto-construct from profile */