From b28a687830483e422747faddca226eadfdf46d2a Mon Sep 17 00:00:00 2001 From: Shakker Date: Fri, 21 Aug 2026 23:19:48 +0100 Subject: [PATCH] fix: scope terminal cleanup by run (#127646) --- .../src/app-server/transcript-mirror.test.ts | 5 ++ .../codex/src/app-server/transcript-mirror.ts | 14 ++++ ...hat-flow.active-run-follow-ups.e2e.test.ts | 43 +++++++---- ui/src/pages/chat/chat-state.test.ts | 70 +++++++++++++++--- ui/src/pages/chat/run-lifecycle.test.ts | 74 +++++++++++++++++++ ui/src/pages/chat/run-lifecycle.ts | 12 ++- ui/src/pages/chat/session-message-apply.ts | 20 +++-- ui/src/pages/chat/stream-causal-boundary.ts | 2 +- ui/src/pages/chat/stream-reconciliation.ts | 60 +++++++++++---- ui/src/pages/chat/tool-stream.ts | 43 +++++++++-- 10 files changed, 292 insertions(+), 51 deletions(-) diff --git a/extensions/codex/src/app-server/transcript-mirror.test.ts b/extensions/codex/src/app-server/transcript-mirror.test.ts index 22d9a4b09f02..b3db829d8652 100644 --- a/extensions/codex/src/app-server/transcript-mirror.test.ts +++ b/extensions/codex/src/app-server/transcript-mirror.test.ts @@ -1061,12 +1061,17 @@ describe("mirrorCodexAppServerTranscript", () => { ), ], idempotencyScope: "codex-app-server:thread-1", + terminalAssistantOwner: { + mirrorIdentity: "turn-1:assistant", + runId: "openclaw-run-1", + }, }); const updates = publishSessionTranscriptUpdateByIdentityMock.mock.calls.map( ([update]) => update as Record & { update?: Record }, ); expect(updates.map((update) => update.update?.messageSeq)).toEqual([1, 2]); + expect(updates.map((update) => update.update?.runId)).toEqual([undefined, "openclaw-run-1"]); expect( updates.map((update) => { const message = update.update?.message as { role?: string } | undefined; diff --git a/extensions/codex/src/app-server/transcript-mirror.ts b/extensions/codex/src/app-server/transcript-mirror.ts index 244dfeac3908..02a94feba24c 100644 --- a/extensions/codex/src/app-server/transcript-mirror.ts +++ b/extensions/codex/src/app-server/transcript-mirror.ts @@ -134,6 +134,10 @@ async function mirrorBestEffort(params: { // identity (not via the scope). Dropping `turnId` from the scope here is // what lets a re-emitted prior-turn entry collide with its existing key. idempotencyScope: `codex-app-server:${params.threadId}`, + terminalAssistantOwner: { + mirrorIdentity: `${params.turnId}:assistant`, + runId: params.params.runId, + }, config: params.params.config, }); for (const receipt of mirrorResult.userMessageReceipts) { @@ -323,6 +327,7 @@ async function mirror(params: { storePath?: string; messages: AgentMessage[]; idempotencyScope?: string; + terminalAssistantOwner?: { mirrorIdentity: string; runId: string }; config?: SessionTranscriptWriteLockParams["config"]; skipBeforeMessageWriteHooks?: boolean; }): Promise { @@ -498,6 +503,14 @@ async function mirror(params: { for (const update of appendedUpdates) { try { + // Commentary and tool rows share the Codex turn but cannot claim terminal run ownership. + const terminalOwner = params.terminalAssistantOwner; + const terminalRunId = + update.message.role === "assistant" && + terminalOwner && + readMirrorIdentity(update.message) === terminalOwner.mirrorIdentity + ? terminalOwner.runId + : undefined; await publishSessionTranscriptUpdateByIdentity({ ...transcriptTarget, update: { @@ -505,6 +518,7 @@ async function mirror(params: { message: update.message, messageId: update.messageId, ...(update.messageSeq !== undefined ? { messageSeq: update.messageSeq } : {}), + ...(terminalRunId ? { runId: terminalRunId } : {}), sessionKey: transcriptTarget.sessionKey, }, }); diff --git a/ui/src/e2e/chat-flow.active-run-follow-ups.e2e.test.ts b/ui/src/e2e/chat-flow.active-run-follow-ups.e2e.test.ts index 3a507f4ebe2f..df7245426e68 100644 --- a/ui/src/e2e/chat-flow.active-run-follow-ups.e2e.test.ts +++ b/ui/src/e2e/chat-flow.active-run-follow-ups.e2e.test.ts @@ -315,15 +315,16 @@ suite.define(() => { expect(steerBounds).not.toBeNull(); expect(streamingBounds).not.toBeNull(); expect(streamingBounds!.y).toBeGreaterThanOrEqual(steerBounds!.y + steerBounds!.height - 1); + const durableFinalMessage = { + role: "assistant", + content: [{ text: finalText, type: "text" }], + __openclaw: { id: "ui4-final", seq: 5 }, + }; await gateway.emitGatewayEvent("session.message", { - activeRunIds: [], + activeRunIds: [runId], clientRunId: runId, - hasActiveRun: false, - message: { - role: "assistant", - content: [{ text: finalText, type: "text" }], - __openclaw: { id: "ui4-final", seq: 5 }, - }, + hasActiveRun: true, + message: durableFinalMessage, messageId: "ui4-final", messageSeq: 5, runId, @@ -344,13 +345,10 @@ suite.define(() => { requestAnimationFrame(wait); }), ); + await streamingBubble.waitFor({ state: "detached" }); expect( - await page - .locator( - "[data-virtual-row-key^='stream-run:'] .chat-group.assistant:not(.chat-group--working)", - ) - .count(), - ).toBe(0); + await page.locator(".chat-thread-inner").getByText(finalText, { exact: true }).count(), + ).toBe(1); const overlaps = await page.locator(".chat-thread").evaluate((thread) => { const rows = Array.from(thread.querySelectorAll(".chat-virtual-row")) .map((row) => { @@ -371,6 +369,25 @@ suite.define(() => { }); }); expect(overlaps).toEqual([]); + await gateway.emitGatewayEvent("session.message", { + activeRunIds: [], + clientRunId: runId, + hasActiveRun: false, + message: durableFinalMessage, + messageId: "ui4-final", + messageSeq: 5, + runId, + sessionKey: "main", + }); + await expect + .poll(() => + page + .locator( + "[data-virtual-row-key^='stream-run:'] .chat-group.assistant:not(.chat-group--working)", + ) + .count(), + ) + .toBe(0); await gateway.emitChatFinal({ runId, text: finalText }); await expect .poll(() => diff --git a/ui/src/pages/chat/chat-state.test.ts b/ui/src/pages/chat/chat-state.test.ts index f2b42a94b00c..b55c5535345f 100644 --- a/ui/src/pages/chat/chat-state.test.ts +++ b/ui/src/pages/chat/chat-state.test.ts @@ -27,6 +27,7 @@ import { buildChatItems } from "./chat-thread-build.ts"; import { getChatSessionProjection, reduceChatSessionProjection } from "./history-merge.ts"; import { scheduleControlUiAfterPaint } from "./performance.ts"; import { applySessionMessagePayload } from "./session-message-apply.ts"; +import { buildToolStreamIdentity } from "./tool-stream-identity.ts"; beforeEach(() => { vi.spyOn(assistantIdentity, "loadLocalAssistantIdentity").mockReturnValue({ @@ -177,19 +178,35 @@ describe("canonical session message recovery", () => { it("retires the complete transient projection when the durable terminal arrives", () => { const runId = "active-run"; + const siblingRunId = "sibling-run"; const finalText = "The durable terminal reply."; const toolMessage = { role: "assistant", runId, toolCallId: "tool-1" }; + const siblingToolMessage = { + role: "assistant", + runId: siblingRunId, + toolCallId: "tool-2", + }; + const toolIdentity = buildToolStreamIdentity(runId, "tool-1"); + const siblingToolIdentity = buildToolStreamIdentity(siblingRunId, "tool-2"); const { state } = createSessionEventState({ connected: false, chatMessages: [], chatRunId: runId, chatStream: finalText, chatStreamStartedAt: 1, - chatStreamSegments: [{ text: "Commentary", ts: 1, runId, itemId: "commentary-1" }], - chatToolMessages: [toolMessage], + chatStreamSegments: [ + { text: "Commentary", ts: 1, runId, itemId: "commentary-1" }, + { + text: "Sibling commentary", + ts: 1, + runId: siblingRunId, + itemId: "commentary-2", + }, + ], + chatToolMessages: [toolMessage, siblingToolMessage], toolStreamById: new Map([ [ - "tool-1", + toolIdentity, { message: toolMessage, name: "exec", @@ -199,8 +216,28 @@ describe("canonical session message recovery", () => { toolCallId: "tool-1", }, ], + [ + siblingToolIdentity, + { + message: siblingToolMessage, + name: "read", + receivedAt: 1, + runId: siblingRunId, + startedAt: 1, + toolCallId: "tool-2", + }, + ], + ]), + toolStreamOrder: [toolIdentity, siblingToolIdentity], + activityEventSeqById: new Map([ + [`tool:${JSON.stringify([runId, "tool-1"])}:result`, 2], + [`tool:${JSON.stringify([siblingRunId, "tool-2"])}:result`, 2], + ]), + knownAgentRunIds: new Set([runId, siblingRunId]), + waitingApprovalStatuses: new Map([ + ["approval-1", { approvalId: "approval-1", toolCallId: "tool-1", runId }], + ["approval-2", { approvalId: "approval-2", toolCallId: "tool-2", runId: siblingRunId }], ]), - toolStreamOrder: ["tool-1"], }); applySessionMessagePayload( @@ -220,12 +257,27 @@ describe("canonical session message recovery", () => { { kind: "live", activeRunId: runId }, ); - expect(renderedTranscript(state)).toEqual([{ role: "assistant", text: finalText }]); + expect(state.chatMessages.filter((message) => extractText(message) === finalText)).toHaveLength( + 1, + ); expect(state.chatStream).toBeNull(); - expect(state.chatStreamSegments).toEqual([]); - expect(state.chatToolMessages).toEqual([]); - expect(state.toolStreamById).toEqual(new Map()); - expect(state.toolStreamOrder).toEqual([]); + expect(state.chatStreamSegments).toEqual([ + { + text: "Sibling commentary", + ts: 1, + runId: siblingRunId, + itemId: "commentary-2", + }, + ]); + expect(state.chatToolMessages).toEqual([siblingToolMessage]); + expect(state.toolStreamById.has(toolIdentity)).toBe(false); + expect(state.toolStreamById.has(siblingToolIdentity)).toBe(true); + expect(state.toolStreamOrder).toEqual([siblingToolIdentity]); + expect(state.knownAgentRunIds).toEqual(new Set([siblingRunId])); + expect([...state.waitingApprovalStatuses.keys()]).toEqual(["approval-2"]); + expect([...(state.activityEventSeqById?.keys() ?? [])]).toEqual([ + `tool:${JSON.stringify([siblingRunId, "tool-2"])}:result`, + ]); }); it("keeps cumulative assistant output split across an authoritative steer", () => { diff --git a/ui/src/pages/chat/run-lifecycle.test.ts b/ui/src/pages/chat/run-lifecycle.test.ts index a6b3c6cc0089..e650107ed1a4 100644 --- a/ui/src/pages/chat/run-lifecycle.test.ts +++ b/ui/src/pages/chat/run-lifecycle.test.ts @@ -16,6 +16,7 @@ import { reconcileStaleChatRunAfterSessionStatePublication, replayPendingChatAbort, } from "./run-lifecycle.ts"; +import { buildToolStreamIdentity } from "./tool-stream-identity.ts"; type ReconcileHost = Parameters[0]; type TestRow = { @@ -349,6 +350,79 @@ describe("reconcileChatRunLifecycle indicators", () => { }); }); +describe("reconcileChatRunFromSessionRow transient projections", () => { + it("clears only the terminal run's tool stream", () => { + const runId = "r1"; + const siblingRunId = "r2"; + const toolIdentity = buildToolStreamIdentity(runId, "tool-1"); + const siblingToolIdentity = buildToolStreamIdentity(siblingRunId, "tool-2"); + const toolMessage = { role: "assistant", runId, toolCallId: "tool-1" }; + const siblingToolMessage = { + role: "assistant", + runId: siblingRunId, + toolCallId: "tool-2", + }; + const host = makeHost({ + chatRunId: runId, + chatStream: "Final reply", + chatStreamSegments: [ + { text: "run one", ts: 1, runId }, + { text: "run two", ts: 2, runId: siblingRunId }, + ], + chatToolMessages: [toolMessage, siblingToolMessage], + toolStreamById: new Map([ + [ + toolIdentity, + { + message: toolMessage, + name: "exec", + receivedAt: 1, + runId, + startedAt: 1, + toolCallId: "tool-1", + }, + ], + [ + siblingToolIdentity, + { + message: siblingToolMessage, + name: "read", + receivedAt: 2, + runId: siblingRunId, + startedAt: 2, + toolCallId: "tool-2", + }, + ], + ]), + toolStreamOrder: [toolIdentity, siblingToolIdentity], + toolStreamSyncTimer: null, + knownAgentRunIds: new Set([runId, siblingRunId]), + waitingApprovalStatuses: new Map([ + ["approval-1", { approvalId: "approval-1", toolCallId: "tool-1", runId }], + ["approval-2", { approvalId: "approval-2", toolCallId: "tool-2", runId: siblingRunId }], + ]), + }); + + expect( + reconcileChatRunFromSessionRow(host, { + key: "s1", + kind: "direct", + updatedAt: 2, + hasActiveRun: false, + status: "done", + }), + ).toBe(true); + + expect(host.chatStreamSegments).toEqual([{ text: "run two", ts: 2, runId: siblingRunId }]); + expect(host.chatToolMessages).toEqual([siblingToolMessage]); + expect(host.toolStreamById?.has(toolIdentity)).toBe(false); + expect(host.toolStreamById?.has(siblingToolIdentity)).toBe(true); + expect(host.toolStreamOrder).toEqual([siblingToolIdentity]); + expect(host.knownAgentRunIds).toEqual(new Set([siblingRunId])); + expect([...host.waitingApprovalStatuses!.keys()]).toEqual(["approval-2"]); + }); +}); + describe("reconcileChatRunFromCurrentSessionRow stale-active suppression (#87875)", () => { it("keeps a local run active when the gateway registry overrides a terminal snapshot", () => { const host = makeHost({ diff --git a/ui/src/pages/chat/run-lifecycle.ts b/ui/src/pages/chat/run-lifecycle.ts index 04ecba2dfdf6..c123d00b7a47 100644 --- a/ui/src/pages/chat/run-lifecycle.ts +++ b/ui/src/pages/chat/run-lifecycle.ts @@ -24,6 +24,7 @@ import { resetChatInputHistoryNavigation, type ChatInputHistoryState } from "./i // Control UI chat module implements run lifecycle behavior. import { resetToolStream, + resetToolStreamRun, type CompactionStatus, type FallbackStatus, type WaitingApprovalStatus, @@ -83,6 +84,7 @@ type ReconcileOptions = { clearChatStream?: boolean; clearIndicators?: boolean; clearToolStream?: boolean; + clearToolStreamForRun?: boolean; clearRunStatus?: boolean; publishRunStatus?: boolean; armLocalTerminalReconcile?: boolean; @@ -452,8 +454,12 @@ export function reconcileChatRunLifecycle(host: RunLifecycleHost, options: Recon if (options.clearLocalRun) { host.chatRunId = null; } - if (options.clearToolStream && canResetToolStream(host)) { - resetToolStream(host); + if (canResetToolStream(host)) { + if (options.clearToolStream) { + resetToolStream(host); + } else if (options.clearToolStreamForRun && runId) { + resetToolStreamRun(host, runId); + } } if (options.outcome) { const status: ChatRunUiStatus = { @@ -606,7 +612,7 @@ export function reconcileChatRunFromSessionRow( sessionKeys: [row.key], clearLocalRun: true, clearChatStream: true, - clearToolStream: true, + clearToolStreamForRun: true, publishRunStatus: options.publishRunStatus, }); return true; diff --git a/ui/src/pages/chat/session-message-apply.ts b/ui/src/pages/chat/session-message-apply.ts index d11fe5407412..c3eafc6da58b 100644 --- a/ui/src/pages/chat/session-message-apply.ts +++ b/ui/src/pages/chat/session-message-apply.ts @@ -18,7 +18,10 @@ import { persistedSteerTargetRunId, rolloverChatStream, } from "./stream-causal-boundary.ts"; -import { maybeResetToolStream } from "./stream-reconciliation.ts"; +import { + assistantMessageReplacesCurrentStream, + maybeResetToolStreamRun, +} from "./stream-reconciliation.ts"; import { prunePersistedAssistantStreamSegments } from "./stream-segment-pruning.ts"; type SessionMessageApplySource = @@ -137,10 +140,17 @@ export function applySessionMessagePayload( ); if (incoming.role === "assistant" && projection.messages.includes(message)) { prunePersistedAssistantStreamSegments(state, message); - if (assistantOwnerRunId && runActive === false) { - state.chatStream = null; - state.chatStreamStartedAt = null; - maybeResetToolStream(state); + if (assistantOwnerRunId) { + if ( + runActive === false || + (state.chatStream !== null && assistantMessageReplacesCurrentStream(state, message)) + ) { + state.chatStream = null; + state.chatStreamStartedAt = null; + } + if (runActive === false) { + maybeResetToolStreamRun(state, assistantOwnerRunId); + } } } const steerTargetRunId = persistedSteerTargetRunId(message); diff --git a/ui/src/pages/chat/stream-causal-boundary.ts b/ui/src/pages/chat/stream-causal-boundary.ts index 2f7089d17f36..dba403d3634f 100644 --- a/ui/src/pages/chat/stream-causal-boundary.ts +++ b/ui/src/pages/chat/stream-causal-boundary.ts @@ -100,7 +100,7 @@ export function indexTurnContinuations( } export function latestPersistedSteerBoundary( - messages: unknown[], + messages: readonly unknown[], activeRunId: string, ): { index: number; runId: string } | null { for (let index = messages.length - 1; index >= 0; index -= 1) { diff --git a/ui/src/pages/chat/stream-reconciliation.ts b/ui/src/pages/chat/stream-reconciliation.ts index 677923c03870..ada697be263b 100644 --- a/ui/src/pages/chat/stream-reconciliation.ts +++ b/ui/src/pages/chat/stream-reconciliation.ts @@ -27,7 +27,7 @@ import { resolveLiveToolStreamRefs, resolveMatchingLiveToolIdentity, } from "./tool-stream-identity.ts"; -import { resetToolStream } from "./tool-stream.ts"; +import { resetToolStream, resetToolStreamRun } from "./tool-stream.ts"; type StreamReconciliationState = StreamCausalBoundaryState & { chatStream: string | null; @@ -64,6 +64,18 @@ type MaterializeVisibleStreamOptions = { isHiddenStreamText: StreamVisibility; }; +function resettableToolStreamHost( + state: StreamReconciliationState, +): Parameters[0] | null { + const toolHost = state as ToolStreamHost & Partial[0]>; + return toolHost.toolStreamById instanceof Map && + Array.isArray(toolHost.toolStreamOrder) && + Array.isArray(toolHost.chatToolMessages) && + Array.isArray(toolHost.chatStreamSegments) + ? (toolHost as Parameters[0]) + : null; +} + export function currentLiveToolCallIds(state: StreamReconciliationState): string[] { const toolHost = state as ToolStreamHost; return Array.isArray(toolHost.toolStreamOrder) @@ -89,20 +101,23 @@ export function maybeResetToolStream( state: StreamReconciliationState, opts?: { preserveStreamSegments?: boolean }, ) { - const toolHost = state as ToolStreamHost & Partial[0]>; - if ( - toolHost.toolStreamById instanceof Map && - Array.isArray(toolHost.toolStreamOrder) && - Array.isArray(toolHost.chatToolMessages) && - Array.isArray(toolHost.chatStreamSegments) - ) { - const preservedStreamSegments = opts?.preserveStreamSegments - ? [...toolHost.chatStreamSegments] - : null; - resetToolStream(toolHost as Parameters[0]); - if (preservedStreamSegments) { - toolHost.chatStreamSegments = preservedStreamSegments; - } + const toolHost = resettableToolStreamHost(state); + if (!toolHost) { + return; + } + const preservedStreamSegments = opts?.preserveStreamSegments + ? [...toolHost.chatStreamSegments] + : null; + resetToolStream(toolHost); + if (preservedStreamSegments) { + toolHost.chatStreamSegments = preservedStreamSegments; + } +} + +export function maybeResetToolStreamRun(state: StreamReconciliationState, runId: string) { + const toolHost = resettableToolStreamHost(state); + if (toolHost) { + resetToolStreamRun(toolHost, runId); } } @@ -266,6 +281,21 @@ function hasAssistantStreamReplacement( }); } +export function assistantMessageReplacesCurrentStream( + state: StreamReconciliationState, + message: unknown, +): boolean { + const currentPart = visibleAssistantStreamParts(state, { + includeCurrent: true, + isHiddenStreamText: () => false, + }).findLast((part) => part.source === "current"); + return Boolean( + currentPart && + (hasAssistantStreamReplacement([message], currentPart.replacementText, () => false, 0) || + hasAssistantStreamReplacement([message], currentPart.text, () => false, 0)), + ); +} + function streamFallbackItemId(message: unknown): string | null { if (!message || typeof message !== "object") { return null; diff --git a/ui/src/pages/chat/tool-stream.ts b/ui/src/pages/chat/tool-stream.ts index 7e542901606e..762d16e2f328 100644 --- a/ui/src/pages/chat/tool-stream.ts +++ b/ui/src/pages/chat/tool-stream.ts @@ -330,11 +330,15 @@ function syncToolStreamMessages(host: ToolStreamHost) { .filter((msg): msg is Record => Boolean(msg)); } -function flushToolStreamSync(host: ToolStreamHost) { +function cancelToolStreamSync(host: ToolStreamHost) { if (host.toolStreamSyncTimer != null) { clearTimeout(host.toolStreamSyncTimer); host.toolStreamSyncTimer = null; } +} + +function flushToolStreamSync(host: ToolStreamHost) { + cancelToolStreamSync(host); syncToolStreamMessages(host); } @@ -354,10 +358,7 @@ function scheduleToolStreamSync(host: ToolStreamHost, force = false) { } export function resetToolStream(host: ToolStreamHost) { - if (host.toolStreamSyncTimer != null) { - clearTimeout(host.toolStreamSyncTimer); - host.toolStreamSyncTimer = null; - } + cancelToolStreamSync(host); host.toolStreamById.clear(); host.toolStreamOrder = []; host.activityEventSeqById?.clear(); @@ -369,6 +370,38 @@ export function resetToolStream(host: ToolStreamHost) { // until snapshot reconciliation observes the approval leaving the queue. } +export function resetToolStreamRun(host: ToolStreamHost, runId: string) { + cancelToolStreamSync(host); + const removedIdentities = new Set(); + for (const identity of host.toolStreamOrder) { + const entry = host.toolStreamById.get(identity); + if (entry?.runId !== runId) { + continue; + } + removedIdentities.add(identity); + } + for (const identity of removedIdentities) { + host.toolStreamById.delete(identity); + } + const activityPrefix = `tool:[${JSON.stringify(runId)},`; + for (const sequenceIdentity of host.activityEventSeqById?.keys() ?? []) { + if (sequenceIdentity.startsWith(activityPrefix)) { + host.activityEventSeqById?.delete(sequenceIdentity); + } + } + host.toolStreamOrder = host.toolStreamOrder.filter( + (identity) => !removedIdentities.has(identity), + ); + syncToolStreamMessages(host); + host.chatStreamSegments = host.chatStreamSegments.filter((segment) => segment.runId !== runId); + host.knownAgentRunIds?.delete(runId); + for (const [approvalId, waitingApproval] of host.waitingApprovalStatuses ?? []) { + if (waitingApproval.runId === runId) { + host.waitingApprovalStatuses?.delete(approvalId); + } + } +} + function toolActivityIdentity(runId: string, toolCallId: string): string { return `tool:${JSON.stringify([runId, toolCallId])}`; }