fix(agents): reset rate-limit retry budget after recovery

lanzhi-lee · lanzhi-lee · commit 0841b694cac7 · 2026-06-10T19:23:15.000+08:00
diff --git a/src/agents/embedded-agent-runner/run.ts b/src/agents/embedded-agent-runner/run.ts
@@ -156,6 +156,7 @@ import {
   resolveOverloadFailoverBackoffMs,
   resolveOverloadProfileRotationLimit,
   resolveRateLimitProfileRotationLimit,
+  resolveNextSameModelRateLimitRetryCount,
   resolveSameModelRateLimitBackoffMs,
   type RuntimeAuthState,
   scrubAnthropicRefusalMagic,
@@ -1214,7 +1215,7 @@ export async function runEmbeddedAgent(
       let lastContextBudgetStatus: EmbeddedAgentMeta["contextBudgetStatus"];
       let runLoopIterations = 0;
       let overloadProfileRotations = 0;
-      let sameModelRateLimitRetries = 0;
+      let consecutiveSameModelRateLimitRetries = 0;
       let planningOnlyRetryAttempts = 0;
       let reasoningOnlyRetryAttempts = 0;
       let emptyResponseRetryAttempts = 0;
@@ -1376,12 +1377,12 @@ export async function runEmbeddedAgent(
         }
       };
       const maybeRetrySameModelRateLimit = async (): Promise<boolean> => {
-        if (sameModelRateLimitRetries >= MAX_SAME_MODEL_RATE_LIMIT_RETRIES) {
+        if (consecutiveSameModelRateLimitRetries >= MAX_SAME_MODEL_RATE_LIMIT_RETRIES) {
           return false;
         }
-        const delayMs = resolveSameModelRateLimitBackoffMs(sameModelRateLimitRetries);
+        const delayMs = resolveSameModelRateLimitBackoffMs(consecutiveSameModelRateLimitRetries);
         log.warn(
-          `rate-limit same-model retry ${sameModelRateLimitRetries + 1}/${MAX_SAME_MODEL_RATE_LIMIT_RETRIES} for ${sanitizeForLog(provider)}/${sanitizeForLog(modelId)}: delayMs=${delayMs}`,
+          `rate-limit same-model retry ${consecutiveSameModelRateLimitRetries + 1}/${MAX_SAME_MODEL_RATE_LIMIT_RETRIES} for ${sanitizeForLog(provider)}/${sanitizeForLog(modelId)}: delayMs=${delayMs}`,
         );
         try {
           await sleepWithAbort(delayMs, params.abortSignal);
@@ -1393,7 +1394,10 @@ export async function runEmbeddedAgent(
           }
           throw err;
         }
-        sameModelRateLimitRetries += 1;
+        consecutiveSameModelRateLimitRetries = resolveNextSameModelRateLimitRetryCount({
+          retriesSoFar: consecutiveSameModelRateLimitRetries,
+          retriedSameModelRateLimit: true,
+        });
         return true;
       };
       // Resolve the context engine once and reuse across retries to avoid
@@ -2926,9 +2930,19 @@ export async function runEmbeddedAgent(
             if (assistantFailoverOutcome.retryKind === "same_model_idle_timeout") {
               sameModelIdleTimeoutRetries += 1;
             }
+            if (assistantFailoverOutcome.retryKind !== "same_model_rate_limit") {
+              consecutiveSameModelRateLimitRetries = resolveNextSameModelRateLimitRetryCount({
+                retriesSoFar: consecutiveSameModelRateLimitRetries,
+                retriedSameModelRateLimit: false,
+              });
+            }
             lastRetryFailoverReason = assistantFailoverOutcome.lastRetryFailoverReason;
             continue;
           }
+          consecutiveSameModelRateLimitRetries = resolveNextSameModelRateLimitRetryCount({
+            retriesSoFar: consecutiveSameModelRateLimitRetries,
+            retriedSameModelRateLimit: false,
+          });
           if (assistantFailoverOutcome.action === "throw") {
             traceAttempts.push({
               provider: activeErrorContext.provider,
diff --git a/src/agents/embedded-agent-runner/run/helpers.test.ts b/src/agents/embedded-agent-runner/run/helpers.test.ts
@@ -7,6 +7,7 @@ import {
   buildErrorAgentMeta,
   resolveFinalAssistantRawText,
   resolveFinalAssistantVisibleText,
+  resolveNextSameModelRateLimitRetryCount,
   resolveSameModelRateLimitBackoffMs,
 } from "./helpers.js";
 
@@ -102,6 +103,34 @@ describe("resolveSameModelRateLimitBackoffMs", () => {
   });
 });
 
+describe("resolveNextSameModelRateLimitRetryCount", () => {
+  it("counts only consecutive same-model rate-limit retries", () => {
+    let retriesSoFar = 0;
+
+    retriesSoFar = resolveNextSameModelRateLimitRetryCount({
+      retriesSoFar,
+      retriedSameModelRateLimit: true,
+    });
+    retriesSoFar = resolveNextSameModelRateLimitRetryCount({
+      retriesSoFar,
+      retriedSameModelRateLimit: true,
+    });
+    expect(retriesSoFar).toBe(2);
+
+    retriesSoFar = resolveNextSameModelRateLimitRetryCount({
+      retriesSoFar,
+      retriedSameModelRateLimit: false,
+    });
+    expect(retriesSoFar).toBe(0);
+
+    retriesSoFar = resolveNextSameModelRateLimitRetryCount({
+      retriesSoFar,
+      retriedSameModelRateLimit: true,
+    });
+    expect(retriesSoFar).toBe(1);
+  });
+});
+
 describe("buildErrorAgentMeta", () => {
   it("preserves active session file for error exits after transcript rotation", () => {
     // Error metadata follows the active session after transcript rotation so
diff --git a/src/agents/embedded-agent-runner/run/helpers.ts b/src/agents/embedded-agent-runner/run/helpers.ts
@@ -70,6 +70,13 @@ export function resolveSameModelRateLimitBackoffMs(retriesSoFar: number): number
   return Math.min(SAME_MODEL_RATE_LIMIT_MAX_BACKOFF_MS, delay);
 }
 
+export function resolveNextSameModelRateLimitRetryCount(params: {
+  retriesSoFar: number;
+  retriedSameModelRateLimit: boolean;
+}): number {
+  return params.retriedSameModelRateLimit ? Math.max(0, params.retriesSoFar) + 1 : 0;
+}
+
 const ANTHROPIC_MAGIC_STRING_TRIGGER_REFUSAL = "ANTHROPIC_MAGIC_STRING_TRIGGER_REFUSAL";
 const ANTHROPIC_MAGIC_STRING_REPLACEMENT = "ANTHROPIC MAGIC STRING TRIGGER REFUSAL (redacted)";
 
