From 1d36dac99d136a551b03fc8db773c0ec6dbd95b6 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 18 Aug 2026 13:33:35 -0700 Subject: [PATCH] fix(gateway): retain chat root through error persistence (#125985) --- .../chat-send-agent-dispatch.ts | 14 ++- .../chat-send-dispatch-errors.test.ts | 40 ++++--- .../chat-send-dispatch-errors.ts | 100 ++++++++---------- .../server.chat.gateway-server-chat.test.ts | 19 ++++ 4 files changed, 95 insertions(+), 78 deletions(-) diff --git a/src/gateway/server-methods/chat-send-agent-dispatch.ts b/src/gateway/server-methods/chat-send-agent-dispatch.ts index f14f6d9b1d56..7571f0d18181 100644 --- a/src/gateway/server-methods/chat-send-agent-dispatch.ts +++ b/src/gateway/server-methods/chat-send-agent-dispatch.ts @@ -213,7 +213,7 @@ export function startChatDispatch(params: StartChatDispatchParams): void { // Reserve the detached dispatch before this request releases its root. Otherwise // its inherited ALS context becomes retired and rejects queued/session work. setReleaseGatewayRootContinuation(retainGatewayRootWorkAdmissionContinuation() ?? undefined); - void replyDispatch + const dispatch = replyDispatch .runAgentMediaTranscript(gatewayWorkAdmission, () => measureDiagnosticsTimelineSpan( "gateway.chat_send.dispatch_inbound", @@ -511,9 +511,12 @@ export function startChatDispatch(params: StartChatDispatchParams): void { }); } }) - .catch(dispatchErrorLifecycle.handleError) - .finally(() => { - dispatchErrorLifecycle.finalize(); + .catch(dispatchErrorLifecycle.handleError); + void (async () => { + try { + await dispatch; + } finally { + await dispatchErrorLifecycle.finalize(); if (userTurnRecorder.isBlocked() && attachments.offloadedRefs.length > 0) { // A blocked turn persists only the redacted block reason — no media // markers — so the prepared inbound media stays unreferenced forever @@ -521,7 +524,8 @@ export function startChatDispatch(params: StartChatDispatchParams): void { // in chat-send-admission.ts: unreferenced staged media is discarded. void discardPreparedInboundMedia(attachments.offloadedRefs); } - }); + } + })(); // Title work starts at turn admission, concurrently with the launched run. It must never run // serially before dispatch (a cold utility runtime can starve the turn) or wait for completion // (long or interrupted first turns would silently remain untitled, and restart loses the chain). diff --git a/src/gateway/server-methods/chat-send-dispatch-errors.test.ts b/src/gateway/server-methods/chat-send-dispatch-errors.test.ts index 8f44ebbb1ae2..0169c90f2698 100644 --- a/src/gateway/server-methods/chat-send-dispatch-errors.test.ts +++ b/src/gateway/server-methods/chat-send-dispatch-errors.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it, vi } from "vitest"; +import { createDeferred } from "../../../test/helpers/promise.js"; import { retainLegacyDefaultAgentId } from "../../config/legacy.default-agent-owner.js"; import { onAgentRuntimeEvent } from "../../infra/agent-events.js"; import { abortChatRunById, registerChatAbortController } from "../chat-abort.js"; @@ -51,7 +52,7 @@ describe("createChatSendDispatchErrorLifecycle", () => { }); await lifecycle.handleError(new Error("late failure")); - lifecycle.finalize(); + await lifecycle.finalize(); expect(warn).toHaveBeenCalledWith( expect.stringContaining("dispatch failed after followup queue admission"), @@ -146,7 +147,7 @@ describe("createChatSendDispatchErrorLifecycle", () => { }); await lifecycle.handleError(new Error("dispatch rejected after explicit abort")); - lifecycle.finalize(); + await lifecycle.finalize(); expect(dedupe.get(`chat:${runId}`)).toMatchObject({ ok: true, @@ -289,7 +290,7 @@ describe("createChatSendDispatchErrorLifecycle", () => { ); }); - it("cleans up a failed non-default global send beside the compatibility owner's run", async () => { + it("keeps a failed non-default global send admitted through lifecycle persistence", async () => { const cfg = retainLegacyDefaultAgentId( { agents: { @@ -298,9 +299,14 @@ describe("createChatSendDispatchErrorLifecycle", () => { }, "main", ); + const persistenceEntered = createDeferred(); + const releasePersistence = createDeferred(); const persistLifecycleEvent = vi .spyOn(sessionLifecycleState, "persistGatewaySessionLifecycleEvent") - .mockResolvedValue(undefined); + .mockImplementation(async () => { + persistenceEntered.resolve(); + await releasePersistence.promise; + }); const cleanupAdmittedRun = vi.fn(); const activeRunCleanup = vi.fn(); const clientRunId = "failed-ops-global-send"; @@ -357,22 +363,24 @@ describe("createChatSendDispatchErrorLifecycle", () => { }); await lifecycle.handleError(new Error("dispatch rejected")); - lifecycle.finalize(); - - await vi.waitFor(() => { - expect(persistLifecycleEvent).toHaveBeenCalledWith({ - sessionKey: "global", - agentId: "ops", - event: expect.objectContaining({ - runId: clientRunId, - sessionId: "sess-ops", - data: expect.objectContaining({ phase: "error" }), - }), - }); + const finalization = lifecycle.finalize(); + await persistenceEntered.promise; + expect(persistLifecycleEvent).toHaveBeenCalledWith({ + sessionKey: "global", + agentId: "ops", + event: expect.objectContaining({ + runId: clientRunId, + sessionId: "sess-ops", + data: expect.objectContaining({ phase: "error" }), + }), }); + expect(cleanupAdmittedRun).not.toHaveBeenCalled(); + releasePersistence.resolve(); + await finalization; expect(activeRunCleanup).toHaveBeenCalledWith({ force: true }); expect(cleanupAdmittedRun).toHaveBeenCalledOnce(); } finally { + releasePersistence.resolve(); persistLifecycleEvent.mockRestore(); } }); diff --git a/src/gateway/server-methods/chat-send-dispatch-errors.ts b/src/gateway/server-methods/chat-send-dispatch-errors.ts index f1af2d3e2ade..1bc2ae34e309 100644 --- a/src/gateway/server-methods/chat-send-dispatch-errors.ts +++ b/src/gateway/server-methods/chat-send-dispatch-errors.ts @@ -1,6 +1,5 @@ import { ErrorCodes, errorShape } from "../../../packages/gateway-protocol/src/index.js"; import { clearAgentRunContext } from "../../infra/agent-run-registry.js"; -import { retainGatewayRootWorkAdmissionContinuation } from "../../process/gateway-work-admission.js"; import type { UserTurnTranscriptRecorder } from "../../sessions/user-turn-transcript.js"; import { setGatewayDedupeEntry } from "../agent-turn/agent-job.js"; import { chatAbortMarkerTimestampMs } from "../server-chat-state.js"; @@ -174,13 +173,6 @@ export function createChatSendDispatchErrorLifecycle(params: { const shouldPersistUserTurn = !userTurnRecorder.hasPersisted() && !userTurnRecorder.isBlocked(); - // Cleanup releases the admitted run; retain its root until the accepted - // user's transcript is settled so shutdown cannot drop that turn. - const releaseAbortTranscriptRoot = shouldPersistUserTurn - ? retainGatewayRootWorkAdmissionContinuation() - : null; - cleanupAdmittedRun(); - clearAgentRunContext(clientRunId, lifecycleGeneration); if (shouldPersistUserTurn) { try { await persistUserTurnTranscript(); @@ -188,8 +180,6 @@ export function createChatSendDispatchErrorLifecycle(params: { context.logGateway.warn( `webchat user transcript update failed after abort: ${formatForLog(transcriptError)}`, ); - } finally { - releaseAbortTranscriptRoot?.(); } } return; @@ -272,19 +262,19 @@ export function createChatSendDispatchErrorLifecycle(params: { } }; - const finalize = () => { + const finalize = async () => { const dispatchError = pendingDispatchLifecycleError; - // Reserve projection before cleanup retires the accepted dispatch root. - const releaseDispatchErrorRoot = dispatchError - ? retainGatewayRootWorkAdmissionContinuation() - : null; - cleanupAdmittedRun(); - clearAgentRunContext(clientRunId, lifecycleGeneration); - context.removeChatRun(clientRunId, clientRunId, sessionKey); if (!dispatchError) { + cleanupAdmittedRun(); + clearAgentRunContext(clientRunId, lifecycleGeneration); + context.removeChatRun(clientRunId, clientRunId, sessionKey); return; } - const persistDispatchLifecycleError = async () => { + // Stop exposing the rejected run before projecting its terminal state, but keep the + // admitted root until persistence settles so restart drain still observes this work. + clearAgentRunContext(clientRunId, lifecycleGeneration); + context.removeChatRun(clientRunId, clientRunId, sessionKey); + try { const hasActiveRun = hasTrackedActiveSessionRun({ context, requestedKey: rawSessionKey, @@ -292,51 +282,47 @@ export function createChatSendDispatchErrorLifecycle(params: { ...(agentId ? { agentId } : {}), defaultAgentId: tryResolveSessionCompatibilityOwnerAgentId(cfg, sessionKey), }); - if (hasActiveRun) { - return; - } - try { - await persistGatewaySessionLifecycleEvent({ - sessionKey, - ...(agentId ? { agentId } : {}), - event: { - runId: clientRunId, - sessionId: dispatchError.sessionId, - lifecycleGeneration, - ts: dispatchError.endedAt, - data: { - phase: "error", - startedAt: dispatchError.startedAt, - endedAt: dispatchError.endedAt, - error: dispatchError.error, + if (!hasActiveRun) { + try { + await persistGatewaySessionLifecycleEvent({ + sessionKey, + ...(agentId ? { agentId } : {}), + event: { + runId: clientRunId, + sessionId: dispatchError.sessionId, + lifecycleGeneration, + ts: dispatchError.endedAt, + data: { + phase: "error", + startedAt: dispatchError.startedAt, + endedAt: dispatchError.endedAt, + error: dispatchError.error, + }, }, - }, - }); - emitSessionsChanged(context, { - sessionKey, - ...(agentId ? { agentId } : {}), - reason: "chat.dispatch-error", - }); - } catch (persistErr: unknown) { - context.logGateway.warn( - `webchat session lifecycle persist failed after error: ${formatForLog(persistErr)}`, - ); + }); + emitSessionsChanged(context, { + sessionKey, + ...(agentId ? { agentId } : {}), + reason: "chat.dispatch-error", + }); + } catch (persistErr: unknown) { + context.logGateway.warn( + `webchat session lifecycle persist failed after error: ${formatForLog(persistErr)}`, + ); + } } - }; - void (async () => { - await persistDispatchLifecycleError(); await persistDispatchErrorUserTurn?.().catch((transcriptErr: unknown) => { context.logGateway.warn( `webchat user transcript update failed after error: ${formatForLog(transcriptErr)}`, ); }); - })() - .catch((continuationErr: unknown) => { - context.logGateway.warn( - `webchat session lifecycle continuation failed: ${formatForLog(continuationErr)}`, - ); - }) - .finally(() => releaseDispatchErrorRoot?.()); + } catch (continuationErr: unknown) { + context.logGateway.warn( + `webchat session lifecycle continuation failed: ${formatForLog(continuationErr)}`, + ); + } finally { + cleanupAdmittedRun(); + } }; return { finalize, handleError }; diff --git a/src/gateway/server.chat.gateway-server-chat.test.ts b/src/gateway/server.chat.gateway-server-chat.test.ts index bcc560e584f8..4f3a6f976ca7 100644 --- a/src/gateway/server.chat.gateway-server-chat.test.ts +++ b/src/gateway/server.chat.gateway-server-chat.test.ts @@ -327,6 +327,10 @@ describe("gateway server chat", () => { expect(res.payload?.status).toBe("ok"); return res; }; + const waitForAgentRunDrained = async (runId: string) => { + await waitForAgentRunOk(runId); + await waitForFast(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); + }; const abortChatRun = async (runId: string) => { const res = await rpcReq(ws, "chat.abort", { sessionKey: "main", @@ -386,6 +390,7 @@ describe("gateway server chat", () => { expect(res.ok).toBe(true); expect(res.payload?.runId).toBe("idem-sessions-send-1"); expect(res.payload?.messageSeq).toBe(1); + await waitForAgentRunDrained("idem-sessions-send-1"); } finally { testState.sessionStorePath = undefined; await removeTempDir(dir); @@ -415,6 +420,7 @@ describe("gateway server chat", () => { storePath: testState.sessionStorePath, })?.sessionId, ).toBeTypeOf("string"); + await waitForAgentRunDrained("idem-sessions-send-orion"); } finally { testState.agentsConfig = undefined; testState.sessionStorePath = undefined; @@ -443,6 +449,7 @@ describe("gateway server chat", () => { expect(res.ok).toBe(true); expect(res.payload?.runId).toBe("idem-sessions-steer-1"); expect(res.payload?.messageSeq).toBe(1); + await waitForAgentRunDrained("idem-sessions-steer-1"); } finally { testState.sessionStorePath = undefined; await removeTempDir(dir); @@ -518,6 +525,7 @@ describe("gateway server chat", () => { } else { expect(abortRes.payload?.abortedRunId).toBeNull(); } + await waitForFast(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); } finally { testState.sessionStorePath = undefined; await removeTempDir(dir); @@ -553,6 +561,7 @@ describe("gateway server chat", () => { if (abortRes.payload?.status === "aborted") { expect(abortRes.payload?.abortedRunId).toBe("idem-sessions-abort-runid-1"); } + await waitForFast(() => expect(getActiveGatewayRootWorkCount()).toBe(0)); } finally { testState.sessionStorePath = undefined; await removeTempDir(dir); @@ -576,6 +585,7 @@ describe("gateway server chat", () => { idempotencyKey: "idem-sanitized-1", }); expect(sanitizedRes.ok).toBe(true); + await waitForAgentRunDrained("idem-sanitized-1"); }); test("handles chat send and history flows", async () => { @@ -605,6 +615,7 @@ describe("gateway server chat", () => { idempotencyKey: "idem-webchat-1", }); expect(webchatRes.ok).toBe(true); + await waitForAgentRunDrained("idem-webchat-1"); webchatWs.close(); webchatWs = undefined; @@ -617,6 +628,7 @@ describe("gateway server chat", () => { }); expect(timeoutRes.ok).toBe(true); expect(timeoutRes.payload?.runId).toBe("idem-timeout-1"); + await waitForAgentRunDrained("idem-timeout-1"); testState.agentConfig = undefined; const sessionRes = await rpcReq(ws, "chat.send", { @@ -626,6 +638,7 @@ describe("gateway server chat", () => { }); expect(sessionRes.ok).toBe(true); expect(sessionRes.payload?.runId).toBe("idem-session-key-1"); + await waitForAgentRunDrained("idem-session-key-1"); const sendPolicyDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-gw-")); tempDirs.push(sendPolicyDir); @@ -719,6 +732,7 @@ describe("gateway server chat", () => { }); expect(imgRes.ok).toBe(true); expectStringRunId(imgRes.payload); + await waitForAgentRunDrained("idem-img"); const imgOnlyRes = await rpcReq(ws, "chat.send", { sessionKey: "main", message: "", @@ -734,6 +748,7 @@ describe("gateway server chat", () => { }); expect(imgOnlyRes.ok).toBe(true); expectStringRunId(imgOnlyRes.payload); + await waitForAgentRunDrained("idem-img-only"); const historyDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-gw-")); tempDirs.push(historyDir); @@ -888,6 +903,7 @@ describe("gateway server chat", () => { await releasePersistence.promise; await persistLifecycleEvent(params); }); + const messagePromises: Promise[] = []; const sessionChanged = await (async () => { try { dispatchInboundMessageMock.mockImplementationOnce(async () => { @@ -904,6 +920,7 @@ describe("gateway server chat", () => { o.payload?.runId === "idem-dispatch-error-1", 8_000, ); + messagePromises.push(errorPromise); const sessionChangedPromise = onceMessage( ws, (o) => @@ -913,6 +930,7 @@ describe("gateway server chat", () => { o.payload?.sessionKey === "agent:main:main", 8_000, ); + messagePromises.push(sessionChangedPromise); const res = await rpcReq(ws, "chat.send", { sessionKey: "main", message: "run: pwd", @@ -936,6 +954,7 @@ describe("gateway server chat", () => { } finally { rejectDispatch.resolve(); releasePersistence.resolve(); + await Promise.allSettled(messagePromises); persistSpy.mockRestore(); resetGatewayWorkAdmission(); }