diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
index f6000cef1d..b39942f182 100644
--- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
+++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
@@ -789,6 +789,101 @@ describe("useChatStore", () => {
).toEqual([]);
});
+ it("closes old WebSocket and resets state when chatID changes", async () => {
+ immediateAnimationFrame();
+
+ const chatID1 = "chat-1";
+ const chatID2 = "chat-2";
+ const msg1 = makeMessage(chatID1, 1, "user", "hello");
+ const msg2 = makeMessage(chatID2, 10, "user", "world");
+
+ const mockSocket1 = createMockSocket();
+ const mockSocket2 = createMockSocket();
+ // Use a fallback so that extra effect re-runs (caused by
+ // dependency changes during rerender) get a valid socket.
+ vi.mocked(watchChat).mockReturnValue(mockSocket2 as never);
+ vi.mocked(watchChat)
+ .mockReturnValueOnce(mockSocket1 as never)
+ .mockReturnValueOnce(mockSocket1 as never)
+ .mockReturnValueOnce(mockSocket1 as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const initialOptions = {
+ chatID: chatID1,
+ chatMessages: [msg1] as TypesGen.ChatMessage[],
+ chatRecord: makeChat(chatID1),
+ chatData: {
+ chat: makeChat(chatID1),
+ messages: [msg1],
+ queued_messages: [] as TypesGen.ChatQueuedMessage[],
+ },
+ chatQueuedMessages: [] as TypesGen.ChatQueuedMessage[],
+ setChatErrorReason,
+ clearChatErrorReason,
+ };
+
+ const { result, rerender } = renderHook(
+ (options: Parameters[0]) => {
+ const { store } = useChatStore(options);
+ return {
+ streamState: useChatSelector(store, selectStreamState),
+ };
+ },
+ { initialProps: initialOptions, wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID1);
+ });
+
+ act(() => {
+ mockSocket1.emitData({
+ type: "message_part",
+ chat_id: chatID1,
+ message_part: {
+ role: "assistant",
+ part: {
+ type: "text",
+ text: "chat1-stream",
+ },
+ },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.streamState?.blocks).toEqual([
+ { type: "response", text: "chat1-stream" },
+ ]);
+ });
+
+ rerender({
+ ...initialOptions,
+ chatID: chatID2,
+ chatMessages: [msg2],
+ chatRecord: makeChat(chatID2),
+ chatData: {
+ chat: makeChat(chatID2),
+ messages: [msg2],
+ queued_messages: [],
+ },
+ });
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID2);
+ });
+
+ // The old WebSocket was closed during effect cleanup.
+ expect(mockSocket1.close).toHaveBeenCalled();
+ // Stream state was reset — no stale stream data from chat-1.
+ expect(result.current.streamState).toBeNull();
+ });
+
it("ignores queue_update events for other chats", async () => {
const chatID = "chat-1";
const otherChatID = "chat-2";
@@ -844,4 +939,277 @@ describe("useChatStore", () => {
).toEqual([queuedMessage.id]);
});
});
+
+ it("filters message events with mismatched chat_id", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-1";
+ const existingMessage = makeMessage(chatID, 1, "user", "hello");
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const { result } = renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [existingMessage],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [existingMessage],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ streamState: useChatSelector(store, selectStreamState),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ // Build up stream state so we can observe whether it gets cleared.
+ act(() => {
+ mockSocket.emitData({
+ type: "message_part",
+ chat_id: chatID,
+ message_part: {
+ role: "assistant",
+ part: {
+ type: "text",
+ text: "streaming",
+ },
+ },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.streamState?.blocks).toEqual([
+ { type: "response", text: "streaming" },
+ ]);
+ });
+
+ // A message event with a mismatched chat_id should be ignored
+ // and should NOT trigger scheduleStreamReset.
+ const mismatchedMessage = makeMessage(
+ "chat-2",
+ 99,
+ "assistant",
+ "wrong chat",
+ );
+ act(() => {
+ mockSocket.emitData({
+ type: "message",
+ chat_id: "chat-2",
+ message: mismatchedMessage,
+ });
+ });
+
+ // Stream state should still be present — the mismatched event
+ // was filtered and did not trigger a stream reset.
+ await waitFor(() => {
+ expect(result.current.streamState?.blocks).toEqual([
+ { type: "response", text: "streaming" },
+ ]);
+ });
+
+ // A message event with the correct chat_id should be processed
+ // and trigger scheduleStreamReset, clearing stream state.
+ const matchingMessage = makeMessage(chatID, 2, "assistant", "correct chat");
+ act(() => {
+ mockSocket.emitData({
+ type: "message",
+ chat_id: chatID,
+ message: matchingMessage,
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.streamState).toBeNull();
+ });
+ });
+
+ it("startTransition deferred parts are discarded after chat switch", async () => {
+ immediateAnimationFrame();
+
+ const chatID1 = "chat-1";
+ const chatID2 = "chat-2";
+ const msg1 = makeMessage(chatID1, 1, "user", "hello");
+ const msg2 = makeMessage(chatID2, 10, "user", "world");
+
+ const mockSocket1 = createMockSocket();
+ const mockSocket2 = createMockSocket();
+ // Use a fallback so that extra effect re-runs get a valid socket.
+ vi.mocked(watchChat).mockReturnValue(mockSocket2 as never);
+ vi.mocked(watchChat)
+ .mockReturnValueOnce(mockSocket1 as never)
+ .mockReturnValueOnce(mockSocket1 as never)
+ .mockReturnValueOnce(mockSocket1 as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const initialOptions = {
+ chatID: chatID1,
+ chatMessages: [msg1] as TypesGen.ChatMessage[],
+ chatRecord: makeChat(chatID1),
+ chatData: {
+ chat: makeChat(chatID1),
+ messages: [msg1],
+ queued_messages: [] as TypesGen.ChatQueuedMessage[],
+ },
+ chatQueuedMessages: [] as TypesGen.ChatQueuedMessage[],
+ setChatErrorReason,
+ clearChatErrorReason,
+ };
+
+ const { result, rerender } = renderHook(
+ (options: Parameters[0]) => {
+ const { store } = useChatStore(options);
+ return {
+ streamState: useChatSelector(store, selectStreamState),
+ };
+ },
+ { initialProps: initialOptions, wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID1);
+ });
+
+ act(() => {
+ mockSocket1.emitData({
+ type: "message_part",
+ chat_id: chatID1,
+ message_part: {
+ role: "assistant",
+ part: {
+ type: "text",
+ text: "stale-part",
+ },
+ },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.streamState?.blocks).toEqual([
+ { type: "response", text: "stale-part" },
+ ]);
+ });
+
+ rerender({
+ ...initialOptions,
+ chatID: chatID2,
+ chatMessages: [msg2],
+ chatRecord: makeChat(chatID2),
+ chatData: {
+ chat: makeChat(chatID2),
+ messages: [msg2],
+ queued_messages: [],
+ },
+ });
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID2);
+ });
+
+ expect(result.current.streamState).toBeNull();
+ });
+
+ it("messages are cleared immediately on chat switch before new query resolves", async () => {
+ immediateAnimationFrame();
+
+ const chatID1 = "chat-1";
+ const chatID2 = "chat-2";
+ const msg1 = makeMessage(chatID1, 1, "user", "first");
+ const queuedMsg = makeQueuedMessage(chatID1, 10, "queued");
+
+ const mockSocket1 = createMockSocket();
+ const mockSocket2 = createMockSocket();
+ // Use a fallback so that extra effect re-runs get a valid socket.
+ vi.mocked(watchChat).mockReturnValue(mockSocket2 as never);
+ vi.mocked(watchChat)
+ .mockReturnValueOnce(mockSocket1 as never)
+ .mockReturnValueOnce(mockSocket1 as never)
+ .mockReturnValueOnce(mockSocket1 as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const initialOptions = {
+ chatID: chatID1,
+ chatMessages: [msg1] as TypesGen.ChatMessage[],
+ chatRecord: makeChat(chatID1),
+ chatData: {
+ chat: makeChat(chatID1),
+ messages: [msg1],
+ queued_messages: [queuedMsg],
+ },
+ chatQueuedMessages: [queuedMsg],
+ setChatErrorReason,
+ clearChatErrorReason,
+ };
+
+ const { result, rerender } = renderHook(
+ (options: Parameters[0]) => {
+ const { store } = useChatStore(options);
+ return {
+ queuedMessages: useChatSelector(store, selectQueuedMessages),
+ };
+ },
+ { initialProps: initialOptions, wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID1);
+ });
+
+ // Verify queued messages from chat-1 are present.
+ expect(result.current.queuedMessages.map((m) => m.id)).toEqual([
+ queuedMsg.id,
+ ]);
+
+ // Switch to chat-2 with no messages and no queued messages
+ // (simulating query not yet resolved for the new chat).
+ rerender({
+ ...initialOptions,
+ chatID: chatID2,
+ chatMessages: [],
+ chatRecord: makeChat(chatID2),
+ chatData: {
+ chat: makeChat(chatID2),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ });
+
+ // After the switch, queued messages from chat-1 should NOT be
+ // visible — the store resets them on chatID change.
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID2);
+ });
+ expect(result.current.queuedMessages).toEqual([]);
+ });
});
diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
index 60bd28d435..8b791e01e5 100644
--- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
+++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
@@ -396,6 +396,8 @@ export const useChatStore = (
const storeRef = useRef(createChatStore());
const streamResetFrameRef = useRef(null);
const queuedMessagesHydratedChatIDRef = useRef(null);
+ const activeChatIDRef = useRef(null);
+ const prevChatIDRef = useRef(chatID);
const store = storeRef.current;
@@ -472,8 +474,15 @@ export const useChatStore = (
);
useEffect(() => {
+ // When the active chat changes, clear stale messages immediately
+ // so the previous chat's messages aren't briefly visible while
+ // the new chat's query resolves.
+ if (prevChatIDRef.current !== chatID) {
+ prevChatIDRef.current = chatID;
+ store.replaceMessages([]);
+ }
store.replaceMessages(chatMessages);
- }, [chatMessages, store]);
+ }, [chatID, chatMessages, store]);
useEffect(() => {
store.setChatStatus(chatRecord?.status ?? null);
@@ -501,6 +510,7 @@ export const useChatStore = (
useEffect(() => {
cancelScheduledStreamReset();
store.resetTransientState();
+ activeChatIDRef.current = chatID ?? null;
if (!chatID) {
return;
@@ -535,7 +545,11 @@ export const useChatStore = (
}
cancelScheduledStreamReset();
const parts = pendingMessageParts.splice(0, pendingMessageParts.length);
+ const currentChatID = chatID;
startTransition(() => {
+ if (activeChatIDRef.current !== currentChatID) {
+ return;
+ }
store.applyMessageParts(parts);
});
};
@@ -563,6 +577,10 @@ export const useChatStore = (
if (!message) {
continue;
}
+ const eventChatID = asString(streamEvent.chat_id);
+ if (eventChatID && eventChatID !== chatID) {
+ continue;
+ }
const { changed } = store.upsertDurableMessage(message);
if (changed) {
scheduleStreamReset();
@@ -657,6 +675,7 @@ export const useChatStore = (
socket.removeEventListener("error", handleError);
socket.close();
cancelScheduledStreamReset();
+ activeChatIDRef.current = null;
};
}, [
cancelScheduledStreamReset,