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

ui/src/components/chat/ChatInterface.tsx

@@ -94,7 +94,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 ?? []);
   };
@@ -164,8 +164,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).
@@ -185,6 +185,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
   }), [selectedNamespace, selectedAgentName]);
 
   useEffect(() => {
+    let cancelled = false;
     async function initializeChat() {
       setSessionStats({ total: 0, prompt: 0, completion: 0 });
       setStreamingMessages([]);
@@ -216,6 +217,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);
@@ -224,6 +226,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);
@@ -232,6 +235,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);
@@ -267,13 +271,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) {
@@ -283,6 +289,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]);
 
@@ -321,16 +330,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;
@@ -375,7 +384,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
@@ -680,13 +689,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();
       }
@@ -772,8 +781,8 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
       const guardResult = await checkAndSyncSessionBeforeAction(currentSessionId, {
         expectedTaskId: approvalTaskId,
         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;
@@ -830,7 +839,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);
@@ -848,7 +857,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
     const reasons = pendingRejectionReasonsRef.current;
 
     if (allApprove) {
-      // Uniform approve — no need for batch
+      // Uniform approve, no need for batch
       await sendApprovalDecision(
         { decision_type: "approve" },
         "Approved",
@@ -860,7 +869,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
         "Rejected",
       );
     } else {
-      // Mixed decisions — use batch mode with per-tool decisions.
+      // Mixed decisions, use batch mode with per-tool decisions.
       // For subagent HITL the keys are inner subagent tool IDs; the backend
       // detects this via hitl_parts in the pending confirmation payload and
       // forwards the batch to the subagent.
@@ -933,8 +942,8 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
       const guardResult = await checkAndSyncSessionBeforeAction(currentSessionId, {
         expectedTaskId: askUserTaskId,
         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;
@@ -1137,7 +1146,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

@@ -75,12 +75,12 @@ 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 count of persisted
 // history messages across all tasks (the high-water mark) against the count this
 // tab last synced. These helpers build tasks whose `history.length` drives that
-// count — the message content is irrelevant to the guard.
+// count, the message content is irrelevant to the guard.
 
 // The backend snapshot the mocked getSessionTasks currently returns. The stream
 // generators advance it to model a turn being persisted after it streams.

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

@@ -0,0 +1,134 @@
+/**
+ * @jest-environment jsdom
+ */
+import { act, render, screen, waitFor } from "@testing-library/react";
+import type { Message, Task } from "@a2a-js/sdk";
+import { checkSessionExists, createSession, getSessionTasks } from "@/app/actions/sessions";
+import { kagentA2AClient } from "@/lib/a2aClient";
+import ChatInterface from "@/components/chat/ChatInterface";
+
+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: Message }) => (
+    <div data-testid={`chat-message-${message.role}`}>
+      {message.parts?.map((part) => (part.kind === "text" ? part.text : "")).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
+>;
+
+function task(sessionId: string, answer: string): Task {
+  return {
+    id: `task-${sessionId}`,
+    contextId: sessionId,
+    status: { state: "completed", timestamp: new Date().toISOString() },
+    history: [
+      {
+        kind: "message",
+        messageId: `${sessionId}-agent`,
+        role: "agent",
+        contextId: sessionId,
+        taskId: `task-${sessionId}`,
+        parts: [{ kind: "text", text: answer }],
+        metadata: { timestamp: Date.now() },
+      } as Message,
+    ],
+  } as Task;
+}
+
+/** 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({ data: true });
+    mockCreateSession.mockResolvedValue({ error: "unexpected createSession call" });
+    mockResubscribeStream.mockReturnValue((async function* () {})());
+  });
+
+  it("does not show a stale session's messages after they arrive out of order", async () => {
+    const sessionA = deferred<{ data: Task[] }>();
+    const sessionB = deferred<{ data: 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({ 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({ 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();
+  });
+});