From ac6009eab195ec09394170c57f8ddef054f6fbac Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 27 Jul 2026 04:39:42 -0400 Subject: [PATCH] fix(tui): keep shared sessions live and bound stream state (#114443) Fix cross-client terminal history synchronization without letting stale snapshots erase in-flight responses. Bound orphaned stream assembly, preserve session and agent ownership, and cover real Gateway delivery, persistence races, reconnects, and stress bursts. Closes #38829. Credits prior issue investigation in #96252 and bounded-stream work in #109492. Co-authored-by: Harjoth Khara <48686985+harjothkhara@users.noreply.github.com> Co-authored-by: Wynne668 <290981215+ZengWen-DT@users.noreply.github.com> --- docs/web/tui.md | 1 + src/tui/tui-event-handlers.test.ts | 337 +++++++++++++++++++++++++++ src/tui/tui-event-handlers.ts | 132 +++++++++-- src/tui/tui-pty-local.e2e.test.ts | 52 +++++ src/tui/tui-stream-assembler.test.ts | 56 +++++ src/tui/tui-stream-assembler.ts | 40 ++-- src/tui/tui-types.ts | 7 + src/tui/tui.ts | 7 + 8 files changed, 604 insertions(+), 28 deletions(-) diff --git a/docs/web/tui.md b/docs/web/tui.md index e39ed8681471..d00ceedc9be6 100644 --- a/docs/web/tui.md +++ b/docs/web/tui.md @@ -219,6 +219,7 @@ Tips: - On connect, the TUI loads the latest history (default 200 messages). - Streaming responses update in place until finalized. +- Messages sent to the same session from another client appear automatically. - The TUI also listens to agent tool events for richer tool cards. ## Connection details diff --git a/src/tui/tui-event-handlers.test.ts b/src/tui/tui-event-handlers.test.ts index 93b64b95a5f0..18c5b517af4e 100644 --- a/src/tui/tui-event-handlers.test.ts +++ b/src/tui/tui-event-handlers.test.ts @@ -8,6 +8,7 @@ import type { BtwEvent, ChatEvent, SessionChangedEvent, + SessionMessageEvent, TuiHistoryLoadResult, TuiStateAccess, } from "./tui-types.js"; @@ -1933,6 +1934,342 @@ describe("tui-event-handlers: handleAgentEvent", () => { expect(loadHistory).toHaveBeenCalledTimes(1); }); + describe("session.message history reload", () => { + it("reloads the current session when another client appends a message", () => { + const { state, loadHistory, handleSessionMessageEvent } = createHandlersHarness({ + state: { + activeChatRunId: null, + currentSessionId: "session-before", + sessionInfo: { verboseLevel: "on", updatedAt: 100 }, + }, + }); + + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + sessionId: "session-after", + updatedAt: 200, + } satisfies SessionMessageEvent); + + expect(state.currentSessionId).toBe("session-after"); + expect(state.sessionInfo.updatedAt).toBe(200); + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + + it("accepts the canonical session's unscoped alias", () => { + const { loadHistory, handleSessionMessageEvent } = createHandlersHarness({ + state: { activeChatRunId: null, currentSessionKey: "agent:main:main" }, + }); + + handleSessionMessageEvent({ sessionKey: "main" } satisfies SessionMessageEvent); + + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + + it("ignores an unscoped alias owned by a different agent", () => { + const { state, loadHistory, handleSessionMessageEvent } = createHandlersHarness({ + state: { + activeChatRunId: null, + currentAgentId: "work", + currentSessionId: "session-work", + currentSessionKey: "agent:work:main", + sessionInfo: { verboseLevel: "on", updatedAt: 100 }, + }, + }); + + handleSessionMessageEvent({ + sessionKey: "main", + agentId: "main", + sessionId: "session-default-agent", + updatedAt: 200, + } satisfies SessionMessageEvent); + + expect(state.currentSessionId).toBe("session-work"); + expect(state.sessionInfo.updatedAt).toBe(100); + expect(loadHistory).not.toHaveBeenCalled(); + }); + + it("does not assign an unqualified default-agent alias to another agent", () => { + const { state, loadHistory, handleSessionMessageEvent } = createHandlersHarness({ + state: { + activeChatRunId: null, + currentAgentId: "work", + currentSessionId: "session-work", + currentSessionKey: "agent:work:main", + }, + }); + + handleSessionMessageEvent({ + sessionKey: "main", + sessionId: "session-default-agent", + } satisfies SessionMessageEvent); + + expect(state.currentSessionId).toBe("session-work"); + expect(loadHistory).not.toHaveBeenCalled(); + }); + + it("ignores messages for another session without changing selected metadata", () => { + const { state, loadHistory, tui, handleSessionMessageEvent } = createHandlersHarness({ + state: { + activeChatRunId: null, + currentAgentId: "work", + currentSessionId: "session-before", + currentSessionKey: "agent:work:main", + sessionInfo: { verboseLevel: "on", updatedAt: 100 }, + }, + }); + + handleSessionMessageEvent({ + sessionKey: "agent:work:other", + agentId: "work", + sessionId: "other-session", + updatedAt: 200, + } satisfies SessionMessageEvent); + + expect(state.currentSessionId).toBe("session-before"); + expect(state.sessionInfo.updatedAt).toBe(100); + expect(loadHistory).not.toHaveBeenCalled(); + expect(tui.requestRender).not.toHaveBeenCalled(); + }); + + it("does not let an older transcript snapshot replace newer session metadata", () => { + const { state, loadHistory, handleSessionMessageEvent } = createHandlersHarness({ + state: { + activeChatRunId: null, + currentSessionId: "session-current", + sessionInfo: { verboseLevel: "on", updatedAt: 200 }, + }, + }); + + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + sessionId: "session-stale", + updatedAt: 100, + } satisfies SessionMessageEvent); + + expect(state.currentSessionId).toBe("session-current"); + expect(state.sessionInfo.updatedAt).toBe(200); + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + + it("reloads a global session only for its selected agent", () => { + const { loadHistory, handleSessionMessageEvent } = createHandlersHarness({ + state: { + agentDefaultId: "main", + activeChatRunId: null, + currentAgentId: "work", + currentSessionKey: "global", + sessionScope: "global", + }, + }); + + handleSessionMessageEvent({ + sessionKey: "global", + agentId: "main", + } satisfies SessionMessageEvent); + expect(loadHistory).not.toHaveBeenCalled(); + + handleSessionMessageEvent({ + sessionKey: "global", + agentId: "work", + } satisfies SessionMessageEvent); + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + + it("coalesces a burst of transcript updates into one follow-up reload", async () => { + let resolveFirstHistory: ((result: TuiHistoryLoadResult) => void) | undefined; + const { state, loadHistory, handleSessionMessageEvent } = createHandlersHarness({ + state: { activeChatRunId: null }, + }); + loadHistory.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveFirstHistory = resolve; + }), + ); + + for (let index = 0; index < 250; index += 1) { + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + updatedAt: index, + } satisfies SessionMessageEvent); + } + + expect(loadHistory).toHaveBeenCalledTimes(1); + expect(state.sessionInfo.updatedAt).toBe(249); + + resolveFirstHistory?.({ loaded: true, inFlightRunId: null }); + await vi.waitFor(() => expect(loadHistory).toHaveBeenCalledTimes(2)); + }); + + it("waits for terminal persistence before refreshing a displayed local final", () => { + const { + state, + chatLog, + loadHistory, + handleChatEvent, + handleSessionsChangedEvent, + handleSessionMessageEvent, + } = createHandlersHarness({ state: { activeChatRunId: "run-active" } }); + + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + } satisfies SessionMessageEvent); + handleChatEvent({ + runId: "run-active", + sessionKey: state.currentSessionKey, + state: "final", + message: { content: [{ type: "text", text: "keep this visible" }] }, + }); + + expect(chatLog.finalizeAssistant).toHaveBeenCalledWith("keep this visible", "run-active"); + expect(loadHistory).not.toHaveBeenCalled(); + + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + updatedAt: 200, + } satisfies SessionMessageEvent); + expect(loadHistory).not.toHaveBeenCalled(); + + handleSessionsChangedEvent({ + sessionKey: state.currentSessionKey, + runId: "run-active", + phase: "end", + } satisfies SessionChangedEvent); + + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + + it("refreshes after a local final when terminal persistence arrives first", () => { + const { + state, + chatLog, + loadHistory, + handleChatEvent, + handleSessionsChangedEvent, + handleSessionMessageEvent, + } = createHandlersHarness({ state: { activeChatRunId: "run-active" } }); + + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + } satisfies SessionMessageEvent); + handleSessionsChangedEvent({ + sessionKey: state.currentSessionKey, + runId: "run-active", + phase: "end", + } satisfies SessionChangedEvent); + expect(loadHistory).not.toHaveBeenCalled(); + + handleChatEvent({ + runId: "run-active", + sessionKey: state.currentSessionKey, + state: "final", + message: { content: [{ type: "text", text: "keep this visible" }] }, + }); + + expect(chatLog.finalizeAssistant).toHaveBeenCalledWith("keep this visible", "run-active"); + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + + it("defers the first external update after a visible final until persistence", () => { + const { + state, + chatLog, + loadHistory, + handleChatEvent, + handleSessionsChangedEvent, + handleSessionMessageEvent, + } = createHandlersHarness({ state: { activeChatRunId: "run-active" } }); + + handleChatEvent({ + runId: "run-active", + sessionKey: state.currentSessionKey, + state: "final", + message: { content: [{ type: "text", text: "keep this visible" }] }, + }); + + expect(chatLog.finalizeAssistant).toHaveBeenCalledWith("keep this visible", "run-active"); + expect(loadHistory).not.toHaveBeenCalled(); + + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + } satisfies SessionMessageEvent); + + expect(loadHistory).not.toHaveBeenCalled(); + + handleSessionsChangedEvent({ + sessionKey: state.currentSessionKey, + runId: "run-active", + phase: "end", + } satisfies SessionChangedEvent); + + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + + it("does not reload until an optimistic submit is resolved", () => { + const { state, loadHistory, handleSessionMessageEvent, flushPendingHistoryRefreshIfIdle } = + createHandlersHarness({ + state: { activeChatRunId: null, pendingSubmit: sendingSubmit("run-pending") }, + }); + + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + } satisfies SessionMessageEvent); + expect(loadHistory).not.toHaveBeenCalled(); + + state.pendingSubmit = null; + flushPendingHistoryRefreshIfIdle(); + + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + + it("waits for persistence when a refresh arrives before an optimistic run is accepted", () => { + const { + state, + chatLog, + loadHistory, + handleAgentEvent, + handleChatEvent, + handleSessionsChangedEvent, + handleSessionMessageEvent, + } = createHandlersHarness({ + state: { activeChatRunId: null, pendingSubmit: sendingSubmit("run-pending") }, + }); + + handleSessionMessageEvent({ + sessionKey: state.currentSessionKey, + } satisfies SessionMessageEvent); + expect(loadHistory).not.toHaveBeenCalled(); + + state.pendingSubmit = acceptedSubmit("run-pending"); + handleAgentEvent({ + runId: "run-pending", + sessionKey: state.currentSessionKey, + stream: "lifecycle", + data: { phase: "start" }, + } satisfies AgentEvent); + handleChatEvent({ + runId: "run-pending", + sessionKey: state.currentSessionKey, + state: "final", + message: { content: [{ type: "text", text: "keep accepted output visible" }] }, + }); + + expect(chatLog.finalizeAssistant).toHaveBeenCalledWith( + "keep accepted output visible", + "run-pending", + ); + expect(loadHistory).not.toHaveBeenCalled(); + + handleSessionsChangedEvent({ + sessionKey: state.currentSessionKey, + runId: "run-pending", + phase: "end", + } satisfies SessionChangedEvent); + + expect(loadHistory).toHaveBeenCalledTimes(1); + }); + }); + describe("sessions.changed history reload", () => { const startRun = ( state: TuiStateAccess, diff --git a/src/tui/tui-event-handlers.ts b/src/tui/tui-event-handlers.ts index 73567203c11a..2cf738e7517c 100644 --- a/src/tui/tui-event-handlers.ts +++ b/src/tui/tui-event-handlers.ts @@ -22,6 +22,7 @@ import type { BtwEvent, ChatEvent, SessionChangedEvent, + SessionMessageEvent, TuiHistoryLoadResult, TuiStateAccess, } from "./tui-types.js"; @@ -119,12 +120,15 @@ export function createEventHandlers(context: EventHandlerContext) { const queuedHistoryReloadRunIds = new Set(); const deferredHistoryRunEvents = new Map(); let historyReloadInFlight = false; + let historyReloadQueued = false; let historyReloadGeneration = 0; const completedRuns = new Map(); const postFinalizingRuns = new Map(); let streamAssembler = new TuiStreamAssembler(); let lastSessionKey = state.currentSessionKey; let pendingHistoryRefresh = false; + let pendingSessionMessageRefresh = false; + let pendingSessionMessageRunId: string | null = null; let reconnectPendingRunId: string | null = null; const pendingTerminalLifecycleErrors = new Map< string, @@ -164,11 +168,19 @@ export function createEventHandlers(context: EventHandlerContext) { }; const flushPendingHistoryRefreshIfIdle = () => { - if (!pendingHistoryRefresh || state.activeChatRunId || hasPendingSubmit(state)) { + if (state.activeChatRunId || hasPendingSubmit(state)) { + return; + } + const canRefreshSessionMessage = + pendingSessionMessageRefresh && pendingSessionMessageRunId === null; + if (!pendingHistoryRefresh && !canRefreshSessionMessage) { return; } pendingHistoryRefresh = false; - void reloadHistoryPreservingTerminalErrors(); + if (canRefreshSessionMessage) { + pendingSessionMessageRefresh = false; + } + queueHistoryReload(); }; const clearStreamingWatchdog = () => { @@ -208,6 +220,7 @@ export function createEventHandlers(context: EventHandlerContext) { historyDisplayedReloadRunIds.clear(); liveTerminalErrorMessages.clear(); queuedHistoryReloadRunIds.clear(); + historyReloadQueued = false; deferredHistoryRunEvents.clear(); finalizedRuns.clear(); finalizedRunsWithDisplay.clear(); @@ -216,6 +229,8 @@ export function createEventHandlers(context: EventHandlerContext) { postFinalizingRuns.clear(); streamAssembler = new TuiStreamAssembler(); pendingHistoryRefresh = false; + pendingSessionMessageRefresh = false; + pendingSessionMessageRunId = null; clearPendingSubmit(state); reconnectPendingRunId = null; clearLocalRunIds?.(); @@ -245,7 +260,7 @@ export function createEventHandlers(context: EventHandlerContext) { state.activityStatus = "idle"; setActivityStatus("idle"); pendingHistoryRefresh = false; - void reloadHistoryPreservingTerminalErrors(); + queueHistoryReload(); tui.requestRender(); return; } @@ -345,6 +360,16 @@ export function createEventHandlers(context: EventHandlerContext) { }; const markSubmittedRunRegistered = (runId: string) => { + if ( + pendingSessionMessageRefresh && + pendingSessionMessageRunId === null && + state.pendingSubmit?.runId === runId && + !persistedTerminalRunIds.has(runId) + ) { + // A transcript invalidation can arrive before Gateway accepts this submit. + // Bind it before clearing the draft so a live final survives until persistence. + pendingSessionMessageRunId = runId; + } clearPendingSubmitDraft(state, runId); }; @@ -569,7 +594,7 @@ export function createEventHandlers(context: EventHandlerContext) { return; } pendingHistoryRefresh = false; - void reloadHistoryPreservingTerminalErrors(); + queueHistoryReload(); }; const messageHasDisplayableNonTextContent = (message: unknown): boolean => { @@ -810,12 +835,13 @@ export function createEventHandlers(context: EventHandlerContext) { }; const drainHistoryReloadQueue = () => { - if (historyReloadInFlight || queuedHistoryReloadRunIds.size === 0 || !loadHistory) { + if (historyReloadInFlight || !historyReloadQueued || !loadHistory) { return; } const reloadGeneration = historyReloadGeneration; const runIds = Array.from(queuedHistoryReloadRunIds); queuedHistoryReloadRunIds.clear(); + historyReloadQueued = false; historyReloadInFlight = true; const finishReload = (result: TuiHistoryLoadResult) => { if (reloadGeneration !== historyReloadGeneration) { @@ -851,15 +877,16 @@ export function createEventHandlers(context: EventHandlerContext) { }); }; - const queueHistoryReload = ( - runIds: Iterable, - historyOwnedRunIds: Iterable, + function queueHistoryReload( + runIds?: Iterable, + historyOwnedRunIds: Iterable = [], displayedRunIds: Iterable = [], - ) => { + ) { const historyOwned = new Set(historyOwnedRunIds); const displayed = new Set(displayedRunIds); + const queuedRunIds = runIds ?? []; if (!loadHistory) { - for (const runId of runIds) { + for (const runId of queuedRunIds) { if (historyOwned.has(runId)) { noteFinalizedRun(runId, { displayedFinal: true }); } @@ -867,7 +894,11 @@ export function createEventHandlers(context: EventHandlerContext) { void refreshSessionInfo?.(); return; } - for (const runId of runIds) { + if (runIds === undefined) { + historyReloadQueued = true; + } + for (const runId of queuedRunIds) { + historyReloadQueued = true; historyReloadRunIds.add(runId); queuedHistoryReloadRunIds.add(runId); if (historyOwned.has(runId)) { @@ -878,7 +909,7 @@ export function createEventHandlers(context: EventHandlerContext) { } } drainHistoryReloadQueue(); - }; + } const collectTrackedSessionRunIds = () => { const runIds = new Set(sessionRuns.keys()); @@ -913,6 +944,9 @@ export function createEventHandlers(context: EventHandlerContext) { if (evt.runId && (evt.phase === "end" || evt.phase === "error")) { persistedTerminalRunIds.set(evt.runId, Date.now()); pruneRunMap(persistedTerminalRunIds); + if (pendingSessionMessageRunId === evt.runId) { + pendingSessionMessageRunId = null; + } if (pendingNewSessionRunIds.delete(evt.runId)) { if (evt.phase === "end") { const displayedRunIds = finalizedRunsWithDisplay.has(evt.runId) ? [evt.runId] : []; @@ -921,6 +955,7 @@ export function createEventHandlers(context: EventHandlerContext) { void refreshSessionInfo?.(); } } + flushPendingHistoryRefreshIfIdle(); return; } if (evt.reason !== "new" && evt.reason !== "reset") { @@ -972,13 +1007,76 @@ export function createEventHandlers(context: EventHandlerContext) { } if (reloadingRunIds.size > 0) { queueHistoryReload(reloadingRunIds, finalizedRunIds, displayedRunIds); - } else if (loadHistory) { - void reloadHistoryPreservingTerminalErrors(); } else { - void refreshSessionInfo?.(); + queueHistoryReload(); } tui.requestRender(); }; + + const handleSessionMessageEvent = (payload: unknown) => { + if (!payload || typeof payload !== "object") { + return; + } + const evt = payload as SessionMessageEvent; + syncSessionKey(); + const eventSessionKey = normalizeLowercaseStringOrEmpty(evt.sessionKey); + const isUnscopedSessionAlias = + eventSessionKey !== "global" && !parseAgentSessionKey(eventSessionKey); + const eventAgentId = normalizeLowercaseStringOrEmpty(evt.agentId); + const selectedAgentId = normalizeLowercaseStringOrEmpty(state.currentAgentId); + const ownsUnscopedSessionAlias = eventAgentId + ? eventAgentId === selectedAgentId + : selectedAgentId === normalizeLowercaseStringOrEmpty(state.agentDefaultId); + if ( + !isSameSessionKey(evt.sessionKey, state.currentSessionKey) || + !isMatchingGlobalAgentEvent(evt.sessionKey, evt.agentId) || + (isUnscopedSessionAlias && !ownsUnscopedSessionAlias) + ) { + return; + } + + const currentUpdatedAt = state.sessionInfo.updatedAt; + const isOlderSnapshot = + typeof evt.updatedAt === "number" && + typeof currentUpdatedAt === "number" && + evt.updatedAt < currentUpdatedAt; + if (!isOlderSnapshot) { + if (typeof evt.sessionId === "string") { + state.currentSessionId = evt.sessionId; + } + if (typeof evt.updatedAt === "number" || evt.updatedAt === null) { + state.sessionInfo.updatedAt = evt.updatedAt; + } + } + + let displayedRunAwaitingPersistence: string | null = null; + for (const runId of finalizedRunsWithDisplay.keys()) { + if (!persistedTerminalRunIds.has(runId)) { + displayedRunAwaitingPersistence = runId; + } + } + + if ( + state.activeChatRunId || + hasPendingSubmit(state) || + pendingSessionMessageRunId || + displayedRunAwaitingPersistence + ) { + pendingSessionMessageRefresh = true; + // Visible chat finals clear the active run before their transcript is committed. + // Keep later client updates behind that commit so history cannot erase the final. + pendingSessionMessageRunId = + state.activeChatRunId ?? pendingSessionMessageRunId ?? displayedRunAwaitingPersistence; + if (pendingSessionMessageRunId && persistedTerminalRunIds.has(pendingSessionMessageRunId)) { + pendingSessionMessageRunId = null; + } + void refreshSessionInfo?.(); + return; + } + + queueHistoryReload(); + }; + const handleAgentEvent = (payload: unknown) => { if (!payload || typeof payload !== "object") { return; @@ -1174,6 +1272,9 @@ export function createEventHandlers(context: EventHandlerContext) { historyDisplayedReloadRunIds.clear(); liveTerminalErrorMessages.clear(); queuedHistoryReloadRunIds.clear(); + historyReloadQueued = false; + pendingSessionMessageRefresh = false; + pendingSessionMessageRunId = null; deferredHistoryRunEvents.clear(); clearStreamingWatchdog(); clearPendingTerminalLifecycleErrors(); @@ -1197,6 +1298,7 @@ export function createEventHandlers(context: EventHandlerContext) { handleAgentEvent, handleBtwEvent, handleSessionsChangedEvent, + handleSessionMessageEvent, pauseStreamingWatchdog, reconnectStreamingWatchdog, consumeCompletedRunForPendingSend, diff --git a/src/tui/tui-pty-local.e2e.test.ts b/src/tui/tui-pty-local.e2e.test.ts index 17a94cc9a44f..6954dc5e24d8 100644 --- a/src/tui/tui-pty-local.e2e.test.ts +++ b/src/tui/tui-pty-local.e2e.test.ts @@ -11,6 +11,7 @@ import { } from "../../test/helpers/openclaw-test-instance.js"; import type { ModelProviderConfig } from "../config/types.models.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; +import { connectGatewayClient } from "../gateway/test-helpers.e2e.js"; import { createDeferred } from "../test-utils/deferred.js"; import { GatewayChatClient } from "./gateway-chat.js"; import { sleep, startPty, waitFor, type PtyRun } from "./tui-pty-test-support.js"; @@ -59,6 +60,14 @@ const GATEWAY_SCENARIOS = { holdFirstResponse: false, followupReplyText: "FOLLOWUP_RUN_COMPLETE", }, + crossClient: { + agentId: "tui-pty-cross-client", + modelId: "tui-pty-cross-client", + toolsProfile: "minimal", + replyText: "FIRST_RUN_ACTIVE", + holdFirstResponse: false, + followupReplyText: "FOLLOWUP_RUN_COMPLETE", + }, followup: { agentId: "tui-pty-followup", modelId: "tui-pty-followup", @@ -1169,6 +1178,49 @@ describe("TUI PTY real backends", () => { LOCAL_TEST_TIMEOUT_MS, ); + registerGatewayTest( + "renders messages sent by another real Gateway client without restarting", + async ({ onTestFinished }) => { + const fixture = await startGatewayModeTui("crossClient", onTestFinished); + let externalClient: Awaited> | undefined; + try { + await fixture.run.waitForOutput("gateway connected", LOCAL_STARTUP_TIMEOUT_MS); + await fixture.run.write("seed cross-client session\r"); + await fixture.run.waitForOutput("FIRST_RUN_ACTIVE", LOCAL_OUTPUT_TIMEOUT_MS); + const firstReplyOffset = fixture.run.output().lastIndexOf("FIRST_RUN_ACTIVE"); + await waitForOutputAfter(fixture.run, "| idle", firstReplyOffset); + externalClient = await connectGatewayClient({ + url: fixture.gateway.url, + token: fixture.gateway.gatewayToken, + scopes: ["operator.read", "operator.write"], + clientDisplayName: "tui-external-session-writer", + }); + + const marker = "EXTERNAL_GATEWAY_SESSION_MESSAGE"; + await externalClient.request("sessions.send", { + key: fixture.sessionKey, + message: marker, + idempotencyKey: `${fixture.sessionKey}:external-message`, + timeoutMs: 30_000, + }); + + await fixture.run.waitForOutput(marker, LOCAL_OUTPUT_TIMEOUT_MS); + await fixture.run.waitForOutput("FOLLOWUP_RUN_COMPLETE", LOCAL_OUTPUT_TIMEOUT_MS); + const followupOffset = fixture.run.output().lastIndexOf("FOLLOWUP_RUN_COMPLETE"); + await waitForOutputAfter(fixture.run, "| idle", followupOffset); + await fixture.run.write("/exit\r", { delay: false }); + expect((await fixture.run.waitForExit()).exitCode).toBe(0); + } finally { + try { + await externalClient?.stopAndWait({ timeoutMs: 1_000 }); + } finally { + await fixture.cleanup(); + } + } + }, + LOCAL_TEST_TIMEOUT_MS, + ); + registerGatewayTest( "creates and adopts a fresh session through the real Gateway backend", async ({ onTestFinished }) => { diff --git a/src/tui/tui-stream-assembler.test.ts b/src/tui/tui-stream-assembler.test.ts index 2767016220c8..893ee4bf5848 100644 --- a/src/tui/tui-stream-assembler.test.ts +++ b/src/tui/tui-stream-assembler.test.ts @@ -125,6 +125,62 @@ describe("TuiStreamAssembler", () => { expect(second).toBeNull(); }); + it("bounds orphaned stream state while preserving recently active runs", () => { + const assembler = new TuiStreamAssembler(); + for (let index = 0; index < 200; index += 1) { + assembler.ingestDelta(`run-${index}`, messageWithContent([text(`Draft ${index}`)]), false); + } + + assembler.ingestDelta("run-0", messageWithContent([text("Recently active")]), false); + assembler.ingestDelta("run-200", messageWithContent([text("Newest")]), false); + + expect(assembler.finalize("run-0", { role: "assistant", content: [] }, false)).toBe( + "Recently active", + ); + expect(assembler.finalize("run-1", { role: "assistant", content: [] }, false)).toBe( + "(no output)", + ); + expect(assembler.finalize("run-200", { role: "assistant", content: [] }, false)).toBe("Newest"); + }); + + it("does not evict an active run when an evicted run finalizes late", () => { + const assembler = new TuiStreamAssembler(); + for (let index = 0; index < 201; index += 1) { + assembler.ingestDelta(`run-${index}`, messageWithContent([text(`Draft ${index}`)]), false); + } + + expect(assembler.finalize("run-0", messageWithContent([text("Late final")]), false)).toBe( + "Late final", + ); + expect(assembler.finalize("run-1", { role: "assistant", content: [] }, false)).toBe("Draft 1"); + }); + + it("keeps a live run available across thousands of orphaned stream updates", () => { + const assembler = new TuiStreamAssembler(); + assembler.ingestDelta("run-live", messageWithContent([text("Still streaming")]), false); + + for (let index = 0; index < 2_000; index += 1) { + assembler.ingestDelta( + `run-orphan-${index}`, + messageWithContent([text(`Draft ${index}`)]), + false, + ); + if (index % 100 === 0) { + assembler.ingestDelta("run-live", messageWithContent([text("Still streaming")]), false); + } + } + + expect(assembler.finalize("run-live", { role: "assistant", content: [] }, false)).toBe( + "Still streaming", + ); + expect(assembler.finalize("run-orphan-0", { role: "assistant", content: [] }, false)).toBe( + "(no output)", + ); + expect(assembler.finalize("run-orphan-1999", { role: "assistant", content: [] }, false)).toBe( + "Draft 1999", + ); + }); + it("keeps streamed delta text when incoming tool boundary drops a block", () => { const assembler = new TuiStreamAssembler(); const first = assembler.ingestDelta("run-delta-boundary", TEXT_ONLY_TWO_BLOCKS, false); diff --git a/src/tui/tui-stream-assembler.ts b/src/tui/tui-stream-assembler.ts index 630532bdb4a9..4e7749b958d7 100644 --- a/src/tui/tui-stream-assembler.ts +++ b/src/tui/tui-stream-assembler.ts @@ -1,4 +1,5 @@ // Assembles streamed backend events into TUI-visible messages. +import { pruneMapToMaxSize } from "../infra/map-size.js"; import { composeThinkingAndContent, extractContentFromMessage, @@ -6,6 +7,8 @@ import { resolveFinalAssistantText, } from "./tui-formatters.js"; +const MAX_TRACKED_STREAM_RUNS = 200; + // Per-run state used to merge streaming deltas with final assistant messages. type RunStreamState = { thinkingText: string; @@ -109,18 +112,28 @@ function shouldPreserveBoundaryDroppedText(params: { export class TuiStreamAssembler { private runs = new Map(); - private getOrCreateRun(runId: string): RunStreamState { - let state = this.runs.get(runId); - if (!state) { - state = { - thinkingText: "", - contentText: "", - contentBlocks: [], - sawNonTextContentBlocks: false, - displayText: "", - }; - this.runs.set(runId, state); + private createRunState(): RunStreamState { + return { + thinkingText: "", + contentText: "", + contentBlocks: [], + sawNonTextContentBlocks: false, + displayText: "", + }; + } + + private getTrackedRun(runId: string): RunStreamState { + const existing = this.runs.get(runId); + if (existing) { + // Keep a still-streaming older run ahead of abandoned runs in eviction order. + this.runs.delete(runId); + this.runs.set(runId, existing); + return existing; } + + const state = this.createRunState(); + this.runs.set(runId, state); + pruneMapToMaxSize(this.runs, MAX_TRACKED_STREAM_RUNS); return state; } @@ -168,7 +181,7 @@ export class TuiStreamAssembler { /** Ingests a streaming delta and returns updated display text only when it changed. */ ingestDelta(runId: string, message: unknown, showThinking: boolean): string | null { - const state = this.getOrCreateRun(runId); + const state = this.getTrackedRun(runId); const previousDisplayText = state.displayText; this.updateRunState(state, message, showThinking, { boundaryDropMode: "streamed-or-incoming", @@ -183,7 +196,8 @@ export class TuiStreamAssembler { /** Finalizes a run, combines any error text, and drops stored stream state. */ finalize(runId: string, message: unknown, showThinking: boolean, errorMessage?: string): string { - const state = this.getOrCreateRun(runId); + // Late finals must not insert an evicted run and displace a live stream. + const state = this.runs.get(runId) ?? this.createRunState(); const streamedDisplayText = state.displayText; const streamedTextBlocks = [...state.contentBlocks]; const streamedSawNonTextContentBlocks = state.sawNonTextContentBlocks; diff --git a/src/tui/tui-types.ts b/src/tui/tui-types.ts index 27135f8619aa..602cd715596b 100644 --- a/src/tui/tui-types.ts +++ b/src/tui/tui-types.ts @@ -65,6 +65,13 @@ export type SessionChangedEvent = { updatedAt?: number | null; }; +export type SessionMessageEvent = { + sessionKey?: string; + agentId?: string; + sessionId?: string; + updatedAt?: number | null; +}; + export type AgentEvent = { runId: string; stream: string; diff --git a/src/tui/tui.ts b/src/tui/tui.ts index 6134d717a23d..ac7ec473db09 100644 --- a/src/tui/tui.ts +++ b/src/tui/tui.ts @@ -1375,11 +1375,13 @@ export async function runTui(opts: RunTuiOptions): Promise { handleAgentEvent, handleBtwEvent, handleSessionsChangedEvent, + handleSessionMessageEvent, pauseStreamingWatchdog, reconnectStreamingWatchdog, consumeCompletedRunForPendingSend, isRunObserved, flushPendingHistoryRefreshIfIdle, + dispose: disposeEventHandlers, } = createEventHandlers({ chatLog, btw, @@ -1417,6 +1419,7 @@ export async function runTui(opts: RunTuiOptions): Promise { exitReason: result?.exitReason ?? "exit", ...(result?.systemAgentMessage ? { systemAgentMessage: result.systemAgentMessage } : {}), }; + disposeEventHandlers(); pluginApprovals?.dispose(); taskSuggestions?.dispose(); beginTuiShutdown({ @@ -1616,6 +1619,9 @@ export async function runTui(opts: RunTuiOptions): Promise { if (evt.event === "sessions.changed") { handleSessionsChangedEvent(evt.payload); } + if (evt.event === "session.message") { + handleSessionMessageEvent(evt.payload); + } }; client.onConnected = () => { @@ -1742,6 +1748,7 @@ export async function runTui(opts: RunTuiOptions): Promise { client.start(); await new Promise((resolve) => { const finish = () => { + disposeEventHandlers(); pluginApprovals?.dispose(); taskSuggestions?.dispose(); if (isLocalMode) {