diff --git a/src/tui/tui-event-handlers.test.ts b/src/tui/tui-event-handlers.test.ts index 352af6af786..3a748ffff41 100644 --- a/src/tui/tui-event-handlers.test.ts +++ b/src/tui/tui-event-handlers.test.ts @@ -364,6 +364,51 @@ describe("tui-event-handlers: handleAgentEvent", () => { expect(loadHistory).toHaveBeenCalledTimes(1); }); + it("does not reload history or clear active run when another run final arrives mid-stream", () => { + const state = makeState({ activeChatRunId: "run-active" }); + const { chatLog, tui, setActivityStatus, loadHistory, isLocalRunId, forgetLocalRunId } = + makeContext(state); + const { handleChatEvent } = createEventHandlers({ + chatLog, + tui, + state, + setActivityStatus, + loadHistory, + isLocalRunId, + forgetLocalRunId, + }); + + handleChatEvent({ + runId: "run-active", + sessionKey: state.currentSessionKey, + state: "delta", + message: { content: "partial" }, + }); + + loadHistory.mockClear(); + setActivityStatus.mockClear(); + + handleChatEvent({ + runId: "run-other", + sessionKey: state.currentSessionKey, + state: "final", + message: { content: [{ type: "text", text: "other final" }] }, + }); + + expect(loadHistory).not.toHaveBeenCalled(); + expect(state.activeChatRunId).toBe("run-active"); + expect(setActivityStatus).not.toHaveBeenCalledWith("idle"); + + handleChatEvent({ + runId: "run-active", + sessionKey: state.currentSessionKey, + state: "delta", + message: { content: "continued" }, + }); + + expect(chatLog.updateAssistant).toHaveBeenLastCalledWith("continued", "run-active"); + }); + it("drops streaming assistant when chat final has no message", () => { const state = makeState({ activeChatRunId: null }); const { chatLog, tui, setActivityStatus } = makeContext(state); diff --git a/src/tui/tui-event-handlers.ts b/src/tui/tui-event-handlers.ts index 743a62a5058..ee3f58b040a 100644 --- a/src/tui/tui-event-handlers.ts +++ b/src/tui/tui-event-handlers.ts @@ -79,6 +79,31 @@ export function createEventHandlers(context: EventHandlerContext) { pruneRunMap(finalizedRuns); }; + const clearActiveRunIfMatch = (runId: string) => { + if (state.activeChatRunId === runId) { + state.activeChatRunId = null; + } + }; + + const hasConcurrentActiveRun = (runId: string) => { + const activeRunId = state.activeChatRunId; + if (!activeRunId || activeRunId === runId) { + return false; + } + return sessionRuns.has(activeRunId); + }; + + const maybeRefreshHistoryForRun = (runId: string) => { + if (isLocalRunId?.(runId)) { + forgetLocalRunId?.(runId); + return; + } + if (hasConcurrentActiveRun(runId)) { + return; + } + void loadHistory?.(); + }; + const handleChatEvent = (payload: unknown) => { if (!payload || typeof payload !== "object") { return; @@ -109,43 +134,36 @@ export function createEventHandlers(context: EventHandlerContext) { setActivityStatus("streaming"); } if (evt.state === "final") { + const wasActiveRun = state.activeChatRunId === evt.runId; if (!evt.message) { - if (isLocalRunId?.(evt.runId)) { - forgetLocalRunId?.(evt.runId); - } else { - void loadHistory?.(); - } + maybeRefreshHistoryForRun(evt.runId); chatLog.dropAssistant(evt.runId); noteFinalizedRun(evt.runId); - state.activeChatRunId = null; - setActivityStatus("idle"); + clearActiveRunIfMatch(evt.runId); + if (wasActiveRun) { + setActivityStatus("idle"); + } void refreshSessionInfo?.(); tui.requestRender(); return; } if (isCommandMessage(evt.message)) { - if (isLocalRunId?.(evt.runId)) { - forgetLocalRunId?.(evt.runId); - } else { - void loadHistory?.(); - } + maybeRefreshHistoryForRun(evt.runId); const text = extractTextFromMessage(evt.message); if (text) { chatLog.addSystem(text); } streamAssembler.drop(evt.runId); noteFinalizedRun(evt.runId); - state.activeChatRunId = null; - setActivityStatus("idle"); + clearActiveRunIfMatch(evt.runId); + if (wasActiveRun) { + setActivityStatus("idle"); + } void refreshSessionInfo?.(); tui.requestRender(); return; } - if (isLocalRunId?.(evt.runId)) { - forgetLocalRunId?.(evt.runId); - } else { - void loadHistory?.(); - } + maybeRefreshHistoryForRun(evt.runId); const stopReason = evt.message && typeof evt.message === "object" && !Array.isArray(evt.message) ? typeof (evt.message as Record).stopReason === "string" @@ -156,36 +174,36 @@ export function createEventHandlers(context: EventHandlerContext) { const finalText = streamAssembler.finalize(evt.runId, evt.message, state.showThinking); chatLog.finalizeAssistant(finalText, evt.runId); noteFinalizedRun(evt.runId); - state.activeChatRunId = null; - setActivityStatus(stopReason === "error" ? "error" : "idle"); + clearActiveRunIfMatch(evt.runId); + if (wasActiveRun) { + setActivityStatus(stopReason === "error" ? "error" : "idle"); + } // Refresh session info to update token counts in footer void refreshSessionInfo?.(); } if (evt.state === "aborted") { + const wasActiveRun = state.activeChatRunId === evt.runId; chatLog.addSystem("run aborted"); streamAssembler.drop(evt.runId); sessionRuns.delete(evt.runId); - state.activeChatRunId = null; - setActivityStatus("aborted"); - void refreshSessionInfo?.(); - if (isLocalRunId?.(evt.runId)) { - forgetLocalRunId?.(evt.runId); - } else { - void loadHistory?.(); + clearActiveRunIfMatch(evt.runId); + if (wasActiveRun) { + setActivityStatus("aborted"); } + void refreshSessionInfo?.(); + maybeRefreshHistoryForRun(evt.runId); } if (evt.state === "error") { + const wasActiveRun = state.activeChatRunId === evt.runId; chatLog.addSystem(`run error: ${evt.errorMessage ?? "unknown"}`); streamAssembler.drop(evt.runId); sessionRuns.delete(evt.runId); - state.activeChatRunId = null; - setActivityStatus("error"); - void refreshSessionInfo?.(); - if (isLocalRunId?.(evt.runId)) { - forgetLocalRunId?.(evt.runId); - } else { - void loadHistory?.(); + clearActiveRunIfMatch(evt.runId); + if (wasActiveRun) { + setActivityStatus("error"); } + void refreshSessionInfo?.(); + maybeRefreshHistoryForRun(evt.runId); } tui.requestRender(); };