From 9fd9aa5fcdd2dec17c5ab892f2938cea347c25cc Mon Sep 17 00:00:00 2001 From: OfflynAI <140015627+joelnishanth@users.noreply.github.com> Date: Tue, 16 Jun 2026 18:04:48 -0700 Subject: [PATCH] fix(agents): defer session suspension across fallback (#93789) Co-authored-by: Vincent Koc <25068+vincentkoc@users.noreply.github.com> --- src/agents/embedded-agent-runner/run.ts | 24 +- src/agents/model-fallback.probe.test.ts | 265 ++++++++++++++++++ .../model-fallback.run-embedded.e2e.test.ts | 91 +++++- src/agents/model-fallback.ts | 159 ++++++++--- src/agents/session-suspension.test.ts | 26 ++ src/agents/session-suspension.ts | 49 +++- 6 files changed, 556 insertions(+), 58 deletions(-) diff --git a/src/agents/embedded-agent-runner/run.ts b/src/agents/embedded-agent-runner/run.ts index 80a66e6a8d92..b630dec97509 100644 --- a/src/agents/embedded-agent-runner/run.ts +++ b/src/agents/embedded-agent-runner/run.ts @@ -121,7 +121,12 @@ import { runAgentCleanupStep } from "../run-cleanup-timeout.js"; import { buildAgentRuntimeAuthPlan } from "../runtime-plan/auth.js"; import { buildAgentRuntimePlan } from "../runtime-plan/build.js"; import { ensureRuntimePluginsLoaded } from "../runtime-plugins.js"; -import { resolveSessionSuspensionReason, suspendSession } from "../session-suspension.js"; +import { + resolveSessionSuspensionReason, + resolveSessionSuspensionTarget, + suspendSession, + type SessionSuspensionParams, +} from "../session-suspension.js"; import { resolveToolLoopDetectionConfig } from "../tool-loop-detection-config.js"; import { derivePromptTokens, normalizeUsage, type UsageLike } from "../usage.js"; import { redactRunIdentifier, resolveRunWorkspaceDir } from "../workspace-run.js"; @@ -619,6 +624,17 @@ async function runEmbeddedAgentInternal( } const sessionLane = resolveSessionLane(params.sessionKey?.trim() || params.sessionId); const globalLane = resolveGlobalLane(params.lane); + // Outer fallback attempts defer session suspension only while another + // candidate remains. Direct and final-candidate runs suspend normally. + const failureSuspension = resolveSessionSuspensionTarget(); + const suspendForFailure = (suspensionParams: Omit) => { + const suspension = { ...suspensionParams, laneId: globalLane }; + if (failureSuspension.mode === "defer") { + failureSuspension.defer(suspension); + return; + } + void suspendSession(suspension); + }; const sessionQueuePriority = resolveEmbeddedRunSessionQueuePriority(params.trigger); const laneTaskTimeoutMs = resolveEmbeddedRunLaneTimeoutMs(params.timeoutMs); let laneTaskProgressAtMs = Date.now(); @@ -2784,11 +2800,10 @@ async function runEmbeddedAgentInternal( ? describeFailoverError(normalizedPromptFailover) : describeFailoverError(promptError); if (normalizedPromptFailover?.suspend) { - void suspendSession({ + suspendForFailure({ cfg: params.config, agentDir, sessionId: activeSessionId ?? params.sessionId, - laneId: globalLane, reason: resolveSessionSuspensionReason(normalizedPromptFailover.reason), failedProvider: normalizedPromptFailover.provider ?? provider, failedModel: normalizedPromptFailover.model ?? modelId, @@ -3258,11 +3273,10 @@ async function runEmbeddedAgentInternal( : {}), }); if (assistantFailoverOutcome.error.suspend) { - void suspendSession({ + suspendForFailure({ cfg: params.config, agentDir, sessionId: activeSessionId ?? params.sessionId, - laneId: globalLane, reason: resolveSessionSuspensionReason(assistantFailoverOutcome.error.reason), failedProvider: assistantFailoverOutcome.error.provider ?? provider, failedModel: assistantFailoverOutcome.error.model ?? modelId, diff --git a/src/agents/model-fallback.probe.test.ts b/src/agents/model-fallback.probe.test.ts index b3d98b69ea08..0895ec1a5dbb 100644 --- a/src/agents/model-fallback.probe.test.ts +++ b/src/agents/model-fallback.probe.test.ts @@ -6,6 +6,7 @@ import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vite import type { OpenClawConfig } from "../config/config.js"; import { createDiagnosticLogRecordCapture } from "../logging/test-helpers/diagnostic-log-capture.js"; import type { AuthProfileStore } from "./auth-profiles.js"; +import type { SessionSuspensionParams } from "./session-suspension.js"; import { makeModelFallbackCfg } from "./test-helpers/model-fallback-config-fixture.js"; // Mock auth-profile submodules before importing model-fallback so the module @@ -29,6 +30,34 @@ vi.mock("./provider-model-normalization.runtime.js", () => ({ normalizeProviderModelIdWithRuntime: () => undefined, })); +const sessionSuspensionMocks = vi.hoisted(() => ({ + suspendSession: vi.fn().mockResolvedValue(undefined), + runWithDeferredSessionSuspension: vi.fn( + (run: () => Promise, onDeferred?: (params: SessionSuspensionParams) => void) => { + onDeferred?.({ + cfg: {}, + sessionId: "test-session", + laneId: "main", + reason: "quota_exhausted", + failedProvider: "openai", + failedModel: "gpt-4.1-mini", + }); + return run(); + }, + ), + resolveSessionSuspensionReason: vi.fn((reason: string) => { + if (reason === "billing") { + return "manual"; + } + if (reason === "rate_limit") { + return "quota_exhausted"; + } + return "circuit_open"; + }), +})); + +vi.mock("./session-suspension.js", () => sessionSuspensionMocks); + const emptyPluginMetadataSnapshot = vi.hoisted(() => ({ policyHash: "model-fallback-probe-test-empty-plugin-policy", configFingerprint: "model-fallback-probe-test-empty-plugin-metadata", @@ -341,6 +370,8 @@ describe("runWithModelFallback – probe logic", () => { cleanupLogCapture = undefined; setLoggerOverride(null); resetLogger(); + sessionSuspensionMocks.suspendSession.mockClear(); + sessionSuspensionMocks.runWithDeferredSessionSuspension.mockClear(); vi.restoreAllMocks(); }); @@ -793,4 +824,238 @@ describe("runWithModelFallback – probe logic", () => { "billing", ); }); + + it("does not lock lane when fallback candidates remain after suspend_lanes decision", async () => { + const cfg = makeCfg({ + agents: { + defaults: { + model: { + primary: "openai/gpt-4.1-mini", + fallbacks: ["anthropic/claude-haiku-3-5"], + }, + }, + }, + } as Partial); + + // Put only OpenAI into cooldown; Anthropic is available + mockedIsProfileInCooldown.mockImplementation((_store: AuthProfileStore, profileId: string) => + profileId.startsWith("openai"), + ); + mockedGetSoonestCooldownExpiry.mockReturnValue(NOW + 30 * 60 * 1000); + mockedResolveProfilesUnavailableReason.mockReturnValue("billing"); + + const run = vi.fn().mockResolvedValue("fallback-ok"); + + await runWithModelFallback({ + cfg, + provider: "openai", + model: "gpt-4.1-mini", + run, + sessionId: "test-session", + lane: "main", + }); + + expect(sessionSuspensionMocks.suspendSession).not.toHaveBeenCalled(); + }); + + it("defers embedded lane suspension only while another candidate remains", async () => { + const cfg = makeCfg({ + agents: { + defaults: { + model: { + primary: "openai/gpt-4.1-mini", + fallbacks: ["anthropic/claude-haiku-3-5"], + }, + }, + }, + } as Partial); + mockedIsProfileInCooldown.mockReturnValue(false); + const run = vi + .fn() + .mockRejectedValueOnce(new Error("primary failed")) + .mockResolvedValueOnce("fallback-ok"); + + const result = await runWithModelFallback({ + cfg, + provider: "openai", + model: "gpt-4.1-mini", + run, + sessionId: "test-session", + lane: "main", + }); + + expect(result.result).toBe("fallback-ok"); + expect(run).toHaveBeenCalledTimes(2); + expect(sessionSuspensionMocks.runWithDeferredSessionSuspension).toHaveBeenCalledOnce(); + }); + + it("discards deferred suspension when the outer run is aborted", async () => { + const cfg = makeCfg({ + agents: { + defaults: { + model: { + primary: "openai/gpt-4.1-mini", + fallbacks: ["anthropic/claude-haiku-3-5"], + }, + }, + }, + } as Partial); + mockedIsProfileInCooldown.mockReturnValue(false); + const controller = new AbortController(); + const disconnect = new Error("client disconnected"); + disconnect.name = "ClientDisconnectError"; + const run = vi.fn().mockImplementation(async () => { + controller.abort(disconnect); + throw disconnect; + }); + + await expect( + runWithModelFallback({ + cfg, + provider: "openai", + model: "gpt-4.1-mini", + run, + sessionId: "test-session", + lane: "main", + abortSignal: controller.signal, + }), + ).rejects.toBe(disconnect); + + expect(run).toHaveBeenCalledOnce(); + expect(sessionSuspensionMocks.runWithDeferredSessionSuspension).toHaveBeenCalledOnce(); + expect(sessionSuspensionMocks.suspendSession).not.toHaveBeenCalled(); + }); + + it("keeps generic no-lane terminal suspension unbound", async () => { + const cfg = makeCfg({ + agents: { + defaults: { + model: { + primary: "openai/gpt-4.1-mini", + fallbacks: ["anthropic/claude-haiku-3-5"], + }, + }, + }, + } as Partial); + + // Both providers in cooldown + mockedIsProfileInCooldown.mockReturnValue(true); + mockedGetSoonestCooldownExpiry.mockReturnValue(NOW + 30 * 60 * 1000); + mockedResolveProfilesUnavailableReason.mockReturnValue("billing"); + mockedResolveAuthProfileOrder.mockImplementation(({ provider }: { provider: string }) => { + if (provider === "openai") { + return ["openai-profile-1"]; + } + if (provider === "anthropic") { + return ["anthropic-profile-1"]; + } + return []; + }); + + // Throttle primary probe so billing goes to suspend_lanes + probeThrottleInternals.lastProbeAttempt.set("openai", NOW - 10_000); + + const run = vi.fn().mockResolvedValue("should-not-run"); + + await expect( + runWithModelFallback({ + cfg, + provider: "openai", + model: "gpt-4.1-mini", + run, + sessionId: "test-session", + }), + ).rejects.toThrow(); + + expect(sessionSuspensionMocks.suspendSession).toHaveBeenCalledWith( + expect.objectContaining({ + laneId: undefined, + failedProvider: "anthropic", + }), + ); + expect(sessionSuspensionMocks.suspendSession).not.toHaveBeenCalledWith( + expect.objectContaining({ failedProvider: "openai" }), + ); + expect( + sessionSuspensionMocks.suspendSession.mock.calls.every( + ([params]) => params.laneId === undefined, + ), + ).toBe(true); + }); + + it("restores a deferred embedded lane when later candidates cannot run", async () => { + const cfg = makeCfg({ + agents: { + defaults: { + model: { + primary: "openai/gpt-4.1-mini", + fallbacks: ["anthropic/claude-haiku-3-5"], + }, + }, + }, + } as Partial); + mockedIsProfileInCooldown.mockImplementation((_store: AuthProfileStore, profileId: string) => + profileId.startsWith("anthropic"), + ); + mockedGetSoonestCooldownExpiry.mockReturnValue(NOW + 30 * 60 * 1000); + mockedResolveProfilesUnavailableReason.mockReturnValue("billing"); + mockedResolveAuthProfileOrder.mockImplementation(({ provider }: { provider: string }) => [ + `${provider}-profile-1`, + ]); + const run = vi.fn().mockRejectedValueOnce(new Error("primary failed")); + + await expect( + runWithModelFallback({ + cfg, + provider: "openai", + model: "gpt-4.1-mini", + run, + sessionId: "test-session", + }), + ).rejects.toThrow(); + + expect(run).toHaveBeenCalledOnce(); + expect(sessionSuspensionMocks.suspendSession).toHaveBeenCalledWith( + expect.objectContaining({ + laneId: "main", + failedProvider: "anthropic", + }), + ); + }); + + it("restores deferred suspension when a later harness precheck fails", async () => { + const cfg = makeCfg({ + agents: { + defaults: { + model: { + primary: "openai/gpt-4.1-mini", + fallbacks: ["anthropic/claude-haiku-3-5"], + }, + }, + }, + } as Partial); + mockedIsProfileInCooldown.mockReturnValue(false); + const run = vi.fn().mockRejectedValueOnce(new Error("primary failed")); + + await expect( + runWithModelFallback({ + cfg, + provider: "openai", + model: "gpt-4.1-mini", + sessionId: "test-session", + resolveAgentHarnessRuntimeOverride: (provider) => + provider === "anthropic" ? "missing-strict-harness" : undefined, + prepareAgentHarnessRuntime: () => undefined, + run, + }), + ).rejects.toThrow('Requested agent harness "missing-strict-harness" is not registered.'); + + expect(run).toHaveBeenCalledOnce(); + expect(sessionSuspensionMocks.suspendSession).toHaveBeenCalledWith( + expect.objectContaining({ + laneId: "main", + failedProvider: "openai", + }), + ); + }); }); diff --git a/src/agents/model-fallback.run-embedded.e2e.test.ts b/src/agents/model-fallback.run-embedded.e2e.test.ts index 8ce01572d3c2..bf55b2651d46 100644 --- a/src/agents/model-fallback.run-embedded.e2e.test.ts +++ b/src/agents/model-fallback.run-embedded.e2e.test.ts @@ -8,7 +8,7 @@ import type { AuthProfileFailureReason } from "./auth-profiles.js"; import { ensureAuthProfileStore, saveAuthProfileStore } from "./auth-profiles/store.js"; import { classifyEmbeddedAgentRunResultForModelFallback } from "./embedded-agent-runner/result-fallback-classifier.js"; import type { EmbeddedRunAttemptResult } from "./embedded-agent-runner/run/types.js"; -import { runWithModelFallback } from "./model-fallback.js"; +import { FailoverError } from "./failover-error.js"; import { buildEmbeddedRunnerAssistant, createResolvedEmbeddedRunnerModel, @@ -21,6 +21,7 @@ import { } from "./test-helpers/embedded-agent-runner-e2e-mocks.js"; const runEmbeddedAttemptMock = vi.fn<(params: unknown) => Promise>(); +const suspendSessionMock = vi.hoisted(() => vi.fn().mockResolvedValue(undefined)); const { computeBackoffMock, sleepWithAbortMock } = vi.hoisted(() => ({ computeBackoffMock: vi.fn( ( @@ -54,18 +55,26 @@ const installRunEmbeddedMocks = () => { resolveModelAsync: async (provider: string, modelId: string) => createResolvedEmbeddedRunnerModel(provider, modelId), })); + vi.doMock("./session-suspension.js", async () => { + const actual = + await vi.importActual("./session-suspension.js"); + return { ...actual, suspendSession: suspendSessionMock }; + }); }; let runEmbeddedAgent: typeof import("./embedded-agent-runner/run.js").runEmbeddedAgent; +let runWithModelFallback: typeof import("./model-fallback.js").runWithModelFallback; beforeAll(async () => { vi.resetModules(); installRunEmbeddedMocks(); ({ runEmbeddedAgent } = await import("./embedded-agent-runner/run.js")); + ({ runWithModelFallback } = await import("./model-fallback.js")); }); beforeEach(() => { runEmbeddedAttemptMock.mockReset(); + suspendSessionMock.mockClear(); computeBackoffMock.mockClear(); sleepWithAbortMock.mockClear(); }); @@ -219,21 +228,26 @@ async function runEmbeddedFallback(params: { workspaceDir: string; sessionKey: string; runId: string; + sessionId?: string; + lane?: string; abortSignal?: AbortSignal; config?: OpenClawConfig; }) { // Runs the same embedded-agent entrypoint that production fallback uses while // keeping provider/model attempts deterministic through mocks. const cfg = params.config ?? makeConfig(); + const sessionId = params.sessionId ?? `session:${params.runId}`; return await runWithModelFallback({ cfg, provider: "openai", model: "mock-1", runId: params.runId, + sessionId: params.sessionId, + lane: params.lane, agentDir: params.agentDir, run: (provider, model, options) => runEmbeddedAgent({ - sessionId: `session:${params.runId}`, + sessionId, sessionKey: params.sessionKey, sessionFile: path.join(params.workspaceDir, `${params.runId}.jsonl`), workspaceDir: params.workspaceDir, @@ -242,6 +256,7 @@ async function runEmbeddedFallback(params: { prompt: "hello", provider, model, + lane: params.lane, authProfileIdSource: "auto", allowTransientCooldownProbe: options?.allowTransientCooldownProbe, timeoutMs: 5_000, @@ -293,6 +308,20 @@ function mockPrimaryPromptErrorThenFallbackSuccess(errorMessage: string) { ); } +function mockPrimarySuspendingPromptErrorThenFallbackSuccess(sessionId: string) { + mockPrimaryFailureThenFallbackSuccess(() => + makeEmbeddedRunnerAttempt({ + sessionIdUsed: sessionId, + promptError: new FailoverError(RATE_LIMIT_ERROR_MESSAGE, { + reason: "rate_limit", + provider: "openai", + model: "mock-1", + suspend: true, + }), + }), + ); +} + function mockPrimaryErrorThenFallbackSuccess(errorMessage: string) { mockPrimaryFailureThenFallbackSuccess(() => makeEmbeddedRunnerAttempt({ @@ -529,6 +558,64 @@ describe("runWithModelFallback + runEmbeddedAgent failover behavior", () => { } }); + it("keeps direct embedded-run lane suspension outside the outer fallback loop", async () => { + await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { + await writeAuthStore(agentDir); + const sessionId = "session:direct-embedded-suspension"; + mockPrimarySuspendingPromptErrorThenFallbackSuccess(sessionId); + + await expect( + runEmbeddedAgent({ + sessionId, + sessionKey: "agent:test:direct-embedded-suspension", + sessionFile: path.join(workspaceDir, "direct-embedded-suspension.jsonl"), + workspaceDir, + agentDir, + config: { + ...makeConfig(), + auth: { cooldowns: { rateLimitedProfileRotations: 0 } }, + }, + prompt: "hello", + provider: "openai", + model: "mock-1", + lane: "direct-lane", + authProfileIdSource: "auto", + timeoutMs: 5_000, + runId: "run:direct-embedded-suspension", + enqueue: async (task) => await task(), + }), + ).rejects.toThrow(); + + expect(suspendSessionMock).toHaveBeenCalledWith( + expect.objectContaining({ laneId: "direct-lane" }), + ); + }); + }); + + it("does not suspend the session while an outer fallback candidate remains", async () => { + await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { + await writeAuthStore(agentDir); + const sessionId = "session:outer-fallback-suspension"; + mockPrimarySuspendingPromptErrorThenFallbackSuccess(sessionId); + + const result = await runEmbeddedFallback({ + agentDir, + workspaceDir, + sessionId, + sessionKey: "agent:test:outer-fallback-suspension", + lane: "outer-fallback-lane", + runId: "run:outer-fallback-suspension", + config: { + ...makeConfig(), + auth: { cooldowns: { rateLimitedProfileRotations: 0 } }, + }, + }); + + expect(result.provider).toBe("groq"); + expect(suspendSessionMock).not.toHaveBeenCalled(); + }); + }); + it("falls back across providers after a bare leading 402 quota-refresh assistant error", async () => { await withAgentWorkspace(async ({ agentDir, workspaceDir }) => { await writeAuthStore(agentDir); diff --git a/src/agents/model-fallback.ts b/src/agents/model-fallback.ts index 293f6058eaad..56ec613ed39f 100644 --- a/src/agents/model-fallback.ts +++ b/src/agents/model-fallback.ts @@ -77,7 +77,12 @@ import { resolveModelRefFromString, } from "./model-selection-resolve.js"; import { isAgentRunRestartAbortReason } from "./run-termination.js"; -import { resolveSessionSuspensionReason, suspendSession } from "./session-suspension.js"; +import { + resolveSessionSuspensionReason, + runWithDeferredSessionSuspension, + suspendSession, + type SessionSuspensionParams, +} from "./session-suspension.js"; const log = createSubsystemLogger("model-fallback"); @@ -337,13 +342,19 @@ async function runFallbackCandidate(params: { provider: string; model: string; options?: ModelFallbackRunOptions; + deferSessionSuspension?: boolean; + onDeferredSessionSuspension?: (params: SessionSuspensionParams) => void; attribution?: FailoverAttribution; abortSignal?: AbortSignal; }): Promise<{ ok: true; result: T } | { ok: false; error: unknown }> { try { - const result = params.options - ? await params.run(params.provider, params.model, params.options) - : await params.run(params.provider, params.model); + const run = () => + params.options + ? params.run(params.provider, params.model, params.options) + : params.run(params.provider, params.model); + const result = params.deferSessionSuspension + ? await runWithDeferredSessionSuspension(run, params.onDeferredSessionSuspension) + : await run(); return { ok: true, result, @@ -379,6 +390,8 @@ async function runFallbackAttempt(params: { model: string; attempts: FallbackAttempt[]; options?: ModelFallbackRunOptions; + deferSessionSuspension?: boolean; + onDeferredSessionSuspension?: (params: SessionSuspensionParams) => void; classifyResult?: ModelFallbackResultClassifier; attempt: number; total: number; @@ -397,6 +410,8 @@ async function runFallbackAttempt(params: { provider: params.provider, model: params.model, options: params.options, + deferSessionSuspension: params.deferSessionSuspension, + onDeferredSessionSuspension: params.onDeferredSessionSuspension, attribution: params.attribution, abortSignal: params.abortSignal, }); @@ -1259,33 +1274,80 @@ function resolveCooldownDecision(params: { }; } -export async function runWithModelFallback( - params: { - cfg: OpenClawConfig | undefined; +type RunWithModelFallbackParams = { + cfg: OpenClawConfig | undefined; + provider: string; + model: string; + runId?: string; + sessionId?: string; + agentId?: string; + sessionKey?: string; + resolveAgentHarnessRuntimeOverride?: (provider: string, model: string) => string | undefined; + prepareAgentHarnessRuntime?: (params: { provider: string; model: string; - runId?: string; - sessionId?: string; - agentId?: string; - sessionKey?: string; - resolveAgentHarnessRuntimeOverride?: (provider: string, model: string) => string | undefined; - prepareAgentHarnessRuntime?: (params: { - provider: string; - model: string; - agentHarnessRuntimeOverride?: string; - }) => Promise | void; - lane?: string; - agentDir?: string; - /** Optional explicit fallbacks list; when provided (even empty), replaces agents.defaults.model.fallbacks. */ - fallbacksOverride?: string[]; - run: ModelFallbackRunFn; - onError?: ModelFallbackErrorHandler; - onFallbackStep?: ModelFallbackStepHandler; - classifyResult?: ModelFallbackResultClassifier; - mergeExhaustedResult?: (params: { latestResult: T; preferredResult: T }) => T; - skipAuthProfileRuntime?: boolean; - abortSignal?: AbortSignal; - } & ModelManifestNormalizationContext, + agentHarnessRuntimeOverride?: string; + }) => Promise | void; + lane?: string; + agentDir?: string; + /** Optional explicit fallbacks list; when provided (even empty), replaces agents.defaults.model.fallbacks. */ + fallbacksOverride?: string[]; + run: ModelFallbackRunFn; + onError?: ModelFallbackErrorHandler; + onFallbackStep?: ModelFallbackStepHandler; + classifyResult?: ModelFallbackResultClassifier; + mergeExhaustedResult?: (params: { latestResult: T; preferredResult: T }) => T; + skipAuthProfileRuntime?: boolean; + abortSignal?: AbortSignal; +} & ModelManifestNormalizationContext; + +type DeferredSessionSuspensionState = { + pending?: SessionSuspensionParams; +}; + +function flushDeferredSessionSuspension(state: DeferredSessionSuspensionState): void { + const pending = state.pending; + if (!pending) { + return; + } + state.pending = undefined; + void suspendSession(pending); +} + +function shouldDiscardDeferredSessionSuspension(params: { + error: unknown; + abortSignal?: AbortSignal; +}): boolean { + return ( + isTerminalAbort(params.abortSignal) || + shouldRethrowAbort(params.error) || + isCommandLaneTaskTimeoutError(params.error) || + isNonProviderRuntimeCoordinationError(params.error) || + isLikelyContextOverflowError(formatErrorMessage(params.error)) + ); +} + +export async function runWithModelFallback( + params: RunWithModelFallbackParams, +): Promise> { + const deferredSuspension: DeferredSessionSuspensionState = {}; + try { + const result = await runWithModelFallbackInternal(params, deferredSuspension); + if (result.outcome === "exhausted") { + flushDeferredSessionSuspension(deferredSuspension); + } + return result; + } catch (err) { + if (!shouldDiscardDeferredSessionSuspension({ error: err, abortSignal: params.abortSignal })) { + flushDeferredSessionSuspension(deferredSuspension); + } + throw err; + } +} + +async function runWithModelFallbackInternal( + params: RunWithModelFallbackParams, + deferredSuspension: DeferredSessionSuspensionState, ): Promise> { const candidates = resolveModelCandidateChain({ cfg: params.cfg, @@ -1311,6 +1373,8 @@ export async function runWithModelFallback( let latestClassifiedResult: ModelFallbackClassifiedResult | undefined; let exhaustionResult: ModelFallbackExhaustionResult | undefined; const cooldownProbeUsedProviders = new Set(); + const resolveTerminalSuspensionLane = () => + deferredSuspension.pending ? deferredSuspension.pending.laneId : params.lane; const observeDecision = async (decision: ModelFallbackDecisionParams) => { if (!params.onFallbackStep && !isModelFallbackDecisionLogEnabled()) { return; @@ -1445,6 +1509,10 @@ export async function runWithModelFallback( authMode, }); + // Only lock the lane when no remaining candidates can serve as + // fallbacks. Per-provider cooldown state already prevents + // re-attempting the failed provider on subsequent turns. + const hasRemainingCandidates = i + 1 < candidates.length; if (params.sessionId) { emitFailoverEvent({ sessionId: params.sessionId, @@ -1452,17 +1520,21 @@ export async function runWithModelFallback( fromProvider: candidate.provider, fromModel: candidate.model, reason: decision.reason, - suspended: true, - }); - void suspendSession({ - cfg: params.cfg, - agentDir: params.agentDir, - sessionId: params.sessionId, - laneId: params.lane, - reason: resolveSessionSuspensionReason(decision.reason), - failedProvider: candidate.provider, - failedModel: candidate.model, + suspended: !hasRemainingCandidates, }); + if (!hasRemainingCandidates) { + const laneId = resolveTerminalSuspensionLane(); + deferredSuspension.pending = undefined; + void suspendSession({ + cfg: params.cfg, + agentDir: params.agentDir, + sessionId: params.sessionId, + laneId, + reason: resolveSessionSuspensionReason(decision.reason), + failedProvider: candidate.provider, + failedModel: candidate.model, + }); + } } await observeDecision({ @@ -1584,6 +1656,13 @@ export async function runWithModelFallback( ...candidate, attempts, options: runOptions, + // Only the outer fallback loop knows another candidate remains. Carry + // that fact through this attempt so the embedded runner does not freeze + // the shared lane before the next candidate can run. + deferSessionSuspension: i + 1 < candidates.length, + onDeferredSessionSuspension: (suspension) => { + deferredSuspension.pending = suspension; + }, classifyResult: params.classifyResult, attempt: i + 1, total: candidates.length, @@ -1795,7 +1874,7 @@ export async function runWithModelFallback( cfg: params.cfg, candidates, }), - attribution: { sessionId: params.sessionId, lane: params.lane }, + attribution: { sessionId: params.sessionId, lane: resolveTerminalSuspensionLane() }, cfg: params.cfg, agentDir: params.agentDir, }); diff --git a/src/agents/session-suspension.test.ts b/src/agents/session-suspension.test.ts index 14774bcf7e50..5f76e4cfc724 100644 --- a/src/agents/session-suspension.test.ts +++ b/src/agents/session-suspension.test.ts @@ -120,6 +120,32 @@ describe("session suspension", () => { expect(patch.quotaSuspension?.expectedResumeBy).toBe(1_000 + MAX_TIMER_TIMEOUT_MS); }); + it("defers session suspension only for the outer fallback candidate run", async () => { + const { resolveSessionSuspensionTarget, runWithDeferredSessionSuspension } = + await import("./session-suspension.js"); + const onDeferred = vi.fn(); + + expect(resolveSessionSuspensionTarget()).toEqual({ mode: "suspend" }); + await runWithDeferredSessionSuspension(async () => { + const target = resolveSessionSuspensionTarget(); + expect(target.mode).toBe("defer"); + if (target.mode === "defer") { + target.defer({ + cfg: {}, + sessionId: "session-1", + laneId: CommandLane.Main, + reason: "quota_exhausted", + failedProvider: "openai", + failedModel: "gpt-5.5", + }); + } + expect(resolveSessionSuspensionTarget()).toEqual({ mode: "suspend" }); + }, onDeferred); + expect(onDeferred).toHaveBeenCalledOnce(); + expect(onDeferred).toHaveBeenCalledWith(expect.objectContaining({ laneId: CommandLane.Main })); + expect(resolveSessionSuspensionTarget()).toEqual({ mode: "suspend" }); + }); + it("maps failover reasons to persisted suspension reasons", async () => { const { testing } = await import("./session-suspension.js"); diff --git a/src/agents/session-suspension.ts b/src/agents/session-suspension.ts index dd55fa3d4398..2bedbc4923a1 100644 --- a/src/agents/session-suspension.ts +++ b/src/agents/session-suspension.ts @@ -3,6 +3,7 @@ * * Records quota/manual/circuit suspensions and temporarily lowers command-lane concurrency. */ +import { AsyncLocalStorage } from "node:async_hooks"; import path from "node:path"; import { resolveAgentMaxConcurrent, resolveSubagentMaxConcurrent } from "../config/agent-limits.js"; import { resolveCronMaxConcurrentRuns } from "../config/cron-limits.js"; @@ -23,8 +24,26 @@ const DEFAULT_CUSTOM_LANE_RESUME_CONCURRENCY = 1; export const DEFAULT_QUOTA_SUSPENSION_RESUME_MS = 30 * 60 * 1000; // 30 min const laneResumeTimers = new Map>(); +const deferredSessionSuspension = new AsyncLocalStorage<{ + claimed: boolean; + onDeferred?: (params: SessionSuspensionParams) => void; +}>(); export type SessionSuspensionReason = "quota_exhausted" | "manual" | "circuit_open"; +export type SessionSuspensionTarget = + | { mode: "defer"; defer: (params: SessionSuspensionParams) => void } + | { mode: "suspend" }; +export type SessionSuspensionParams = { + cfg: OpenClawConfig | undefined; + agentDir?: string; + sessionId: string; + laneId?: string; + reason: SessionSuspensionReason; + failedProvider: string; + failedModel: string; + summary?: string; + ttlMs?: number; +}; function resolveLaneResumeConcurrency(cfg: OpenClawConfig | undefined, laneId: string): number { switch (laneId) { @@ -50,6 +69,24 @@ export function resolveSessionSuspensionReason(reason: FailoverReason): SessionS return "circuit_open"; } +export function runWithDeferredSessionSuspension( + run: () => Promise, + onDeferred?: (params: SessionSuspensionParams) => void, +): Promise { + return deferredSessionSuspension.run({ claimed: false, onDeferred }, run); +} + +export function resolveSessionSuspensionTarget(): SessionSuspensionTarget { + const scope = deferredSessionSuspension.getStore(); + if (!scope || scope.claimed) { + return { mode: "suspend" }; + } + // One candidate callback may launch nested direct embedded runs. Only its + // first embedded run inherits the outer fallback's remaining-candidate fact. + scope.claimed = true; + return { mode: "defer", defer: (params) => scope.onDeferred?.(params) }; +} + function scheduleLaneAutoResume(laneId: string, delayMs: number, resumeConcurrency: number) { const existing = laneResumeTimers.get(laneId); if (existing) { @@ -78,17 +115,7 @@ export function cancelLaneAutoResume(laneId: string) { } } -export async function suspendSession(params: { - cfg: OpenClawConfig | undefined; - agentDir?: string; - sessionId: string; - laneId?: string; - reason: SessionSuspensionReason; - failedProvider: string; - failedModel: string; - summary?: string; - ttlMs?: number; -}) { +export async function suspendSession(params: SessionSuspensionParams) { if (!params.cfg) { return; }