Skip to content

Commit b3dc274

Browse files
scotthuangscotthuangvincentkoc
authored
fix(gateway): preserve active runs during plugin finalization (#92746)
* fix(gateway): preserve active run during plugin finalization * fix(ui): skip session.message history reload while gateway reports active run * fix(ui): remove unused eslint-disable directive * fix(ui): preserve active runs through finalization --------- Co-authored-by: scotthuang <[email protected]> Co-authored-by: Vincent Koc <[email protected]>
1 parent 1acca03 commit b3dc274

8 files changed

Lines changed: 213 additions & 17 deletions

‎src/gateway/server-runtime-subscriptions.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,7 @@ export function startGatewayEventSubscriptions(params: {
156156
broadcastToConnIds: params.broadcastToConnIds,
157157
sessionEventSubscribers: params.sessionEventSubscribers,
158158
sessionMessageSubscribers: params.sessionMessageSubscribers,
159+
chatAbortControllers: params.chatAbortControllers,
159160
}),
160161
);
161162
return transcriptUpdateHandlerPromise;

‎src/gateway/server-session-events.test.ts‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
import { beforeEach, describe, expect, it, vi } from "vitest";
2+
import type { ChatAbortControllerEntry } from "./chat-abort.js";
3+
4+
const sessionRow = vi.hoisted(() => ({
5+
key: "agent:main:main",
6+
kind: "direct",
7+
status: "done",
8+
updatedAt: 1,
9+
}));
10+
11+
vi.mock("../config/io.js", () => ({ getRuntimeConfig: () => ({}) }));
12+
vi.mock("./chat-display-projection.js", () => ({
13+
projectChatDisplayMessage: (message: unknown) => message,
14+
}));
15+
vi.mock("./session-utils.js", () => ({
16+
attachOpenClawTranscriptMeta: (message: unknown) => message,
17+
loadGatewaySessionRow: () => sessionRow,
18+
loadSessionEntry: () => ({ entry: undefined, storePath: "" }),
19+
readSessionMessageCountAsync: vi.fn(),
20+
}));
21+
22+
const { createTranscriptUpdateBroadcastHandler } = await import("./server-session-events.js");
23+
24+
function createActiveRun(projectSessionActive: boolean): ChatAbortControllerEntry {
25+
return {
26+
controller: new AbortController(),
27+
sessionId: "sess-main",
28+
sessionKey: "agent:main:main",
29+
startedAtMs: Date.now(),
30+
expiresAtMs: Date.now() + 60_000,
31+
projectSessionActive,
32+
};
33+
}
34+
35+
function createHandler(projectSessionActive: boolean) {
36+
const broadcastToConnIds = vi.fn();
37+
const handler = createTranscriptUpdateBroadcastHandler({
38+
broadcastToConnIds,
39+
sessionEventSubscribers: { getAll: () => new Set(["conn-1"]) },
40+
sessionMessageSubscribers: { get: () => new Set<string>() },
41+
chatAbortControllers: new Map([["run-before-finalize", createActiveRun(projectSessionActive)]]),
42+
});
43+
return { broadcastToConnIds, handler };
44+
}
45+
46+
async function emitAssistantTranscriptUpdate(projectSessionActive: boolean) {
47+
const { broadcastToConnIds, handler } = createHandler(projectSessionActive);
48+
handler({
49+
sessionFile: "/tmp/sess-main.jsonl",
50+
sessionKey: "agent:main:main",
51+
message: { role: "assistant", content: [{ type: "text", text: "Final answer" }] },
52+
messageId: "message-1",
53+
messageSeq: 1,
54+
});
55+
await vi.waitFor(() => expect(broadcastToConnIds).toHaveBeenCalledTimes(1));
56+
return broadcastToConnIds.mock.calls[0]?.[1];
57+
}
58+
59+
describe("createTranscriptUpdateBroadcastHandler", () => {
60+
beforeEach(() => {
61+
vi.clearAllMocks();
62+
});
63+
64+
it("keeps transcript snapshots active while plugin finalization delays the terminal event", async () => {
65+
// before_agent_finalize hooks run after the assistant transcript write but
66+
// before terminal delivery. The active-run registry remains authoritative
67+
// during that interval even when the persisted session row says done.
68+
await expect(emitAssistantTranscriptUpdate(true)).resolves.toMatchObject({
69+
sessionKey: "agent:main:main",
70+
hasActiveRun: true,
71+
session: { key: "agent:main:main", status: "done", hasActiveRun: true },
72+
});
73+
});
74+
75+
it("keeps stale-run recovery when terminal lifecycle has cleared active projection", async () => {
76+
await expect(emitAssistantTranscriptUpdate(false)).resolves.toMatchObject({
77+
sessionKey: "agent:main:main",
78+
hasActiveRun: false,
79+
session: { hasActiveRun: false },
80+
});
81+
});
82+
});

‎src/gateway/server-session-events.ts‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,14 @@ import { getRuntimeConfig } from "../config/io.js";
77
import { normalizeAgentId } from "../routing/session-key.js";
88
import type { SessionLifecycleEvent } from "../sessions/session-lifecycle-events.js";
99
import type { SessionTranscriptUpdate } from "../sessions/transcript-events.js";
10+
import type { ChatAbortControllerEntry } from "./chat-abort.js";
1011
import { projectChatDisplayMessage } from "./chat-display-projection.js";
1112
import type { GatewayBroadcastToConnIdsFn } from "./server-broadcast-types.js";
1213
import type {
1314
SessionEventSubscriberRegistry,
1415
SessionMessageSubscriberRegistry,
1516
} from "./server-chat.js";
17+
import { hasTrackedActiveSessionRun } from "./server-methods/session-active-runs.js";
1618
import { resolveSessionKeyForTranscriptFile } from "./session-transcript-key.js";
1719
import {
1820
attachOpenClawTranscriptMeta,
@@ -49,6 +51,7 @@ function buildGatewaySessionSnapshot(params: {
4951
label?: string;
5052
displayName?: string;
5153
parentSessionKey?: string;
54+
hasActiveRun?: boolean;
5255
}): Record<string, unknown> {
5356
const { sessionRow } = params;
5457
if (!sessionRow) {
@@ -61,6 +64,9 @@ function buildGatewaySessionSnapshot(params: {
6164
if (session && omitUnscopedGlobalGoal) {
6265
delete session.goal;
6366
}
67+
if (session && params.hasActiveRun !== undefined) {
68+
session.hasActiveRun = params.hasActiveRun;
69+
}
6470
return {
6571
...(session ? { session } : {}),
6672
updatedAt: sessionRow.updatedAt ?? undefined,
@@ -107,6 +113,7 @@ function buildGatewaySessionSnapshot(params: {
107113
modelProvider: sessionRow.modelProvider,
108114
model: sessionRow.model,
109115
status: sessionRow.status,
116+
...(params.hasActiveRun === undefined ? {} : { hasActiveRun: params.hasActiveRun }),
110117
subagentRunState: sessionRow.subagentRunState,
111118
hasActiveSubagentRun: sessionRow.hasActiveSubagentRun,
112119
startedAt: sessionRow.startedAt,
@@ -122,6 +129,7 @@ export function createTranscriptUpdateBroadcastHandler(params: {
122129
broadcastToConnIds: GatewayBroadcastToConnIdsFn;
123130
sessionEventSubscribers: SessionEventSubscribers;
124131
sessionMessageSubscribers: SessionMessageSubscribers;
132+
chatAbortControllers: Map<string, ChatAbortControllerEntry>;
125133
}) {
126134
let broadcastQueue = Promise.resolve();
127135
return (update: SessionTranscriptUpdate): void => {
@@ -138,6 +146,7 @@ async function handleTranscriptUpdateBroadcast(
138146
broadcastToConnIds: GatewayBroadcastToConnIdsFn;
139147
sessionEventSubscribers: SessionEventSubscribers;
140148
sessionMessageSubscribers: SessionMessageSubscribers;
149+
chatAbortControllers: Map<string, ChatAbortControllerEntry>;
141150
},
142151
update: SessionTranscriptUpdate,
143152
): Promise<void> {
@@ -176,13 +185,24 @@ async function handleTranscriptUpdateBroadcast(
176185
)
177186
: undefined;
178187
}
188+
const sessionRow = loadGatewaySessionRow(sessionKey, {
189+
agentId: visibleAgentId,
190+
transcriptUsageMaxBytes: 64 * 1024,
191+
});
192+
const hasActiveRun = sessionRow
193+
? hasTrackedActiveSessionRun({
194+
context: params,
195+
requestedKey: sessionKey,
196+
canonicalKey: sessionRow.key,
197+
...(sessionRow.key === "global" && visibleAgentId ? { agentId: visibleAgentId } : {}),
198+
defaultAgentId: normalizeAgentId(resolveDefaultAgentId(getRuntimeConfig())),
199+
})
200+
: false;
179201
const sessionSnapshot = buildGatewaySessionSnapshot({
180-
sessionRow: loadGatewaySessionRow(sessionKey, {
181-
agentId: visibleAgentId,
182-
transcriptUsageMaxBytes: 64 * 1024,
183-
}),
202+
sessionRow,
184203
agentId: visibleAgentId,
185204
includeSession: true,
205+
hasActiveRun,
186206
});
187207
const rawMessage = attachOpenClawTranscriptMeta(update.message, {
188208
...(typeof update.messageId === "string" ? { id: update.messageId } : {}),

‎ui/src/ui/app-gateway.sessions.node.test.ts‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ const loadChatHistoryMock = vi.fn();
66
const applySessionsChangedEventMock = vi.fn();
77
const clearPendingQueueItemsForRunMock = vi.fn();
88
const flushChatQueueForEventMock = vi.fn();
9-
const handleChatEventMock = vi.fn(() => "idle");
9+
const handleChatEventMock = vi.fn((_state?: any) => "idle");
1010
const handleSessionOperationEventMock = vi.fn();
1111
const recordFirstAssistantChatTimingMock = vi.fn();
1212

@@ -716,6 +716,29 @@ describe("handleGatewayEvent session.message", () => {
716716
expect(loadChatHistoryMock).not.toHaveBeenCalled();
717717
});
718718

719+
it("defers history reload when a session message confirms the chat run is still active", () => {
720+
loadChatHistoryMock.mockReset();
721+
applySessionsChangedEventMock.mockReset().mockReturnValue({ applied: false });
722+
loadSessionsMock.mockReset();
723+
const host = createHost();
724+
host.sessionKey = "agent:qa:main";
725+
host.chatRunId = "run-123";
726+
727+
handleGatewayEvent(host, {
728+
type: "event",
729+
event: "session.message",
730+
payload: { sessionKey: "agent:qa:main", hasActiveRun: true },
731+
seq: 1,
732+
});
733+
734+
expect(loadChatHistoryMock).not.toHaveBeenCalled();
735+
expect(loadSessionsMock).not.toHaveBeenCalled();
736+
expect(
737+
(host as typeof host & { pendingSessionMessageReloadSessionKey?: string | null })
738+
.pendingSessionMessageReloadSessionKey,
739+
).toBe("agent:qa:main");
740+
});
741+
719742
it("scopes selected-global session.message refreshes while a chat run is active", async () => {
720743
loadChatHistoryMock.mockReset();
721744
applySessionsChangedEventMock.mockReset().mockReturnValue({ applied: false });

‎ui/src/ui/app-gateway.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1176,6 +1176,12 @@ function handleSessionMessageGatewayEvent(
11761176
// chatStream, which delays the user message card from appearing until the
11771177
// first LLM delta arrives.
11781178
if (host.chatRunId) {
1179+
// Gateway confirms the run is still active (plugin hook window, etc.).
1180+
// Skip reload — the pending chat terminal owns history reconciliation.
1181+
if ((payload as Record<string, unknown> | null)?.hasActiveRun === true) {
1182+
deferredReloadHost.pendingSessionMessageReloadSessionKey = sessionKey;
1183+
return;
1184+
}
11791185
deferredReloadHost.pendingSessionMessageReloadSessionKey = sessionKey;
11801186
const refreshStartedAt = Date.now();
11811187
const runIdBeforeRefresh = host.chatRunId;

‎ui/src/ui/chat/run-lifecycle.test.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,25 @@ function rowActive(host: ReconcileHost): boolean {
3333
}
3434

3535
describe("reconcileChatRunFromCurrentSessionRow stale-active suppression (#87875)", () => {
36+
it("keeps a local run active when the gateway registry overrides a terminal snapshot", () => {
37+
const host = makeHost({
38+
chatRunId: "run-before-finalize",
39+
chatStream: "final answer",
40+
});
41+
42+
expect(
43+
reconcileChatRunFromSessionRow(host, {
44+
key: "s1",
45+
kind: "direct",
46+
updatedAt: 1,
47+
hasActiveRun: true,
48+
status: "done",
49+
}),
50+
).toBe(false);
51+
expect(host.chatRunId).toBe("run-before-finalize");
52+
expect(host.chatStream).toBe("final answer");
53+
});
54+
3655
it("suppresses a stale active row after a recent local completion", () => {
3756
const host = makeHost({
3857
lastLocalTerminalReconcile: {

‎ui/src/ui/chat/run-lifecycle.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,9 @@ export function reconcileChatRunFromSessionRow(
296296
if (!host.chatRunId && host.chatStream == null) {
297297
return false;
298298
}
299+
if (row.hasActiveRun === true) {
300+
return false;
301+
}
299302
if (isSessionRunActive(row)) {
300303
return false;
301304
}

‎ui/src/ui/controllers/sessions.test.ts‎

Lines changed: 54 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -655,7 +655,7 @@ describe("loadSessions", () => {
655655
expect(state.sessionsResult?.count).toBe(2);
656656
});
657657

658-
it("uses session list terminal state to clear stale local run tracking", async () => {
658+
it("keeps local run tracking while the session list reports an active terminal snapshot", async () => {
659659
vi.useFakeTimers();
660660
try {
661661
const request = vi.fn(async (method: string) => {
@@ -710,18 +710,17 @@ describe("loadSessions", () => {
710710

711711
await loadSessions(state);
712712

713-
expect(state.chatRunId).toBeNull();
714-
expect(state.chatStream).toBeNull();
715-
expect(state.chatStreamStartedAt).toBeNull();
716-
expect(state.compactionStatus).toBeNull();
717-
expect(state.compactionClearTimer).toBeNull();
718-
expect(state.fallbackStatus).toBeNull();
719-
expect(state.fallbackClearTimer).toBeNull();
720-
expect(state.chatRunStatus).toMatchObject({
721-
phase: "done",
722-
runId: "run-1",
723-
sessionKey: "main",
713+
expect(state.chatRunId).toBe("run-1");
714+
expect(state.chatStream).toBe("Visible answer");
715+
expect(state.chatStreamStartedAt).toBe(123);
716+
expect(state.compactionStatus).toMatchObject({ phase: "active", runId: "run-1" });
717+
expect(state.compactionClearTimer).not.toBeNull();
718+
expect(state.fallbackStatus).toMatchObject({
719+
selected: "openai/gpt-5.5",
720+
active: "anthropic/claude-sonnet-4-6",
724721
});
722+
expect(state.fallbackClearTimer).not.toBeNull();
723+
expect(state.chatRunStatus).toBeUndefined();
725724
} finally {
726725
vi.useRealTimers();
727726
}
@@ -1574,6 +1573,49 @@ describe("applySessionsChangedEvent", () => {
15741573
});
15751574
});
15761575

1576+
it("keeps the local run active when a transcript snapshot reports plugin finalization pending", () => {
1577+
const state = {
1578+
...createState(async () => undefined, {
1579+
sessionsResult: {
1580+
ts: 1,
1581+
path: "(multiple)",
1582+
count: 1,
1583+
defaults: { modelProvider: null, model: null, contextTokens: null },
1584+
sessions: [
1585+
{
1586+
key: "agent:main:main",
1587+
kind: "direct",
1588+
updatedAt: 1,
1589+
hasActiveRun: true,
1590+
status: "running",
1591+
},
1592+
],
1593+
},
1594+
}),
1595+
sessionKey: "agent:main:main",
1596+
chatRunId: "run-before-finalize",
1597+
} as SessionsState & { sessionKey: string; chatRunId: string | null };
1598+
1599+
const applied = applySessionsChangedEvent(state, {
1600+
sessionKey: "agent:main:main",
1601+
session: {
1602+
key: "agent:main:main",
1603+
kind: "direct",
1604+
updatedAt: 2,
1605+
status: "done",
1606+
hasActiveRun: true,
1607+
},
1608+
ts: 2,
1609+
});
1610+
1611+
expect(applied).toEqual({ applied: true, change: "updated" });
1612+
expect(state.chatRunId).toBe("run-before-finalize");
1613+
expect(state.sessionsResult?.sessions[0]).toMatchObject({
1614+
status: "done",
1615+
hasActiveRun: true,
1616+
});
1617+
});
1618+
15771619
it("clears the local chat run when an applied websocket patch makes the current session terminal", () => {
15781620
const requestUpdate = vi.fn();
15791621
const state: SessionsState & {

0 commit comments

Comments
 (0)