From 986cf4bbc3fdf1feb854aaba8c23954af28c06b1 Mon Sep 17 00:00:00 2001 From: mesutoezdil Date: Wed, 26 Aug 2026 09:17:41 +0200 Subject: [PATCH] fix(ui): stop stale session load from overwriting the active chat ChatInterface is not remounted when the session id prop changes, so the message-load effect for the old session keeps running after a new session is picked. If the old fetch resolves after the new one (plausible under normal network jitter), its response silently overwrites the new session's messages on screen, no error shown, wrong conversation visible. Added a cancelled flag to the effect, same pattern already used by the neighboring effect in this file, so a stale load exits before touching state. Added a test that reproduces the race with out of order responses, confirmed it fails against the old code and passes with the fix. Signed-off-by: mesutoezdil --- ui/src/components/chat/ChatInterface.tsx | 39 +++-- .../ChatInterface.sendGuard.test.tsx | 2 +- .../ChatInterface.sessionSwitch.test.tsx | 139 ++++++++++++++++++ 3 files changed, 164 insertions(+), 16 deletions(-) create mode 100644 ui/src/components/chat/__tests__/ChatInterface.sessionSwitch.test.tsx diff --git a/ui/src/components/chat/ChatInterface.tsx b/ui/src/components/chat/ChatInterface.tsx index e48ea09e0..e8ae0a814 100644 --- a/ui/src/components/chat/ChatInterface.tsx +++ b/ui/src/components/chat/ChatInterface.tsx @@ -115,7 +115,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se // Single place that computes the high-water mark, so every update site stays // consistent. Accepts the raw server Task[] (artifacts/synthetic cards are - // intentionally ignored — only persisted history counts). + // intentionally ignored, only persisted history counts). const setServerMark = (tasks: Task[] | undefined) => { syncedServerMsgCountRef.current = countServerMessages(tasks ?? []); }; @@ -185,8 +185,8 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se [isStandaloneToolName, pendingDecisions, pendingApprovalIds], ); // Group over the COMBINED transcript (stored + streaming) so a run that - // spans the boundary — e.g. an approval request persisted at - // input_required and its tool result arriving on the post-approval stream — + // spans the boundary (e.g. an approval request persisted at + // input_required and its tool result arriving on the post-approval stream) // folds into a single group instead of two. const renderItems = useMemo(() => groupToolCallMessages(allMessages, groupingOptions), [allMessages, groupingOptions]); // Shared call_id -> is_error lookup so each group summary is O(group size). @@ -224,6 +224,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se }), [selectedNamespace, selectedAgentName, onTerminalTask]); useEffect(() => { + let cancelled = false; async function initializeChat() { setSessionStats({ total: 0, prompt: 0, completion: 0 }); setTerminalTaskIds(new Set()); @@ -257,6 +258,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se if (shareToken) { // Fetch session info to get authoritative read_only status from the server. const sessionInfoResponse = await getSessionWithEvents(sessionId, shareToken); + if (cancelled) return; if (sessionInfoResponse.error || !sessionInfoResponse.data) { setSessionNotFound(true); setIsLoading(false); @@ -265,6 +267,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se setShareReadOnly(sessionInfoResponse.data.read_only === true); } else { const sessionExistsResponse = await checkSessionExists(sessionId); + if (cancelled) return; if (sessionExistsResponse.error || !sessionExistsResponse.data) { setSessionNotFound(true); setIsLoading(false); @@ -273,6 +276,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se } const messagesResponse = await getSessionTasks(sessionId, shareToken); + if (cancelled) return; if (messagesResponse.error) { toast.error("Failed to load messages"); setIsLoading(false); @@ -314,6 +318,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se } setServerMark(messagesResponse.data); } catch (error) { + if (cancelled) return; console.error("Error loading messages:", error); toast.error("Error loading messages"); setSessionNotFound(true); @@ -321,6 +326,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se return; } + if (cancelled) return; setIsLoading(false); if (activeTask) { @@ -330,6 +336,9 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se } initializeChat(); + return () => { + cancelled = true; + }; // eslint-disable-next-line react-hooks/exhaustive-deps }, [sessionId, selectedAgentName, selectedNamespace, isFirstMessage, shareToken]); @@ -368,16 +377,16 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se } // Cross-tab guard: fetch the latest session state before mutating anything. - // Two cases: (1) another tab is still streaming — reconnect instead of sending; - // (2) another tab completed a turn we haven't loaded — reload so the user sees + // Two cases: (1) another tab is still streaming, reconnect instead of sending; + // (2) another tab completed a turn we haven't loaded, reload so the user sees // the full context before their next message goes out. const guardSessionId = session?.id || sessionId; if (guardSessionId) { const guardResult = await checkAndSyncSessionBeforeAction(guardSessionId, { messages: { - inFlight: "This session is already being processed — reconnecting to live updates", - inputRequired: "Session is awaiting your input — please review before sending", - staleOrChanged: "New messages loaded — please review before sending", + inFlight: "This session is already being processed, reconnecting to live updates", + inputRequired: "Session is awaiting your input, please review before sending", + staleOrChanged: "New messages loaded, please review before sending", }, }); if (guardResult === "blocked") return; @@ -415,7 +424,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se // rename block below must be skipped: the title was already set at creation // time, and session React state hasn't yet re-rendered (so session?.name // is still null, which would make isPlaceholderSessionTitle return true - // incorrectly and queue a redundant — potentially hanging — POST /sessions). + // incorrectly and queue a redundant, potentially hanging, POST /sessions). let justCreatedSession = false; // If there's no session, create one @@ -725,13 +734,13 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se await consumeStream(stream); - // Stream ended cleanly — reload final state from DB and settle. + // Stream ended cleanly, reload final state from DB and settle. await reloadSessionFromDB(); } catch (error: unknown) { if (error instanceof Error && error.name !== "AbortError" && !isTerminalError(error)) { console.error("Resubscribe failed:", error); } - // Terminal, AbortError, or unexpected error — reload whatever state we have. + // Terminal, AbortError, or unexpected error, reload whatever state we have. if (!(error instanceof Error && error.name === "AbortError")) { await reloadSessionFromDB(); } @@ -798,8 +807,8 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se const guardResult = await checkAndSyncSessionBeforeAction(currentSessionId, { expectedTaskId: taskId, messages: { - inFlight: "Another tab already responded — reconnecting to live updates", - staleOrChanged: "Session state changed — please review", + inFlight: "Another tab already responded, reconnecting to live updates", + staleOrChanged: "Session state changed, please review", }, }); if (guardResult === "blocked") return; @@ -832,7 +841,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se setPendingDecisions({}); pendingDecisionsRef.current = {}; pendingRejectionReasonsRef.current = {}; - // Only reset "thinking" → "ready". Do NOT reset "input_required" — + // Only reset "thinking" to "ready". Do NOT reset "input_required", // handleMessageEvent may have already set it for the next HITL cycle // during this same stream. setChatStatus(prev => prev === "thinking" ? "ready" : prev); @@ -1058,7 +1067,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se ? voiceError : isListening ? "Stop listening" - : "Voice input — click and speak"} + : "Voice input, click and speak"} diff --git a/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx b/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx index 756c68594..7d36026d0 100644 --- a/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx +++ b/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx @@ -86,7 +86,7 @@ const mockGetSessionTasks = getSessionTasks as jest.MockedFunction; const mockToastInfo = toast.info as jest.MockedFunction; -const staleToastMessage = "New messages loaded — please review before sending"; +const staleToastMessage = "New messages loaded, please review before sending"; // The send guard is server-authoritative: it compares the persisted user-message // high-water mark against what this tab last synced. In A2A v1, assistant output diff --git a/ui/src/components/chat/__tests__/ChatInterface.sessionSwitch.test.tsx b/ui/src/components/chat/__tests__/ChatInterface.sessionSwitch.test.tsx new file mode 100644 index 000000000..fe049d37c --- /dev/null +++ b/ui/src/components/chat/__tests__/ChatInterface.sessionSwitch.test.tsx @@ -0,0 +1,139 @@ +/** + * @jest-environment jsdom + */ +import { act, render, screen, waitFor } from "@testing-library/react"; +import { Role, type Task } from "@a2a-js/sdk"; +import { checkSessionExists, createSession, getSessionTasks } from "@/app/actions/sessions"; +import { kagentA2AClient } from "@/lib/a2aClient"; +import ChatInterface from "@/components/chat/ChatInterface"; +import { createMockTask, createMockTextMessage, createTextPart } from "@/mocks/factories"; +import type { BaseResponse } from "@/types"; + +jest.mock("@/app/actions/sessions", () => ({ + checkSessionExists: jest.fn(), + createSession: jest.fn(), + getSessionTasks: jest.fn(), +})); + +jest.mock("@/app/actions/agents", () => ({ + getAgentWithResolvedKind: jest.fn(), + waitForSandboxAgentReady: jest.fn(), +})); + +jest.mock("@/lib/a2aClient", () => ({ + kagentA2AClient: { + sendMessageStream: jest.fn(), + resubscribeStream: jest.fn(), + }, +})); + +jest.mock("sonner", () => ({ + toast: { info: jest.fn(), error: jest.fn(), loading: jest.fn(), dismiss: jest.fn() }, +})); + +jest.mock("@/hooks/useSpeechRecognition", () => ({ + useSpeechRecognition: () => ({ + isListening: false, + isSupported: false, + startListening: jest.fn(), + stopListening: jest.fn(), + error: null, + }), +})); + +jest.mock("@/components/chat/ChatAgentContext", () => ({ + useChatRunInSandbox: () => false, + useChatSubstrateSandbox: () => false, + useCurrentChatAgent: () => ({ deploymentReady: true }), +})); + +jest.mock("@/components/chat/ChatMessage", () => ({ + __esModule: true, + default: ({ message }: { message: import("@a2a-js/sdk").Message }) => ( +
+ {message.parts + ?.map((part) => (part.content?.$case === "text" ? part.content.value : "")) + .join("")} +
+ ), +})); + +jest.mock("@/components/chat/StreamingMessage", () => ({ + __esModule: true, + default: ({ content }: { content: string }) =>
{content}
, +})); + +const mockCheckSessionExists = checkSessionExists as jest.MockedFunction; +const mockCreateSession = createSession as jest.MockedFunction; +const mockGetSessionTasks = getSessionTasks as jest.MockedFunction; +const mockResubscribeStream = kagentA2AClient.resubscribeStream as jest.MockedFunction< + typeof kagentA2AClient.resubscribeStream +>; + +/** A completed A2A v1 turn: user request in history, agent output in an artifact. */ +function task(sessionId: string, answer: string): Task { + const taskId = `task-${sessionId}`; + const result = createMockTask(taskId, sessionId, [ + createMockTextMessage(`${sessionId}-user`, Role.ROLE_USER, "hi", { + contextId: sessionId, + taskId, + }), + ]); + result.artifacts = [{ + artifactId: `${taskId}-answer`, + name: "", + description: "", + parts: [createTextPart(answer)], + extensions: [], + metadata: undefined, + }]; + return result; +} + +/** A promise this test can resolve on demand, to control arrival order. */ +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { promise, resolve }; +} + +describe("ChatInterface session switch", () => { + beforeEach(() => { + jest.clearAllMocks(); + mockCheckSessionExists.mockResolvedValue({ message: "ok", data: true }); + mockCreateSession.mockResolvedValue({ message: "unexpected createSession call", error: "unexpected createSession call" }); + mockResubscribeStream.mockReturnValue(Promise.resolve((async function* () {})())); + }); + + it("does not show a stale session's messages after they arrive out of order", async () => { + const sessionA = deferred>(); + const sessionB = deferred>(); + mockGetSessionTasks.mockImplementation(async (sessionId: string) => + sessionId === "session-a" ? sessionA.promise : sessionB.promise, + ); + + const { rerender } = render(); + await waitFor(() => expect(mockGetSessionTasks).toHaveBeenCalledWith("session-a", undefined)); + + rerender(); + await waitFor(() => expect(mockGetSessionTasks).toHaveBeenCalledWith("session-b", undefined)); + + // session-b's fetch resolves first, then session-a's late response arrives. + sessionB.resolve({ message: "ok", data: [task("session-b", "answer b")] }); + await screen.findByText("answer b"); + + // Let the late session-a response run its full async continuation past + // the awaited getSessionTasks call before asserting on the DOM. + await act(async () => { + sessionA.resolve({ message: "ok", data: [task("session-a", "answer a")] }); + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); + }); + + expect(screen.queryByText("answer a")).not.toBeInTheDocument(); + expect(screen.getByText("answer b")).toBeInTheDocument(); + }); +});