fix: improve auth profile failover

src/agents/auth-profiles.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest";
 
 import {
   type AuthProfileStore,
+  calculateAuthProfileCooldownMs,
   resolveAuthProfileOrder,
 } from "./auth-profiles.js";
 
@@ -105,4 +106,87 @@ describe("resolveAuthProfileOrder", () => {
     });
     expect(order).toEqual(["anthropic:oauth", "anthropic:default"]);
   });
+
+  it("orders by lastUsed when no explicit order exists", () => {
+    const order = resolveAuthProfileOrder({
+      store: {
+        version: 1,
+        profiles: {
+          "anthropic:a": {
+            type: "oauth",
+            provider: "anthropic",
+            access: "access-token",
+            refresh: "refresh-token",
+            expires: Date.now() + 60_000,
+          },
+          "anthropic:b": {
+            type: "api_key",
+            provider: "anthropic",
+            key: "sk-b",
+          },
+          "anthropic:c": {
+            type: "api_key",
+            provider: "anthropic",
+            key: "sk-c",
+          },
+        },
+        usageStats: {
+          "anthropic:a": { lastUsed: 200 },
+          "anthropic:b": { lastUsed: 100 },
+          "anthropic:c": { lastUsed: 300 },
+        },
+      },
+      provider: "anthropic",
+    });
+    expect(order).toEqual(["anthropic:b", "anthropic:a", "anthropic:c"]);
+  });
+
+  it("pushes cooldown profiles to the end, ordered by cooldown expiry", () => {
+    const now = Date.now();
+    const order = resolveAuthProfileOrder({
+      store: {
+        version: 1,
+        profiles: {
+          "anthropic:ready": {
+            type: "api_key",
+            provider: "anthropic",
+            key: "sk-ready",
+          },
+          "anthropic:cool1": {
+            type: "oauth",
+            provider: "anthropic",
+            access: "access-token",
+            refresh: "refresh-token",
+            expires: now + 60_000,
+          },
+          "anthropic:cool2": {
+            type: "api_key",
+            provider: "anthropic",
+            key: "sk-cool",
+          },
+        },
+        usageStats: {
+          "anthropic:ready": { lastUsed: 50 },
+          "anthropic:cool1": { cooldownUntil: now + 5_000 },
+          "anthropic:cool2": { cooldownUntil: now + 1_000 },
+        },
+      },
+      provider: "anthropic",
+    });
+    expect(order).toEqual([
+      "anthropic:ready",
+      "anthropic:cool2",
+      "anthropic:cool1",
+    ]);
+  });
+});
+
+describe("auth profile cooldowns", () => {
+  it("applies exponential backoff with a 1h cap", () => {
+    expect(calculateAuthProfileCooldownMs(1)).toBe(60_000);
+    expect(calculateAuthProfileCooldownMs(2)).toBe(5 * 60_000);
+    expect(calculateAuthProfileCooldownMs(3)).toBe(25 * 60_000);
+    expect(calculateAuthProfileCooldownMs(4)).toBe(60 * 60_000);
+    expect(calculateAuthProfileCooldownMs(5)).toBe(60 * 60_000);
+  });
 });

src/agents/auth-profiles.ts

@@ -368,6 +368,14 @@ export function markAuthProfileUsed(params: {
   saveAuthProfileStore(store);
 }
 
+export function calculateAuthProfileCooldownMs(errorCount: number): number {
+  const normalized = Math.max(1, errorCount);
+  return Math.min(
+    60 * 60 * 1000, // 1 hour max
+    60 * 1000 * 5 ** Math.min(normalized - 1, 3),
+  );
+}
+
 /**
  * Mark a profile as failed/rate-limited. Applies exponential backoff cooldown.
  * Cooldown times: 1min, 5min, 25min, max 1 hour.
@@ -384,10 +392,7 @@ export function markAuthProfileCooldown(params: {
   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)),
-  );
+  const backoffMs = calculateAuthProfileCooldownMs(errorCount);
 
   store.usageStats[profileId] = {
     ...existing,

src/agents/pi-embedded-runner.ts

@@ -24,7 +24,11 @@ import {
 } from "../process/command-queue.js";
 import { resolveUserPath } from "../utils.js";
 import { resolveClawdbotAgentDir } from "./agent-paths.js";
-import { markAuthProfileGood, markAuthProfileUsed, markAuthProfileCooldown } from "./auth-profiles.js";
+import {
+  markAuthProfileCooldown,
+  markAuthProfileGood,
+  markAuthProfileUsed,
+} from "./auth-profiles.js";
 import type { BashElevatedDefaults } from "./bash-tools.js";
 import { DEFAULT_MODEL, DEFAULT_PROVIDER } from "./defaults.js";
 import {
@@ -955,14 +959,18 @@ export async function runEmbeddedPiAgent(params: {
             (params.config?.agent?.model?.fallbacks?.length ?? 0) > 0;
           const authFailure = isAuthAssistantError(lastAssistant);
           const rateLimitFailure = isRateLimitAssistantError(lastAssistant);
-          
+
           // Treat timeout as potential rate limit (Antigravity hangs on rate limit)
-          const shouldRotate = (!aborted && (authFailure || rateLimitFailure)) || timedOut;
-          
+          const shouldRotate =
+            (!aborted && (authFailure || rateLimitFailure)) || timedOut;
+
           if (shouldRotate) {
             // Mark current profile for cooldown before rotating
             if (lastProfileId) {
-              markAuthProfileCooldown({ store: authStore, profileId: lastProfileId });
+              markAuthProfileCooldown({
+                store: authStore,
+                profileId: lastProfileId,
+              });
               if (timedOut) {
                 log.warn(
                   `Profile ${lastProfileId} timed out (possible rate limit). Trying next account...`,
@@ -973,15 +981,17 @@ export async function runEmbeddedPiAgent(params: {
             if (rotated) {
               continue;
             }
-            if (fallbackConfigured && !timedOut) {
+            if (fallbackConfigured) {
               const message =
                 lastAssistant?.errorMessage?.trim() ||
                 (lastAssistant
                   ? formatAssistantErrorText(lastAssistant)
                   : "") ||
-                (rateLimitFailure
-                  ? "LLM request rate limited."
-                  : "LLM request unauthorized.");
+                (timedOut
+                  ? "LLM request timed out."
+                  : rateLimitFailure
+                    ? "LLM request rate limited."
+                    : "LLM request unauthorized.");
               throw new Error(message);
             }
           }

src/gateway/protocol/index.ts

@@ -394,4 +394,5 @@ export type {
   CronRunParams,
   CronRunsParams,
   CronRunLogEntry,
+  PollParams,
 };