From 1a87e74574c2429bf87b5529bd8e899b57fbd8a2 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Fri, 27 Feb 2026 15:33:16 -0500 Subject: [PATCH] fix(site): prevent race conditions when switching chats on agents page (#22404) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem When switching between chats on the agents page, stream parts could be lost or applied to the wrong chat due to several race conditions in `ChatContext.ts`: 1. **`startTransition` deferred parts escape cleanup** — `startTransition(() => store.applyMessageParts(parts))` defers the state update. If a chat switch happens between `flushMessageParts` being called and the transition executing, old-chat parts could apply after `resetTransientState()` has already cleared stream state for the new chat. 2. **`message` event has no `chat_id` filter** — Unlike `message_part`, `queue_update`, and `status` events, the `message` event handler did not check `streamEvent.chat_id`. While the server-scoped WebSocket makes this safe in practice, it's an inconsistency in defensive programming. 3. **Brief stale message window on switch** — Between `chatID` changing and `replaceMessages()` firing (after the query resolves), the store held old-chat messages while the new WebSocket was already connected. ## Changes ### `ChatContext.ts` - Added `activeChatIDRef` to track the currently active chat ID - Guard `startTransition` callback: check `activeChatIDRef` before applying message parts, discarding them if the chat has switched - Added `chat_id` filter to `message` event handler, matching the pattern used by all other event types - Added `store.replaceMessages([])` to the chatID-change effect so messages are cleared immediately on switch ### `ChatContext.test.tsx` Four new tests covering the chat-switch lifecycle: - WebSocket closure and state reset when chatID changes - `message` event filtering by `chat_id` - `startTransition` deferred parts discarded after switch - Messages cleared immediately before new query resolves All 13 tests pass (8 existing + 4 new + 1 existing). --- .../AgentDetail/ChatContext.test.tsx | 368 ++++++++++++++++++ .../AgentsPage/AgentDetail/ChatContext.ts | 21 +- 2 files changed, 388 insertions(+), 1 deletion(-) 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,