diff --git a/src/agents/embedded-agent-runner/run-loop.ts b/src/agents/embedded-agent-runner/run-loop.ts index d2a614f03f8e..16e9cbe4128a 100644 --- a/src/agents/embedded-agent-runner/run-loop.ts +++ b/src/agents/embedded-agent-runner/run-loop.ts @@ -115,7 +115,6 @@ export async function runPreparedEmbeddedLoop( profileFailureStore, pluginHarnessOwnsAuthBootstrap, attemptedThinking, - advanceAttemptAuthProfile, maybeRefreshRuntimeAuthForAuthError, stopRuntimeAuthRefreshTimer, getApiKeyInfo, @@ -266,6 +265,7 @@ export async function runPreparedEmbeddedLoop( harnessOwnsTransport: () => preparedRuntime.snapshot().pluginHarnessOwnsTransport, getRuntimeAuthOwnerId: () => preparedRuntime.snapshot().agentHarness.id, getApiKeyInfo, + advanceAuthProfile: preparedRuntime.advanceAttemptAuthProfile, }); const ownsContextEngineLogicalTurnLease = params.contextEngineLogicalTurnLease === undefined; const contextEngineLogicalTurnLease = @@ -489,17 +489,14 @@ export async function runPreparedEmbeddedLoop( emptyErrorRetries, overloadProfileRotations, overloadProfileRotationLimit: failoverRetryController.overloadProfileRotationLimit, - rateLimitProfileRotations: failoverRetryController.rateLimitProfileRotations, - rateLimitProfileRotationLimit: failoverRetryController.rateLimitProfileRotationLimit, sameModelIdleTimeoutRetries, previousRetryFailoverReason: lastRetryFailoverReason, maybeMarkAuthProfileFailure: failoverRetryController.maybeMarkAuthProfileFailure, - maybeEscalateRateLimitProfileFallback: - failoverRetryController.maybeEscalateRateLimitProfileFallback, maybeRetrySameModelRateLimit: failoverRetryController.maybeRetrySameModelRateLimit, maybeBackoffBeforeOverloadFailover: failoverRetryController.maybeBackoffBeforeOverloadFailover, - advanceAttemptAuthProfile, + advanceAuthProfile: failoverRetryController.advanceAuthProfile, + advanceRateLimitAuthProfile: failoverRetryController.advanceRateLimitAuthProfile, traceAttempts, suspendForFailure, suspensionSessionId: sessionPromptState.sessionId ?? params.sessionId, diff --git a/src/agents/embedded-agent-runner/run.empty-error-retry.owner-test-support.ts b/src/agents/embedded-agent-runner/run.empty-error-retry.owner-test-support.ts index a75e20f6c87e..73d4d32ee73f 100644 --- a/src/agents/embedded-agent-runner/run.empty-error-retry.owner-test-support.ts +++ b/src/agents/embedded-agent-runner/run.empty-error-retry.owner-test-support.ts @@ -70,15 +70,13 @@ function makeInput( emptyErrorRetries: options.emptyErrorRetries ?? 0, overloadProfileRotations: 0, overloadProfileRotationLimit: 1, - rateLimitProfileRotations: 0, - rateLimitProfileRotationLimit: 1, sameModelIdleTimeoutRetries: 0, previousRetryFailoverReason: null, maybeMarkAuthProfileFailure: vi.fn(async () => {}), - maybeEscalateRateLimitProfileFallback: vi.fn(), maybeRetrySameModelRateLimit: vi.fn(async () => false), maybeBackoffBeforeOverloadFailover: vi.fn(async () => {}), - advanceAttemptAuthProfile: vi.fn(async () => false), + advanceAuthProfile: vi.fn(async () => false), + advanceRateLimitAuthProfile: vi.fn(async () => false), traceAttempts: [], suspendForFailure: vi.fn(), suspensionSessionId: "session:empty-error", diff --git a/src/agents/embedded-agent-runner/run/assistant-failover.test.ts b/src/agents/embedded-agent-runner/run/assistant-failover.test.ts index 6897db0df0be..bd163e1fff2e 100644 --- a/src/agents/embedded-agent-runner/run/assistant-failover.test.ts +++ b/src/agents/embedded-agent-runner/run/assistant-failover.test.ts @@ -21,7 +21,6 @@ function makeParams(overrides: Partial = {}): Params { failoverReason: "billing", harnessOwnsTransport: false, allowSameModelIdleTimeoutRetry: false, - allowSameModelRateLimitRetry: true, assistantProfileFailureReason: null, lastProfileId: undefined, modelId: model, @@ -41,10 +40,10 @@ function makeParams(overrides: Partial = {}): Params { logAssistantFailoverDecision: vi.fn(), warn: vi.fn(), maybeMarkAuthProfileFailure: vi.fn(async () => {}), - maybeEscalateRateLimitProfileFallback: vi.fn(), maybeRetrySameModelRateLimit: vi.fn(async () => false), maybeBackoffBeforeOverloadFailover: vi.fn(async () => {}), advanceAuthProfile: vi.fn(async () => false), + advanceRateLimitAuthProfile: vi.fn(async () => false), }; return { ...defaults, ...overrides }; } @@ -88,7 +87,7 @@ describe("handleAssistantFailover", () => { billingFailure: false, rateLimitFailure: true, maybeMarkAuthProfileFailure, - advanceAuthProfile: vi.fn(async () => { + advanceRateLimitAuthProfile: vi.fn(async () => { events.push("advance"); return true; }), @@ -108,8 +107,7 @@ describe("handleAssistantFailover", () => { it("retries the same model before spending a rate-limit profile rotation", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -121,8 +119,7 @@ describe("handleAssistantFailover", () => { errorMessage: "HTTP 429 Too Many Requests: requests per minute exceeded", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -133,14 +130,12 @@ describe("handleAssistantFailover", () => { expect(outcome.retryKind).toBe("same_model_rate_limit"); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledTimes(1); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledWith({}); - expect(maybeEscalateRateLimitProfileFallback).not.toHaveBeenCalled(); - expect(advanceAuthProfile).not.toHaveBeenCalled(); + expect(advanceRateLimitAuthProfile).not.toHaveBeenCalled(); }); - it("honors disabled rate-limit profile rotations before same-model retry", async () => { - const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + it("rotates when the rate-limit controller denies a same-model retry", async () => { + const maybeRetrySameModelRateLimit = vi.fn(async () => false); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -148,13 +143,11 @@ describe("handleAssistantFailover", () => { failoverReason: "rate_limit", billingFailure: false, rateLimitFailure: true, - allowSameModelRateLimitRetry: false, lastAssistant: { errorMessage: "HTTP 429 Too Many Requests: requests per minute exceeded", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -163,15 +156,13 @@ describe("handleAssistantFailover", () => { return; } expect(outcome.retryKind).toBe("profile_rotation"); - expect(maybeRetrySameModelRateLimit).not.toHaveBeenCalled(); - expect(maybeEscalateRateLimitProfileFallback).toHaveBeenCalledTimes(1); - expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(maybeRetrySameModelRateLimit).toHaveBeenCalledTimes(1); + expect(advanceRateLimitAuthProfile).toHaveBeenCalledTimes(1); }); it("does not spend same-model retry budget on quota-style rate limits", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -184,8 +175,7 @@ describe("handleAssistantFailover", () => { "You exceeded your current quota, please check your plan and billing details.", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -195,14 +185,12 @@ describe("handleAssistantFailover", () => { } expect(outcome.retryKind).toBe("profile_rotation"); expect(maybeRetrySameModelRateLimit).not.toHaveBeenCalled(); - expect(maybeEscalateRateLimitProfileFallback).toHaveBeenCalledTimes(1); - expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(advanceRateLimitAuthProfile).toHaveBeenCalledTimes(1); }); it("does not treat bare 429 quota_exceeded as a short-window throttle", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -214,8 +202,7 @@ describe("handleAssistantFailover", () => { errorMessage: "Provider API error (429): Quota exceeded [code=quota_exceeded]", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -225,14 +212,12 @@ describe("handleAssistantFailover", () => { } expect(outcome.retryKind).toBe("profile_rotation"); expect(maybeRetrySameModelRateLimit).not.toHaveBeenCalled(); - expect(maybeEscalateRateLimitProfileFallback).toHaveBeenCalledTimes(1); - expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(advanceRateLimitAuthProfile).toHaveBeenCalledTimes(1); }); it("does not treat generic rate-limit text as a short-window throttle", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -244,8 +229,7 @@ describe("handleAssistantFailover", () => { errorMessage: "rate limit exceeded", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -255,14 +239,12 @@ describe("handleAssistantFailover", () => { } expect(outcome.retryKind).toBe("profile_rotation"); expect(maybeRetrySameModelRateLimit).not.toHaveBeenCalled(); - expect(maybeEscalateRateLimitProfileFallback).toHaveBeenCalledTimes(1); - expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(advanceRateLimitAuthProfile).toHaveBeenCalledTimes(1); }); it("retries the same model on a status-prefixed 429 with no window wording", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -274,8 +256,7 @@ describe("handleAssistantFailover", () => { errorMessage: "429 Provider returned error", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -286,14 +267,12 @@ describe("handleAssistantFailover", () => { expect(outcome.retryKind).toBe("same_model_rate_limit"); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledTimes(1); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledWith({}); - expect(maybeEscalateRateLimitProfileFallback).not.toHaveBeenCalled(); - expect(advanceAuthProfile).not.toHaveBeenCalled(); + expect(advanceRateLimitAuthProfile).not.toHaveBeenCalled(); }); it("does not spend same-model retry budget when Retry-After is long", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -305,8 +284,7 @@ describe("handleAssistantFailover", () => { errorMessage: "429 rate_limit_exceeded; Retry-After: 3600", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -316,8 +294,7 @@ describe("handleAssistantFailover", () => { } expect(outcome.retryKind).toBe("profile_rotation"); expect(maybeRetrySameModelRateLimit).not.toHaveBeenCalled(); - expect(maybeEscalateRateLimitProfileFallback).toHaveBeenCalledTimes(1); - expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(advanceRateLimitAuthProfile).toHaveBeenCalledTimes(1); }); it("does not spend same-model retry budget when Retry-After date is beyond the retry budget", async () => { @@ -325,8 +302,7 @@ describe("handleAssistantFailover", () => { vi.setSystemTime(new Date("2026-06-11T00:00:00.000Z")); try { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -338,8 +314,7 @@ describe("handleAssistantFailover", () => { errorMessage: "429 rate_limit_exceeded; Retry-After: Thu, 11 Jun 2026 01:05:00 GMT", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -349,8 +324,7 @@ describe("handleAssistantFailover", () => { } expect(outcome.retryKind).toBe("profile_rotation"); expect(maybeRetrySameModelRateLimit).not.toHaveBeenCalled(); - expect(maybeEscalateRateLimitProfileFallback).toHaveBeenCalledTimes(1); - expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(advanceRateLimitAuthProfile).toHaveBeenCalledTimes(1); } finally { vi.useRealTimers(); } @@ -358,8 +332,7 @@ describe("handleAssistantFailover", () => { it("allows short Retry-After intervals to use same-model retry", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -371,8 +344,7 @@ describe("handleAssistantFailover", () => { errorMessage: "429 rate_limit_exceeded; Retry-After: 30 seconds", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -383,14 +355,12 @@ describe("handleAssistantFailover", () => { expect(outcome.retryKind).toBe("same_model_rate_limit"); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledTimes(1); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledWith({ retryAfterSeconds: 30 }); - expect(maybeEscalateRateLimitProfileFallback).not.toHaveBeenCalled(); - expect(advanceAuthProfile).not.toHaveBeenCalled(); + expect(advanceRateLimitAuthProfile).not.toHaveBeenCalled(); }); it("allows RESOURCE_EXHAUSTED messages with short-window 429 hints", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -402,8 +372,7 @@ describe("handleAssistantFailover", () => { errorMessage: "429 RESOURCE_EXHAUSTED: tokens per minute limit exceeded", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -413,14 +382,12 @@ describe("handleAssistantFailover", () => { } expect(outcome.retryKind).toBe("same_model_rate_limit"); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledTimes(1); - expect(maybeEscalateRateLimitProfileFallback).not.toHaveBeenCalled(); - expect(advanceAuthProfile).not.toHaveBeenCalled(); + expect(advanceRateLimitAuthProfile).not.toHaveBeenCalled(); }); it("allows quota wording when it points at a per-minute throttle", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => true); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -433,8 +400,7 @@ describe("handleAssistantFailover", () => { "Quota exceeded for quota metric 'Generate requests per minute' and limit 'Generate requests per minute per project'.", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -444,14 +410,12 @@ describe("handleAssistantFailover", () => { } expect(outcome.retryKind).toBe("same_model_rate_limit"); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledTimes(1); - expect(maybeEscalateRateLimitProfileFallback).not.toHaveBeenCalled(); - expect(advanceAuthProfile).not.toHaveBeenCalled(); + expect(advanceRateLimitAuthProfile).not.toHaveBeenCalled(); }); it("falls back to profile rotation after the same-model rate-limit budget is exhausted", async () => { const maybeRetrySameModelRateLimit = vi.fn(async () => false); - const maybeEscalateRateLimitProfileFallback = vi.fn(); - const advanceAuthProfile = vi.fn(async () => true); + const advanceRateLimitAuthProfile = vi.fn(async () => true); const outcome = await handleAssistantFailover( makeParams({ @@ -463,8 +427,7 @@ describe("handleAssistantFailover", () => { errorMessage: "429 rate_limit_exceeded: too many requests per minute", } as Params["lastAssistant"], maybeRetrySameModelRateLimit, - maybeEscalateRateLimitProfileFallback, - advanceAuthProfile, + advanceRateLimitAuthProfile, }), ); @@ -474,8 +437,7 @@ describe("handleAssistantFailover", () => { } expect(outcome.retryKind).toBe("profile_rotation"); expect(maybeRetrySameModelRateLimit).toHaveBeenCalledTimes(1); - expect(maybeEscalateRateLimitProfileFallback).toHaveBeenCalledTimes(1); - expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(advanceRateLimitAuthProfile).toHaveBeenCalledTimes(1); }); it("does not log profile-specific warnings without a failed profile id", async () => { diff --git a/src/agents/embedded-agent-runner/run/assistant-failover.ts b/src/agents/embedded-agent-runner/run/assistant-failover.ts index 7741ed25bb04..96ebfd7fd41f 100644 --- a/src/agents/embedded-agent-runner/run/assistant-failover.ts +++ b/src/agents/embedded-agent-runner/run/assistant-failover.ts @@ -73,7 +73,6 @@ export async function handleAssistantFailover(params: { failoverReason: FailoverReason | null; harnessOwnsTransport: boolean; allowSameModelIdleTimeoutRetry: boolean; - allowSameModelRateLimitRetry: boolean; assistantProfileFailureReason: AuthProfileFailureReason | null; lastProfileId?: string; modelId: string; @@ -102,14 +101,14 @@ export async function handleAssistantFailover(params: { reason?: AuthProfileFailureReason | null; modelId?: string; }) => Promise; - maybeEscalateRateLimitProfileFallback: (params: { - failoverProvider: string; - failoverModel: string; - logFallbackDecision: (decision: "fallback_model", extra?: { status?: number }) => void; - }) => void; maybeRetrySameModelRateLimit: (retry?: ShortWindowRateLimitRetry) => Promise; maybeBackoffBeforeOverloadFailover: (reason: FailoverReason | null) => Promise; advanceAuthProfile: () => Promise; + advanceRateLimitAuthProfile: (context: { + failoverProvider: string; + failoverModel: string; + logFallbackDecision: (decision: "fallback_model", extra?: { status?: number }) => void; + }) => Promise; }): Promise { const terminal = projectAgentRunAttemptTerminal(params.terminal); const externalAbort = terminal.externalAbort || params.signalOwnedInterruption; @@ -190,26 +189,24 @@ export async function handleAssistantFailover(params: { } } + let rotated: boolean; if (params.failoverReason === "rate_limit") { // Minute-scale RPM windows can clear without spending a profile rotation // or model fallback. Keep the retry bounded; once exhausted, continue // through the existing rate-limit escalation path. const shortWindowRetry = resolveShortWindowRateLimitRetry(params.lastAssistant?.errorMessage); - if ( - params.allowSameModelRateLimitRetry && - shortWindowRetry && - (await params.maybeRetrySameModelRateLimit(shortWindowRetry)) - ) { + if (shortWindowRetry && (await params.maybeRetrySameModelRateLimit(shortWindowRetry))) { return sameModelRateLimitRetry(); } - params.maybeEscalateRateLimitProfileFallback({ + rotated = await params.advanceRateLimitAuthProfile({ failoverProvider: params.activeErrorContext.provider, failoverModel: params.activeErrorContext.model, logFallbackDecision: params.logAssistantFailoverDecision, }); + } else { + rotated = await params.advanceAuthProfile(); } - const rotated = await params.advanceAuthProfile(); const markFailedProfilePromise = markFailedProfile(); if (timeoutFailure && !params.isProbeSession && failedProfileId) { const timeoutLabel = terminal.idleTimedOut ? "idle timeout (model silent)" : "timed out"; diff --git a/src/agents/embedded-agent-runner/run/assistant-failure.test.ts b/src/agents/embedded-agent-runner/run/assistant-failure.test.ts index edfb4a80a784..e2ba8d8fa432 100644 --- a/src/agents/embedded-agent-runner/run/assistant-failure.test.ts +++ b/src/agents/embedded-agent-runner/run/assistant-failure.test.ts @@ -27,7 +27,7 @@ function makeExhaustedCredentialFailureInput(options?: { replaySafe?: boolean }) currentAttemptAssistant: assistant, toolMetas: replaySafe ? [] : [{ toolName: "write", replaySafe: false }], }); - const advanceAttemptAuthProfile = vi.fn(async () => true); + const advanceAuthProfile = vi.fn(async () => true); const maybeMarkAuthProfileFailure = vi.fn(async () => {}); const traceAttempts: AssistantFailureInput["traceAttempts"] = []; const input: AssistantFailureInput = { @@ -79,15 +79,13 @@ function makeExhaustedCredentialFailureInput(options?: { replaySafe?: boolean }) emptyErrorRetries: 3, overloadProfileRotations: 0, overloadProfileRotationLimit: 1, - rateLimitProfileRotations: 0, - rateLimitProfileRotationLimit: 1, sameModelIdleTimeoutRetries: 0, previousRetryFailoverReason: null, maybeMarkAuthProfileFailure, - maybeEscalateRateLimitProfileFallback: vi.fn(), maybeRetrySameModelRateLimit: vi.fn(async () => false), maybeBackoffBeforeOverloadFailover: vi.fn(async () => {}), - advanceAttemptAuthProfile, + advanceAuthProfile, + advanceRateLimitAuthProfile: vi.fn(async () => true), traceAttempts, suspendForFailure: vi.fn(), suspensionSessionId: "session:credential-enoent", @@ -95,7 +93,7 @@ function makeExhaustedCredentialFailureInput(options?: { replaySafe?: boolean }) isProbeSession: false, }; return { - advanceAttemptAuthProfile, + advanceAuthProfile, input, maybeMarkAuthProfileFailure, traceAttempts, @@ -113,7 +111,7 @@ describe("handleEmbeddedAssistantFailure", () => { rawError: CREDENTIAL_FILE_ENOENT_MESSAGE, }); - expect(fixture.advanceAttemptAuthProfile).not.toHaveBeenCalled(); + expect(fixture.advanceAuthProfile).not.toHaveBeenCalled(); expect(fixture.maybeMarkAuthProfileFailure).not.toHaveBeenCalled(); expect(fixture.input.authProfileStore.usageStats).toEqual({ "anthropic:p1": { lastUsed: 1 }, @@ -136,7 +134,7 @@ describe("handleEmbeddedAssistantFailure", () => { const outcome = await handleEmbeddedAssistantFailure(fixture.input); expect(outcome.action).toBe("proceed"); - expect(fixture.advanceAttemptAuthProfile).not.toHaveBeenCalled(); + expect(fixture.advanceAuthProfile).not.toHaveBeenCalled(); expect(fixture.maybeMarkAuthProfileFailure).not.toHaveBeenCalled(); expect(fixture.traceAttempts).toEqual([]); }); diff --git a/src/agents/embedded-agent-runner/run/assistant-failure.ts b/src/agents/embedded-agent-runner/run/assistant-failure.ts index 38ade3e55778..6de73b3eb6e4 100644 --- a/src/agents/embedded-agent-runner/run/assistant-failure.ts +++ b/src/agents/embedded-agent-runner/run/assistant-failure.ts @@ -74,8 +74,6 @@ export async function handleEmbeddedAssistantFailure(input: { emptyErrorRetries: number; overloadProfileRotations: number; overloadProfileRotationLimit: number; - rateLimitProfileRotations: number; - rateLimitProfileRotationLimit: number; sameModelIdleTimeoutRetries: number; previousRetryFailoverReason: FailoverReason | null; maybeMarkAuthProfileFailure: (failure: { @@ -83,12 +81,12 @@ export async function handleEmbeddedAssistantFailure(input: { reason?: AuthProfileFailureReason | null; modelId?: string; }) => Promise; - maybeEscalateRateLimitProfileFallback: Parameters< - typeof handleAssistantFailover - >[0]["maybeEscalateRateLimitProfileFallback"]; maybeRetrySameModelRateLimit: (retry?: { retryAfterSeconds?: number }) => Promise; maybeBackoffBeforeOverloadFailover: (reason: FailoverReason | null) => Promise; - advanceAttemptAuthProfile: () => Promise; + advanceAuthProfile: Parameters[0]["advanceAuthProfile"]; + advanceRateLimitAuthProfile: Parameters< + typeof handleAssistantFailover + >[0]["advanceRateLimitAuthProfile"]; traceAttempts: TraceAttempt[]; suspendForFailure: (params: Omit) => void; suspensionSessionId: string; @@ -268,8 +266,6 @@ export async function handleEmbeddedAssistantFailure(input: { !input.fallbackConfigured && input.canRestartForLiveSwitch && input.sameModelIdleTimeoutRetries < MAX_SAME_MODEL_IDLE_TIMEOUT_RETRIES, - allowSameModelRateLimitRetry: - input.rateLimitProfileRotations < input.rateLimitProfileRotationLimit, assistantProfileFailureReason, lastProfileId: input.authProfileId, modelId: input.modelId, @@ -292,10 +288,10 @@ export async function handleEmbeddedAssistantFailure(input: { logAssistantFailoverDecision: logFailoverDecision, warn: (message) => log.warn(message), maybeMarkAuthProfileFailure: input.maybeMarkAuthProfileFailure, - maybeEscalateRateLimitProfileFallback: input.maybeEscalateRateLimitProfileFallback, maybeRetrySameModelRateLimit: input.maybeRetrySameModelRateLimit, maybeBackoffBeforeOverloadFailover: input.maybeBackoffBeforeOverloadFailover, - advanceAuthProfile: input.advanceAttemptAuthProfile, + advanceAuthProfile: input.advanceAuthProfile, + advanceRateLimitAuthProfile: input.advanceRateLimitAuthProfile, }); if (outcome.action === "retry") { const retryTraceResult = diff --git a/src/agents/embedded-agent-runner/run/attempt-recovery.ts b/src/agents/embedded-agent-runner/run/attempt-recovery.ts index 3d429b4f0f8c..ae0ddca38193 100644 --- a/src/agents/embedded-agent-runner/run/attempt-recovery.ts +++ b/src/agents/embedded-agent-runner/run/attempt-recovery.ts @@ -351,9 +351,8 @@ export async function recoverEmbeddedRunAttempt(input: { pluginHarnessOwnsTransport: runtime.pluginHarnessOwnsTransport, timedOutByRunBudget, resolveAuthProfileFailureReason: failoverRetryController.resolveAuthProfileFailureReason, - maybeEscalateRateLimitProfileFallback: - failoverRetryController.maybeEscalateRateLimitProfileFallback, - advanceAttemptAuthProfile: preparedRuntime.advanceAttemptAuthProfile, + advanceAuthProfile: failoverRetryController.advanceAuthProfile, + advanceRateLimitAuthProfile: failoverRetryController.advanceRateLimitAuthProfile, maybeMarkAuthProfileFailure: failoverRetryController.maybeMarkAuthProfileFailure, maybeBackoffBeforeOverloadFailover: failoverRetryController.maybeBackoffBeforeOverloadFailover, diff --git a/src/agents/embedded-agent-runner/run/failover-retry-controller.test.ts b/src/agents/embedded-agent-runner/run/failover-retry-controller.test.ts index 3acf2d329b6b..86dbd84ddc52 100644 --- a/src/agents/embedded-agent-runner/run/failover-retry-controller.test.ts +++ b/src/agents/embedded-agent-runner/run/failover-retry-controller.test.ts @@ -1,76 +1,110 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { FailoverError } from "../../failover-error.js"; + +const mocks = vi.hoisted(() => ({ + sleepWithAbort: vi.fn(async () => {}), +})); + +vi.mock("../../../infra/backoff.js", async () => { + const actual = await vi.importActual( + "../../../infra/backoff.js", + ); + return { ...actual, sleepWithAbort: mocks.sleepWithAbort }; +}); + import { createEmbeddedRunFailoverRetryController } from "./failover-retry-controller.js"; -const { sleepWithAbortMock } = vi.hoisted(() => ({ - sleepWithAbortMock: vi.fn(async () => {}), -})); +type ControllerInput = Parameters[0]; -vi.mock("../../../infra/backoff.js", () => ({ - sleepWithAbort: sleepWithAbortMock, -})); - -function createController(fallbackConfigured: boolean) { +function createController( + advanceAuthProfile: ControllerInput["advanceAuthProfile"], + fallbackConfigured = false, +) { return createEmbeddedRunFailoverRetryController({ runParams: { - sessionId: "session:rate-limit-controller", - runId: "run:rate-limit-controller", - } as never, + runId: "run:failover-retry-controller-test", + } as ControllerInput["runParams"], provider: "openai", - modelId: "mock-1", + modelId: "gpt-5.6-luna", globalLane: "test", - agentDir: "/tmp/openclaw-rate-limit-controller-test", + agentDir: "/tmp/openclaw-failover-retry-controller-test", fallbackConfigured, - profileFailureStore: { version: 1, profiles: {} } as never, + profileFailureStore: { version: 1, profiles: {} }, getLastProfileId: () => "openai:p1", - getSessionId: () => "session:rate-limit-controller", + getSessionId: () => "session:failover-retry-controller-test", harnessOwnsTransport: () => false, - getRuntimeAuthOwnerId: () => "pi", + getRuntimeAuthOwnerId: () => "embedded", getApiKeyInfo: () => null, + advanceAuthProfile, }); } +const rateLimitContext = { + failoverProvider: "openai", + failoverModel: "gpt-5.6-luna", + logFallbackDecision: vi.fn(), +}; + describe("createEmbeddedRunFailoverRetryController", () => { beforeEach(() => { - sleepWithAbortMock.mockClear(); + mocks.sleepWithAbort.mockClear(); + rateLimitContext.logFallbackDecision.mockClear(); }); - it("keeps the full same-model retry budget when no fallback rotation is configured", async () => { - const controller = createController(false); + it("preserves the full same-model retry budget when rate-limit rotation does not advance", async () => { + const advanceAuthProfile = vi.fn(async () => false); + const controller = createController(advanceAuthProfile); - controller.maybeEscalateRateLimitProfileFallback({ - failoverProvider: "openai", - failoverModel: "mock-1", - logFallbackDecision: vi.fn(), - }); - - expect(controller.rateLimitProfileRotations).toBe(0); - expect(controller.rateLimitProfileRotations).toBeLessThan( - controller.rateLimitProfileRotationLimit, - ); + await expect(controller.advanceRateLimitAuthProfile(rateLimitContext)).resolves.toBe(false); await expect(controller.maybeRetrySameModelRateLimit()).resolves.toBe(true); await expect(controller.maybeRetrySameModelRateLimit()).resolves.toBe(true); await expect(controller.maybeRetrySameModelRateLimit()).resolves.toBe(true); await expect(controller.maybeRetrySameModelRateLimit()).resolves.toBe(false); + + expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(mocks.sleepWithAbort).toHaveBeenCalledTimes(3); }); - it("counts one actual rotation and does not increment while enforcing the cap", () => { - const controller = createController(true); - const logFallbackDecision = vi.fn(); - const escalation = { - failoverProvider: "groq", - failoverModel: "mock-2", - logFallbackDecision, - }; + it("consumes same-model retry eligibility after a successful rate-limit rotation", async () => { + const advanceAuthProfile = vi.fn(async () => true); + const controller = createController(advanceAuthProfile); - controller.maybeEscalateRateLimitProfileFallback(escalation); - expect(controller.rateLimitProfileRotations).toBe(1); + await expect(controller.advanceRateLimitAuthProfile(rateLimitContext)).resolves.toBe(true); + await expect(controller.maybeRetrySameModelRateLimit()).resolves.toBe(false); - expect(() => controller.maybeEscalateRateLimitProfileFallback(escalation)).toThrow( + expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(mocks.sleepWithAbort).not.toHaveBeenCalled(); + }); + + it("does not spend rate-limit rotation eligibility on an ordinary profile advance", async () => { + const advanceAuthProfile = vi.fn(async () => true); + const controller = createController(advanceAuthProfile); + + await expect(controller.advanceAuthProfile()).resolves.toBe(true); + await expect(controller.maybeRetrySameModelRateLimit()).resolves.toBe(true); + + expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(mocks.sleepWithAbort).toHaveBeenCalledWith(10_000, undefined); + }); + + it("escalates after one successful rate-limit rotation without advancing again", async () => { + const advanceAuthProfile = vi.fn(async () => true); + const controller = createController(advanceAuthProfile, true); + + await expect(controller.advanceRateLimitAuthProfile(rateLimitContext)).resolves.toBe(true); + await expect(controller.advanceRateLimitAuthProfile(rateLimitContext)).rejects.toMatchObject({ + name: "FailoverError", + reason: "rate_limit", + status: 429, + } satisfies Partial); + await expect(controller.advanceRateLimitAuthProfile(rateLimitContext)).rejects.toBeInstanceOf( FailoverError, ); - expect(controller.rateLimitProfileRotations).toBe(1); - expect(logFallbackDecision).toHaveBeenCalledOnce(); - expect(logFallbackDecision).toHaveBeenCalledWith("fallback_model", { status: 429 }); + + expect(advanceAuthProfile).toHaveBeenCalledTimes(1); + expect(rateLimitContext.logFallbackDecision).toHaveBeenCalledTimes(2); + expect(rateLimitContext.logFallbackDecision).toHaveBeenNthCalledWith(1, "fallback_model", { + status: 429, + }); }); }); diff --git a/src/agents/embedded-agent-runner/run/failover-retry-controller.ts b/src/agents/embedded-agent-runner/run/failover-retry-controller.ts index 38f7925d4752..7cd674ae60b5 100644 --- a/src/agents/embedded-agent-runner/run/failover-retry-controller.ts +++ b/src/agents/embedded-agent-runner/run/failover-retry-controller.ts @@ -24,6 +24,12 @@ import type { prepareEmbeddedRunRuntime } from "./runtime-preparation.js"; type PreparedRuntime = Awaited>; +type RateLimitAuthProfileContext = { + failoverProvider: string; + failoverModel: string; + logFallbackDecision: (decision: "fallback_model", extra?: { status?: number }) => void; +}; + export function createEmbeddedRunFailoverRetryController(input: { runParams: PreparedEmbeddedRunInput["runParams"]; provider: string; @@ -37,6 +43,7 @@ export function createEmbeddedRunFailoverRetryController(input: { harnessOwnsTransport: () => boolean; getRuntimeAuthOwnerId: () => string; getApiKeyInfo: () => ResolvedProviderAuth | null; + advanceAuthProfile: PreparedRuntime["advanceAttemptAuthProfile"]; }) { const { runParams: params, @@ -68,10 +75,6 @@ export function createEmbeddedRunFailoverRetryController(input: { return { overloadProfileRotationLimit, - rateLimitProfileRotationLimit, - get rateLimitProfileRotations() { - return rateLimitProfileRotations; - }, get consecutiveSameModelRateLimitRetries() { return consecutiveSameModelRateLimitRetries; }, @@ -81,37 +84,32 @@ export function createEmbeddedRunFailoverRetryController(input: { retriedSameModelRateLimit: false, }); }, - maybeEscalateRateLimitProfileFallback: (paramsLocal: { - failoverProvider: string; - failoverModel: string; - logFallbackDecision: (decision: "fallback_model", extra?: { status?: number }) => void; - }) => { - if (!fallbackConfigured) { - return; + advanceAuthProfile: input.advanceAuthProfile, + advanceRateLimitAuthProfile: async (context: RateLimitAuthProfileContext): Promise => { + if (rateLimitProfileRotations >= rateLimitProfileRotationLimit && fallbackConfigured) { + const status = resolveFailoverStatus("rate_limit"); + log.warn( + `rate-limit profile rotation cap reached for ${sanitizeForLog(provider)}/${sanitizeForLog(modelId)} after ${rateLimitProfileRotations} rotations; escalating to model fallback`, + ); + context.logFallbackDecision("fallback_model", { status }); + throw new FailoverError( + "The AI service is temporarily rate-limited. Please try again in a moment.", + { + reason: "rate_limit", + provider: context.failoverProvider, + model: context.failoverModel, + profileId: input.getLastProfileId(), + sessionId: input.getSessionId(), + lane: globalLane, + status, + }, + ); } - if (rateLimitProfileRotations < rateLimitProfileRotationLimit) { - // This state gates same-model retries, so skipped rotations must not consume it; - // otherwise one rate-limit response can disable the remaining retry budget. + const rotated = await input.advanceAuthProfile(); + if (rotated) { rateLimitProfileRotations += 1; - return; } - const status = resolveFailoverStatus("rate_limit"); - log.warn( - `rate-limit profile rotation cap reached for ${sanitizeForLog(provider)}/${sanitizeForLog(modelId)} after ${rateLimitProfileRotations} rotations; escalating to model fallback`, - ); - paramsLocal.logFallbackDecision("fallback_model", { status }); - throw new FailoverError( - "The AI service is temporarily rate-limited. Please try again in a moment.", - { - reason: "rate_limit", - provider: paramsLocal.failoverProvider, - model: paramsLocal.failoverModel, - profileId: input.getLastProfileId(), - sessionId: input.getSessionId(), - lane: globalLane, - status, - }, - ); + return rotated; }, maybeMarkAuthProfileFailure: async (failure: { profileId?: string; @@ -195,7 +193,10 @@ export function createEmbeddedRunFailoverRetryController(input: { maybeRetrySameModelRateLimit: async (retry?: { retryAfterSeconds?: number; }): Promise => { - if (consecutiveSameModelRateLimitRetries >= MAX_SAME_MODEL_RATE_LIMIT_RETRIES) { + if ( + rateLimitProfileRotations >= rateLimitProfileRotationLimit || + consecutiveSameModelRateLimitRetries >= MAX_SAME_MODEL_RATE_LIMIT_RETRIES + ) { return false; } const delayMs = resolveSameModelRateLimitRetryDelayMs({ diff --git a/src/agents/embedded-agent-runner/run/prompt-failure.test.ts b/src/agents/embedded-agent-runner/run/prompt-failure.test.ts index 995a042b5ab8..1fc17b62446c 100644 --- a/src/agents/embedded-agent-runner/run/prompt-failure.test.ts +++ b/src/agents/embedded-agent-runner/run/prompt-failure.test.ts @@ -46,8 +46,8 @@ function makeParams(overrides: Partial = {}): Params { resolveAuthProfileFailureReason: vi.fn( () => "rate_limit", ), - maybeEscalateRateLimitProfileFallback: vi.fn(), - advanceAttemptAuthProfile: vi.fn(async () => true), + advanceAuthProfile: vi.fn(async () => true), + advanceRateLimitAuthProfile: vi.fn(async () => true), maybeMarkAuthProfileFailure: vi.fn(async () => {}), maybeBackoffBeforeOverloadFailover: vi.fn(async () => {}), attemptedThinking: new Set(), @@ -75,7 +75,7 @@ describe("handleEmbeddedPromptFailure", () => { try { const outcome = await handleEmbeddedPromptFailure( makeParams({ - advanceAttemptAuthProfile: vi.fn(async () => { + advanceRateLimitAuthProfile: vi.fn(async () => { events.push("advance"); return true; }), diff --git a/src/agents/embedded-agent-runner/run/prompt-failure.ts b/src/agents/embedded-agent-runner/run/prompt-failure.ts index 961b1fbb5efa..b15b69256ea3 100644 --- a/src/agents/embedded-agent-runner/run/prompt-failure.ts +++ b/src/agents/embedded-agent-runner/run/prompt-failure.ts @@ -67,12 +67,12 @@ export async function handleEmbeddedPromptFailure(input: { reason: FailoverReason | null, options?: { providerStarted?: boolean; transientRateLimit?: boolean }, ) => AuthProfileFailureReason | null; - maybeEscalateRateLimitProfileFallback: (params: { + advanceAuthProfile: () => Promise; + advanceRateLimitAuthProfile: (context: { failoverProvider: string; failoverModel: string; logFallbackDecision: ReturnType; - }) => void; - advanceAttemptAuthProfile: () => Promise; + }) => Promise; maybeMarkAuthProfileFailure: (failure: { profileId?: string; reason?: AuthProfileFailureReason | null; @@ -156,13 +156,6 @@ export async function handleEmbeddedPromptFailure(input: { fallbackConfigured: input.fallbackConfigured, aborted: input.aborted, }); - if (promptFailoverReason === "rate_limit") { - input.maybeEscalateRateLimitProfileFallback({ - failoverProvider: input.provider, - failoverModel: input.modelId, - logFallbackDecision: logFailoverDecision, - }); - } let failoverDecision = resolveRunFailoverDecision({ stage: "prompt", aborted: input.aborted, @@ -176,7 +169,19 @@ export async function handleEmbeddedPromptFailure(input: { timedOutByRunBudget: input.timedOutByRunBudget, profileRotated: false, }); - if (failoverDecision.action === "rotate_profile" && (await input.advanceAttemptAuthProfile())) { + let rotated = false; + if (failoverDecision.action === "rotate_profile") { + if (promptFailoverReason === "rate_limit") { + rotated = await input.advanceRateLimitAuthProfile({ + failoverProvider: input.provider, + failoverModel: input.modelId, + logFallbackDecision: logFailoverDecision, + }); + } else { + rotated = await input.advanceAuthProfile(); + } + } + if (rotated) { if (promptProfileFailureReason) { void input .maybeMarkAuthProfileFailure({