diff --git a/ui/src/e2e/new-session-page.prompt-attachments.e2e.test.ts b/ui/src/e2e/new-session-page.prompt-attachments.e2e.test.ts index 3061d64c28ed..48a6e678da6c 100644 --- a/ui/src/e2e/new-session-page.prompt-attachments.e2e.test.ts +++ b/ui/src/e2e/new-session-page.prompt-attachments.e2e.test.ts @@ -3,6 +3,7 @@ import { mkdir } from "node:fs/promises"; import path from "node:path"; import type { Page } from "playwright"; import { expect, it } from "vitest"; +import { waitForControlUiGatewayReady } from "../test-helpers/control-ui-e2e-readiness.ts"; import { ONE_PIXEL_PNG_B64, SESSION_LIST_DEFAULTS, @@ -50,6 +51,7 @@ suite.define(() => { await firstMessage.fill(text); await waitForCommittedNewSessionDraft(firstPage, text, 0); await firstPage.reload(); + await waitForControlUiGatewayReady(firstPage); await expect.poll(() => firstMessage.inputValue()).toBe(text); await firstPage.close(); diff --git a/ui/src/pages/chat/chat-gateway.test.ts b/ui/src/pages/chat/chat-gateway.test.ts index 9a45f8559850..d4255624276e 100644 --- a/ui/src/pages/chat/chat-gateway.test.ts +++ b/ui/src/pages/chat/chat-gateway.test.ts @@ -106,7 +106,13 @@ function seedChatSnapshot( type SessionTestState = ChatState & { [key: string]: unknown; chatRunStatus?: { phase: string; runId: string | null; sessionKey: string } | null; - lastLocalTerminalReconcile?: { sessionStatus: string } | null; + knownAgentRunIds: Set; + lastLocalTerminalReconcile?: { + phase: string; + runId: string | null; + sessionKey: string; + sessionStatus: string; + } | null; sessionsResult: { [key: string]: unknown; sessions: Array>; @@ -830,6 +836,54 @@ describe("handleChatGatewayEvent", () => { }); }); + it("settles an active run when its terminal reply was already accepted", () => { + const runId = "run-1"; + const final = createTextChatMessage("assistant", "Delivered answer"); + const state = createStateWithRunningSession({ + sessionKey: "main", + chatRunId: runId, + chatStream: "Delivered answer", + chatStreamStartedAt: 123, + chatRunStartup: { state: "activity", runId }, + }) as SessionTestState & { chatStreamSegments: HistoryToolSegment[] }; + state.chatStreamSegments = [{ text: "Retained commentary", ts: 122, toolCallId: "call-1" }]; + state.knownAgentRunIds = new Set([runId]); + const scope = { sessionKey: state.sessionKey }; + const projection = reduceSessionProjection( + getChatSessionProjection(state, state.chatMessages, scope), + { type: "runTerminal", runId, status: "completed", message: final, scope }, + ); + setChatSessionProjection(state, projection); + state.chatMessages = [final]; + + expect( + handleChatGatewayEvent(state, { + runId, + sessionKey: "main", + state: "final", + message: final, + }), + ).toBe("final"); + + expect(state.chatMessages).toEqual([final]); + expect(state.chatRunId).toBeNull(); + expect(state.chatStream).toBeNull(); + expect(state.chatStreamSegments).toEqual([]); + expect(state.chatRunStartup).toBeNull(); + expect(state.knownAgentRunIds.has(runId)).toBe(false); + expect(state.sessionsResult.sessions[0]).toMatchObject({ + status: "done", + hasActiveRun: false, + activeRunIds: [], + }); + expect(state.lastLocalTerminalReconcile).toMatchObject({ + runId, + sessionKey: "main", + phase: "done", + sessionStatus: "done", + }); + }); + it("persists keyed commentary with the final answer by default", () => { const user = { role: "user", content: [{ type: "text", text: "Ask" }], timestamp: 1 }; const state = createState({ diff --git a/ui/src/pages/chat/chat-gateway.ts b/ui/src/pages/chat/chat-gateway.ts index 00e3bd13f806..21fc12fe5db5 100644 --- a/ui/src/pages/chat/chat-gateway.ts +++ b/ui/src/pages/chat/chat-gateway.ts @@ -264,6 +264,49 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { if (projectedRun) { setChatSessionProjection(state, projectedRun.projection); } + const terminalRunId = payload.runId ?? state.chatRunId; + const reconcileOwnedTerminalRun = () => { + const terminalStatus = projectedRun?.currentRun?.status; + if ( + !payload.runId || + payload.runId !== state.chatRunId || + !terminalStatus || + terminalStatus === "streaming" + ) { + return; + } + clearToolStreamSegments(state); + const sessionKeys = sessionMatches ? [state.sessionKey, payload.sessionKey] : []; + if (terminalStatus === "yielded") { + reconcileChatRunLifecycle(state, { + yielded: true, + runId: terminalRunId, + sessionKey: state.sessionKey, + sessionKeys, + clearLocalRun: true, + clearChatStream: true, + }); + return; + } + const sessionStatus = + terminalStatus === "completed" + ? ("done" as const) + : terminalStatus === "aborted" + ? ("killed" as const) + : terminalStatus === "timeout" + ? ("timeout" as const) + : ("failed" as const); + reconcileChatRunLifecycle(state, { + outcome: terminalStatus === "completed" ? "done" : "interrupted", + sessionStatus, + runId: terminalRunId, + sessionKey: state.sessionKey, + sessionKeys, + clearLocalRun: true, + clearChatStream: true, + armLocalTerminalReconcile: hadActiveRunBeforeEvent && activeRunMatches, + }); + }; const previousTerminalRun = projectedRun?.previousRun; if (previousTerminalRun && previousTerminalRun.status !== "streaming") { if (payload.state === "delta") { @@ -285,6 +328,7 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { setChatRunError(state, resolveGatewayErrorText(payload, null)); } if (payload.state === "error") { + reconcileOwnedTerminalRun(); return "error"; } } @@ -296,6 +340,7 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { shouldHideAssistantChatMessage(incomingFinal) || hasSessionProjectionAcceptedFinal(previousTerminalRun, incomingFinal))) ) { + reconcileOwnedTerminalRun(); return payload.state; } } @@ -325,7 +370,6 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { return null; } - const terminalRunId = payload.runId ?? state.chatRunId; const terminalAfterBoundaryRunId = latestStreamBoundaryRunId(state); const materializeVisibleStream = ( materializeOpts: Parameters[2] = {}, @@ -333,21 +377,6 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { materializeVisibleAssistantStreamMessages(state.chatMessages, state, { ...materializeOpts, }); - const reconcileTerminalRun = ( - outcome: "done" | "interrupted", - sessionStatus: "done" | "failed" | "killed" | "timeout", - ) => - reconcileChatRunLifecycle(state, { - outcome, - sessionStatus, - runId: terminalRunId, - sessionKey: state.sessionKey, - sessionKeys: sessionMatches ? [state.sessionKey, payload.sessionKey] : [], - clearLocalRun: true, - clearChatStream: true, - armLocalTerminalReconcile: hadActiveRunBeforeEvent && activeRunMatches, - }); - if (payload.state === "status") { if (!payload.runId || payload.runId !== state.chatRunId) { return null; @@ -424,18 +453,7 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { state.chatMessages = materializeVisibleStream(); } } - if (payload.yielded === true && payload.stopReason === "end_turn") { - reconcileChatRunLifecycle(state, { - yielded: true, - runId: terminalRunId, - sessionKey: state.sessionKey, - sessionKeys: sessionMatches ? [state.sessionKey, payload.sessionKey] : [], - clearLocalRun: true, - clearChatStream: true, - }); - } else { - reconcileTerminalRun("done", "done"); - } + reconcileOwnedTerminalRun(); } else if (payload.state === "aborted") { const normalizedMessage = normalizeAbortedAssistantMessage(payload.message); if (normalizedMessage && !shouldHideAssistantChatMessage(normalizedMessage)) { @@ -460,7 +478,7 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { if (payload.errorMessage?.trim()) { setChatRunError(state, resolveGatewayErrorText(payload, null)); } - reconcileTerminalRun("interrupted", "killed"); + reconcileOwnedTerminalRun(); } else if (payload.state === "error") { const payloadMessage = normalizeFinalAssistantMessage(payload.message); const visiblePayloadMessage = @@ -523,10 +541,7 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { } // The shared Gateway projection owns timeout classification; preserve it // when publishing selected-session and sidebar terminal status. - reconcileTerminalRun( - "interrupted", - projectedRun?.currentRun?.status === "timeout" ? "timeout" : "failed", - ); + reconcileOwnedTerminalRun(); setChatRunError( state, resolveGatewayErrorText(payload, projectedErrorMessage ? visiblePayloadMessage : null),