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 */