diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx
index 59a46c953c..692bb3543b 100644
--- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx
+++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx
@@ -963,6 +963,112 @@ describe("useChatStore", () => {
expect(cachedData?.pages[0]?.queued_messages).toEqual([]);
});
+ it("writes WebSocket message events into the chat query cache", async () => {
+ const chatID = "chat-1";
+ const existingMessage = makeMessage(chatID, 1, "user", "hello");
+ const mockSocket = createMockSocket();
+ mockWatchChatReturn(mockSocket);
+
+ const queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ gcTime: Number.POSITIVE_INFINITY,
+ refetchOnWindowFocus: false,
+ networkMode: "offlineFirst",
+ },
+ },
+ });
+ const initialChatMessagesData: TypesGen.ChatMessagesResponse = {
+ messages: [existingMessage],
+ queued_messages: [],
+ has_more: false,
+ };
+ queryClient.setQueryData(chatMessagesKey(chatID), {
+ pages: [initialChatMessagesData],
+ pageParams: [undefined],
+ });
+
+ 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),
+ chatMessagesData: initialChatMessagesData,
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ orderedIDs: useChatSelector(store, selectOrderedMessageIDs),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID, 1);
+ });
+
+ const newMessage = makeMessage(chatID, 2, "assistant", "hi there");
+ act(() => {
+ mockSocket.emitData({
+ type: "message",
+ chat_id: chatID,
+ message: newMessage,
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.orderedIDs).toContain(2);
+ });
+
+ // The React Query cache should also contain the new message.
+ const cachedData = queryClient.getQueryData<{
+ pages: TypesGen.ChatMessagesResponse[];
+ pageParams: unknown[];
+ }>(chatMessagesKey(chatID));
+ const cachedMessages = cachedData?.pages[0]?.messages ?? [];
+ // Verifies insertion, preservation, and DESC order.
+ expect(cachedMessages.map((m) => m.id)).toEqual([2, 1]);
+ // Emitting the same message again should not change the
+ // cache reference (reference stability).
+ const refBefore = queryClient.getQueryData(chatMessagesKey(chatID));
+ act(() => {
+ mockSocket.emitData({
+ type: "message",
+ chat_id: chatID,
+ message: newMessage,
+ });
+ });
+ const refAfter = queryClient.getQueryData(chatMessagesKey(chatID));
+ expect(refAfter).toBe(refBefore);
+
+ // Emitting the same message ID with different content should
+ // update the cached entry (content-update path).
+ const revised = makeMessage(chatID, 2, "assistant", "revised");
+ act(() => {
+ mockSocket.emitData({
+ type: "message",
+ chat_id: chatID,
+ message: revised,
+ });
+ });
+ const updatedCache = queryClient.getQueryData<{
+ pages: TypesGen.ChatMessagesResponse[];
+ pageParams: unknown[];
+ }>(chatMessagesKey(chatID));
+ const updatedFirst = updatedCache?.pages[0]?.messages[0];
+ expect(updatedFirst?.content).toEqual([{ type: "text", text: "revised" }]);
+ });
+
it("closes old WebSocket and resets state when chatID changes", async () => {
immediateAnimationFrame();
diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts
index 978ffd87b6..0558d68245 100644
--- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts
+++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts
@@ -64,7 +64,7 @@ const jsonValuesEqual = (left: unknown, right: unknown): boolean => {
}
};
-const chatMessagesEqualByValue = (
+export const chatMessagesEqualByValue = (
left: TypesGen.ChatMessage,
right: TypesGen.ChatMessage,
): boolean =>
diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts
index 93117ccc59..cac0a21dab 100644
--- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts
+++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts
@@ -11,6 +11,7 @@ import { asNumber, asString } from "../ChatElements/runtimeTypeUtils";
import {
type ChatStore,
type ChatStoreState,
+ chatMessagesEqualByValue,
chatQueuedMessagesEqualByID,
createChatStore,
isActiveChatStatus,
@@ -286,6 +287,55 @@ export const useChatStore = (
};
});
};
+
+ // Write WebSocket-delivered durable messages into the React
+ // Query infinite cache so that navigating away and back
+ // serves up-to-date data instead of the stale REST snapshot.
+ // Without this, the cache only contains messages from the
+ // last REST fetch, and structural sharing can suppress the
+ // refetch-driven store update when no new durable messages
+ // have been committed to the DB yet.
+ const upsertCacheMessages = (messages: readonly TypesGen.ChatMessage[]) => {
+ if (!chatID || messages.length === 0) {
+ return;
+ }
+ queryClient.setQueryData<
+ InfiniteData | undefined
+ >(chatMessagesKey(chatID), (currentData) => {
+ if (!currentData?.pages?.length) {
+ return currentData;
+ }
+ const firstPage = currentData.pages[0];
+ const existingByID = new Map(firstPage.messages.map((m) => [m.id, m]));
+
+ let changed = false;
+ for (const msg of messages) {
+ const existing = existingByID.get(msg.id);
+ if (!existing || !chatMessagesEqualByValue(existing, msg)) {
+ changed = true;
+ existingByID.set(msg.id, msg);
+ }
+ }
+
+ if (!changed) {
+ return currentData;
+ }
+
+ // Sort descending to match the API page order
+ // (newest first).
+ const updatedMessages = Array.from(existingByID.values());
+ updatedMessages.sort((a, b) => b.id - a.id);
+
+ return {
+ ...currentData,
+ pages: [
+ { ...firstPage, messages: updatedMessages },
+ ...currentData.pages.slice(1),
+ ],
+ };
+ });
+ };
+
cancelScheduledStreamReset();
store.resetTransientState();
activeChatIDRef.current = chatID ?? null;
@@ -499,6 +549,7 @@ export const useChatStore = (
// pass: one Map copy + one sort instead of N each.
if (pendingMessages.length > 0) {
store.upsertDurableMessages(pendingMessages);
+ upsertCacheMessages(pendingMessages);
}
});
if (needsStreamReset) {