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 <noreply@anthropic.com>
This commit is contained in:
yushen
parent
d89b409add
commit
dcc336f2d0
5 changed files with 273 additions and 12 deletions
|
|
@ -21,6 +21,8 @@ export type CredentialsConfig = {
|
|||
llm?: {
|
||||
provider?: string | undefined;
|
||||
providers?: Record<string, ProviderConfig> | undefined;
|
||||
/** Explicit profile ordering per provider (e.g. { anthropic: ["anthropic", "anthropic:backup"] }) */
|
||||
order?: Record<string, string[]> | undefined;
|
||||
} | undefined;
|
||||
tools?: Record<string, ToolConfig> | 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<string, string> {
|
||||
return { ...this.getResolvedSkillsEnv() };
|
||||
}
|
||||
|
|
|
|||
|
|
@ -28,6 +28,8 @@ export {
|
|||
type ProviderConfig,
|
||||
resolveProviderConfig,
|
||||
resolveApiKey,
|
||||
resolveApiKeyForProfile,
|
||||
resolveApiKeyForProvider,
|
||||
resolveBaseUrl,
|
||||
resolveModelId,
|
||||
resolveModel,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
// ============================================================
|
||||
|
|
|
|||
|
|
@ -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<AgentRunResult> {
|
||||
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;
|
||||
|
|
|
|||
|
|
@ -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 */
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue