diff --git a/ui/src/ui/app-gateway.node.test.ts b/ui/src/ui/app-gateway.node.test.ts index 8cd7377b8bc..7045ddb8e10 100644 --- a/ui/src/ui/app-gateway.node.test.ts +++ b/ui/src/ui/app-gateway.node.test.ts @@ -622,6 +622,38 @@ describe("connectGateway", () => { }, ); + it.each(["aborted", "error"] as const)( + "replays deferred session.message reloads after %s clears the active run", + (terminalState) => { + const { host, client } = connectHostGateway(); + host.chatRunId = "main-run-3"; + loadChatHistoryMock.mockClear(); + + client.emitEvent({ + event: "session.message", + payload: { + sessionKey: "main", + }, + }); + + expect(loadChatHistoryMock).not.toHaveBeenCalled(); + + client.emitEvent({ + event: "chat", + payload: { + runId: "main-run-3", + sessionKey: "main", + state: terminalState, + errorMessage: terminalState === "error" ? "chat failed" : undefined, + }, + }); + + expect(host.chatRunId).toBeNull(); + expect(loadChatHistoryMock).toHaveBeenCalledTimes(1); + expect(loadChatHistoryMock).toHaveBeenCalledWith(host); + }, + ); + it("clears tracked BTW terminal runs after reconnect hello", () => { const host = createHost(); diff --git a/ui/src/ui/app-gateway.ts b/ui/src/ui/app-gateway.ts index 24774177dd8..e9f12390319 100644 --- a/ui/src/ui/app-gateway.ts +++ b/ui/src/ui/app-gateway.ts @@ -97,6 +97,10 @@ type GatewayHost = { updateAvailable: UpdateAvailable | null; }; +type GatewayHostWithDeferredSessionMessageReload = GatewayHost & { + pendingSessionMessageReloadSessionKey?: string | null; +}; + type SessionDefaultsSnapshot = { defaultAgentId?: string; mainKey?: string; @@ -409,8 +413,26 @@ function handleChatGatewayEvent(host: GatewayHost, payload: ChatEventPayload | u } const state = handleChatEvent(host as unknown as ChatState, payload); const historyReloaded = handleTerminalChatEvent(host, payload, state); + const deferredReloadHost = host as GatewayHostWithDeferredSessionMessageReload; + const deferredSessionKey = deferredReloadHost.pendingSessionMessageReloadSessionKey?.trim(); + const payloadSessionKey = payload?.sessionKey?.trim(); + const shouldReplayDeferredSessionMessageReload = Boolean( + deferredSessionKey && + payloadSessionKey && + deferredSessionKey === payloadSessionKey && + isTerminalChatState(state) && + payloadSessionKey === host.sessionKey && + !host.chatRunId, + ); + if (deferredSessionKey && payloadSessionKey && deferredSessionKey === payloadSessionKey) { + deferredReloadHost.pendingSessionMessageReloadSessionKey = null; + } if (state === "final" && !historyReloaded && shouldReloadHistoryForFinalEvent(payload)) { void loadChatHistory(host as unknown as ChatState); + return; + } + if (shouldReplayDeferredSessionMessageReload && !historyReloaded) { + void loadChatHistory(host as unknown as ChatState); } } @@ -418,6 +440,7 @@ function handleSessionMessageGatewayEvent( host: GatewayHost, payload: { sessionKey?: string } | undefined, ) { + const deferredReloadHost = host as GatewayHostWithDeferredSessionMessageReload; const sessionKey = payload?.sessionKey?.trim(); if (!sessionKey || sessionKey !== host.sessionKey) { return; @@ -428,8 +451,10 @@ function handleSessionMessageGatewayEvent( // chatStream, which delays the user message card from appearing until the // first LLM delta arrives. if (host.chatRunId) { + deferredReloadHost.pendingSessionMessageReloadSessionKey = sessionKey; return; } + deferredReloadHost.pendingSessionMessageReloadSessionKey = null; void loadChatHistory(host as unknown as ChatState); }