diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
index b4244bb406..7a737e433d 100644
--- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
+++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
@@ -797,6 +797,80 @@ describe("useChatStore", () => {
});
});
+ it("corrects stale queued messages from cache when switching back to a chat", async () => {
+ const chatID = "chat-1";
+ const existingMessage = makeMessage(chatID, 1, "user", "hello");
+ const queuedMessage = makeQueuedMessage(chatID, 10, "queued");
+ 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();
+
+ // Start with queued messages from a stale React Query cache.
+ // This simulates coming back to a chat whose queue was drained
+ // server-side while the user was viewing a different chat.
+ const staleOptions = {
+ chatID,
+ chatMessages: [existingMessage],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [existingMessage],
+ queued_messages: [queuedMessage],
+ },
+ chatQueuedMessages: [queuedMessage],
+ setChatErrorReason,
+ clearChatErrorReason,
+ };
+
+ const { result, rerender } = renderHook(
+ (options: Parameters[0]) => {
+ const { store } = useChatStore(options);
+ return {
+ queuedMessages: useChatSelector(store, selectQueuedMessages),
+ };
+ },
+ {
+ initialProps: staleOptions,
+ wrapper,
+ },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID, 1);
+ });
+ // Initially shows the stale queued message from cache.
+ expect(result.current.queuedMessages.map((m) => m.id)).toEqual([
+ queuedMessage.id,
+ ]);
+
+ // Simulate the REST query refetching and returning fresh
+ // data with an empty queue (no queue_update from WS yet).
+ rerender({
+ ...staleOptions,
+ chatData: {
+ chat: {
+ ...makeChat(chatID),
+ updated_at: "2025-01-01T00:00:02.000Z",
+ },
+ messages: [existingMessage],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ });
+
+ // The store should accept the fresh REST data because the
+ // WebSocket hasn't sent a queue_update yet.
+ await waitFor(() => {
+ expect(result.current.queuedMessages).toEqual([]);
+ });
+ });
+
it("writes queue_update snapshots into the chat query cache", async () => {
const chatID = "chat-1";
const existingMessage = makeMessage(chatID, 1, "user", "hello");
diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
index 6c571d9b73..6217cd76f4 100644
--- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
+++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
@@ -454,6 +454,13 @@ export const useChatStore = (
const storeRef = useRef(createChatStore());
const streamResetFrameRef = useRef(null);
const queuedMessagesHydratedChatIDRef = useRef(null);
+ // Tracks whether the WebSocket has delivered a queue_update for the
+ // current chat. When true, the stream is the authoritative source
+ // and REST re-fetches must not overwrite the store. When false,
+ // REST data is allowed to re-hydrate so stale cached queued
+ // messages are corrected when switching back to a chat whose
+ // queue was drained while the user was away.
+ const wsQueueUpdateReceivedRef = useRef(false);
const activeChatIDRef = useRef(null);
const prevChatIDRef = useRef(chatID);
@@ -553,6 +560,7 @@ export const useChatStore = (
useEffect(() => {
queuedMessagesHydratedChatIDRef.current = null;
+ wsQueueUpdateReceivedRef.current = false;
store.setQueuedMessages([]);
if (!chatID) {
return;
@@ -563,7 +571,15 @@ export const useChatStore = (
if (!chatID || !chatData) {
return;
}
- if (queuedMessagesHydratedChatIDRef.current === chatID) {
+ // Allow re-hydration from REST as long as the WebSocket hasn't
+ // delivered a queue_update yet (which would be fresher). This
+ // ensures that when the user navigates back to a chat whose
+ // queued messages were drained server-side while they were
+ // away, the REST refetch corrects the stale cached state.
+ if (
+ queuedMessagesHydratedChatIDRef.current === chatID &&
+ wsQueueUpdateReceivedRef.current
+ ) {
return;
}
queuedMessagesHydratedChatIDRef.current = chatID;
@@ -688,6 +704,7 @@ export const useChatStore = (
continue;
}
}
+ wsQueueUpdateReceivedRef.current = true;
store.setQueuedMessages(streamEvent.queued_messages);
updateChatQueuedMessages(streamEvent.queued_messages);
continue;