From 18c7795ee06b2c1441a73e4cc90c90902b34eb91 Mon Sep 17 00:00:00 2001
From: Muhammed Mukhthar CM <mukhtharcm@gmail.com>
Date: Tue, 6 Jan 2026 04:48:34 +0000
Subject: [PATCH] feat: treat timeout as rate limit for profile rotation

Antigravity rate limits cause requests to hang indefinitely rather than
returning 429 errors. This change detects timeouts and treats them as
potential rate limits:

- Added timedOut flag to track timeout-triggered aborts
- Timeout now triggers profile cooldown + rotation
- Logs: "Profile X timed out (possible rate limit). Trying next account..."

This ensures automatic failover when Antigravity hangs due to rate limiting.
---
 src/agents/pi-embedded-runner.ts | 19 +++++++++++++++----
 src/gateway/protocol/index.ts    |  1 -
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/agents/pi-embedded-runner.ts b/src/agents/pi-embedded-runner.ts
index 079b43eba..eb2bb78f9 100644
--- a/src/agents/pi-embedded-runner.ts
+++ b/src/agents/pi-embedded-runner.ts
@@ -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,16 +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
diff --git a/src/gateway/protocol/index.ts b/src/gateway/protocol/index.ts
index 6dad0d7e3..6fc3d7ecd 100644
--- a/src/gateway/protocol/index.ts
+++ b/src/gateway/protocol/index.ts
@@ -395,5 +395,4 @@ export type {
   CronRunParams,
   CronRunsParams,
   CronRunLogEntry,
-  PollParams,
 };