fix(ai): Stop-button / stream-ownership machine (stacked on #2018)

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<void>) | 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 };
 }

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<string | null>(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:<convId>` (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<string | null>(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<string | null>(null);
+  heldStreamConvIdRef.current = holdForStream({
+    current: heldStreamConvIdRef.current,
+    isStreaming,
+    liveValue: currentConversationId,
+  });
+  const heldStreamMsgIdRef = useRef<string | null>(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,40 +1073,82 @@ 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:<convId>`, 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:<convId>` (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 });
       }
     }
   }, [
     selectedAgent,
     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,
   ]);
 

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,
       });
     });