diff --git a/webui/src/components/thread/ThreadShell.tsx b/webui/src/components/thread/ThreadShell.tsx index 65ffd1e0d..7dc2afaec 100644 --- a/webui/src/components/thread/ThreadShell.tsx +++ b/webui/src/components/thread/ThreadShell.tsx @@ -44,6 +44,7 @@ export function ThreadShell({ const [booting, setBooting] = useState(false); const pendingFirstRef = useRef(null); const messageCacheRef = useRef>(new Map()); + const lastCachedChatIdRef = useRef(null); const initial = useMemo(() => { if (!chatId) return historical; @@ -91,6 +92,13 @@ export function ThreadShell({ useEffect(() => { if (!chatId) return; + // Skip the first cache write after a chat switch. During that render, + // `messages` can still belong to the previous chat until the stream hook + // resets its local state for the new session. + if (lastCachedChatIdRef.current !== chatId) { + lastCachedChatIdRef.current = chatId; + return; + } messageCacheRef.current.set(chatId, messages); }, [chatId, messages]); diff --git a/webui/src/tests/thread-shell.test.tsx b/webui/src/tests/thread-shell.test.tsx index d134fcce2..f5dea5960 100644 --- a/webui/src/tests/thread-shell.test.tsx +++ b/webui/src/tests/thread-shell.test.tsx @@ -267,6 +267,93 @@ describe("ThreadShell", () => { expect(screen.queryByText("old answer")).not.toBeInTheDocument(); }); + it("does not cache optimistic messages under the next chat during a session switch", async () => { + const client = makeClient(); + const onNewChat = vi.fn().mockResolvedValue("chat-b"); + + const { rerender } = render( + wrap( + client, + {}} + onGoHome={() => {}} + onNewChat={onNewChat} + />, + ), + ); + + fireEvent.change(screen.getByLabelText("Message input"), { + target: { value: "only in chat a" }, + }); + fireEvent.click(screen.getByRole("button", { name: "Send message" })); + + await waitFor(() => + expect(client.sendMessage).toHaveBeenCalledWith( + "chat-a", + "only in chat a", + undefined, + ), + ); + expect(screen.getByText("only in chat a")).toBeInTheDocument(); + + await act(async () => { + rerender( + wrap( + client, + {}} + onGoHome={() => {}} + onNewChat={onNewChat} + />, + ), + ); + }); + + await waitFor(() => { + expect(screen.queryByText("only in chat a")).not.toBeInTheDocument(); + }); + + await act(async () => { + rerender( + wrap( + client, + {}} + onGoHome={() => {}} + onNewChat={onNewChat} + />, + ), + ); + }); + + expect(screen.getByText("only in chat a")).toBeInTheDocument(); + + await act(async () => { + rerender( + wrap( + client, + {}} + onGoHome={() => {}} + onNewChat={onNewChat} + />, + ), + ); + }); + + await waitFor(() => { + expect(screen.queryByText("only in chat a")).not.toBeInTheDocument(); + }); + }); + it("surfaces a dismissible banner when the stream reports message_too_big", async () => { const client = makeClient(); const onNewChat = vi.fn().mockResolvedValue("chat-a");