diff --git a/src/agents/auth-profiles.ts b/src/agents/auth-profiles.ts index f390061dc..8fa3080d5 100644 --- a/src/agents/auth-profiles.ts +++ b/src/agents/auth-profiles.ts @@ -32,10 +32,19 @@ export type OAuthCredential = OAuthCredentials & { export type AuthProfileCredential = ApiKeyCredential | OAuthCredential; +/** Per-profile usage statistics for round-robin and cooldown tracking */ +export type ProfileUsageStats = { + lastUsed?: number; + cooldownUntil?: number; + errorCount?: number; +}; + export type AuthProfileStore = { version: number; profiles: Record; lastGood?: Record; + /** Usage statistics per profile for round-robin rotation */ + usageStats?: Record; }; type LegacyAuthStore = Record; @@ -183,6 +192,10 @@ function coerceAuthStore(raw: unknown): AuthProfileStore | null { record.lastGood && typeof record.lastGood === "object" ? (record.lastGood as Record) : undefined, + usageStats: + record.usageStats && typeof record.usageStats === "object" + ? (record.usageStats as Record) + : undefined, }; } @@ -300,6 +313,7 @@ export function saveAuthProfileStore(store: AuthProfileStore): void { version: AUTH_STORE_VERSION, profiles: store.profiles, lastGood: store.lastGood ?? undefined, + usageStats: store.usageStats ?? undefined, } satisfies AuthProfileStore; saveJsonFile(authPath, payload); } @@ -322,6 +336,85 @@ export function listProfilesForProvider( .map(([id]) => id); } +/** + * Check if a profile is currently in cooldown (due to rate limiting or errors). + */ +export function isProfileInCooldown( + store: AuthProfileStore, + profileId: string, +): boolean { + const stats = store.usageStats?.[profileId]; + if (!stats?.cooldownUntil) return false; + return Date.now() < stats.cooldownUntil; +} + +/** + * Mark a profile as successfully used. Resets error count and updates lastUsed. + */ +export function markAuthProfileUsed(params: { + store: AuthProfileStore; + profileId: string; +}): void { + const { store, profileId } = params; + if (!store.profiles[profileId]) return; + + store.usageStats = store.usageStats ?? {}; + store.usageStats[profileId] = { + ...store.usageStats[profileId], + lastUsed: Date.now(), + errorCount: 0, + cooldownUntil: undefined, + }; + saveAuthProfileStore(store); +} + +/** + * Mark a profile as failed/rate-limited. Applies exponential backoff cooldown. + * Cooldown times: 1min, 5min, 25min, max 1 hour. + */ +export function markAuthProfileCooldown(params: { + store: AuthProfileStore; + profileId: string; +}): void { + const { store, profileId } = params; + if (!store.profiles[profileId]) return; + + store.usageStats = store.usageStats ?? {}; + const existing = store.usageStats[profileId] ?? {}; + const errorCount = (existing.errorCount ?? 0) + 1; + + // Exponential backoff: 1min, 5min, 25min, capped at 1h + const backoffMs = Math.min( + 60 * 60 * 1000, // 1 hour max + 60 * 1000 * Math.pow(5, Math.min(errorCount - 1, 3)), + ); + + store.usageStats[profileId] = { + ...existing, + errorCount, + cooldownUntil: Date.now() + backoffMs, + }; + saveAuthProfileStore(store); +} + +/** + * Clear cooldown for a profile (e.g., manual reset). + */ +export function clearAuthProfileCooldown(params: { + store: AuthProfileStore; + profileId: string; +}): void { + const { store, profileId } = params; + if (!store.usageStats?.[profileId]) return; + + store.usageStats[profileId] = { + ...store.usageStats[profileId], + errorCount: 0, + cooldownUntil: undefined, + }; + saveAuthProfileStore(store); +} + export function resolveAuthProfileOrder(params: { cfg?: ClawdbotConfig; store: AuthProfileStore; @@ -376,14 +469,50 @@ function orderProfilesByMode( order: string[], store: AuthProfileStore, ): string[] { - const scored = order.map((profileId) => { + const now = Date.now(); + + // Partition into available and in-cooldown + const available: string[] = []; + const inCooldown: string[] = []; + + for (const profileId of order) { + if (isProfileInCooldown(store, profileId)) { + inCooldown.push(profileId); + } else { + available.push(profileId); + } + } + + // Sort available profiles by lastUsed (oldest first = round-robin) + // Then by type (oauth preferred over api_key) + const scored = available.map((profileId) => { const type = store.profiles[profileId]?.type; - const score = type === "oauth" ? 0 : type === "api_key" ? 1 : 2; - return { profileId, score }; + const typeScore = type === "oauth" ? 0 : type === "api_key" ? 1 : 2; + const lastUsed = store.usageStats?.[profileId]?.lastUsed ?? 0; + return { profileId, typeScore, lastUsed }; }); - return scored - .sort((a, b) => a.score - b.score) + + // Primary sort: lastUsed (oldest first for round-robin) + // Secondary sort: type preference (oauth > api_key) + const sorted = scored + .sort((a, b) => { + // First by lastUsed (oldest first) + if (a.lastUsed !== b.lastUsed) return a.lastUsed - b.lastUsed; + // Then by type + return a.typeScore - b.typeScore; + }) .map((entry) => entry.profileId); + + // Append cooldown profiles at the end (sorted by cooldown expiry, soonest first) + const cooldownSorted = inCooldown + .map((profileId) => ({ + profileId, + cooldownUntil: store.usageStats?.[profileId]?.cooldownUntil ?? now, + })) + .sort((a, b) => a.cooldownUntil - b.cooldownUntil) + .map((entry) => entry.profileId); + + return [...sorted, ...cooldownSorted]; } export async function resolveApiKeyForProfile(params: { diff --git a/src/agents/pi-embedded-runner.ts b/src/agents/pi-embedded-runner.ts index 5666e85ea..eb2bb78f9 100644 --- a/src/agents/pi-embedded-runner.ts +++ b/src/agents/pi-embedded-runner.ts @@ -24,7 +24,7 @@ import { } from "../process/command-queue.js"; import { resolveUserPath } from "../utils.js"; import { resolveClawdbotAgentDir } from "./agent-paths.js"; -import { markAuthProfileGood } from "./auth-profiles.js"; +import { markAuthProfileGood, markAuthProfileUsed, markAuthProfileCooldown } from "./auth-profiles.js"; import type { BashElevatedDefaults } from "./bash-tools.js"; import { DEFAULT_MODEL, DEFAULT_PROVIDER } from "./defaults.js"; import { @@ -808,8 +808,10 @@ export async function runEmbeddedPiAgent(params: { session.agent.replaceMessages(prior); } let aborted = Boolean(params.abortSignal?.aborted); - const abortRun = () => { + let timedOut = false; + const abortRun = (isTimeout = false) => { aborted = true; + if (isTimeout) timedOut = true; void session.abort(); }; const subscription = subscribeEmbeddedPiSession({ @@ -848,7 +850,7 @@ export async function runEmbeddedPiAgent(params: { log.warn( `embedded run timeout: runId=${params.runId} sessionId=${params.sessionId} timeoutMs=${params.timeoutMs}`, ); - abortRun(); + abortRun(true); if (!abortWarnTimer) { abortWarnTimer = setTimeout(() => { if (!session.isStreaming) return; @@ -953,12 +955,25 @@ export async function runEmbeddedPiAgent(params: { (params.config?.agent?.model?.fallbacks?.length ?? 0) > 0; const authFailure = isAuthAssistantError(lastAssistant); const rateLimitFailure = isRateLimitAssistantError(lastAssistant); - if (!aborted && (authFailure || rateLimitFailure)) { + + // Treat timeout as potential rate limit (Antigravity hangs on rate limit) + const shouldRotate = (!aborted && (authFailure || rateLimitFailure)) || timedOut; + + if (shouldRotate) { + // Mark current profile for cooldown before rotating + if (lastProfileId) { + markAuthProfileCooldown({ store: authStore, profileId: lastProfileId }); + if (timedOut) { + log.warn( + `Profile ${lastProfileId} timed out (possible rate limit). Trying next account...`, + ); + } + } const rotated = await advanceAuthProfile(); if (rotated) { continue; } - if (fallbackConfigured) { + if (fallbackConfigured && !timedOut) { const message = lastAssistant?.errorMessage?.trim() || (lastAssistant @@ -1040,6 +1055,8 @@ export async function runEmbeddedPiAgent(params: { provider, profileId: lastProfileId, }); + // Track usage for round-robin rotation + markAuthProfileUsed({ store: authStore, profileId: lastProfileId }); } return { payloads: payloads.length ? payloads : undefined, diff --git a/src/commands/configure.ts b/src/commands/configure.ts index e2b8b6451..d65908f5a 100644 --- a/src/commands/configure.ts +++ b/src/commands/configure.ts @@ -323,7 +323,7 @@ async function promptAuthConfig( if (oauthCreds) { await writeOAuthCredentials("google-antigravity", oauthCreds); next = applyAuthProfileConfig(next, { - profileId: "google-antigravity:default", + profileId: `google-antigravity:${oauthCreds.email ?? "default"}`, provider: "google-antigravity", mode: "oauth", }); diff --git a/src/commands/onboard-auth.ts b/src/commands/onboard-auth.ts index 3da496b34..8151bfca3 100644 --- a/src/commands/onboard-auth.ts +++ b/src/commands/onboard-auth.ts @@ -7,7 +7,7 @@ export async function writeOAuthCredentials( creds: OAuthCredentials, ): Promise { upsertAuthProfile({ - profileId: `${provider}:default`, + profileId: `${provider}:${creds.email ?? "default"}`, credential: { type: "oauth", provider, diff --git a/src/wizard/onboarding.ts b/src/wizard/onboarding.ts index 35af38387..d157c7cf3 100644 --- a/src/wizard/onboarding.ts +++ b/src/wizard/onboarding.ts @@ -448,7 +448,7 @@ export async function runOnboardingWizard( if (oauthCreds) { await writeOAuthCredentials("google-antigravity", oauthCreds); nextConfig = applyAuthProfileConfig(nextConfig, { - profileId: "google-antigravity:default", + profileId: `google-antigravity:${oauthCreds.email ?? "default"}`, provider: "google-antigravity", mode: "oauth", });