From 6e8067e3a94d1e1dbe181505189a6a6f536021b0 Mon Sep 17 00:00:00 2001 From: HeyItsChloe Date: Thu, 28 May 2026 15:25:52 -0700 Subject: [PATCH] fix(conversation): fix messages disappearing and cross-conversation event bleed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Skip clearEvents on same-conversation remount (Settings → back) - Guard REST seed against stale React Query cache to prevent event mixing - Re-fire REST seed after active conv is stamped to fix empty chat state - Reject stale WebSocket messages from previous conversation's socket - Flush in-progress draft on unmount before debounce fires - Wire up clearPendingMessages on conversation switch (was never called) --- .../conversation-websocket-context.tsx | 81 ++++++++++++++++++- src/hooks/chat/use-draft-persistence.ts | 51 +++++++++--- src/routes/conversation.tsx | 42 ++++++++++ src/stores/use-event-store.ts | 47 +++++++++-- 4 files changed, 205 insertions(+), 16 deletions(-) diff --git a/src/contexts/conversation-websocket-context.tsx b/src/contexts/conversation-websocket-context.tsx index e7ea2c008..92ed83221 100644 --- a/src/contexts/conversation-websocket-context.tsx +++ b/src/contexts/conversation-websocket-context.tsx @@ -129,6 +129,14 @@ export function ConversationWebSocketProvider({ const hasConnectedRefMain = React.useRef(false); const hasConnectedRefPlanning = React.useRef(false); + // Track which conversationId each socket was opened for. + // useWebSocket routes all onMessage calls through optionsRef.current, which + // always points to the latest handler. Without this guard, a stale socket + // from ConvA that fires after a switch to ConvB would call ConvB's handler + // and inject ConvA's events into ConvB's event store. + const mainSocketConvIdRef = React.useRef(null); + const planningSocketConvIdRef = React.useRef(null); + const posthog = usePostHog(); const queryClient = useQueryClient(); const addEvent = useEventStore((state) => state.addEvent); @@ -219,12 +227,48 @@ export function ConversationWebSocketProvider({ const isLoadingHistoryMain = !!conversationId && isPreloadingHistory; + // Subscribe to activeConversationId so this effect re-runs after the + // parent's useEffect stamps the new conversation. Without this subscription, + // if the REST seed fires before the parent clears/stamps (stale cache race), + // the seed is skipped and never retried — leaving the chat permanently empty. + const activeConversationId = useEventStore( + (state) => state.activeConversationId, + ); + useLayoutEffect(() => { if (!preloadedHistory || preloadedHistory.events.length === 0) { return; } + + console.log( + "[OH-DEBUG][ws] REST seed: seeding %d events for conv=%s (storeActiveConv=%s)", + preloadedHistory.events.length, + conversationId, + activeConversationId, + ); + + // Guard against stale React Query cache: when the user switches conversations + // quickly, the incoming conv's cached preloadedHistory fires this layout + // effect before the parent's useEffect has cleared the previous conv's events + // and stamped the new active conv. If the store is already active for a + // *different* conversation, seeding here would mix two convs' events in the + // same store. We skip here and re-run once activeConversationId updates + // (because it is in the deps array below). Null means very first page load + // — allow through. + if ( + activeConversationId !== null && + activeConversationId !== conversationId + ) { + console.log( + "[OH-DEBUG][ws] REST seed SKIPPED — storeActive=%s differs from conv=%s (stale cache)", + activeConversationId, + conversationId, + ); + return; + } + addEvents(preloadedHistory.events); - }, [preloadedHistory, addEvents]); + }, [preloadedHistory, activeConversationId, conversationId, addEvents]); /** * Timestamp of the latest event we already have from REST. Used as @@ -392,6 +436,19 @@ export function ConversationWebSocketProvider({ // Separate message handlers for each connection const handleMainMessage = useCallback( (messageEvent: MessageEvent) => { + // Guard against stale sockets: useWebSocket always calls the latest + // onMessage handler via optionsRef, so a socket opened for ConvA can + // fire this callback after we've switched to ConvB. Reject any message + // whose socket was opened for a different conversation. + if (mainSocketConvIdRef.current !== conversationId) { + console.log( + "[OH-DEBUG][ws:main] STALE REJECT — socket opened for conv=%s but current conv=%s", + mainSocketConvIdRef.current, + conversationId, + ); + return; + } + try { const event = JSON.parse(messageEvent.data); @@ -566,6 +623,16 @@ export function ConversationWebSocketProvider({ const handlePlanningMessage = useCallback( (messageEvent: MessageEvent) => { + // Same stale-socket guard as handleMainMessage. + if (planningSocketConvIdRef.current !== conversationId) { + console.log( + "[OH-DEBUG][ws:planning] STALE REJECT — socket opened for conv=%s but current conv=%s", + planningSocketConvIdRef.current, + conversationId, + ); + return; + } + try { const event = JSON.parse(messageEvent.data); @@ -762,6 +829,11 @@ export function ConversationWebSocketProvider({ queryParams, reconnect: { enabled: true }, onOpen: () => { + mainSocketConvIdRef.current = conversationId ?? null; // Stamp which conv this socket serves + console.log( + "[OH-DEBUG][ws:main] onOpen — socket stamped for conv=%s", + conversationId, + ); setMainConnectionState("OPEN"); hasConnectedRefMain.current = true; // Mark that we've successfully connected clearConnectionError(); // Clear a previous connection error; keep sticky conversation errors @@ -779,6 +851,7 @@ export function ConversationWebSocketProvider({ onMessage: handleMainMessage, }; }, [ + conversationId, handleMainMessage, setErrorMessage, clearConnectionError, @@ -803,6 +876,11 @@ export function ConversationWebSocketProvider({ queryParams, reconnect: { enabled: true }, onOpen: async () => { + planningSocketConvIdRef.current = conversationId ?? null; // Stamp which conv this socket serves + console.log( + "[OH-DEBUG][ws:planning] onOpen — socket stamped for conv=%s", + conversationId, + ); setPlanningConnectionState("OPEN"); hasConnectedRefPlanning.current = true; // Mark that we've successfully connected clearConnectionError(); // Clear a previous connection error; keep sticky conversation errors @@ -843,6 +921,7 @@ export function ConversationWebSocketProvider({ onMessage: handlePlanningMessage, }; }, [ + conversationId, handlePlanningMessage, setErrorMessage, clearConnectionError, diff --git a/src/hooks/chat/use-draft-persistence.ts b/src/hooks/chat/use-draft-persistence.ts index 0a7152400..61176abfb 100644 --- a/src/hooks/chat/use-draft-persistence.ts +++ b/src/hooks/chat/use-draft-persistence.ts @@ -55,6 +55,10 @@ export const useDraftPersistence = ( // React 18 (refs are cleared during the synchronous commit phase, before // passive effects fire), so we can't read from the DOM there. const lastHomeTextRef = useRef(""); + // Same pattern for conversation routes: updated synchronously on every + // saveDraft call so the unmount flush captures the very latest text even + // when the 500 ms debounce hasn't fired yet (e.g. rapid navigate-away). + const lastConvTextRef = useRef(""); // IMPORTANT: This effect must run FIRST when conversation changes. // It handles three concerns: @@ -159,6 +163,12 @@ export const useDraftPersistence = ( // The hook's state may not have synced yet after conversationId change const { draftMessage } = getConversationState(conversationId); + console.log( + "[OH-DEBUG][draft] restore: conv=%s draftMessage=%s", + conversationId, + draftMessage ? `"${draftMessage.slice(0, 40)}…"` : "null", + ); + // Only restore if there's a saved draft and the input is empty if (draftMessage && getTextContent(element).trim() === "") { element.textContent = draftMessage; @@ -199,6 +209,13 @@ export const useDraftPersistence = ( // Capture the conversationId at the time of input const capturedConversationId = conversationId; + // Track latest text synchronously (mirrors lastHomeTextRef) so the unmount + // flush can use it even if the debounce timer hasn't fired yet. + const elementNow = chatInputRef.current; + if (elementNow) { + lastConvTextRef.current = getTextContent(elementNow).trim(); + } + saveTimeoutRef.current = setTimeout(() => { // Verify we're still on the same conversation before saving // This prevents saving draft to wrong conversation if user switched quickly @@ -238,17 +255,19 @@ export const useDraftPersistence = ( setDraftMessage(null); }, [conversationId, setDraftMessage]); - // Cleanup on unmount: cancel any pending debounce timer and, for the home - // page, flush the last-tracked text to sessionStorage. We read from - // lastHomeTextRef rather than chatInputRef because React clears ref.current - // during the synchronous commit phase — before async useEffect cleanups run - // — so the DOM ref is null by the time this function executes. + // Cleanup on unmount: cancel any pending debounce timer and flush the + // latest draft text. We read from refs rather than chatInputRef because + // React clears ref.current during the synchronous commit phase — before + // async useEffect cleanups run — so the DOM ref is null by then. + // + // Home page: only write back if the sessionStorage key already exists. + // saveDraft writes synchronously on every keystroke, so the key is present + // whenever there is unsaved text. If absent it was intentionally removed + // (e.g. by HomeChatLauncher.onSuccess) and we must not restore it. // - // We only write text back if the key already exists in sessionStorage. - // saveDraft writes synchronously on every keystroke, so the key is present - // whenever there is unsaved text. If the key is absent it was intentionally - // removed — most importantly by HomeChatLauncher.onSuccess after a - // successful conversation start — and we must not restore it here. + // Conversation route: always flush lastConvTextRef to localStorage so a + // draft in progress is never lost when the user navigates to Settings or + // another conversation before the 500 ms debounce fires. useEffect( () => () => { if (saveTimeoutRef.current) { @@ -267,6 +286,18 @@ export const useDraftPersistence = ( } catch { // sessionStorage not available } + } else if (!isTaskId(currentConversationIdRef.current)) { + // Flush the last-tracked text synchronously to localStorage so + // in-progress drafts survive navigation before the debounce fires. + const text = lastConvTextRef.current; + console.log( + "[OH-DEBUG][draft] unmount flush: conv=%s text=%s", + currentConversationIdRef.current, + text ? `"${text.slice(0, 40)}…"` : "(empty)", + ); + setConversationState(currentConversationIdRef.current, { + draftMessage: text || null, + }); } }, [], diff --git a/src/routes/conversation.tsx b/src/routes/conversation.tsx index e210ea879..f323ddd95 100644 --- a/src/routes/conversation.tsx +++ b/src/routes/conversation.tsx @@ -30,6 +30,7 @@ import { WebSocketProviderWrapper } from "#/contexts/websocket-provider-wrapper" import { useErrorMessageStore } from "#/stores/error-message-store"; import { I18nKey } from "#/i18n/declaration"; import { useEventStore } from "#/stores/use-event-store"; +import { useOptimisticUserMessageStore } from "#/stores/optimistic-user-message-store"; import { resumeCloudSandbox } from "#/api/cloud/conversation-service.api"; function AppContent() { @@ -37,6 +38,9 @@ function AppContent() { const { conversationId } = useConversationId(); const panelViewMatch = useMatch("/conversations/:conversationId/panel"); const clearEvents = useEventStore((state) => state.clearEvents); + const clearPendingMessages = useOptimisticUserMessageStore( + (state) => state.clearPendingMessages, + ); const { isTask, taskStatus, taskDetail } = useTaskPolling(); @@ -71,12 +75,49 @@ function AppContent() { ); React.useEffect(() => { + // Read directly from the store (not via a hook subscription) so this value + // doesn't need to be a dep and can't cause the effect to loop. + // The store survives component unmount (Zustand is module-level), so on + // remount the store still holds the id that was active when we left. + const { activeConversationId, setActiveConversationId } = + useEventStore.getState(); + + const isSameConversation = activeConversationId === conversationId; + console.log( + "[OH-DEBUG][conversation] mount effect — convId=%s storeActiveId=%s isSameConversation=%s", + conversationId, + activeConversationId, + isSameConversation, + ); + + // Mark this conversation as the currently active one before any early return + // so subsequent effects always see an up-to-date value. + setActiveConversationId(conversationId); + + if (isSameConversation) { + // Returning to the same conversation (e.g. navigated to Settings and back). + // The event store already holds the correct history — clearing it would + // wipe messages that the REST cache seeded and that the WebSocket won't + // re-deliver because it uses resend_mode=since. + console.log( + "[OH-DEBUG][conversation] SKIP reset — same conversation remount, events preserved", + ); + return; + } + + // Genuine conversation switch: tear down stale state from the previous one. + console.log( + "[OH-DEBUG][conversation] RESET — switching from %s to %s", + activeConversationId, + conversationId, + ); clearTerminal(); resetConversationState(); resetConversationRuntimeState(); setCurrentAgentState(AgentState.LOADING); removeErrorMessage(); clearEvents(); + clearPendingMessages(); }, [ conversationId, clearTerminal, @@ -85,6 +126,7 @@ function AppContent() { setCurrentAgentState, removeErrorMessage, clearEvents, + clearPendingMessages, ]); React.useEffect(() => { diff --git a/src/stores/use-event-store.ts b/src/stores/use-event-store.ts index 6abd6c81d..5093609ca 100644 --- a/src/stores/use-event-store.ts +++ b/src/stores/use-event-store.ts @@ -48,9 +48,17 @@ const needsSorting = (events: OHEvent[], newEvent: OHEvent): boolean => { }; export interface EventState { + /** + * The conversation ID whose events are currently held in the store. + * Survives component unmount (Zustand is module-level) so conversation.tsx + * can compare on remount and skip a redundant clearEvents when the user + * navigates away to Settings and back to the same conversation. + */ + activeConversationId: string | null; events: OHEvent[]; eventIds: Set; uiEvents: OHEvent[]; + setActiveConversationId: (id: string | null) => void; addEvent: (event: OHEvent) => void; /** * Bulk-insert events. Used for the initial REST history load and for @@ -105,12 +113,28 @@ const applyAddEvent = (state: EventState, event: OHEvent): EventState => { }; export const useEventStore = create()((set) => ({ + activeConversationId: null, events: [], eventIds: new Set(), uiEvents: [], + setActiveConversationId: (id) => + set((state) => { + console.log( + "[OH-DEBUG][event-store] setActiveConversationId: %s → %s", + state.activeConversationId, + id, + ); + return { activeConversationId: id }; + }), addEvent: (event: OHEvent) => set((state) => applyAddEvent(state, event)), addEvents: (incoming: OHEvent[]) => set((state) => { + console.log( + "[OH-DEBUG][event-store] addEvents: +%d events (activeConv=%s, storeSize=%d)", + incoming.length, + state.activeConversationId, + state.events.length, + ); if (incoming.length === 0) return state; const eventIds = new Set(state.eventIds); @@ -144,11 +168,24 @@ export const useEventStore = create()((set) => ({ }); }), clearEvents: () => - set(() => ({ - events: [], - eventIds: new Set(), - uiEvents: [], - })), + set((state) => { + console.log( + "[OH-DEBUG][event-store] clearEvents: wiping %d events (wasActiveConv=%s)", + state.events.length, + state.activeConversationId, + ); + // activeConversationId is intentionally NOT reset here. + // It must survive clearEvents so that the next mount of conversation.tsx + // can compare the incoming convId against the last active one and detect + // a same-conversation remount (e.g. Settings → back). Resetting to null + // here caused every mount to see storeActiveId=null, making + // isSameConversation always false and defeating the whole guard. + return { + events: [], + eventIds: new Set(), + uiEvents: [], + }; + }), })); // In dev builds, expose the store on `window` so that fixture/preview