diff --git a/apps/web/src/components/layout/middle-content/page-views/dashboard/GlobalAssistantView.tsx b/apps/web/src/components/layout/middle-content/page-views/dashboard/GlobalAssistantView.tsx index e485f86775..6767c22f4e 100644 --- a/apps/web/src/components/layout/middle-content/page-views/dashboard/GlobalAssistantView.tsx +++ b/apps/web/src/components/layout/middle-content/page-views/dashboard/GlobalAssistantView.tsx @@ -50,7 +50,9 @@ import { useDriveStore } from '@/hooks/useDrive'; import { fetchWithAuth } from '@/lib/auth/auth-fetch'; import { useAssistantSettingsStore } from '@/stores/useAssistantSettingsStore'; import { useGlobalChatConfig, useGlobalChatStream, useGlobalChatConversation } from '@/contexts/GlobalChatContext'; -import { usePageAgentDashboardStore } from '@/stores/page-agents'; +import { usePageAgentDashboardStore, agentStreamKey, selectIsAgentStreaming, selectAgentStop, type AgentStreamKey } from '@/stores/page-agents'; +import { holdForStream } from '@/lib/ai/streams/holdForStream'; +import { selectLiveAssistantIds } from '@/lib/ai/streams/selectLiveAssistantIds'; import { useVoiceModeStore, type VoiceModeOwner } from '@/stores/useVoiceModeStore'; import { VoiceCallPanel } from '@/components/ai/voice/VoiceCallPanel'; import { useDisplayPreferences } from '@/hooks/useDisplayPreferences'; @@ -72,7 +74,7 @@ import { buildGlobalChatRequestBody, } from '@/lib/ai/shared'; import { AskUserAnswerProvider } from '@/components/ai/shared/chat/ask-user/AskUserAnswerContext'; -import { abortActiveStream, clearActiveStreamId } from '@/lib/ai/core/client'; +import { abortActiveStream, abortActiveStreamByMessageId, clearActiveStreamId } from '@/lib/ai/core/client'; import { useAppStateRecovery } from '@/hooks/useAppStateRecovery'; import { useEditingStore } from '@/stores/useEditingStore'; import { useAgentChannelMultiplayer } from '@/hooks/useAgentChannelMultiplayer'; @@ -125,7 +127,7 @@ const GlobalAssistantView: React.FC = () => { const createAgentConversation = usePageAgentDashboardStore((state) => state.createNewConversation); const loadMostRecentConversation = usePageAgentDashboardStore((state) => state.loadMostRecentConversation); const setAgentStreaming = usePageAgentDashboardStore((state) => state.setAgentStreaming); - const setAgentStopStreaming = usePageAgentDashboardStore((state) => state.setAgentStopStreaming); + const setAgentStop = usePageAgentDashboardStore((state) => state.setAgentStop); const setActiveTab = usePageAgentDashboardStore((state) => state.setActiveTab); const loadAgentConversation = usePageAgentDashboardStore((state) => state.loadConversation); @@ -365,30 +367,146 @@ const GlobalAssistantView: React.FC = () => { useEffect(() => { latestGlobalMessagesRef.current = globalLocalMessages; }, [globalLocalMessages]); - const stop = useChatStop(currentConversationId, rawStop); + // The STREAM's conversation, not the surface's. `heldStopConversationId` is computed below from + // the hold-refs; before a stream exists it falls back to the live id, which is correct because + // there is nothing else to name. See useChatStop for why the chatId map alone is not enough. + // PER-MODE, and that is the whole point. + // + // This surface hosts TWO independent chats (agent and global), and both can be in flight at + // once — switching mode does not abort the running POST, because useChat's id is constant. + // Deriving one id from the MODE-SELECTED `status`/`messages` and feeding it to BOTH hold-refs + // let the IDLE mode's ref latch the ACTIVE mode's messageId — and holdForStream then pinned it + // for the rest of that stream. Stop, back in the other mode, aborted the WRONG stream: the + // agent's answer died mid-sentence while the global generation kept running its write tools and + // kept billing, its own Stop permanently wired to an id that was never its. + // + // A stream's identity comes from ITS OWN chat, never from whichever one the surface happens to + // be rendering. (Both gate on 'streaming', never 'submitted' — see holdForStream's contract: + // during submitted the array's last assistant message is the PREVIOUS turn's.) + const { agentLiveId, globalLiveId } = useMemo( + () => selectLiveAssistantIds({ + agent: { status: agentStatus, messages: agentMessages }, + global: { status: globalStatus, messages: globalLocalMessages }, + }), + [agentStatus, agentMessages, globalStatus, globalLocalMessages], + ); + + // The conversation the CURRENT local stream belongs to, captured when it starts and held + // until it ends — the stream's identity, not the surface's. See the flag effect below. + // Assigned during render, and it only ever changes when `agentStatus` does (a stream + // starting or ending) — which every effect below already depends on. + const streamConvIdRef = useRef(null); + // Whether OUR local stream set the flag — the multiplayer hook may hold the same key for a + // different, still-live stream, and we must not clear its flag. + const ownsFlagRef = useRef(false); + const isAgentStreamingNow = agentStatus === 'submitted' || agentStatus === 'streaming'; + streamConvIdRef.current = holdForStream({ + current: streamConvIdRef.current, + isStreaming: isAgentStreamingNow, + liveValue: agentConversationId, + }); + // The stream's assistant messageId, captured when the first chunk arrives and held for the + // rest of the stream. THIS is what we abort by. + // + // Aborting by chatId cannot work here: `abortActiveStream` is a lookup in the client-side + // activeStreams map, and the conversation-change cleanup below DELETES this stream's entry + // the instant the surface switches conversation. So the chatId abort became a map miss — + // the local fetch stopped and the SERVER KEPT GENERATING AND KEPT BILLING. The multiplayer + // hook was always immune because it aborts by messageId, which needs no map. + const streamMsgIdRef = useRef(null); + streamMsgIdRef.current = holdForStream({ + current: streamMsgIdRef.current, + isStreaming: isAgentStreamingNow, + // The AGENT chat's own id — never the mode-selected one. See above. + liveValue: agentLiveId, + }); + const globalStreamingNow = globalStatus === 'submitted' || globalStatus === 'streaming'; + const globalStreamConvIdRef = useRef(null); + globalStreamConvIdRef.current = holdForStream({ + current: globalStreamConvIdRef.current, + isStreaming: globalStreamingNow, + liveValue: globalConversationId, + }); + const globalStreamMsgIdRef = useRef(null); + globalStreamMsgIdRef.current = holdForStream({ + current: globalStreamMsgIdRef.current, + isStreaming: globalStreamingNow, + // The GLOBAL chat's own id — never the mode-selected one. See above. + liveValue: globalLiveId, + }); + + // THE STOP BUTTON THE USER ACTUALLY CLICKS aborts by the HELD id, not the live one. + // + // `liveAssistantMessageId` is derived from the live `messages` array, and "New Chat" (and + // history-select) empties that array mid-stream with no streaming guard. The id therefore + // vanished at exactly the moment the user most needed it: Stop fell through to the chatId + // fallback, whose map entry the conversation-change cleanup had already deleted. The local + // fetch stopped, the button looked like it worked, and the SERVER KEPT GENERATING — running + // write tools and billing — against a conversation the user had already navigated away from. + // + // The held ref survives the array being cleared, because it names the STREAM and not the + // surface. The stop functions published to the OTHER surfaces already used it; this one, + // purely by declaration order, did not. + const heldStreamMsgId = (selectedAgent ? streamMsgIdRef.current : globalStreamMsgIdRef.current) ?? undefined; + + // The conversation the Stop button must NAME. The held one while a stream is running — never + // the live one, or a mid-stream conversation switch would abort the wrong generation (or + // none). Falls back to the live id before any stream exists, which is correct: there is + // nothing else to name, and the server simply reports that nothing was in flight. + const heldStopConversationId = + (selectedAgent ? streamConvIdRef.current : globalStreamConvIdRef.current) ?? currentConversationId; + + const stop = useChatStop(currentConversationId, rawStop, heldStopConversationId); // The stable assistant messageId of the live stream (the rendered streaming // bubble === serverAssistantMessageId). Used to abort authoritatively by // messageId rather than the fragile chatId→streamId map. - const liveAssistantMessageId = useMemo(() => { - if (!isStreaming) return undefined; - for (let i = messages.length - 1; i >= 0; i--) { - if (messages[i].role === 'assistant') return messages[i].id; - } - return undefined; - }, [isStreaming, messages]); + // + // Resolved ONLY during 'streaming' — never 'submitted'. This is the difference between + // aborting THIS stream and aborting the previous turn's finished reply. + // + // useChat sets `status: 'submitted'` BEFORE it issues the request, and only pushes the new + // assistant message inside `write()`, which flips the status to 'streaming' in the same + // synchronous job (see ai/dist/index.mjs: setStatus('submitted') at the top of sendMessage; + // the pushMessage + setStatus('streaming') together in write()). So for the whole submitted + // window the array's last assistant message is THE PREVIOUS TURN'S. + // + // That matters because `holdForStream` below latches this value on the first render where the + // stream is live — which is a 'submitted' render. Gated on the looser `isStreaming` (which + // includes 'submitted'), it therefore captured and held the id of a reply that finished + // minutes ago, on every turn after the first: Stop aborted a messageId the server registry no + // longer knew, the local fetch stopped, the button looked like it worked, and the real + // generation kept running its write tools and kept billing. + // + // At the first 'streaming' render the push has already happened, so the last assistant IS the + // stream's. Before that we return undefined, and callers correctly fall back to the chatId map. // After a refresh mid-stream, useChat starts at idle — but the // GlobalChatContext bootstrap may have detected an own in-flight stream // and registered a stop function. Surface either source so the UI shows // a stop button + streaming indicator from both bootstrap and live paths. + // A bootstrap-restored AGENT stream (after a refresh mid-stream). useAgentChannelMultiplayer + // claims this slot; the sidebar has always read it, and the DASHBOARD never did — so the surface + // that started the stream rendered Send while the sidebar showed a working Stop. Keyed by the + // STREAM's conversation, not the surface's. + const agentBootstrapKey: AgentStreamKey = { + agentId: selectedAgent?.id ?? null, + conversationId: streamConvIdRef.current ?? agentConversationId, + }; + const agentBootstrapIsStreaming = usePageAgentDashboardStore( + selectIsAgentStreaming(agentBootstrapKey), + ); + const agentBootstrapStop = usePageAgentDashboardStore(selectAgentStop(agentBootstrapKey)); + const { effectiveIsStreaming, effectiveStop } = useGlobalEffectiveStream({ localIsStreaming: isStreaming, rawStop: stop, selectedAgent, contextIsStreaming, contextStopStreaming, - activeMessageId: liveAssistantMessageId, + activeMessageId: heldStreamMsgId, + agentBootstrapIsStreaming, + agentBootstrapStop, }); const remoteStreamingUser = !effectiveIsStreaming @@ -583,14 +701,24 @@ const GlobalAssistantView: React.FC = () => { enabled: !isStreaming && currentConversationId !== null && !useEditingStore.getState().isAnyEditing(), }); - // Clean up stream tracking on unmount + // Clean up stream tracking on unmount / conversation change. + // + // Keyed by `agentConversationId` — THE ONLY KEY THIS SURFACE REGISTERS (its agent transport, + // `useChatTransport(agentConversationId, …)`). It used to clear `currentConversationId`, which + // in GLOBAL mode is `globalConversationId` — and that is GlobalChatContext's transport key + // (`useChatTransport(currentConversationId, …)`), not ours. So navigating away from the + // dashboard mid-global-stream deleted the CONTEXT's activeStreams entry, and the context (which + // outlives this component) was left unable to abort by chatId in the pre-first-chunk window: + // the server kept generating and kept billing. + // + // Same rule as the sidebar's cleanup: a surface may only free what it allocated. useEffect(() => { return () => { - if (currentConversationId) { - clearActiveStreamId({ chatId: currentConversationId }); + if (agentConversationId) { + clearActiveStreamId({ chatId: agentConversationId }); } }; - }, [currentConversationId]); + }, [agentConversationId]); // ============================================ // GLOBAL MODE SYNC EFFECTS @@ -629,9 +757,13 @@ const GlobalAssistantView: React.FC = () => { const isCurrentlyStreaming = globalStatus === 'submitted' || globalStatus === 'streaming'; const wasStreaming = prevStatusRef.current === 'submitted' || prevStatusRef.current === 'streaming'; - if (isCurrentlyStreaming && !wasStreaming) { + // Level-triggered set, edge-triggered clear — see the agent-mode twin below. An + // edge-guarded set left the flag FALSE for the whole streaming phase, because the + // level-triggered cleanup clears it on the submitted -> streaming transition and the + // body then declined to re-assert it. + if (isCurrentlyStreaming) { setGlobalIsStreaming(true); - } else if (!isCurrentlyStreaming && wasStreaming) { + } else if (wasStreaming) { setGlobalIsStreaming(false); } prevStatusRef.current = globalStatus; @@ -646,28 +778,67 @@ const GlobalAssistantView: React.FC = () => { // Register stop function to global context (global mode only) // Combined function calls both abort endpoint (server-side) and useChat stop (client-side) // Use try/finally to guarantee client-side stop runs even if server abort fails + // The stop slot is SHARED, and this component is not its only writer: the stream socket + // claims it on bootstrap for a stream restored after a refresh (GlobalChatContext's + // onOwnStreamBootstrap). This effect used to null it UNCONDITIONALLY — on its + // else-branch and on its cleanup, both of which fire whenever globalStatus is 'ready', + // which it is for the ENTIRE life of a bootstrapped stream, and whose deps + // (globalConversationId) resolve asynchronously right after the claim by design. + // + // So it destroyed a live Stop button belonging to someone else, leaving + // `isStreaming: true` with `stopStreaming: null` — the Stop renders and does nothing + // while the stream keeps generating and keeps billing. Only clear what we installed. + const ownedGlobalStopFnRef = useRef<(() => void | Promise) | null>(null); + const contextStopStreamingRef = useRef(contextStopStreaming); + contextStopStreamingRef.current = contextStopStreaming; + + const clearGlobalStopIfOurs = useCallback(() => { + if (ownedGlobalStopFnRef.current === null) return; + const stillOurs = contextStopStreamingRef.current === ownedGlobalStopFnRef.current; + ownedGlobalStopFnRef.current = null; + if (!stillOurs) return; + setGlobalStopStreaming(null); + }, [setGlobalStopStreaming]); + useEffect(() => { if (selectedAgent) return; if (globalStatus === 'submitted' || globalStatus === 'streaming') { - setGlobalStopStreaming(() => async () => { + // Same as agent mode: GLOBAL_CHAT_ID is a constant too, so useChat never recreates the + // global Chat and a mid-stream conversation switch does NOT abort the POST. Name the + // STREAM, not the surface — and abort by messageId, because the conversation-change + // cleanup deletes this stream's entry from the client-side chatId map. + const stopFn = async () => { try { - // Call abort endpoint to stop server-side processing - if (globalConversationId) { - await abortActiveStream({ chatId: globalConversationId }); + const messageId = globalStreamMsgIdRef.current; + const convId = globalStreamConvIdRef.current; + if (messageId) { + await abortActiveStreamByMessageId({ messageId }); + } else if (convId) { + // Pre-first-chunk: no assistant id yet — and the chatId map is EMPTY here, not stale + // (setActiveStreamId only runs once the response headers land, 0.5-3s into a real + // send). Name the conversation too, or this abort is a guaranteed no-op while the + // server keeps generating and keeps billing. The agent-mode stop above already did + // this; this one was missed. + await abortActiveStream({ chatId: convId, conversationId: convId }); } } finally { // Call useChat's stop to abort client-side fetch globalStop(); } - }); + }; + ownedGlobalStopFnRef.current = stopFn; + // setGlobalStopStreaming IS a useState dispatch, so a function argument is an + // UPDATER — the wrapper is required here. (Contrast setAgentStopStreaming below, + // a plain zustand value setter, where the wrapper would store the wrapper itself.) + setGlobalStopStreaming(() => stopFn); } else { - setGlobalStopStreaming(null); + clearGlobalStopIfOurs(); } return () => { - setGlobalStopStreaming(null); + clearGlobalStopIfOurs(); }; - }, [selectedAgent, globalStatus, globalStop, globalConversationId, setGlobalStopStreaming]); + }, [selectedAgent, globalStatus, globalStop, setGlobalStopStreaming, clearGlobalStopIfOurs]); // ============================================ // AGENT MODE SYNC EFFECTS @@ -678,45 +849,110 @@ const GlobalAssistantView: React.FC = () => { if (!selectedAgent) return; const isCurrentlyStreaming = agentStatus === 'submitted' || agentStatus === 'streaming'; const wasStreaming = prevAgentStatusRef.current === 'submitted' || prevAgentStatusRef.current === 'streaming'; - if (isCurrentlyStreaming && !wasStreaming) { - setAgentStreaming(true); - } else if (!isCurrentlyStreaming && wasStreaming) { - setAgentStreaming(false); + // Set is LEVEL-triggered, clear is edge-triggered. The cleanup below is level-triggered + // and `agentStatus` is a dep, so on the submitted -> streaming transition React runs the + // previous cleanup (which sets false) and then this body. With an edge-guarded set + // (`&& !wasStreaming`) the body then refused to re-assert it — and the flag stayed FALSE + // for the entire streaming phase, killing the cross-surface sync this state exists for + // and dropping SWR protection mid-stream. + // + // Keyed by the conversation the stream STARTED in — captured at the transition and held + // in a ref — NOT by the surface's live `agentConversationId`. + // + // `useChat` only recreates its Chat when its `id` changes, and ours is a constant + // (AGENT_CHAT_ID). So switching conversation mid-stream does NOT abort the POST: the + // stream keeps running while `agentConversationId` moves. Keying off the live value + // MIGRATED ownership — the cleanup cleared the running stream's key and the body + // installed a fresh claim under a conversation with NO stream. The abandoned stream lost + // its Stop and its SWR protection while still generating; the new key showed a spinner + // and a Stop that aborted nothing. (History-select and New Chat both do this with no + // streaming guard at all.) + const streamConvId = streamConvIdRef.current; + const flagKey = { agentId: selectedAgent.id, conversationId: streamConvId }; + // Ownership-guarded, like the stop below. The multiplayer hook can hold this same key for + // a DIFFERENT, still-live stream (a bootstrap-restored one, or a cross-instance stream + // takeover could not abort). Clearing the flag unconditionally would strip that stream's + // Stop affordance and its SWR protection while it is still generating. + if (isCurrentlyStreaming) { + ownsFlagRef.current = true; + setAgentStreaming(flagKey, true); + } else if (wasStreaming && ownsFlagRef.current) { + ownsFlagRef.current = false; + setAgentStreaming(flagKey, false); } prevAgentStatusRef.current = agentStatus; return () => { - if (isCurrentlyStreaming) { - setAgentStreaming(false); + if (isCurrentlyStreaming && ownsFlagRef.current) { + ownsFlagRef.current = false; + setAgentStreaming(flagKey, false); } }; }, [selectedAgent, agentStatus, setAgentStreaming]); + // Same shared-slot discipline as the global stop above: useAgentChannelMultiplayer claims + // this slot on bootstrap for a stream restored after a refresh, and nulling it + // unconditionally from here destroyed that live Stop button while isAgentStreaming stayed + // true. Only clear what we installed. + const ownedAgentStopFnRef = useRef<(() => void | Promise) | null>(null); + // Still identity-guarded WITHIN the agent: useAgentChannelMultiplayer claims this same + // agent's stop on bootstrap (a stream restored after a refresh), and clearing that would + // destroy a live Stop button. Cross-AGENT collisions are now impossible by construction. + const clearAgentStopIfOurs = useCallback((key: AgentStreamKey) => { + if (ownedAgentStopFnRef.current === null) return; + const k = agentStreamKey(key); + const current = k === null ? undefined : usePageAgentDashboardStore.getState().agentStops[k]; + const stillOurs = current === ownedAgentStopFnRef.current; + ownedAgentStopFnRef.current = null; + if (!stillOurs) return; + setAgentStop(key, null); + }, [setAgentStop]); + // Register stop function to dashboard store (agent mode only) // Combined function calls both abort endpoint (server-side) and useChat stop (client-side) // Use try/finally to guarantee client-side stop runs even if server abort fails useEffect(() => { if (!selectedAgent) return; + // Named by (agent, conversation) — an agent id alone cannot say WHICH conversation, and + // the sidebar keeps its own conversation for the same agent. + // Keyed by the conversation the stream STARTED in — see the flag effect. Keying off the + // live `agentConversationId` migrated ownership on a mid-stream conversation switch, and + // aborted the WRONG conversation (a server no-op) while the real stream kept billing. + const streamConvId = streamConvIdRef.current; + const stopKey: AgentStreamKey = { agentId: selectedAgent.id, conversationId: streamConvId }; if (agentStatus === 'submitted' || agentStatus === 'streaming') { - setAgentStopStreaming(() => async () => { + // setAgentStop is a plain zustand VALUE setter, NOT a useState dispatch — so pass the + // fn itself, never the `() => fn` updater form (which would be stored verbatim, and + // calling it would merely return the inner fn: a Stop button that does nothing). + const stopFn = async () => { try { - // Call abort endpoint to stop server-side processing - if (agentConversationId) { - await abortActiveStream({ chatId: agentConversationId }); + // Read at CALL time: the messageId only exists once the first chunk lands. + const messageId = streamMsgIdRef.current; + if (messageId) { + await abortActiveStreamByMessageId({ messageId }); + } else if (streamConvId) { + // Pre-first-chunk: no assistant id yet. The chatId map is NOT enough — it is empty + // until the response headers land, and the conversation-change cleanup deletes the + // running stream's entry on a mid-stream switch. Naming the conversation is what + // makes this abort actually reach the server instead of silently no-opping while the + // generation keeps running and keeps billing. + await abortActiveStream({ chatId: streamConvId, conversationId: streamConvId }); } } finally { // Call useChat's stop to abort client-side fetch agentStop(); } - }); + }; + ownedAgentStopFnRef.current = stopFn; + setAgentStop(stopKey, stopFn); } else { - setAgentStopStreaming(null); + clearAgentStopIfOurs(stopKey); } return () => { - setAgentStopStreaming(null); + clearAgentStopIfOurs(stopKey); }; - }, [selectedAgent, agentStatus, agentStop, agentConversationId, setAgentStopStreaming]); + }, [selectedAgent, agentStatus, agentStop, setAgentStop, clearAgentStopIfOurs]); // Agent-mode load-on-select guarantee: the store's conversationLoadSignal // fires on explicit load/create (not on streaming updates). We use it rather diff --git a/apps/web/src/components/layout/middle-content/page-views/dashboard/useGlobalEffectiveStream.ts b/apps/web/src/components/layout/middle-content/page-views/dashboard/useGlobalEffectiveStream.ts index ea57da953f..ce28b068c2 100644 --- a/apps/web/src/components/layout/middle-content/page-views/dashboard/useGlobalEffectiveStream.ts +++ b/apps/web/src/components/layout/middle-content/page-views/dashboard/useGlobalEffectiveStream.ts @@ -14,6 +14,16 @@ interface UseGlobalEffectiveStreamArgs { * path is still served by `contextStopStreaming`. */ activeMessageId?: string; + /** + * A bootstrap-restored AGENT stream, from the dashboard store (keyed by (agentId, + * conversationId)). The global mode had this all along via the context; agent mode did not — so + * after a refresh mid-agent-stream, `useAgentChannelMultiplayer` claimed the store slot, the + * SIDEBAR read it and showed a working Stop, and the DASHBOARD — the surface that started the + * stream — rendered Send. The user had no way to stop their own generation, and it kept running + * and kept billing. + */ + agentBootstrapIsStreaming?: boolean; + agentBootstrapStop?: (() => void | Promise) | null; } interface GlobalEffectiveStream { @@ -37,11 +47,15 @@ export function useGlobalEffectiveStream({ contextIsStreaming, contextStopStreaming, activeMessageId, + agentBootstrapIsStreaming = false, + agentBootstrapStop = null, }: UseGlobalEffectiveStreamArgs): GlobalEffectiveStream { const inGlobalMode = !selectedAgent; + // Both modes now surface a bootstrap-restored stream, not just global. See + // agentBootstrapIsStreaming: agent mode used to show Send after a refresh mid-stream. const effectiveIsStreaming = inGlobalMode ? localIsStreaming || contextIsStreaming - : localIsStreaming; + : localIsStreaming || agentBootstrapIsStreaming; const effectiveStop = useCallback(() => { // Authoritative: abort by the stable assistant messageId when the live stream @@ -58,12 +72,25 @@ export function useGlobalEffectiveStream({ rawStop(); return; } - // Idle locally but resumed via bootstrap after refresh: use the context's - // messageId-based stop registered at bootstrap. + // Idle locally but resumed via bootstrap after refresh: use the messageId-based stop the + // bootstrap registered. Global mode reads it from the context; agent mode from the dashboard + // store — and agent mode never used to, which is how the dashboard ended up with no Stop + // button at all for a stream it had started itself. if (inGlobalMode && contextStopStreaming) { contextStopStreaming(); + return; + } + if (!inGlobalMode && agentBootstrapStop) { + void agentBootstrapStop(); } - }, [activeMessageId, localIsStreaming, rawStop, inGlobalMode, contextStopStreaming]); + }, [ + activeMessageId, + localIsStreaming, + rawStop, + inGlobalMode, + contextStopStreaming, + agentBootstrapStop, + ]); return { effectiveIsStreaming, effectiveStop }; } diff --git a/apps/web/src/components/layout/right-sidebar/ai-assistant/SidebarChatTab.tsx b/apps/web/src/components/layout/right-sidebar/ai-assistant/SidebarChatTab.tsx index 4906b7e9d1..1032349fef 100644 --- a/apps/web/src/components/layout/right-sidebar/ai-assistant/SidebarChatTab.tsx +++ b/apps/web/src/components/layout/right-sidebar/ai-assistant/SidebarChatTab.tsx @@ -21,7 +21,7 @@ import { useAssistantSettingsStore } from '@/stores/useAssistantSettingsStore'; import { useVoiceModeStore, type VoiceModeOwner } from '@/stores/useVoiceModeStore'; import { useGlobalChatConversation, useGlobalChatConfig, useGlobalChatStream } from '@/contexts/GlobalChatContext'; import { usePageAgentSidebarState, usePageAgentSidebarChat, type SidebarAgentInfo } from '@/hooks/page-agents'; -import { usePageAgentDashboardStore } from '@/stores/page-agents'; +import { usePageAgentDashboardStore, selectIsAgentStreaming, selectAgentStop } from '@/stores/page-agents'; import { usePendingStreamsStore, type PendingStream } from '@/stores/usePendingStreamsStore'; import { useShallow } from 'zustand/react/shallow'; import { useAuth } from '@/hooks/useAuth'; @@ -35,6 +35,7 @@ import { LocationContext } from '@/lib/ai/shared'; import { parseTabPath, getStaticTabMeta } from '@/lib/tabs/tab-title'; import { abortActiveStream, abortActiveStreamByMessageId, clearActiveStreamId } from '@/lib/ai/core/client'; import { resolveActiveAssistantMessageId } from '@/lib/ai/streams/resolveActiveAssistantMessageId'; +import { holdForStream } from '@/lib/ai/streams/holdForStream'; import { useChatTransport, useStreamingRegistration, useSendHandoff, useMessageActions, useStreamRecovery, useAskUserAnswering, buildChatConfig, SIDEBAR_AGENT_CHAT_ID, buildGlobalChatRequestBody } from '@/lib/ai/shared'; import { AskUserAnswerProvider } from '@/components/ai/shared/chat/ask-user/AskUserAnswerContext'; import { useMobileKeyboard } from '@/hooks/useMobileKeyboard'; @@ -262,8 +263,19 @@ const SidebarChatTab: React.FC = () => { // ============================================ // Dashboard Streaming State (for agent mode sync) // ============================================ - const dashboardIsStreaming = usePageAgentDashboardStore(state => state.isAgentStreaming); - const dashboardStopStreaming = usePageAgentDashboardStore(state => state.agentStopStreaming); + // Scoped to THIS surface's agent. The dashboard holds a different one (its agent comes + // from usePageAgentDashboardStore; ours comes from useSidebarAgentStore), and + // GlobalAssistantView never unmounts — CenterPanel only hides it — so after one dashboard + // visit we are co-mounted with it on every page. Reading the slot unscoped meant a stream + // on the dashboard's agent B lit up OUR Stop button for agent A, and clicking it aborted + // B while A kept generating and kept billing. + // Named by (agent, conversation). The dashboard holds a different agent — and, for the + // SAME agent, a different conversation (each surface keeps its own; "New Chat" in either + // diverges them). With either half missing the store answers a question we did not ask: + // a dashboard stream on conv X2 lighting up OUR Stop while we are showing conv X1. + const dashboardStreamKey = { agentId: selectedAgent?.id, conversationId: agentConversationId }; + const dashboardIsStreaming = usePageAgentDashboardStore(selectIsAgentStreaming(dashboardStreamKey)); + const dashboardStopStreaming = usePageAgentDashboardStore(selectAgentStop(dashboardStreamKey)); // ============================================ // Derived State @@ -681,22 +693,34 @@ const SidebarChatTab: React.FC = () => { enabled: !isStreaming && currentConversationId !== null && !useEditingStore.getState().isAnyEditing(), }); - // Clean up stream tracking on unmount or conversation change - // Use ref to capture current ID so cleanup clears the correct stream - const prevConversationIdRef = useRef(null); + // Clean up stream tracking on unmount or conversation change. + // + // Keyed by `sidebarChatId` — THE KEY THIS SURFACE ACTUALLY REGISTERED. It used to clear the bare + // `currentConversationId`, which this surface never writes: the sidebar's transport registers + // under the namespaced `sidebar:` (see sidebarChatId — the namespace exists precisely so + // the sidebar and the dashboard can view the same conversation without colliding). So the old + // cleanup did both halves of the wrong thing at once: it leaked its own `sidebar:` entry forever, + // and the bare id it *did* delete belongs to ANOTHER surface — in agent mode, the dashboard's + // transport (`useChatTransport(agentConversationId, …)`). + // + // Concretely: dashboard streaming on agent A / conversation C, sidebar open on the same agent and + // conversation. Collapse the sidebar → this cleanup ran → the DASHBOARD's streamId entry vanished + // → its pre-first-chunk Stop became a map miss and the server kept generating. + // + // A surface may only free what it allocated. + const prevSidebarChatIdRef = useRef(null); useEffect(() => { - // Clear previous conversation's stream ID when switching conversations - if (prevConversationIdRef.current && prevConversationIdRef.current !== currentConversationId) { - clearActiveStreamId({ chatId: prevConversationIdRef.current }); + if (prevSidebarChatIdRef.current && prevSidebarChatIdRef.current !== sidebarChatId) { + clearActiveStreamId({ chatId: prevSidebarChatIdRef.current }); } - prevConversationIdRef.current = currentConversationId; + prevSidebarChatIdRef.current = sidebarChatId; return () => { - if (currentConversationId) { - clearActiveStreamId({ chatId: currentConversationId }); + if (sidebarChatId) { + clearActiveStreamId({ chatId: sidebarChatId }); } }; - }, [currentConversationId]); + }, [sidebarChatId]); // ============================================ // Effects: Initialize Settings Store @@ -928,6 +952,41 @@ const SidebarChatTab: React.FC = () => { regenerate, }); + // The live stream's assistant messageId, captured when the first chunk lands and HELD for the + // rest of the stream — the STREAM's identity, not the surface's. + // + // `lastAssistantMessageId` is derived from the live `messages` array, and `handleNewConversation` + // below calls `setMessages([])` outright with no streaming guard. So the id vanished at exactly + // the moment Stop needed it: the abort fell through to the chatId fallback, keyed by the + // conversation the surface had just switched TO — a map miss. The local fetch stopped, the button + // looked like it worked, and the SERVER KEPT GENERATING (write tools, billing) against the + // conversation the user had already left. Streams deliberately survive a client disconnect + // (see the abort registry), so only an explicit, correctly-keyed abort can stop one. + // + // The live value is read ONLY during 'streaming', never 'submitted'. useChat sets + // status='submitted' BEFORE issuing the request and only pushes the new assistant message + // inside write(), which flips the status to 'streaming' in the same job. So for the whole + // submitted window `lastAssistantMessageId` (which has no streaming guard of its own — see + // useMessageActions) is THE PREVIOUS TURN'S reply. Latching that as "the stream's id" made + // Stop abort a message that finished minutes ago while the real generation kept running and + // kept billing — on every turn after the first. + const isActuallyStreaming = status === 'streaming'; + // The conversation the CURRENT stream belongs to, held from when it starts. The surface moves + // independently of the stream — switching conversation mid-stream does NOT abort the POST — so + // the abort must name the conversation the generation is actually running on. + const heldStreamConvIdRef = useRef(null); + heldStreamConvIdRef.current = holdForStream({ + current: heldStreamConvIdRef.current, + isStreaming, + liveValue: currentConversationId, + }); + const heldStreamMsgIdRef = useRef(null); + heldStreamMsgIdRef.current = holdForStream({ + current: heldStreamMsgIdRef.current, + isStreaming, + liveValue: isActuallyStreaming ? (lastAssistantMessageId ?? null) : null, + }); + // Rejoin-first recovery probe for useStreamRecovery. // On a network error (e.g. iOS backgrounding kills the fetch): // 1. Check /api/ai/chat/active-streams — if the original run is still live, rejoin it. @@ -1014,31 +1073,66 @@ const SidebarChatTab: React.FC = () => { // Stop handler that uses appropriate stop function based on mode // All stop functions call both abort endpoint (server-side) and useChat stop (client-side) const handleStop = useCallback(async () => { - // Use the appropriate stop function based on mode - if (!selectedAgent && contextStopStreaming) { - // Global mode: use context stop function (already calls abort endpoint) + // OUR OWN live stream wins. The shared stop (context / dashboard store) belongs to + // whichever surface installed it, and the two surfaces register under DIFFERENT chatIds + // — ours is `sidebar:`, the dashboard's is the bare convId — so the shared stop + // literally cannot abort a stream we started. Reaching for it first meant clicking Stop + // aborted the DASHBOARD's stream while ours was never stopped at all: it kept + // generating, and kept billing. (The dashboard's own dispatcher, useGlobalEffectiveStream, + // already gets this order right; this one was inverted.) + // + // The shared stop is for a stream this surface does NOT locally own — one restored by + // the bootstrap after a refresh, where there is no local fetch to stop. + if (isStreaming) { + // Fall through to the local path below. + } else if (!selectedAgent && contextStopStreaming) { + // Global mode, no local stream: a bootstrap-restored stream owns the context stop. contextStopStreaming(); + return; } else if (selectedAgent && dashboardStopStreaming) { - // Agent mode: use dashboard store stop function (already calls abort endpoint) + // Agent mode, no local stream: a bootstrap-restored stream owns the dashboard stop. dashboardStopStreaming(); - } else { + return; + } + { // Fallback (live stream, no bootstrap-registered stop): stop the local fetch // first, then abort authoritatively by the stable assistant messageId — this // reaches the server registry even if the conversation id shifted mid-stream // and tears down any multicast SSE join. Fall back to the chatId map only when // no assistant id exists yet (submitted, before the first chunk). stop(); + // Read the HELD id at call time — see heldStreamMsgIdRef. `lastAssistantMessageId` is the + // live array's, and it is gone the moment the surface switches conversation mid-stream. const messageId = resolveActiveAssistantMessageId({ - ownStreamMessageId: undefined, - isStreaming, + ownStreamMessageId: heldStreamMsgIdRef.current ?? undefined, + // 'streaming', NOT the looser isStreaming (which includes 'submitted'). During submitted + // the array's last assistant message is the previous turn's — see isActuallyStreaming. + isStreaming: isActuallyStreaming, lastAssistantMessageId, }); if (messageId) { void abortActiveStreamByMessageId({ messageId }); return; } - if (currentConversationId) { - await abortActiveStream({ chatId: currentConversationId }); + // Key by the TRANSPORT's chatId, not the bare conversation id. In agent mode the + // transport registers the streamId under `sidebar:` (see sidebarChatId — the + // namespace exists so the sidebar and the dashboard can view the same conversation + // without colliding in the activeStreams map). Aborting under the bare id was a map + // miss: the local fetch stopped, but the SERVER kept generating and kept billing, + // because the abort registry deliberately lets streams survive a client disconnect. + // Reachable in the pre-first-chunk window, where there is no assistant messageId yet + // and this fallback is the only route to a server-side abort. + // Name the CONVERSATION as well as the transport key. The chatId map is empty until the + // response headers land (0.5-3s into a real send) and is torn down by the conversation-change + // cleanup on a mid-stream switch — so on both of the paths a user actually takes, the chatId + // abort was a guaranteed no-op. It cancelled the local fetch and returned, while the server + // (which deliberately survives client disconnect) kept generating and kept billing. + const abortChatId = selectedAgent ? sidebarChatId : currentConversationId; + const abortConversationId = heldStreamConvIdRef.current ?? currentConversationId; + if (abortChatId) { + await abortActiveStream({ chatId: abortChatId, conversationId: abortConversationId }); + } else if (abortConversationId) { + await abortActiveStream({ chatId: abortConversationId, conversationId: abortConversationId }); } } }, [ @@ -1046,8 +1140,15 @@ const SidebarChatTab: React.FC = () => { contextStopStreaming, dashboardStopStreaming, currentConversationId, + sidebarChatId, stop, isStreaming, + // The callback READS this (it is what keeps Stop from resolving the previous turn's + // messageId during the submitted window). Omitting it meant the memo only happened to stay + // fresh because lastAssistantMessageId co-varies on the submitted -> streaming transition — + // an accident, not a guarantee. It is one refactor away from Stop silently capturing a stale + // value and aborting the wrong message while the real generation keeps billing. + isActuallyStreaming, lastAssistantMessageId, ]); diff --git a/apps/web/src/contexts/__tests__/GlobalChatContext.test.tsx b/apps/web/src/contexts/__tests__/GlobalChatContext.test.tsx index 931971f541..c7110dbc9e 100644 --- a/apps/web/src/contexts/__tests__/GlobalChatContext.test.tsx +++ b/apps/web/src/contexts/__tests__/GlobalChatContext.test.tsx @@ -908,6 +908,12 @@ describe('GlobalChatProvider — global channel stream socket', () => { mockSocket._trigger('chat:stream_complete', { messageId: 'msg-live', pageId: GLOBAL_CHANNEL_ID, + // The server ALWAYS sends this (broadcastAiStreamComplete passes it unconditionally); the + // type merely marks it optional. Omitting it modelled a payload production never emits — + // and it matters now: this join never delivered a part, so the hook correctly reports the + // stream as non-authoritative and drops it, leaving the conversationId as the only way to + // route the reload-from-DB. + conversationId: CONV_ID, }); }); diff --git a/apps/web/src/hooks/__tests__/useAgentChannelMultiplayer.test.ts b/apps/web/src/hooks/__tests__/useAgentChannelMultiplayer.test.ts index d1de9a38c1..10c2979b9a 100644 --- a/apps/web/src/hooks/__tests__/useAgentChannelMultiplayer.test.ts +++ b/apps/web/src/hooks/__tests__/useAgentChannelMultiplayer.test.ts @@ -64,18 +64,34 @@ vi.mock('@/lib/ai/core/stream-abort-client', () => ({ })); // Real Zustand store imported AFTER mocks -import { usePageAgentDashboardStore } from '@/stores/page-agents'; +import { usePageAgentDashboardStore, selectIsAgentStreaming, selectAgentStop } from '@/stores/page-agents'; import { useAgentChannelMultiplayer } from '../useAgentChannelMultiplayer'; +const CONV_ID = 'conv-1'; const AGENT = { id: 'agent-1' }; -const seedDashboard = (overrides: Partial<{ agentStopStreaming: (() => void) | null; isAgentStreaming: boolean }> = {}) => { +type StopSlot = { agentId: string; stop: () => void | Promise } | null; + +/** Seed the (agent-keyed) store the way production writes it. */ +const seedDashboard = (overrides: Partial<{ agentStopStreaming: StopSlot; isAgentStreaming: boolean }> = {}) => { + const slot = overrides.agentStopStreaming ?? null; usePageAgentDashboardStore.setState({ - isAgentStreaming: overrides.isAgentStreaming ?? false, - agentStopStreaming: overrides.agentStopStreaming !== undefined ? overrides.agentStopStreaming : null, + streamingAgentIds: overrides.isAgentStreaming ? { [key(AGENT.id, CONV_ID)]: true } : {}, + agentStops: slot ? { [key(slot.agentId, CONV_ID)]: slot.stop } : {}, }); }; +// Read the slot through the PRODUCTION selectors, not a reimplementation of them — the +// whole bug was that a reader could get an un-scoped answer, so a test that does its own +// agent comparison would prove nothing about the code that ships. +const dashIsStreaming = (agentId: string = AGENT.id, conversationId: string = CONV_ID) => + selectIsAgentStreaming({ agentId, conversationId })(usePageAgentDashboardStore.getState()); +const dashStop = (agentId: string = AGENT.id, conversationId: string = CONV_ID) => + selectAgentStop({ agentId, conversationId })(usePageAgentDashboardStore.getState()); + +/** Store key, the way production builds it. */ +const key = (agentId: string, conversationId: string) => `${agentId}::${conversationId}`; + const baseOptions = ( overrides: Partial[0]>, ): Parameters[0] => ({ @@ -108,7 +124,7 @@ describe('useAgentChannelMultiplayer', () => { it('given a selected agent, the channel-stream socket should be subscribed to the agent id', () => { renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); expect(capturedChannel.channelId).toBe(AGENT.id); @@ -122,88 +138,456 @@ describe('useAgentChannelMultiplayer', () => { }); describe('own-stream slot ownership', () => { + // A channel can carry two of this user's own streams at once (send, then New Chat + // while it runs), and only one holds the Stop slot. A bare boolean let the OTHER + // one's finalize release the slot out from under the stream that actually owns it — + // killing a live Stop button. The conversation guard makes this MORE reachable: a + // declined stream is still registered for finalization. + it('given a DIFFERENT own stream finalizes, the slot claimed by another stream should survive', () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-claimer', conversationId: CONV_ID }); + }); + expect(dashIsStreaming()).toBe(true); + + // A second own stream on the channel finishes — it never claimed the slot. + act(() => { + capturedChannel.options?.onOwnStreamFinalize?.({ messageId: 'msg-someone-else' }); + }); + + expect(dashIsStreaming()).toBe(true); + expect(dashStop()).not.toBeNull(); + }); + + // The claim is only ever released by onOwnStreamFinalize — and there are paths where + // that event CANNOT fire: the socket effect tears down without finalizing on a + // socket-instance swap (an auth:refreshed reconnect builds a brand-new io()), and if + // the stream ended during that gap nothing is left to announce it. The flag would + // strand true: a Stop button over a dead stream, plus permanent SWR suppression via + // useStreamingRegistration. Bootstrap is the server's word on what is still running. + it('given a claimed stream is absent from the next active-streams snapshot, the claim should be released', () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-gone', conversationId: CONV_ID }); + }); + expect(dashIsStreaming()).toBe(true); + + // The next bootstrap says that stream is no longer running. + act(() => { + capturedChannel.options?.onActiveStreamsSnapshot?.(new Set()); + }); + + expect(dashIsStreaming()).toBe(false); + expect(dashStop()).toBeNull(); + }); + + it('given the claimed stream is STILL in the snapshot, the claim should survive', () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-live', conversationId: CONV_ID }); + }); + + act(() => { + capturedChannel.options?.onActiveStreamsSnapshot?.(new Set(['msg-live'])); + }); + + expect(dashIsStreaming()).toBe(true); + expect(dashStop()).not.toBeNull(); + }); + + // The dashboard store's setAgentStopStreaming is a plain zustand VALUE setter, not a + // useState dispatch — so an updater-shaped argument (`() => fn`) is stored VERBATIM. + // GlobalAssistantView was passing exactly that, so the slot held the outer wrapper and + // SidebarChatTab's `dashboardStopStreaming()` merely RETURNED the inner fn instead of + // running it: a Stop button that silently did nothing while the stream kept generating + // and kept billing. Typechecked because `() => (() => Promise)` is assignable to + // `() => void`. Guard the invariant: whatever is in the slot must ACT when called. + it('given a claimed stop fn in the slot, calling it must actually abort (not merely return another function)', () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-stop', conversationId: CONV_ID }); + }); + + const stop = dashStop(); + expect(typeof stop).toBe('function'); + + const returned = stop!() as unknown; + + expect(mockAbortByMessageId).toHaveBeenCalledWith({ messageId: 'msg-stop' }); + expect(typeof returned).not.toBe('function'); + }); + + // BUG EIGHT. The dashboard and the sidebar do NOT share an agent — GlobalAssistantView's + // comes from usePageAgentDashboardStore, SidebarChatTab's from useSidebarAgentStore + // (independent, localStorage-persisted). And GlobalAssistantView never unmounts; + // CenterPanel only HIDES it. So after one dashboard visit the two are co-mounted on + // every page, holding different agents, reading one slot. + // + // With no identity on the slot, a stream on the dashboard's agent B lit up the + // sidebar's Stop for agent A — and clicking it aborted B while A kept generating and + // kept billing. The slot now names its agent, and a reader asking about a different one + // gets nothing. + it("given ANOTHER agent's stream owns the slot, this surface must not see it as its own", () => { + const otherAgentsStop = vi.fn(); + usePageAgentDashboardStore.setState({ + streamingAgentIds: { [key('agent-OTHER', CONV_ID)]: true }, + agentStops: { [key('agent-OTHER', CONV_ID)]: otherAgentsStop }, + }); + + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + // Our agent is AGENT.id; the slot belongs to agent-OTHER. + expect(dashIsStreaming(AGENT.id)).toBe(false); + expect(dashStop(AGENT.id)).toBeNull(); + + // And the other agent's own view of it is intact. + expect(dashIsStreaming('agent-OTHER')).toBe(true); + expect(dashStop('agent-OTHER')).toBe(otherAgentsStop); + }); + + it("given another agent owns the slot, finalizing OUR stream must not clear THEIR state", () => { + const otherAgentsStop = vi.fn(); + usePageAgentDashboardStore.setState({ + streamingAgentIds: { [key('agent-OTHER', CONV_ID)]: true }, + agentStops: { [key('agent-OTHER', CONV_ID)]: otherAgentsStop }, + }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-ours', conversationId: CONV_ID }); + }); + act(() => { + capturedChannel.options?.onOwnStreamFinalize?.({ messageId: 'msg-ours' }); + }); + + expect(dashIsStreaming('agent-OTHER')).toBe(true); + expect(dashStop('agent-OTHER')).toBe(otherAgentsStop); + }); + + // The single-slot shape had a second consequence beyond the cross-wire: whichever agent + // claimed first held the ONLY slot, so a second agent streaming at the same time could + // never get a Stop button at all. Keying by agent removes the artificial contention — + // two agents streaming concurrently each keep their own control. + it('given another agent is already streaming, this agent should still get its own Stop', () => { + const otherAgentsStop = vi.fn(); + usePageAgentDashboardStore.setState({ + streamingAgentIds: { [key('agent-OTHER', CONV_ID)]: true }, + agentStops: { [key('agent-OTHER', CONV_ID)]: otherAgentsStop }, + }); + + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-ours', conversationId: CONV_ID }); + }); + + // Ours was claimed... + expect(dashIsStreaming(AGENT.id)).toBe(true); + expect(typeof dashStop(AGENT.id)).toBe('function'); + // ...and theirs is untouched. + expect(dashIsStreaming('agent-OTHER')).toBe(true); + expect(dashStop('agent-OTHER')).toBe(otherAgentsStop); + }); + + // The last coarseness in the key. The dashboard and the sidebar keep INDEPENDENT + // conversations for the SAME agent ("New Chat" in either diverges them). Keyed by agent + // alone, a dashboard stream on conversation X2 still lit up the sidebar's Stop while it + // was showing X1 — and clicking it aborted X2. Ownership is per stream, so the key is. + it("given the same agent but a DIFFERENT conversation, this surface must not see that stream as its own", () => { + const otherConvStop = vi.fn(); + usePageAgentDashboardStore.setState({ + streamingAgentIds: { [key(AGENT.id, 'conv-OTHER')]: true }, + agentStops: { [key(AGENT.id, 'conv-OTHER')]: otherConvStop }, + }); + + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + // Same agent, our conversation: nothing. + expect(dashIsStreaming(AGENT.id, CONV_ID)).toBe(false); + expect(dashStop(AGENT.id, CONV_ID)).toBeNull(); + + // The other conversation's own state is intact. + expect(dashIsStreaming(AGENT.id, 'conv-OTHER')).toBe(true); + expect(dashStop(AGENT.id, 'conv-OTHER')).toBe(otherConvStop); + }); + + // A stream in a conversation this surface is not showing is not ours to control — the + // surface showing THAT conversation claims it. What must never happen is it leaking into + // OUR key, which is what an agent-only key allowed. + it("given a bootstrap for another conversation, it must not claim anything under OUR key", () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ + messageId: 'msg-other-conv', + conversationId: 'conv-OTHER', + }); + }); + + expect(dashIsStreaming(AGENT.id, CONV_ID)).toBe(false); + expect(dashStop(AGENT.id, CONV_ID)).toBeNull(); + }); + + // The claim is keyed by the STREAM's own conversation, so the DB bootstrap landing + // before this surface has resolved its identity (which it is designed to tolerate) + // cannot claim under the wrong conversation. + it("given the surface's conversation is not resolved yet, the claim should be keyed to the STREAM's conversation", () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: null })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ + messageId: 'msg-boot', + conversationId: CONV_ID, + }); + }); + + expect(dashIsStreaming(AGENT.id, CONV_ID)).toBe(true); + expect(typeof dashStop(AGENT.id, CONV_ID)).toBe('function'); + }); + + // BUG TWELVE, from the hook's side. GlobalAssistantView used to key its flag/stop by the + // surface's LIVE conversation. Because `useChat` only recreates its Chat when its `id` + // changes (and GAV's is a constant), switching conversation mid-stream does NOT abort the + // POST — so the key migrated: the running stream's entry was cleared and a fresh claim was + // installed under a conversation with NO stream. The abandoned stream lost its Stop and + // its SWR protection while still generating; the new key showed a Stop that aborted + // nothing. GAV now keys by the conversation the stream STARTED in. + // + // From here, the invariant to hold is that a claim made for conversation X stays under X + // even as the surface moves on. + it('given the surface switches conversation while our claimed stream runs, the claim must stay under the STREAM\'s conversation', () => { + seedDashboard({ agentStopStreaming: null }); + const { rerender } = renderHook(({ opts }) => useAgentChannelMultiplayer(opts), { + initialProps: { opts: baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID }) }, + }); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-x', conversationId: CONV_ID }); + }); + expect(dashIsStreaming(AGENT.id, CONV_ID)).toBe(true); + + // The user switches to another conversation on the same agent. The stream keeps running. + act(() => { + rerender({ opts: baseOptions({ selectedAgent: AGENT, agentConversationId: 'conv-NEW' }) }); + }); + + // The claim is still where the stream is... + expect(dashIsStreaming(AGENT.id, CONV_ID)).toBe(true); + expect(typeof dashStop(AGENT.id, CONV_ID)).toBe('function'); + // ...and nothing was fabricated under the conversation we moved to. + expect(dashIsStreaming(AGENT.id, 'conv-NEW')).toBe(false); + expect(dashStop(AGENT.id, 'conv-NEW')).toBeNull(); + }); + + // A NULL slot is FREE, not foreign. GlobalAssistantView nulls the stop fn on ordinary + // paths (its effect's else-branch and cleanup fire whenever agentStatus is 'ready' — + // which it is for the whole life of a BOOTSTRAPPED stream) without touching + // isAgentStreaming. Treating that as "not ours" would skip setAgentStreaming(false) + // and strand the surface showing "streaming" with no way out. + it('given the stop fn was nulled by another surface, finalizing must still clear isAgentStreaming', () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-boot', conversationId: CONV_ID }); + }); + expect(dashIsStreaming()).toBe(true); + + // Another surface's effect cleanup nulls the stop fn — it never touches the flag. + act(() => { usePageAgentDashboardStore.setState({ agentStops: {} }); }); + + act(() => { + capturedChannel.options?.onOwnStreamFinalize?.({ messageId: 'msg-boot' }); + }); + + expect(dashIsStreaming()).toBe(false); + }); + + // The stop slot is a single shared singleton, and GlobalAssistantView writes it + // DIRECTLY from its own local status without going through the claim protocol. So by + // the time this surface releases, the slot may belong to somebody else — and nulling + // it then kills THEIR live Stop button and their isAgentStreaming flag mid-stream. + it('given the slot has since been taken over by another surface, switching agents should NOT clobber it', () => { + seedDashboard({ agentStopStreaming: null }); + const { rerender } = renderHook(({ opts }) => useAgentChannelMultiplayer(opts), { + initialProps: { opts: baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID }) }, + }); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-ours', conversationId: CONV_ID }); + }); + + // Another surface (e.g. GlobalAssistantView) overwrites the shared slot. + const someoneElsesStop = vi.fn(); + act(() => { + usePageAgentDashboardStore.setState({ agentStops: { [key(AGENT.id, CONV_ID)]: someoneElsesStop } }); + usePageAgentDashboardStore.setState({ streamingAgentIds: { [key(AGENT.id, CONV_ID)]: true } }); + }); + + // We switch agents. Our claim is stale — we must not touch their state. + act(() => { + rerender({ opts: baseOptions({ selectedAgent: { id: 'agent-b' }, agentConversationId: 'conv-b' }) }); + }); + + expect(dashStop()).toBe(someoneElsesStop); + expect(dashIsStreaming()).toBe(true); + }); + + // The sidebar swaps agents IN PLACE (no unmount), and useChannelStreamSocket tears + // down its effect on that channelId change WITHOUT firing onOwnStreamFinalize (by + // design). So the claim must be released by the channel-keyed cleanup — otherwise the + // next agent's own stream finds the slot occupied, declines it, and its finalize does + // not match the stale messageId: the dashboard is stuck streaming with a Stop button + // wired to the PREVIOUS agent's dead message. (The old boolean ref recovered from this + // by accident, so scoping the release by messageId without this is a regression.) + it('given the agent is switched mid-stream, the stale claim should be released rather than stranding a dead Stop button', () => { + seedDashboard({ agentStopStreaming: null }); + const { rerender } = renderHook(({ opts }) => useAgentChannelMultiplayer(opts), { + initialProps: { opts: baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID }) }, + }); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-agent-a', conversationId: CONV_ID }); + }); + expect(dashIsStreaming()).toBe(true); + + // Switch to a different agent — same component, new channel. + act(() => { + rerender({ + opts: baseOptions({ selectedAgent: { id: 'agent-b' }, agentConversationId: 'conv-b' }), + }); + }); + + expect(dashIsStreaming()).toBe(false); + expect(dashStop()).toBeNull(); + }); + + it('given the CLAIMING stream finalizes, the slot should be released', () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ selectedAgent: AGENT, agentConversationId: CONV_ID })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-claimer', conversationId: CONV_ID }); + }); + act(() => { + capturedChannel.options?.onOwnStreamFinalize?.({ messageId: 'msg-claimer' }); + }); + + expect(dashIsStreaming()).toBe(false); + expect(dashStop()).toBeNull(); + }); + + // Conversation-scoped: an own stream in ANOTHER conversation on this agent channel + // must not light up the Stop button for the conversation on screen (it would abort + // the wrong stream). Same class of bug as the remote-stream conversation filter. + it('given onOwnStreamBootstrap fires for a DIFFERENT conversation, the slot should NOT be claimed', () => { + seedDashboard({ agentStopStreaming: null }); + renderWiring(baseOptions({ + selectedAgent: AGENT, + agentConversationId: CONV_ID, + })); + + act(() => { + capturedChannel.options?.onOwnStreamBootstrap?.({ + messageId: 'msg-other-conv', + conversationId: 'a-different-conversation', + }); + }); + + expect(dashStop()).toBeNull(); + expect(dashIsStreaming()).toBe(false); + }); + it('given onOwnStreamBootstrap fires while the dashboard stop slot is empty, the slot should be claimed', () => { seedDashboard({ agentStopStreaming: null }); renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); act(() => { - capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: 'conv-1' }); + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: CONV_ID }); }); - const state = usePageAgentDashboardStore.getState(); - expect(typeof state.agentStopStreaming).toBe('function'); - expect(state.isAgentStreaming).toBe(true); + expect(typeof dashStop()).toBe('function'); + expect(dashIsStreaming()).toBe(true); }); it('given onOwnStreamBootstrap fires while the slot is already populated, the existing slot should be preserved', () => { const existingStop = vi.fn(); - seedDashboard({ agentStopStreaming: existingStop }); + seedDashboard({ agentStopStreaming: { agentId: AGENT.id, stop: existingStop } }); renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); act(() => { - capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: 'conv-1' }); + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: CONV_ID }); }); - expect(usePageAgentDashboardStore.getState().agentStopStreaming).toBe(existingStop); + expect(dashStop()).toBe(existingStop); }); it('given onOwnStreamFinalize fires after this surface claimed the slot, the slot should be cleared', () => { seedDashboard({ agentStopStreaming: null }); renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); act(() => { - capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: 'conv-1' }); + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: CONV_ID }); }); act(() => { capturedChannel.options?.onOwnStreamFinalize?.({ messageId: 'msg-own' }); }); - const state = usePageAgentDashboardStore.getState(); - expect(state.agentStopStreaming).toBeNull(); - expect(state.isAgentStreaming).toBe(false); + expect(dashStop()).toBeNull(); + expect(dashIsStreaming()).toBe(false); }); it('given onOwnStreamFinalize fires when this surface never claimed the slot, the existing stop should remain', () => { const otherSurfaceStop = vi.fn(); - seedDashboard({ agentStopStreaming: otherSurfaceStop }); + seedDashboard({ agentStopStreaming: { agentId: AGENT.id, stop: otherSurfaceStop } }); renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); // Bootstrap fires but slot is occupied; this surface declines to claim. act(() => { - capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: 'conv-1' }); + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: CONV_ID }); }); // Finalize arrives — must NOT clear the other surface's slot. act(() => { capturedChannel.options?.onOwnStreamFinalize?.({ messageId: 'msg-own' }); }); - expect(usePageAgentDashboardStore.getState().agentStopStreaming).toBe(otherSurfaceStop); + expect(dashStop()).toBe(otherSurfaceStop); }); it("given the slot is claimed, calling the stop function in the slot should invoke the abort endpoint", () => { seedDashboard({ agentStopStreaming: null }); renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); act(() => { - capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: 'conv-1' }); + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: CONV_ID }); }); - const stop = usePageAgentDashboardStore.getState().agentStopStreaming; + const stop = dashStop(); stop?.(); expect(mockAbortByMessageId).toHaveBeenCalledWith({ messageId: 'msg-own' }); @@ -386,7 +770,7 @@ describe('useAgentChannelMultiplayer', () => { mockSocketStatus.current = status; return useAgentChannelMultiplayer(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, loadConversation: surfaceLoadConversation, })); }, @@ -411,7 +795,7 @@ describe('useAgentChannelMultiplayer', () => { mockSocketStatus.current = status; return useAgentChannelMultiplayer(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, loadConversation: surfaceLoadConversation, })); }, @@ -446,7 +830,7 @@ describe('useAgentChannelMultiplayer', () => { mockSocketStatus.current = status; return useAgentChannelMultiplayer(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, loadConversation: surfaceLoadConversation, })); }, @@ -496,42 +880,41 @@ describe('useAgentChannelMultiplayer', () => { seedDashboard({ agentStopStreaming: null }); const { unmount } = renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); // Bootstrap claims the slot. act(() => { - capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: 'conv-1' }); + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: CONV_ID }); }); - expect(typeof usePageAgentDashboardStore.getState().agentStopStreaming).toBe('function'); + expect(typeof dashStop()).toBe('function'); // Surface unmounts before stream finalizes (e.g. user navigates away). // useChannelStreamSocket aborts the controller but does NOT fire // onOwnStreamFinalize, so the hook's own cleanup must clear the slot. unmount(); - const state = usePageAgentDashboardStore.getState(); - expect(state.agentStopStreaming).toBeNull(); - expect(state.isAgentStreaming).toBe(false); + expect(dashStop()).toBeNull(); + expect(dashIsStreaming()).toBe(false); }); it('given this surface never claimed the slot, unmount should NOT clear another writer\'s slot', () => { const otherSurfaceStop = vi.fn(); - seedDashboard({ agentStopStreaming: otherSurfaceStop, isAgentStreaming: true }); + seedDashboard({ agentStopStreaming: { agentId: AGENT.id, stop: otherSurfaceStop }, isAgentStreaming: true }); const { unmount } = renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); // Bootstrap fires but slot is occupied — this surface declines. act(() => { - capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: 'conv-1' }); + capturedChannel.options?.onOwnStreamBootstrap?.({ messageId: 'msg-own', conversationId: CONV_ID }); }); // Surface unmounts. The other writer's slot must survive. unmount(); - expect(usePageAgentDashboardStore.getState().agentStopStreaming).toBe(otherSurfaceStop); - expect(usePageAgentDashboardStore.getState().isAgentStreaming).toBe(true); + expect(dashStop()).toBe(otherSurfaceStop); + expect(dashIsStreaming()).toBe(true); }); }); @@ -539,7 +922,7 @@ describe('useAgentChannelMultiplayer', () => { it('given a selected agent, the registration key should be `ai-channel-${agent.id}`', () => { renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, isLocallyStreaming: true, })); @@ -561,7 +944,7 @@ describe('useAgentChannelMultiplayer', () => { seedDashboard({ isAgentStreaming: true }); renderWiring(baseOptions({ selectedAgent: AGENT, - agentConversationId: 'conv-1', + agentConversationId: CONV_ID, })); const lastCall = mockUseStreamingRegistration.mock.calls.at(-1); diff --git a/apps/web/src/hooks/__tests__/useChannelStreamSocket.test.ts b/apps/web/src/hooks/__tests__/useChannelStreamSocket.test.ts index a5f06b136d..ae1fd6c0a4 100644 --- a/apps/web/src/hooks/__tests__/useChannelStreamSocket.test.ts +++ b/apps/web/src/hooks/__tests__/useChannelStreamSocket.test.ts @@ -244,8 +244,18 @@ describe('useChannelStreamSocket', () => { it('should call removeStream and onStreamComplete after consumeStreamJoin resolves', async () => { let resolveJoin!: () => void; - mockConsumeStreamJoin.mockReturnValue( - new Promise<{ aborted: boolean }>((res) => { resolveJoin = () => res({ aborted: false }); }), + // A successful join to a live stream always delivers at least the buffered prefix (the SSE + // route replays it synchronously on subscribe). A join that delivers NOTHING means our copy + // is not authoritative — the hook now reports joinFailed so consumers reload from the DB + // rather than rendering the stale, debounced seed. Model the delivery. + mockConsumeStreamJoin.mockImplementation( + (_id: string, _signal: AbortSignal, onChunk: (part: unknown) => void) => + new Promise<{ aborted: boolean }>((res) => { + resolveJoin = () => { + onChunk({ type: 'text', text: 'hello' }); + res({ aborted: false }); + }; + }), ); const onStreamComplete = vi.fn(); @@ -260,8 +270,18 @@ describe('useChannelStreamSocket', () => { it('should call onStreamComplete before removeStream so stream data is available in the callback (SSE path)', async () => { let resolveJoin!: () => void; - mockConsumeStreamJoin.mockReturnValue( - new Promise<{ aborted: boolean }>((res) => { resolveJoin = () => res({ aborted: false }); }), + // A real join DELIVERS at least one part — that is how the store gets the data this test is + // about. A join that delivers nothing leaves only the seeded (debounced, possibly shorter) + // DB snapshot, which is not authoritative, so the hook now drops it and lets consumers + // reload from the DB instead of rendering a truncated bubble. + mockConsumeStreamJoin.mockImplementation( + (_id: string, _signal: AbortSignal, onChunk: (part: unknown) => void) => + new Promise<{ aborted: boolean }>((res) => { + resolveJoin = () => { + onChunk({ type: 'text', text: 'hello' }); + res({ aborted: false }); + }; + }), ); const callOrder: string[] = []; mockRemoveStream.mockImplementation(() => { callOrder.push('removeStream'); }); @@ -730,23 +750,39 @@ describe('useChannelStreamSocket', () => { }); describe('chat:stream_complete', () => { - it('should call removeStream and onStreamComplete', () => { + it('should call removeStream and onStreamComplete, reporting joinFailed for a stream we never joined', () => { const onStreamComplete = vi.fn(); renderHook(() => useChannelStreamSocket('page-a', { onStreamComplete })); + // No stream_start, no join: we hold no authoritative copy of this stream, so `joinFailed` + // is TRUE — that flag means "our parts are not the source of truth, reload from the DB", + // not "the network failed". Consumers behave identically either way here (there is no store + // entry, so they take the reload branch regardless); the flag now simply tells the truth. act(() => { mockSocket._trigger('chat:stream_complete', COMPLETE_PAYLOAD); }); expect(mockRemoveStream).toHaveBeenCalledWith('msg-1'); - expect(onStreamComplete).toHaveBeenCalledWith('msg-1', 'conv-1', { joinFailed: false }); + expect(onStreamComplete).toHaveBeenCalledWith('msg-1', 'conv-1', { joinFailed: true }); }); - it('should call onStreamComplete before removeStream so stream data is available in the callback (socket path)', () => { + it('should call onStreamComplete before removeStream so stream data is available in the callback (socket path)', async () => { const callOrder: string[] = []; - mockRemoveStream.mockImplementation(() => { callOrder.push('removeStream'); }); + // The stream must actually have been joined AND delivered content — otherwise there is no + // "stream data" for the callback to see, and the hook deliberately drops the entry so + // consumers reload the authoritative message from the DB. + let deliver!: () => void; + mockConsumeStreamJoin.mockImplementation( + (_id: string, _signal: AbortSignal, onChunk: (part: unknown) => void) => + new Promise<{ aborted: boolean }>(() => { + deliver = () => onChunk({ type: 'text', text: 'hello' }); + }), + ); const onStreamComplete = vi.fn(() => { callOrder.push('onStreamComplete'); }); renderHook(() => useChannelStreamSocket('page-a', { onStreamComplete })); + act(() => { mockSocket._trigger('chat:stream_start', START_PAYLOAD); }); + await act(async () => { await Promise.resolve(); deliver(); }); + mockRemoveStream.mockImplementation(() => { callOrder.push('removeStream'); }); act(() => { mockSocket._trigger('chat:stream_complete', COMPLETE_PAYLOAD); }); expect(callOrder).toEqual(['onStreamComplete', 'removeStream']); diff --git a/apps/web/src/hooks/useAgentChannelMultiplayer.ts b/apps/web/src/hooks/useAgentChannelMultiplayer.ts index 011ca8139c..ea98d8820c 100644 --- a/apps/web/src/hooks/useAgentChannelMultiplayer.ts +++ b/apps/web/src/hooks/useAgentChannelMultiplayer.ts @@ -4,7 +4,12 @@ import { useChannelStreamSocket } from './useChannelStreamSocket'; import { usePageSocketRoom } from './usePageSocketRoom'; import { useSocketStore } from '@/stores/useSocketStore'; import { usePendingStreamsStore } from '@/stores/usePendingStreamsStore'; -import { usePageAgentDashboardStore } from '@/stores/page-agents'; +import { + usePageAgentDashboardStore, + selectIsAgentStreaming, + agentStreamKey, + type AgentStreamKey, +} from '@/stores/page-agents'; import { useStreamingRegistration } from '@/lib/ai/shared'; import { abortActiveStreamByMessageId } from '@/lib/ai/core/stream-abort-client'; import { synthesizeAssistantMessage } from '@/lib/ai/streams/synthesizeAssistantMessage'; @@ -81,16 +86,77 @@ export function useAgentChannelMultiplayer({ // this surface unmounts mid-stream — useChannelStreamSocket intentionally // does not fire onOwnStreamFinalize on teardown, so without the cleanup // a navigate-away mid-stream would leave the slot stuck. - const ownedStopSlotRef = useRef(false); + // The messageId this surface claimed the slot FOR — not a bare boolean. A channel can + // carry two of this user's own streams at once (send, then New Chat while it runs), and + // only one of them holds the slot. With a boolean, the OTHER one's finalize would + // release the slot out from under the stream that actually owns it, killing a live + // Stop button. `onOwnStreamBootstrap` can also decline to claim (conversation + // mismatch), while `ownStreamIds` still registers the stream for finalization — so a + // declined stream must never be able to release anything. + const ownedStopSlotRef = useRef(null); + // The exact stop function we installed. Releasing on "I claimed once" is not enough: + // the slot is a single shared singleton, and GlobalAssistantView writes it directly + // from its own local status without going through the claim protocol. So by the time we + // release, the slot may belong to somebody else — and nulling it then would kill THEIR + // live Stop button and their isAgentStreaming flag. Only release what is still ours. + const ownedStopFnRef = useRef<(() => void) | null>(null); + // The agent we claimed FOR — not the agent we happen to be on now. The release can run + // from a cleanup triggered by an agent switch, where `channelId` has already moved on; + // clearing the flag against the new agent would leave the OLD one's flag set forever. + // The KEY we claimed for — (agent, conversation), not just the agent, and not the agent we + // happen to be on now (a cleanup fired by an agent switch has already moved on). + const ownedKeyRef = useRef(null); + + // KNOWN GAP (pre-existing, unchanged by this hook's messageId scoping): the claim + // protocol has no HANDOFF. If a co-mounted surface was declined the slot (first writer + // wins) and the claimant later releases it, the declined surface does not re-claim — so + // it can render a live own stream with no Stop button until it remounts. Closing this + // needs a re-claim protocol on the dashboard store (subscribe to the slot, re-claim when + // it frees and our stream is still in flight), which is a change to that store's + // contract rather than to stream ownership. Deliberately left for a follow-up. + const releaseStopSlotIfStillOurs = () => { + if (ownedStopSlotRef.current === null) return; + const dashboard = usePageAgentDashboardStore.getState(); + // A null slot is FREE, not foreign. GlobalAssistantView nulls the stop fn on ordinary + // paths (its effect's else-branch and cleanup, whenever agentStatus is 'ready' — which + // it is for the whole life of a BOOTSTRAPPED stream) without touching isAgentStreaming. + // Treating that as "not ours" would skip setAgentStreaming(false) and strand it true. + // A DIFFERENT fn means another, live stream owns the slot — leave both halves alone. + // A NULL slot means nobody owns it: it is free, and still ours to clear. + const claimedKey = ownedKeyRef.current; + const k = claimedKey === null ? null : agentStreamKey(claimedKey); + const ownedStopFn = ownedStopFnRef.current; + // Clear our refs FIRST, unconditionally — whatever we decide below, we no longer hold a + // claim, and an early return must never leave them dangling. + ownedStopSlotRef.current = null; + ownedStopFnRef.current = null; + ownedKeyRef.current = null; + if (claimedKey === null || k === null) return; + // Identity-guarded WITHIN the key: GlobalAssistantView installs its own local stop for + // this same (agent, conversation) and would otherwise have ours clobbered. Collisions + // across different streams are impossible by construction — the state is keyed per stream. + const current = dashboard.agentStops[k]; + const stillOurs = current === ownedStopFn || current === undefined; + if (!stillOurs) return; + dashboard.setAgentStreaming(claimedKey, false); + dashboard.setAgentStop(claimedKey, null); + }; + const releaseStopSlotRef = useRef(releaseStopSlotIfStillOurs); + releaseStopSlotRef.current = releaseStopSlotIfStillOurs; + // Keyed on channelId, NOT []. `channelId` is `selectedAgent?.id`, and the sidebar swaps + // agents in place without unmounting — while useChannelStreamSocket's effect (deps + // [socket, channelId]) tears down on that same change and, by design, does NOT fire + // onOwnStreamFinalize. A []-keyed cleanup would therefore never run, leaving the claim + // behind: the next agent's own stream would find the slot occupied and decline it, its + // finalize would not match the stale messageId, and the dashboard would be stuck + // `isAgentStreaming: true` with a Stop button wired to the previous agent's dead + // message. (The old boolean ref accidentally recovered from this — any own finalize + // released it — so scoping the release by messageId without this would be a regression.) useEffect(() => { return () => { - if (!ownedStopSlotRef.current) return; - ownedStopSlotRef.current = false; - const dashboard = usePageAgentDashboardStore.getState(); - dashboard.setAgentStreaming(false); - dashboard.setAgentStopStreaming(null); + releaseStopSlotRef.current(); }; - }, []); + }, [channelId]); const { rejoinActiveStreams } = useChannelStreamSocket(channelId, { onUserMessage: (message, payload) => { @@ -116,31 +182,76 @@ export function useAgentChannelMultiplayer({ onStreamComplete: (messageId, completedConvId) => { const stream = usePendingStreamsStore.getState().streams.get(messageId); if (stream && stream.parts.length > 0 && stream.conversationId === agentConversationIdRef.current) { - setLocalMessagesRef.current((prev) => [ - ...prev, - synthesizeAssistantMessage(messageId, stream.parts, stream.startedAt), - ]); + // REPLACE by id — do not skip. An existing message with this id is NOT proof we already + // have the content. + // + // The server names the assistant message (`generateId: () => serverAssistantMessageId`), + // so useChat's copy and the stream's `messageId` are the SAME id. And useChat does not + // roll back on error: a mid-stream network drop leaves its HALF-STREAMED message sitting + // in the array. That is exactly the path the rejoin machinery exists for — recovery + // rejoins the multicast, `stream.parts` accumulates the FULL reply, and this fires on + // completion. + // + // A skip-if-present guard therefore threw the complete reply away and left the user + // staring at the truncated one, with the real text stranded in the DB until they + // navigated away and back. Replacing is right in both cases: same id means same message, + // and `stream.parts` is the authoritative, complete version of it. + const synthesized = synthesizeAssistantMessage(messageId, stream.parts, stream.startedAt); + setLocalMessagesRef.current((prev) => { + const i = prev.findIndex((m) => m.id === messageId); + return i === -1 + ? [...prev, synthesized] + : prev.map((m, j) => (j === i ? synthesized : m)); + }); return; } if (shouldReloadOnComountComplete(stream, completedConvId, agentConversationIdRef.current)) { loadConversationRef.current(completedConvId!); } }, - onOwnStreamBootstrap: ({ messageId }) => { + onOwnStreamBootstrap: ({ messageId, conversationId }) => { + // Conversation-scoped, same reasoning as the remote-stream filter: a channel + // carries every conversation's streams, so an own stream in conversation X must + // not make conversation Y render as streaming with a Stop button that aborts X. + // Only reject a KNOWN mismatch — the DB bootstrap can land before the surface has + // resolved its conversation, and rejecting on a null id there would drop the very + // stream it is about to render. + const activeId = agentConversationIdRef.current; + if (activeId !== null && conversationId !== activeId) return; const dashboard = usePageAgentDashboardStore.getState(); - if (!shouldClaimAgentStopSlot(dashboard.agentStopStreaming)) return; - ownedStopSlotRef.current = true; - dashboard.setAgentStreaming(true); - dashboard.setAgentStopStreaming(() => { + if (!channelId) return; + // Keyed by the STREAM's OWN conversation — not the surface's active one. That is what + // lets the bootstrap claim land before identity resolves without owning the wrong + // thing: the claim names its own stream, and a reader asking about a different + // conversation simply never sees it. + const claimKey: AgentStreamKey = { agentId: channelId, conversationId }; + const claimK = agentStreamKey(claimKey); + if (claimK === null) return; + if (!shouldClaimAgentStopSlot(dashboard.agentStops[claimK] ?? null)) return; + const stopFn = () => { abortActiveStreamByMessageId({ messageId }); - }); + }; + ownedStopSlotRef.current = messageId; + ownedStopFnRef.current = stopFn; + ownedKeyRef.current = claimKey; + dashboard.setAgentStreaming(claimKey, true); + dashboard.setAgentStop(claimKey, stopFn); }, - onOwnStreamFinalize: () => { - if (!ownedStopSlotRef.current) return; - ownedStopSlotRef.current = false; - const dashboard = usePageAgentDashboardStore.getState(); - dashboard.setAgentStreaming(false); - dashboard.setAgentStopStreaming(null); + // A claim is only ever released by onOwnStreamFinalize — and there are paths where + // that event never fires (the socket effect tears down without finalizing on a + // socket-instance swap; if the stream ended in that gap, nothing announces it). The + // flag would strand true: a Stop button over a dead stream, plus permanent SWR + // suppression via useStreamingRegistration. Bootstrap knows the truth, so reconcile. + onActiveStreamsSnapshot: (liveMessageIds) => { + const claimed = ownedStopSlotRef.current; + if (claimed === null || liveMessageIds.has(claimed)) return; + releaseStopSlotRef.current(); + }, + onOwnStreamFinalize: ({ messageId }) => { + // Only the stream that actually claimed the slot may release it — and only if the + // slot still holds our stop function (see releaseStopSlotIfStillOurs). + if (ownedStopSlotRef.current !== messageId) return; + releaseStopSlotRef.current(); }, }); @@ -149,7 +260,10 @@ export function useAgentChannelMultiplayer({ // dedups same-key writes naturally. The flag ORs the surface's local // streaming flag with the dashboard store's bootstrap-driven flag so a // mid-stream refresh stays SWR-protected before useChat re-engages. - const dashboardAgentStreaming = usePageAgentDashboardStore((s) => s.isAgentStreaming); + // Scoped to OUR agent — the dashboard slot is shared with a surface that may hold another. + const dashboardAgentStreaming = usePageAgentDashboardStore( + selectIsAgentStreaming({ agentId: channelId, conversationId: agentConversationId }), + ); useStreamingRegistration( selectedAgent ? `ai-channel-${selectedAgent.id}` : 'ai-channel-no-agent', Boolean(selectedAgent) && (isLocallyStreaming || dashboardAgentStreaming), diff --git a/apps/web/src/hooks/useChannelStreamSocket.ts b/apps/web/src/hooks/useChannelStreamSocket.ts index e6de811a1d..b7c77517e7 100644 --- a/apps/web/src/hooks/useChannelStreamSocket.ts +++ b/apps/web/src/hooks/useChannelStreamSocket.ts @@ -205,6 +205,18 @@ export function useChannelStreamSocket( // chat:stream_complete and reload the (durably persisted) message from the // DB, which is what this set enables in handleStreamComplete below. const joinFailed = new Set(); + // Streams whose SSE join actually DELIVERED at least one part. + // + // `joinFailed` alone is not enough, because it is populated in an ASYNC `.catch`. A join + // commonly 404s *because* the stream just ended — and `chat:stream_complete` can land before + // that rejection settles. In that window the store still holds only the seeded DB snapshot, + // which is debounced and can be SHORTER than what useChat already has. Consumers would then + // replace a longer visible reply with a truncated one AND skip the reload-from-DB branch, so + // the full persisted message was never fetched. Content loss. + // + // A join that delivered nothing is not authoritative, whatever the catch has managed to + // record yet. This is the synchronous fact that says so. + const joinDelivered = new Set(); const { addStream, appendPart, removeStream, clearPageStreams } = usePendingStreamsStore.getState(); @@ -242,6 +254,7 @@ export function useChannelStreamSocket( chunksToSkip -= 1; return; } + joinDelivered.add(messageId); appendPart(messageId, part); }) .then(() => { @@ -449,11 +462,15 @@ export function useChannelStreamSocket( // — the authoritative message is in the DB. Dropping the store entry BEFORE // firing makes consumers take their reload-from-DB branch // (shouldReloadOnComountComplete) instead of synthesizing a truncated bubble. - const didJoinFail = joinFailed.has(payload.messageId); + // `|| !joinDelivered` closes the race described at joinDelivered: the catch that would have + // set joinFailed may not have settled yet, but a join that delivered nothing cannot be the + // authoritative copy either way. Treat both as "reload from the DB". + const didJoinFail = joinFailed.has(payload.messageId) || !joinDelivered.has(payload.messageId); if (didJoinFail) { joinFailed.delete(payload.messageId); removeStream(payload.messageId); } + joinDelivered.delete(payload.messageId); try { fireComplete(payload.messageId, payload.conversationId, { joinFailed: didJoinFail }); } finally { diff --git a/apps/web/src/lib/ai/shared/hooks/useChatStop.ts b/apps/web/src/lib/ai/shared/hooks/useChatStop.ts index cb366840e8..c43edc0778 100644 --- a/apps/web/src/lib/ai/shared/hooks/useChatStop.ts +++ b/apps/web/src/lib/ai/shared/hooks/useChatStop.ts @@ -6,18 +6,42 @@ import { abortActiveStream } from '@/lib/ai/core/client'; * then stops the client-side fetch via useChat's stop. * * Uses try/finally to guarantee client-side stop runs even if server abort fails. + * + * `conversationId` is what makes this able to stop anything at all in the windows that matter. + * + * The `chatId` lookup is a client-side map (`activeStreams`) populated only once the response + * HEADERS land — and a real agent send spends 0.5-3 seconds before that (auth, rate limit, DB + * reads, context assembly, connecting to the provider). The same map is also torn down by the + * surfaces' own conversation-change cleanups, which fire on a mid-stream switch. So on the two + * paths a user is most likely to take — Stop right after Send, and Stop after switching + * conversation — the map was EMPTY, the abort was a guaranteed no-op, and this cancelled the + * local fetch and returned. + * + * Cancelling the fetch stops NOTHING. Streams are deliberately server-owned and survive a client + * disconnect — that is the architecture. So the generation kept running, kept calling write + * tools, and kept BILLING, while the button flipped back to Send and the user believed it had + * stopped. + * + * The conversationId is the one name the client holds from t=0, and the server can abort by it + * (the caller's OWN streams only — see abort-conversation-streams.ts). Pass the STREAM's + * conversation, held from when it started — never the live one — so a mid-stream switch still + * names the generation that is actually running. */ export function useChatStop( chatId: string | null, - chatStop: () => void + chatStop: () => void, + conversationId?: string | null, ): () => Promise { return useCallback(async () => { try { if (chatId) { - await abortActiveStream({ chatId }); + await abortActiveStream({ chatId, conversationId }); + } else if (conversationId) { + // No transport key, but we still know which conversation is generating. + await abortActiveStream({ chatId: conversationId, conversationId }); } } finally { chatStop(); } - }, [chatId, chatStop]); + }, [chatId, chatStop, conversationId]); } diff --git a/apps/web/src/lib/ai/streams/__tests__/holdForStream.test.ts b/apps/web/src/lib/ai/streams/__tests__/holdForStream.test.ts new file mode 100644 index 0000000000..7c5c5641ef --- /dev/null +++ b/apps/web/src/lib/ai/streams/__tests__/holdForStream.test.ts @@ -0,0 +1,143 @@ +import { describe, it, expect } from 'vitest'; +import { holdForStream } from '../holdForStream'; + +describe('holdForStream', () => { + it('given no stream, should hold nothing', () => { + expect(holdForStream({ current: null, isStreaming: false, liveValue: 'X' })) + .toBeNull(); + }); + + it('given a stream starting, should capture the surface\'s conversation', () => { + expect(holdForStream({ current: null, isStreaming: true, liveValue: 'X' })) + .toBe('X'); + }); + + // THE point of this module. useChat does not recreate its Chat when the conversation + // changes (its id is a constant), so a mid-stream switch does NOT abort the POST — the + // stream keeps running while the surface moves on. Following the surface here migrated + // ownership: the running stream's entry was cleared and a fresh claim installed under a + // conversation with no stream, so the real stream lost its Stop and kept billing. + it('given the surface switches conversation MID-STREAM, should keep holding the conversation the stream started in', () => { + expect(holdForStream({ current: 'X', isStreaming: true, liveValue: 'Y' })) + .toBe('X'); + }); + + it('given the stream ends, should release', () => { + expect(holdForStream({ current: 'X', isStreaming: false, liveValue: 'Y' })) + .toBeNull(); + }); + + it('given the next stream starts after a release, should capture the NEW conversation', () => { + const afterRelease = holdForStream({ current: 'X', isStreaming: false, liveValue: 'Y' }); + expect(holdForStream({ current: afterRelease, isStreaming: true, liveValue: 'Y' })) + .toBe('Y'); + }); + + // A stream that starts before the surface has resolved its conversation cannot be named, + // and an un-named stream owns nothing — the store warns rather than silently claiming. + it('given a stream starts with no conversation resolved yet, should hold nothing', () => { + expect(holdForStream({ current: null, isStreaming: true, liveValue: null })) + .toBeNull(); + expect(holdForStream({ current: null, isStreaming: true, liveValue: undefined })) + .toBeNull(); + }); + + // ...and must still be able to capture once it resolves, rather than being stuck at null. + it('given the conversation resolves after the stream started, should capture it then', () => { + expect(holdForStream({ current: null, isStreaming: true, liveValue: 'X' })) + .toBe('X'); + }); + + // ── The submitted-window latch ────────────────────────────────────────────────────────── + // + // This module latches on the FIRST render where the stream is live — and with the callers' + // `isStreaming` (which is `status === 'submitted' || status === 'streaming'`), that first + // render is a SUBMITTED render. + // + // useChat sets status='submitted' BEFORE issuing the request, and pushes the new assistant + // message only inside write(), which flips the status to 'streaming' in the same job. So at + // the submitted render the array's last assistant message is THE PREVIOUS TURN'S reply. + // + // Latching it meant Stop aborted a message that finished minutes ago: the registry no longer + // knew it, the local fetch stopped, the button LOOKED like it worked, and the real generation + // kept running its write tools and kept billing. On every turn after the first. + // + // The fix is at the callers — they must feed `liveValue: status === 'streaming' ? id : null` + // — so what this pins is the property that makes that fix WORK: a null liveValue while the + // stream is live must not latch, and must stay capturable once the real id arrives. + describe('the submitted-window latch', () => { + it('given the stream is live but the live id is not yet knowable, should hold nothing rather than latch a wrong id', () => { + expect(holdForStream({ current: null, isStreaming: true, liveValue: null })) + .toBeNull(); + }); + + it('THE BUG: after holding nothing through submitted, should capture THIS stream\'s id on the first streaming render', () => { + // submitted: caller passes null (status !== 'streaming'), so nothing is latched... + const held = holdForStream({ current: null, isStreaming: true, liveValue: null }); + expect(held).toBeNull(); + // ...streaming: the real assistant message has now been pushed. Capture THAT. + expect(holdForStream({ current: held, isStreaming: true, liveValue: 'M_new' })) + .toBe('M_new'); + }); + + it('given the id was captured, should keep holding it for the rest of the stream', () => { + expect(holdForStream({ current: 'M_new', isStreaming: true, liveValue: 'M_new' })) + .toBe('M_new'); + }); + + it('given the stream ends, should release so the NEXT turn captures its own id and not this one', () => { + expect(holdForStream({ current: 'M_new', isStreaming: false, liveValue: 'M_new' })) + .toBeNull(); + }); + }); + + + // ── Cross-contamination between two concurrent streams on one surface ──────────────────── + // + // The dashboard hosts TWO independent chats (agent and global) and both can be in flight at + // once — switching mode does NOT abort the running POST, because useChat's id is constant. + // + // Feeding ONE mode-selected liveValue into BOTH hold-refs let the IDLE mode's ref latch the + // ACTIVE mode's messageId and pin it. Stop, back in the other mode, then aborted the WRONG + // stream: one answer died mid-sentence while the other kept running its write tools and kept + // billing, its Stop permanently wired to an id that was never its own. + // + // The fix is at the caller (each ref is fed its OWN chat's id). What this pins is the reducer + // property that makes it work: a ref whose stream is live but whose own id is not yet known + // must latch NOTHING, so it stays capturable — it must never be fillable from elsewhere. + describe('two concurrent streams on one surface', () => { + it("THE BUG: a live stream with no id of its own yet must not latch, so it cannot be filled with the OTHER stream's id", () => { + // Global stream is live but still in 'submitted' — its own id is unknowable, so the caller + // passes null. If this latched anything here, the next render (when the AGENT's id becomes + // the mode-selected value) would have pinned the agent's id onto the global ref. + const globalRef = holdForStream({ current: null, isStreaming: true, liveValue: null }); + expect(globalRef).toBeNull(); + + // The agent's stream starts and reaches 'streaming'. Its ref captures ITS id. + const agentRef = holdForStream({ current: null, isStreaming: true, liveValue: 'M_agent' }); + expect(agentRef).toBe('M_agent'); + + // The global ref, still live, still without its own id, must STILL hold nothing. + expect(holdForStream({ current: globalRef, isStreaming: true, liveValue: null })).toBeNull(); + + // ...and when the global stream finally produces its own id, it captures THAT. + expect(holdForStream({ current: globalRef, isStreaming: true, liveValue: 'M_global' })) + .toBe('M_global'); + }); + + it('given both streams have captured, each holds its own id independently', () => { + expect(holdForStream({ current: 'M_agent', isStreaming: true, liveValue: 'M_agent' })) + .toBe('M_agent'); + expect(holdForStream({ current: 'M_global', isStreaming: true, liveValue: 'M_global' })) + .toBe('M_global'); + }); + + it('given one stream ends, it releases without disturbing the other', () => { + expect(holdForStream({ current: 'M_agent', isStreaming: false, liveValue: 'M_agent' })) + .toBeNull(); + expect(holdForStream({ current: 'M_global', isStreaming: true, liveValue: 'M_global' })) + .toBe('M_global'); + }); + }); + +}); diff --git a/apps/web/src/lib/ai/streams/__tests__/selectLiveAssistantIds.test.ts b/apps/web/src/lib/ai/streams/__tests__/selectLiveAssistantIds.test.ts new file mode 100644 index 0000000000..4786cca5fd --- /dev/null +++ b/apps/web/src/lib/ai/streams/__tests__/selectLiveAssistantIds.test.ts @@ -0,0 +1,84 @@ +import { describe, it, expect } from 'vitest'; +import { selectLiveAssistantIds } from '../selectLiveAssistantIds'; + +const user = (id: string) => ({ id, role: 'user' }); +const assistant = (id: string) => ({ id, role: 'assistant' }); + +describe('selectLiveAssistantIds', () => { + it('resolves each chat\'s live assistant id from its OWN status and messages', () => { + expect( + selectLiveAssistantIds({ + agent: { status: 'streaming', messages: [user('u1'), assistant('M_agent')] }, + global: { status: 'streaming', messages: [user('u2'), assistant('M_global')] }, + }), + ).toEqual({ agentLiveId: 'M_agent', globalLiveId: 'M_global' }); + }); + + // THE BUG. Both chats can stream at once (switching mode does not abort the running POST), and + // the component used to feed ONE mode-selected id into both hold-refs. While the global stream + // sat in 'submitted' with no id of its own, the AGENT's id was the mode-selected value — so the + // global ref latched it and pinned it. Stop, back in global mode, aborted the AGENT's stream: + // that answer died mid-sentence while the global generation kept billing. + it('THE BUG: an agent stream in flight must NOT leak its id into the global slot', () => { + const { agentLiveId, globalLiveId } = selectLiveAssistantIds({ + agent: { status: 'streaming', messages: [user('u1'), assistant('M_agent')] }, + // Global is live, but still submitted — it has no id of its own yet. The array's last + // assistant message is the PREVIOUS turn's reply, and must not be mistaken for this one. + global: { status: 'submitted', messages: [assistant('M_global_prev'), user('u2')] }, + }); + + expect(agentLiveId).toBe('M_agent'); + expect(globalLiveId).toBeUndefined(); + expect(globalLiveId).not.toBe('M_agent'); + }); + + it('THE BUG, mirrored: a global stream in flight must NOT leak its id into the agent slot', () => { + const { agentLiveId, globalLiveId } = selectLiveAssistantIds({ + agent: { status: 'submitted', messages: [assistant('M_agent_prev'), user('u1')] }, + global: { status: 'streaming', messages: [user('u2'), assistant('M_global')] }, + }); + + expect(globalLiveId).toBe('M_global'); + expect(agentLiveId).toBeUndefined(); + expect(agentLiveId).not.toBe('M_global'); + }); + + it('given a chat is submitted, returns undefined rather than the PREVIOUS turn\'s assistant id', () => { + // The submitted-window latch: useChat pushes this turn's assistant message only on reaching + // 'streaming'. Reading the array any earlier yields the last completed reply — aborting that + // id is a no-op the server registry no longer knows, while the real stream keeps billing. + const { agentLiveId } = selectLiveAssistantIds({ + agent: { status: 'submitted', messages: [assistant('M_prev'), user('u_new')] }, + global: { status: 'ready', messages: [] }, + }); + expect(agentLiveId).toBeUndefined(); + }); + + it('given a chat is idle, returns undefined even though a completed assistant message exists', () => { + const { agentLiveId, globalLiveId } = selectLiveAssistantIds({ + agent: { status: 'ready', messages: [user('u1'), assistant('M_done')] }, + global: { status: 'error', messages: [user('u2'), assistant('M_partial')] }, + }); + expect(agentLiveId).toBeUndefined(); + expect(globalLiveId).toBeUndefined(); + }); + + it('given a streaming chat with no assistant message yet, returns undefined', () => { + const { agentLiveId } = selectLiveAssistantIds({ + agent: { status: 'streaming', messages: [user('u1')] }, + global: { status: 'ready', messages: [] }, + }); + expect(agentLiveId).toBeUndefined(); + }); + + it('picks the LAST assistant message, not the first', () => { + const { agentLiveId } = selectLiveAssistantIds({ + agent: { + status: 'streaming', + messages: [assistant('M_old'), user('u1'), assistant('M_current')], + }, + global: { status: 'ready', messages: [] }, + }); + expect(agentLiveId).toBe('M_current'); + }); +}); diff --git a/apps/web/src/lib/ai/streams/__tests__/shouldClaimAgentStopSlot.test.ts b/apps/web/src/lib/ai/streams/__tests__/shouldClaimAgentStopSlot.test.ts index 073f9356f0..6256dafc84 100644 --- a/apps/web/src/lib/ai/streams/__tests__/shouldClaimAgentStopSlot.test.ts +++ b/apps/web/src/lib/ai/streams/__tests__/shouldClaimAgentStopSlot.test.ts @@ -6,12 +6,14 @@ describe('shouldClaimAgentStopSlot', () => { expect(shouldClaimAgentStopSlot(null)).toBe(true); }); - it('given a populated slot, should not claim (so the first writer is preserved)', () => { - const existingStop = () => {}; - expect(shouldClaimAgentStopSlot(existingStop)).toBe(false); + it('given a populated slot, should not claim (so the first writer for THIS agent is preserved)', () => { + expect(shouldClaimAgentStopSlot(() => {})).toBe(false); }); - it('given a populated slot whose function is a no-op, should still not claim', () => { - expect(shouldClaimAgentStopSlot(() => undefined)).toBe(false); + // The store is keyed by agent, so the caller looks up its OWN agent's stop before asking. + // Another agent's stop is not reachable here at all — which is what stops one agent's + // stream from blocking (or hijacking) another's Stop button. + it('given an async stop fn, should still not claim (the void-return trap must not hide it)', () => { + expect(shouldClaimAgentStopSlot(async () => {})).toBe(false); }); }); diff --git a/apps/web/src/lib/ai/streams/holdForStream.ts b/apps/web/src/lib/ai/streams/holdForStream.ts new file mode 100644 index 0000000000..20821550ea --- /dev/null +++ b/apps/web/src/lib/ai/streams/holdForStream.ts @@ -0,0 +1,63 @@ +/** + * Capture a fact about the CURRENT local stream when it appears, and hold it until the + * stream ends — rather than following the surface, which moves independently of it. + * + * A stream's identity is fixed when it starts. The surface's conversation is not: `useChat` + * only recreates its Chat when its `id` changes, and the agent/global surfaces use a constant + * id — so switching conversation mid-stream does NOT abort the POST. The stream keeps running + * while the surface's conversation moves out from under it. + * + * Keying stream ownership off the *live* conversation therefore migrated it: the running + * stream's entry was cleared and a fresh claim was installed under a conversation with no + * stream at all. The abandoned stream lost its Stop and its SWR protection while still + * generating; the new key showed a Stop that aborted nothing — and the abort it did issue + * named the wrong conversation, so the real stream kept billing. + * + * So: capture on the way in, hold for the stream's life, release on the way out. + * + * ── CALLER CONTRACT: `liveValue` MUST BE null UNTIL IT NAMES *THIS* STREAM ────────────────── + * + * This captures on the FIRST render where `isStreaming` is true. Callers derive `isStreaming` + * from `status === 'submitted' || status === 'streaming'`, so that first render is a SUBMITTED + * render — and at that moment useChat has NOT yet pushed this stream's assistant message. It + * sets status='submitted' before issuing the request and pushes the message only inside + * `write()`, which flips the status to 'streaming' in the same job. + * + * So during the submitted window, "the last assistant message in the array" is the PREVIOUS + * TURN'S reply. Feeding that in as `liveValue` latched a message that had finished minutes ago + * and held it for the whole stream: Stop then aborted an id the server registry no longer knew, + * the local fetch stopped, the button LOOKED like it worked, and the real generation kept + * running its write tools and kept billing — on every turn after the first. + * + * Hence the contract. For a messageId, pass `status === 'streaming' ? id : null` — never the id + * derived from the looser `isStreaming`. Holding nothing is safe (callers fall back to the + * chatId map, which is correct in that window); holding the WRONG thing is not. + * + * (An earlier version of this doc said the messageId is "captured when the first chunk + * arrives". It is not, and that sentence is what the bug above was made of.) + */ +export const holdForStream = ({ + current, + isStreaming, + liveValue, +}: { + /** What we captured previously (null when no stream is running). */ + current: string | null; + /** Is a local stream in flight right now? */ + isStreaming: boolean; + /** + * The value to capture — but ONLY once it actually names THIS stream. See the caller + * contract above: pass null while it does not, because we latch the first non-null value we + * see and hold it for the whole stream. + * + * The conversationId is knowable immediately (the surface's own conversation at send time). + * The assistant messageId is NOT: it does not exist until useChat reaches 'streaming', and + * reading it any earlier yields the previous turn's reply. + */ + liveValue: string | null | undefined; +}): string | null => { + if (!isStreaming) return null; + // Already captured: hold it, even if the surface has since moved on. + if (current !== null) return current; + return liveValue ?? null; +}; diff --git a/apps/web/src/lib/ai/streams/selectLiveAssistantIds.ts b/apps/web/src/lib/ai/streams/selectLiveAssistantIds.ts new file mode 100644 index 0000000000..a864c1284a --- /dev/null +++ b/apps/web/src/lib/ai/streams/selectLiveAssistantIds.ts @@ -0,0 +1,52 @@ +/** + * The live assistant messageId of EACH of a surface's chats, kept strictly apart. + * + * GlobalAssistantView hosts two independent chats — the agent one and the global one — and both + * can be in flight at the same time: switching mode does NOT abort the running POST, because + * useChat's `id` is a constant and the Chat is never recreated. + * + * The component used to derive ONE id from the mode-selected `status`/`messages` (correct for + * what it renders) and feed that single value into BOTH stream-identity hold-refs. So while the + * global stream sat in 'submitted' with no id of its own, the agent's id became the mode-selected + * value — and the global ref latched it and pinned it for the rest of the stream. Stop, back in + * global mode, then aborted the AGENT's stream: that answer died mid-sentence while the global + * generation kept running its write tools and kept billing, its own Stop permanently wired to an + * id that was never its. + * + * This exists as a pure function because the bug was in the WIRING, not in the reducer it fed. + * A test of `holdForStream` alone passes happily while the caller hands it the wrong chat's id — + * which is exactly what happened. Making the wiring itself the unit under test is the only way to + * pin it. + * + * Both ids resolve ONLY at `status === 'streaming'`, never 'submitted': useChat sets 'submitted' + * before issuing the request and pushes the new assistant message inside write(), which flips the + * status in the same job — so during 'submitted' the array's last assistant message is the + * PREVIOUS turn's reply. Returning undefined there is correct; callers fall back to the chatId + * map. See holdForStream's caller contract. + */ + +interface ChatSnapshot { + status: string; + messages: ReadonlyArray<{ id: string; role: string }>; +} + +const lastAssistantIdWhileStreaming = ({ status, messages }: ChatSnapshot): string | undefined => { + if (status !== 'streaming') return undefined; + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].role === 'assistant') return messages[i].id; + } + return undefined; +}; + +export const selectLiveAssistantIds = ({ + agent, + global, +}: { + agent: ChatSnapshot; + global: ChatSnapshot; +}): { agentLiveId: string | undefined; globalLiveId: string | undefined } => ({ + // Each id is computed from ITS OWN chat's status and messages. Neither function argument can + // reach the other's result — that separation is the entire contract. + agentLiveId: lastAssistantIdWhileStreaming(agent), + globalLiveId: lastAssistantIdWhileStreaming(global), +}); diff --git a/apps/web/src/lib/ai/streams/shouldClaimAgentStopSlot.ts b/apps/web/src/lib/ai/streams/shouldClaimAgentStopSlot.ts index 31d1d8f150..b58e57dd20 100644 --- a/apps/web/src/lib/ai/streams/shouldClaimAgentStopSlot.ts +++ b/apps/web/src/lib/ai/streams/shouldClaimAgentStopSlot.ts @@ -5,5 +5,12 @@ * other's stop function — first writer wins, second is a no-op. */ export const shouldClaimAgentStopSlot = ( - currentStop: (() => void) | null, + // THIS AGENT's current stop, already looked up by key. The store is keyed by agent now, so + // "is the slot free" is a question about one agent and cannot be answered on behalf of + // another — which is exactly what used to cross-wire the dashboard and the sidebar. + // + // The union return type is deliberate: `(() => void)` would let TypeScript's void-return + // rule accept a function returning ANYTHING, which is how an updater-shaped value slipped + // into a plain value setter and shipped a Stop button that did nothing. + currentStop: (() => void | Promise) | null, ): boolean => currentStop === null; diff --git a/apps/web/src/stores/page-agents/index.ts b/apps/web/src/stores/page-agents/index.ts index b01fc00d77..550dd994a0 100644 --- a/apps/web/src/stores/page-agents/index.ts +++ b/apps/web/src/stores/page-agents/index.ts @@ -7,6 +7,10 @@ export { usePageAgentDashboardStore, + selectIsAgentStreaming, + agentStreamKey, + type AgentStreamKey, + selectAgentStop, type AgentInfo, type SidebarTab, } from './usePageAgentDashboardStore'; diff --git a/apps/web/src/stores/page-agents/usePageAgentDashboardStore.ts b/apps/web/src/stores/page-agents/usePageAgentDashboardStore.ts index 0e03a2aa95..0f9ec0a6b6 100644 --- a/apps/web/src/stores/page-agents/usePageAgentDashboardStore.ts +++ b/apps/web/src/stores/page-agents/usePageAgentDashboardStore.ts @@ -59,9 +59,56 @@ interface AgentState { * even when the conversation ID doesn't change (clicking the same conversation). */ conversationLoadSignal: number; - // Streaming state (for agent mode sync between GlobalAssistantView and sidebar) - isAgentStreaming: boolean; - agentStopStreaming: (() => void) | null; + // Agent-mode streaming state, shared between GlobalAssistantView and the sidebar. + // + // BOTH CARRY AN AGENT ID, and that is the whole point. + // + // These used to be a bare `boolean` and a bare function. But the two surfaces that read + // this slot do NOT share an agent: GlobalAssistantView's agent comes from THIS store, + // while SidebarChatTab's comes from useSidebarAgentStore (independent, localStorage- + // persisted). And GlobalAssistantView never unmounts — CenterPanel only HIDES it — so + // after one dashboard visit the two are co-mounted on every page, with different agents. + // + // An identity-less slot therefore cross-wired them: a stream on the dashboard's agent B + // lit up the sidebar's Stop button for agent A, and clicking it aborted B while A's own + // stream was never stopped — and kept generating, and kept billing. A boolean cannot + // express "whose", so it lied. Now a reader must name the agent it is asking about, and + // a slot belonging to someone else simply doesn't answer. + /** + * Which agents are currently streaming. Ask via `selectIsAgentStreaming(myAgentId)`. + * + * KEYED BY (AGENT, CONVERSATION) — see agentStreamKey. + * + * Not a single slot, and not agent alone. The surfaces that read this hold different + * agents AND different conversations within the same agent, and can stream at the same + * time. Every coarser key forced them to share one answer, and every bug this state has + * shed came from that: a boolean cannot say WHOSE, and an agent id cannot say WHICH + * CONVERSATION. Ownership is per stream, so the key is per stream. + * + * WITHIN a key, at most one stream is *expected* — a new send takes over the conversation + * (stream-takeover.ts) — but that is NOT a guarantee, and this state must not pretend it + * is. Takeover's abort registry is in-process, so a stream running on another web instance + * cannot be stopped from here ("this send runs alongside it", stream-takeover.ts), and the + * whole takeover is best-effort on a DB error. PageSpace runs multi-instance, so two live + * streams in one conversation is rare but real. + * + * When that happens this state holds ONE entry for the key: the first claimant wins + * (shouldClaimAgentStopSlot) and the second stream runs without a Stop of its own. That is + * a known, accepted limitation of the same family as the cross-instance rejoin gap — not + * something the key can fix, because the second stream is unreachable from this process + * anyway (its abort registry entry lives elsewhere). Adding a messageId axis would let the + * state DESCRIBE it without being able to ACT on it. + */ + streamingAgentIds: Readonly>; + /** + * Stop controls, keyed by the agent each belongs to. Ask via `selectAgentStop(myAgentId)`. + * + * The union return type is deliberate: `(() => void)` would let TypeScript's void-return + * rule accept a function returning ANYTHING — which is how an updater-shaped value + * (`() => async () => {}`) slipped into a plain VALUE setter and stored the outer wrapper. + * Calling it merely returned the inner fn: a Stop button that did nothing. + */ + agentStops: Readonly void | Promise>>; // Sidebar tab state (for dashboard context only - GlobalAssistantView <-> RightPanel sync) activeTab: SidebarTab; @@ -78,9 +125,9 @@ interface AgentState { clearConversation: () => void; loadMostRecentConversation: () => Promise; - // Streaming methods - setAgentStreaming: (isStreaming: boolean) => void; - setAgentStopStreaming: (stop: (() => void) | null) => void; + // Both name the stream — this state has no meaning without (agent, conversation). + setAgentStreaming: (key: AgentStreamKey, isStreaming: boolean) => void; + setAgentStop: (key: AgentStreamKey, stop: (() => void | Promise) | null) => void; } const IDLE_IDENTITY: ConversationIdentityState = { status: 'idle' }; @@ -116,8 +163,8 @@ export const usePageAgentDashboardStore = create()((set, get) => { isConversationMessagesLoading: false, conversationAgentId: null, conversationLoadSignal: 0, - isAgentStreaming: false, - agentStopStreaming: null, + streamingAgentIds: {}, + agentStops: {}, activeTab: 'history', // Default for dashboard (no chat tab in dashboard context) /** @@ -392,18 +439,75 @@ export const usePageAgentDashboardStore = create()((set, get) => { } }, - /** - * Set agent streaming state (for sync between GlobalAssistantView and sidebar) - */ - setAgentStreaming: (isStreaming: boolean) => { - set({ isAgentStreaming: isStreaming }); + /** Declare whether THIS stream is running. Never touches another stream's state. */ + setAgentStreaming: (key: AgentStreamKey, isStreaming: boolean) => { + const k = agentStreamKey(key); + if (k === null) { + // A write we cannot key is a claim silently dropped — the stream would run with no + // flag and no Stop, and nothing would ever say so. Be loud instead of quiet. + if (isStreaming) { + console.warn('[dashboardStore] setAgentStreaming with an unidentified stream — dropped', key); + } + return; + } + set((state) => { + const next = { ...state.streamingAgentIds }; + if (isStreaming) next[k] = true; + else delete next[k]; + return { streamingAgentIds: next }; + }); }, - /** - * Set agent stop streaming function (for sync between GlobalAssistantView and sidebar) - */ - setAgentStopStreaming: (stop: (() => void) | null) => { - set({ agentStopStreaming: stop }); + /** Install (or clear) THIS stream's stop control. Never touches another stream's. */ + setAgentStop: (key: AgentStreamKey, stop: (() => void | Promise) | null) => { + const k = agentStreamKey(key); + if (k === null) { + if (stop) { + console.warn('[dashboardStore] setAgentStop with an unidentified stream — dropped', key); + } + return; + } + set((state) => { + const next = { ...state.agentStops }; + if (stop) next[k] = stop; + else delete next[k]; + return { agentStops: next }; + }); }, }; }); + +/** + * Identifies ONE stream: an agent, and the conversation within it. + * + * Both halves are required. The dashboard and the sidebar hold different agents AND, for + * the same agent, different conversations (each keeps its own; "New Chat" in either + * diverges them). A key missing either half answers a question the reader did not ask. + */ +export interface AgentStreamKey { + agentId: string | null | undefined; + conversationId: string | null | undefined; +} + +/** null when the stream is not yet fully identified — an un-named stream owns nothing. */ +export const agentStreamKey = ({ agentId, conversationId }: AgentStreamKey): string | null => + agentId != null && conversationId != null ? `${agentId}::${conversationId}` : null; + +/** + * Ask "is MY stream running?" — never "is something running?". + * + * The distinction is the bug this slot's identity exists to prevent: the dashboard and the + * sidebar hold different agents, so an un-named answer belongs to somebody else's stream. + */ +export const selectIsAgentStreaming = (key: AgentStreamKey) => + (state: AgentState): boolean => { + const k = agentStreamKey(key); + return k !== null && state.streamingAgentIds[k] === true; + }; + +/** The stop control, but only if it is MY agent's. Otherwise null — it is not mine to call. */ +export const selectAgentStop = (key: AgentStreamKey) => + (state: AgentState): (() => void | Promise) | null => { + const k = agentStreamKey(key); + return (k !== null ? state.agentStops[k] : undefined) ?? null; + };