diff --git a/apps/web/src/domains/chat/chat-layout.tsx b/apps/web/src/domains/chat/chat-layout.tsx index d8ec9f0728c..861c3018be2 100644 --- a/apps/web/src/domains/chat/chat-layout.tsx +++ b/apps/web/src/domains/chat/chat-layout.tsx @@ -17,6 +17,7 @@ import type { AssistantContextValue } from "@/domains/chat/assistant-context.js" import { useConversationListStore } from "@/domains/conversations/conversation-list-store.js"; import { useConversationListInit } from "@/domains/conversations/use-conversation-list-init.js"; +import { useAttentionTracking } from "@/domains/chat/hooks/use-attention-tracking.js"; import { useFeatureFlagStore } from "@/lib/feature-flags/feature-flag-store.js"; import { useViewerStore } from "@/stores/viewer-store.js"; import { useSubagentStore } from "@/domains/subagents/subagent-store.js"; @@ -129,6 +130,17 @@ export function ChatLayout() { conversationGroupsUI, }); + // Track processing/attention indicators for every conversation in the + // sidebar, on every chat-layout child route. Mounted here (not ChatPage) + // so the 10s polling loop and graduation logic stay live when the user is + // on home/library/contacts/identity. Pass lifecycle values directly — + // `useAssistantContext()` would crash here since this hook runs inside + // the layout that PROVIDES that context (no parent outlet to read from). + useAttentionTracking({ + assistantId: lifecycle.assistantId, + assistantStateKind: lifecycle.assistantState.kind, + }); + // --- Layout slot state for child route content --- const [topBarCenter, setTopBarCenter] = useState(null); const [topBarRightSlot, setTopBarRightSlot] = useState(null); diff --git a/apps/web/src/domains/chat/chat-page.tsx b/apps/web/src/domains/chat/chat-page.tsx index 042b82001df..a78e3bf5389 100644 --- a/apps/web/src/domains/chat/chat-page.tsx +++ b/apps/web/src/domains/chat/chat-page.tsx @@ -144,7 +144,6 @@ export function ChatPage() { const activeConversationKey = useConversationListStore.use.activeConversationKey(); const editingConversationKey = useConversationListStore.use.editingConversationKey(); const processingKeys = useConversationListStore.use.processingKeys(); - const attentionKeys = useConversationListStore.use.attentionKeys(); const viewerState = useViewerStore(useShallow((s) => ({ mainView: s.mainView, activeAppId: s.activeAppId, @@ -176,15 +175,12 @@ export function ChatPage() { const assistantIdRef = useRef(assistantId); useEffect(() => { assistantIdRef.current = assistantId; }, [assistantId]); - const conversationsRef = useRef(conversations); - conversationsRef.current = conversations; const streamRef = useRef(null); const streamEpochRef = useRef(0); const streamContextRef = useRef<{ assistantId: string; conversationKey: string } | null>(null); const reconcileAfterNextStreamOpenRef = useRef(false); const needsNewBubbleRef = useRef(true); - const processingSnapshotsRef = useRef>(new Map()); const dismissedSurfaceIdsRef = useRef>(new Set()); const pendingOnboardingContextRef = useRef(null); const onboardingDraftConversationKeyRef = useRef(null); @@ -319,9 +315,6 @@ export function ChatPage() { searchParams, navigate, conversations, - activeConversation, - processingKeys, - attentionKeys, transcriptPagination, conversationGroupsUI, refreshEpoch, @@ -335,7 +328,6 @@ export function ChatPage() { inputRef, draftsRef, messagesRef, - conversationsRef, contextWindowUsageByConversationRef, dismissedSurfaceIdsRef, needsNewBubbleRef, @@ -344,7 +336,6 @@ export function ChatPage() { requestIdToStableIdRef, pendingLocalDeletionsRef, confirmationToolCallMapRef, - processingSnapshotsRef, refreshSettleRef, lastSuggestionMsgIdRef, autoGreetRef, @@ -469,7 +460,6 @@ export function ChatPage() { setMessages, messagesRef, needsNewBubbleRef, - processingSnapshotsRef, setError, streamRef, cancelReconciliation, @@ -508,12 +498,10 @@ export function ChatPage() { assistantIdRef, activeConversationKeyRef, messagesRef, - conversationsRef, streamRef, streamContextRef, streamEpochRef, needsNewBubbleRef, - processingSnapshotsRef, dismissedSurfaceIdsRef, pendingOnboardingContextRef, onboardingDraftConversationKeyRef, @@ -580,7 +568,6 @@ export function ChatPage() { reachabilityProbe: reachability.probe, reachabilityPhase: reachability.state.phase, reachabilityReset: reachability.reset, - processingSnapshotsRef, setMessages, setError, streamRetryNonce, @@ -1087,7 +1074,6 @@ export function ChatPage() { refreshSettleRef, streamRef, streamEpochRef, - processingSnapshotsRef, historyLoadedRef, pendingQueuedStableIdsRef, requestIdToStableIdRef, diff --git a/apps/web/src/domains/chat/components/chat-route-content.tsx b/apps/web/src/domains/chat/components/chat-route-content.tsx index 5ab84b23127..f1ec0c264d8 100644 --- a/apps/web/src/domains/chat/components/chat-route-content.tsx +++ b/apps/web/src/domains/chat/components/chat-route-content.tsx @@ -217,7 +217,6 @@ export interface ChatRouteRefs { refreshSettleRef: MutableRefObject; streamRef: MutableRefObject; streamEpochRef: MutableRefObject; - processingSnapshotsRef: MutableRefObject>; historyLoadedRef: MutableRefObject; pendingQueuedStableIdsRef: MutableRefObject; requestIdToStableIdRef: MutableRefObject>; @@ -489,7 +488,6 @@ export function ChatRouteContent({ refreshSettleRef, streamRef: _streamRef, streamEpochRef: _streamEpochRef, - processingSnapshotsRef: _processingSnapshotsRef, historyLoadedRef: _historyLoadedRef, pendingQueuedStableIdsRef: _pendingQueuedStableIdsRef, requestIdToStableIdRef: _requestIdToStableIdRef, diff --git a/apps/web/src/domains/chat/hooks/use-attention-tracking.ts b/apps/web/src/domains/chat/hooks/use-attention-tracking.ts index 2af0b4a41eb..8b4968c7eb3 100644 --- a/apps/web/src/domains/chat/hooks/use-attention-tracking.ts +++ b/apps/web/src/domains/chat/hooks/use-attention-tracking.ts @@ -1,31 +1,17 @@ import * as Sentry from "@sentry/react"; -import { - type MutableRefObject, - useEffect, - useRef, -} from "react"; +import { useEffect, useRef } from "react"; import { useConversationListStore } from "@/domains/conversations/conversation-list-store.js"; -import type { AssistantStateKind } from "@/domains/chat/types.js"; -import { type Conversation, markConversationSeen } from "@/domains/chat/api/conversations.js"; +import { markConversationSeen } from "@/domains/chat/api/conversations.js"; import { listConversationKeysWithPendingInteractions } from "@/domains/chat/api/interactions.js"; +import type { AssistantState } from "@/domains/chat/hooks/use-assistant-lifecycle.js"; interface UseAttentionTrackingParams { + /** From `useAssistantLifecycle` in `ChatLayout`. */ assistantId: string | null; - assistantStateKind: AssistantStateKind; - activeConversationKey: string | null; - - // Collections - conversations: Conversation[]; - activeConversation: Conversation | undefined; - processingKeys: Set; - attentionKeys: Set; - - // Refs - conversationsRef: MutableRefObject; - processingSnapshotsRef: MutableRefObject>; - + /** From `useAssistantLifecycle` in `ChatLayout`. */ + assistantStateKind: AssistantState["kind"]; } // --------------------------------------------------------------------------- @@ -36,6 +22,11 @@ interface UseAttentionTrackingParams { * Tracks which conversations need user attention (pending interactions) * and manages processing-key lifecycle for background conversations. * + * Reads conversations, processingKeys, attentionKeys, and processingSnapshots + * directly from `useConversationListStore`. Mounted in `ChatLayout` so the + * sidebar's processing/attention indicators stay live on every chat-layout + * route (home, library, contacts, identity, chat) — not only `/assistant`. + * * Handles: * - Marking conversations as seen when opened * - Graduating processing keys when the assistant finishes responding @@ -82,14 +73,16 @@ export function decideGraduationDispatches( export function useAttentionTracking({ assistantId, assistantStateKind, - activeConversationKey, - conversations, - activeConversation, - processingKeys, - attentionKeys, - conversationsRef, - processingSnapshotsRef, }: UseAttentionTrackingParams) { + const conversations = useConversationListStore.use.conversations(); + const activeConversationKey = useConversationListStore.use.activeConversationKey(); + const processingKeys = useConversationListStore.use.processingKeys(); + const attentionKeys = useConversationListStore.use.attentionKeys(); + + const activeConversation = conversations.find( + (c) => c.conversationKey === activeConversationKey, + ); + const lastSeenOnOpenConversationKeyRef = useRef(null); const initialAttentionSweepDoneRef = useRef(false); @@ -135,12 +128,13 @@ export function useAttentionTracking({ // ------------------------------------------------------------------------- useEffect(() => { if (processingKeys.size === 0) return; + const snapshots = useConversationListStore.getState().processingSnapshots; const graduatingKeys: string[] = []; for (const key of processingKeys) { if (key === activeConversationKey) continue; const conv = conversations.find((c) => c.conversationKey === key); if (!conv) continue; - const snapshot = processingSnapshotsRef.current.get(key); + const snapshot = snapshots.get(key); if (conv.latestAssistantMessageAt && conv.latestAssistantMessageAt !== snapshot) { graduatingKeys.push(key); } @@ -164,13 +158,12 @@ export function useAttentionTracking({ useConversationListStore.getState().addAttentionKey(action.key); } else { useConversationListStore.getState().removeProcessingKey(action.key); - processingSnapshotsRef.current.delete(action.key); } } })(); return () => { cancelled = true; }; - }, [conversations, processingKeys, activeConversationKey, assistantId, processingSnapshotsRef]); + }, [conversations, processingKeys, activeConversationKey, assistantId]); // ------------------------------------------------------------------------- // Poll processing + attention conversations every 10s. @@ -198,27 +191,35 @@ export function useAttentionTracking({ } if (cancelled) return; + // Read latest store values inside the tick — the effect captured the + // sets at scheduling time, which would be stale ten seconds later. + const state = useConversationListStore.getState(); + const currentProcessingKeys = state.processingKeys; + const currentAttentionKeys = state.attentionKeys; + const currentConversations = state.conversations; + const currentSnapshots = state.processingSnapshots; + const currentActiveKey = state.activeConversationKey; + // Graduate processing keys that are now pending; drop ones the // assistant has finished responding to without raising anything. - for (const key of processingKeys) { - if (key === activeConversationKey) continue; - if (attentionKeys.has(key)) continue; + for (const key of currentProcessingKeys) { + if (key === currentActiveKey) continue; + if (currentAttentionKeys.has(key)) continue; if (pendingKeys.has(key)) { useConversationListStore.getState().addAttentionKey(key); useConversationListStore.getState().removeProcessingKey(key); continue; } - const conv = conversationsRef.current.find((c) => c.conversationKey === key); - const snapshot = processingSnapshotsRef.current.get(key); + const conv = currentConversations.find((c) => c.conversationKey === key); + const snapshot = currentSnapshots.get(key); if (conv?.latestAssistantMessageAt && conv.latestAssistantMessageAt !== snapshot) { useConversationListStore.getState().removeProcessingKey(key); - processingSnapshotsRef.current.delete(key); } } // Clear attention keys whose interaction has been resolved. - for (const key of attentionKeys) { - if (key === activeConversationKey) continue; + for (const key of currentAttentionKeys) { + if (key === currentActiveKey) continue; if (!pendingKeys.has(key)) { useConversationListStore.getState().removeAttentionKey(key); } @@ -229,14 +230,7 @@ export function useAttentionTracking({ cancelled = true; clearInterval(pollInterval); }; - }, [ - assistantId, - processingKeys, - attentionKeys, - activeConversationKey, - conversationsRef, - processingSnapshotsRef, - ]); + }, [assistantId, processingKeys, attentionKeys]); // ------------------------------------------------------------------------- // One-time sweep on mount: seed attention keys for every non-active diff --git a/apps/web/src/domains/chat/hooks/use-conversation-loader.ts b/apps/web/src/domains/chat/hooks/use-conversation-loader.ts index 613eda2b22a..2831c4fc7ee 100644 --- a/apps/web/src/domains/chat/hooks/use-conversation-loader.ts +++ b/apps/web/src/domains/chat/hooks/use-conversation-loader.ts @@ -38,7 +38,6 @@ import { useConversationHistory, type HistoryPaginationSnapshot, } from "@/domains/chat/hooks/use-conversation-history.js"; -import { useAttentionTracking } from "@/domains/chat/hooks/use-attention-tracking.js"; import { useQueryClient } from "@tanstack/react-query"; import { getChatContext } from "@/domains/chat/api/assistant.js"; @@ -79,9 +78,6 @@ interface UseConversationLoaderParams { // Collections conversations: Conversation[]; - activeConversation: Conversation | undefined; - processingKeys: Set; - attentionKeys: Set; transcriptPagination: Omit; // Feature flags / epochs @@ -101,7 +97,6 @@ interface UseConversationLoaderParams { inputRef: MutableRefObject; draftsRef: MutableRefObject>; messagesRef: MutableRefObject; - conversationsRef: MutableRefObject; contextWindowUsageByConversationRef: MutableRefObject>; dismissedSurfaceIdsRef: MutableRefObject>; needsNewBubbleRef: MutableRefObject; @@ -110,7 +105,6 @@ interface UseConversationLoaderParams { requestIdToStableIdRef: MutableRefObject>; pendingLocalDeletionsRef: MutableRefObject>; confirmationToolCallMapRef: MutableRefObject>; - processingSnapshotsRef: MutableRefObject>; refreshSettleRef: MutableRefObject; lastSuggestionMsgIdRef: MutableRefObject; autoGreetRef: MutableRefObject; @@ -167,7 +161,10 @@ interface UseConversationLoaderParams { * * Delegates to: * - `useConversationHistory` -- conversation switch, cache, and history loading - * - `useAttentionTracking` -- processing/attention key lifecycle and polling + * + * Attention/processing-key tracking is now owned by `useAttentionTracking`, + * mounted in `ChatLayout` so its 10s polling loop covers every chat-layout + * route (home/library/contacts/identity), not only `/assistant`. */ export function useConversationLoader({ assistantId, @@ -177,9 +174,6 @@ export function useConversationLoader({ searchParams, navigate, conversations, - activeConversation, - processingKeys, - attentionKeys, transcriptPagination, conversationGroupsUI, refreshEpoch, @@ -193,7 +187,6 @@ export function useConversationLoader({ inputRef, draftsRef, messagesRef, - conversationsRef, contextWindowUsageByConversationRef, dismissedSurfaceIdsRef, needsNewBubbleRef, @@ -202,7 +195,6 @@ export function useConversationLoader({ requestIdToStableIdRef, pendingLocalDeletionsRef, confirmationToolCallMapRef, - processingSnapshotsRef, refreshSettleRef, lastSuggestionMsgIdRef, autoGreetRef, @@ -484,21 +476,6 @@ export function useConversationLoader({ shouldSuppressGenericChatErrorNotice, }); - // ------------------------------------------------------------------------- - // Delegate: attention tracking and processing key lifecycle - // ------------------------------------------------------------------------- - useAttentionTracking({ - assistantId, - assistantStateKind, - activeConversationKey, - conversations, - activeConversation, - processingKeys, - attentionKeys, - conversationsRef, - processingSnapshotsRef, - }); - // ------------------------------------------------------------------------- // switchConversation // ------------------------------------------------------------------------- diff --git a/apps/web/src/domains/chat/hooks/use-event-stream.ts b/apps/web/src/domains/chat/hooks/use-event-stream.ts index b93fd81d980..a6f39c4eecd 100644 --- a/apps/web/src/domains/chat/hooks/use-event-stream.ts +++ b/apps/web/src/domains/chat/hooks/use-event-stream.ts @@ -88,9 +88,6 @@ export interface UseEventStreamParams { reachabilityPhase: string; reachabilityReset: () => void; - // Conversation list - processingSnapshotsRef: MutableRefObject>; - // Messages setMessages: Dispatch>; @@ -140,7 +137,6 @@ export function useEventStream({ reachabilityProbe, reachabilityPhase, reachabilityReset, - processingSnapshotsRef, setMessages, setError, streamRetryNonce, @@ -249,8 +245,8 @@ export function useEventStream({ { const convKey = streamContextRef.current?.conversationKey; if (convKey) { + // `removeProcessingKey` clears the matching snapshot atomically. useConversationListStore.getState().removeProcessingKey(convKey); - processingSnapshotsRef.current.delete(convKey); } } reachabilityProbeRef.current(); diff --git a/apps/web/src/domains/chat/hooks/use-send-message.ts b/apps/web/src/domains/chat/hooks/use-send-message.ts index f450ec08d30..abf14ef83f6 100644 --- a/apps/web/src/domains/chat/hooks/use-send-message.ts +++ b/apps/web/src/domains/chat/hooks/use-send-message.ts @@ -82,7 +82,6 @@ interface UseSendMessageParams { activeConversationKeyRef: MutableRefObject; messagesRef: MutableRefObject; - conversationsRef: MutableRefObject; streamRef: MutableRefObject; streamContextRef: MutableRefObject<{ assistantId: string; @@ -90,7 +89,6 @@ interface UseSendMessageParams { } | null>; streamEpochRef: MutableRefObject; needsNewBubbleRef: MutableRefObject; - processingSnapshotsRef: MutableRefObject>; dismissedSurfaceIdsRef: MutableRefObject>; pendingOnboardingContextRef: MutableRefObject; onboardingDraftConversationKeyRef: MutableRefObject; @@ -128,12 +126,10 @@ export function useSendMessage({ assistantIdRef, activeConversationKeyRef, messagesRef, - conversationsRef, streamRef, streamContextRef, streamEpochRef, needsNewBubbleRef, - processingSnapshotsRef, dismissedSurfaceIdsRef, pendingOnboardingContextRef, onboardingDraftConversationKeyRef, @@ -483,14 +479,19 @@ export function useSendMessage({ const fallbackTurnId = newTurnId(); useTurnStore.getState().requestSend(fallbackTurnId); useTurnStore.getState().acceptSend(fallbackTurnId); - useConversationListStore.getState().addProcessingKey(activeConversationKey); - const currentConv = conversationsRef.current.find( - (c) => c.conversationKey === activeConversationKey, - ); - processingSnapshotsRef.current.set( - activeConversationKey, - currentConv?.latestAssistantMessageAt as string | undefined, - ); + { + const currentConv = useConversationListStore + .getState() + .conversations.find( + (c) => c.conversationKey === activeConversationKey, + ); + useConversationListStore + .getState() + .addProcessingKey( + activeConversationKey, + currentConv?.latestAssistantMessageAt as string | undefined, + ); + } return; } } catch { @@ -503,12 +504,15 @@ export function useSendMessage({ const turnId = newTurnId(); useTurnStore.getState().requestSend(turnId); - useConversationListStore.getState().addProcessingKey(activeConversationKey); - const currentConv = conversationsRef.current.find(c => c.conversationKey === activeConversationKey); - processingSnapshotsRef.current.set( - activeConversationKey, - currentConv?.latestAssistantMessageAt as string | undefined, - ); + const currentConv = useConversationListStore + .getState() + .conversations.find((c) => c.conversationKey === activeConversationKey); + useConversationListStore + .getState() + .addProcessingKey( + activeConversationKey, + currentConv?.latestAssistantMessageAt as string | undefined, + ); // Optimistically add a stub conversation to the sidebar for draft // conversations that don't exist on the server yet. @@ -533,12 +537,9 @@ export function useSendMessage({ // Resolve draft key -> server-assigned conversation ID. if (resolvedId && resolvedId !== activeConversationKey) { const newKey = resolvedId; - useConversationListStore.getState().transferProcessingKey(activeConversationKey, newKey); - const snapshot = processingSnapshotsRef.current.get(activeConversationKey); - processingSnapshotsRef.current.delete(activeConversationKey); - if (snapshot !== undefined) { - processingSnapshotsRef.current.set(newKey, snapshot); - } + useConversationListStore + .getState() + .transferProcessingKey(activeConversationKey, newKey); useConversationListStore.getState().resolveDraftKey(activeConversationKey, newKey); resolveEditChatDraftKey(activeConversationKey, newKey); @@ -562,7 +563,6 @@ export function useSendMessage({ setError({ message: "Something went wrong. Please try again." }); useTurnStore.getState().onStreamError(); const keysToClean = [activeConversationKey, resolvedId].filter(Boolean) as string[]; - for (const k of keysToClean) processingSnapshotsRef.current.delete(k); if (keysToClean.length > 0) { useConversationListStore.getState().removeMultipleProcessingKeys(keysToClean); } @@ -595,7 +595,6 @@ export function useSendMessage({ useSubagentStore.getState().reset(); confirmationToolCallMapRef.current.clear(); useConversationListStore.getState().removeProcessingKey(activeConversationKey); - processingSnapshotsRef.current.delete(activeConversationKey); try { await cancelGeneration(assistantId, activeConversationKey); } catch { diff --git a/apps/web/src/domains/chat/hooks/use-stream-event-handler.ts b/apps/web/src/domains/chat/hooks/use-stream-event-handler.ts index 9affc49693f..9467ff883b4 100644 --- a/apps/web/src/domains/chat/hooks/use-stream-event-handler.ts +++ b/apps/web/src/domains/chat/hooks/use-stream-event-handler.ts @@ -94,11 +94,6 @@ export interface UseStreamEventHandlerParams { messagesRef: MutableRefObject; needsNewBubbleRef: MutableRefObject; - // --- Processing --- - processingSnapshotsRef: MutableRefObject< - Map - >; - // --- Error & stream lifecycle --- setError: Dispatch>; streamRef: MutableRefObject; @@ -174,7 +169,6 @@ export function useStreamEventHandler( setMessages, messagesRef, needsNewBubbleRef, - processingSnapshotsRef, setError, streamRef, cancelReconciliation, @@ -212,13 +206,10 @@ export function useStreamEventHandler( invalidateAvatarRef.current = invalidateAvatar; /** Remove a conversation key from the processing set and snapshots map. */ - const clearProcessingKey = useCallback( - (convKey: string) => { - useConversationListStore.getState().removeProcessingKey(convKey); - processingSnapshotsRef.current.delete(convKey); - }, - [processingSnapshotsRef], - ); + const clearProcessingKey = useCallback((convKey: string) => { + // `removeProcessingKey` clears the matching snapshot in the same set call. + useConversationListStore.getState().removeProcessingKey(convKey); + }, []); // --- Main event handler --- diff --git a/apps/web/src/domains/conversations/conversation-list-store.ts b/apps/web/src/domains/conversations/conversation-list-store.ts index d28e50e507a..657c4b398a3 100644 --- a/apps/web/src/domains/conversations/conversation-list-store.ts +++ b/apps/web/src/domains/conversations/conversation-list-store.ts @@ -96,6 +96,15 @@ export interface ConversationListState { activeConversationKey: string | null; editingConversationKey: string | null; processingKeys: Set; + /** + * Per-conversation snapshot of `latestAssistantMessageAt` at the moment the + * key was added to `processingKeys`. The attention-tracking graduation logic + * compares the current `latestAssistantMessageAt` against this snapshot to + * detect when the assistant has finished responding. Entries are added by + * `addProcessingKey` and cleared by every action that removes from + * `processingKeys`, so the two collections stay in sync. + */ + processingSnapshots: Map; attentionKeys: Set; } @@ -121,7 +130,7 @@ export interface ConversationListActions { setEditingKey: (key: string | null) => void; // --- Processing keys --- - addProcessingKey: (key: string) => void; + addProcessingKey: (key: string, snapshot?: string) => void; removeProcessingKey: (key: string) => void; removeMultipleProcessingKeys: (keys: string[]) => void; transferProcessingKey: (oldKey: string, newKey: string) => void; @@ -145,9 +154,22 @@ const INITIAL_STATE: ConversationListState = { activeConversationKey: null, editingConversationKey: null, processingKeys: new Set(), + processingSnapshots: new Map(), attentionKeys: new Set(), }; +/** + * Return a new Map with the given key removed, or the same reference if the + * key wasn't present — lets Zustand's shallow equality bail out of + * unnecessary re-renders. + */ +function deleteFromMap(prev: Map, key: K): Map { + if (!prev.has(key)) return prev; + const next = new Map(prev); + next.delete(key); + return next; +} + // --------------------------------------------------------------------------- // Store // --------------------------------------------------------------------------- @@ -272,27 +294,46 @@ export const useConversationListStore = createSelectors( // --- Processing keys --- - addProcessingKey: (key) => { - set({ processingKeys: addToSet(get().processingKeys, key) }); + addProcessingKey: (key, snapshot) => { + const { processingKeys, processingSnapshots } = get(); + const nextSnapshots = new Map(processingSnapshots); + nextSnapshots.set(key, snapshot); + set({ + processingKeys: addToSet(processingKeys, key), + processingSnapshots: nextSnapshots, + }); }, removeProcessingKey: (key) => { - set({ processingKeys: removeFromSet(get().processingKeys, key) }); + set({ + processingKeys: removeFromSet(get().processingKeys, key), + processingSnapshots: deleteFromMap(get().processingSnapshots, key), + }); }, removeMultipleProcessingKeys: (keys) => { + const { processingKeys, processingSnapshots } = get(); + let nextSnapshots = processingSnapshots; + for (const key of keys) { + nextSnapshots = deleteFromMap(nextSnapshots, key); + } set({ - processingKeys: removeMultipleFromSet(get().processingKeys, keys), + processingKeys: removeMultipleFromSet(processingKeys, keys), + processingSnapshots: nextSnapshots, }); }, transferProcessingKey: (oldKey, newKey) => { - const { processingKeys } = get(); + const { processingKeys, processingSnapshots } = get(); if (!processingKeys.has(oldKey)) return; - const next = new Set(processingKeys); - next.delete(oldKey); - next.add(newKey); - set({ processingKeys: next }); + const nextKeys = new Set(processingKeys); + nextKeys.delete(oldKey); + nextKeys.add(newKey); + const nextSnapshots = new Map(processingSnapshots); + const snapshot = nextSnapshots.get(oldKey); + nextSnapshots.delete(oldKey); + nextSnapshots.set(newKey, snapshot); + set({ processingKeys: nextKeys, processingSnapshots: nextSnapshots }); }, // --- Attention keys --- @@ -310,6 +351,7 @@ export const useConversationListStore = createSelectors( graduateProcessingKey: (key, hasPendingInteraction) => { set((state) => ({ processingKeys: removeFromSet(state.processingKeys, key), + processingSnapshots: deleteFromMap(state.processingSnapshots, key), attentionKeys: hasPendingInteraction ? addToSet(state.attentionKeys, key) : state.attentionKeys, @@ -319,7 +361,12 @@ export const useConversationListStore = createSelectors( // --- Reset --- reset: () => { - set({ ...INITIAL_STATE, processingKeys: new Set(), attentionKeys: new Set() }); + set({ + ...INITIAL_STATE, + processingKeys: new Set(), + processingSnapshots: new Map(), + attentionKeys: new Set(), + }); }, })), );