fix(gateway): retain chat root through error persistence (#125985)

This commit is contained in:
Peter Steinberger
2026-08-18 13:33:35 -07:00
committed by GitHub
parent a57d9b5060
commit 1d36dac99d
4 changed files with 95 additions and 78 deletions
@@ -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).
@@ -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();
}
});
@@ -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 };
@@ -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<unknown>[] = [];
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();
}