fix(ui): stop stale session load from overwriting the active chat

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,13 +318,15 @@ 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);
         setIsLoading(false);
         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"}
                       </TooltipContent>
                     </Tooltip>
                   </TooltipProvider>

ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx

@@ -86,7 +86,7 @@ const mockGetSessionTasks = getSessionTasks as jest.MockedFunction<typeof getSes
 const mockSendMessageStream = kagentA2AClient.sendMessageStream as jest.MockedFunction<typeof kagentA2AClient.sendMessageStream>;
 const mockToastInfo = toast.info as jest.MockedFunction<typeof toast.info>;
 
-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

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 }) => (
+    <div data-testid={`chat-message-${message.role}`}>
+      {message.parts
+        ?.map((part) => (part.content?.$case === "text" ? part.content.value : ""))
+        .join("")}
+    </div>
+  ),
+}));
+
+jest.mock("@/components/chat/StreamingMessage", () => ({
+  __esModule: true,
+  default: ({ content }: { content: string }) => <div>{content}</div>,
+}));
+
+const mockCheckSessionExists = checkSessionExists as jest.MockedFunction<typeof checkSessionExists>;
+const mockCreateSession = createSession as jest.MockedFunction<typeof createSession>;
+const mockGetSessionTasks = getSessionTasks as jest.MockedFunction<typeof getSessionTasks>;
+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<T>() {
+  let resolve!: (value: T) => void;
+  const promise = new Promise<T>((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<BaseResponse<Task[]>>();
+    const sessionB = deferred<BaseResponse<Task[]>>();
+    mockGetSessionTasks.mockImplementation(async (sessionId: string) =>
+      sessionId === "session-a" ? sessionA.promise : sessionB.promise,
+    );
+
+    const { rerender } = render(<ChatInterface selectedAgentName="test-agent" selectedNamespace="kagent" sessionId="session-a" />);
+    await waitFor(() => expect(mockGetSessionTasks).toHaveBeenCalledWith("session-a", undefined));
+
+    rerender(<ChatInterface selectedAgentName="test-agent" selectedNamespace="kagent" sessionId="session-b" />);
+    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();
+  });
+});