Merge pull request #269 from mukhtharcm/feat/multi-account-roundrobin

feat: Multi-account OAuth with round-robin rotation

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<string, AuthProfileCredential>;
   lastGood?: Record<string, string>;
+  /** Usage statistics per profile for round-robin rotation */
+  usageStats?: Record<string, ProfileUsageStats>;
 };
 
 type LegacyAuthStore = Record<string, AuthProfileCredential>;
@@ -183,6 +192,10 @@ function coerceAuthStore(raw: unknown): AuthProfileStore | null {
       record.lastGood && typeof record.lastGood === "object"
         ? (record.lastGood as Record<string, string>)
         : undefined,
+    usageStats:
+      record.usageStats && typeof record.usageStats === "object"
+        ? (record.usageStats as Record<string, ProfileUsageStats>)
+        : 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: {

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,

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",
         });

src/commands/onboard-auth.ts

@@ -7,7 +7,7 @@ export async function writeOAuthCredentials(
   creds: OAuthCredentials,
 ): Promise<void> {
   upsertAuthProfile({
-    profileId: `${provider}:default`,
+    profileId: `${provider}:${creds.email ?? "default"}`,
     credential: {
       type: "oauth",
       provider,

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",
         });