Merge pull request #430 from Mng-dev-ai/fix/streaming-projection-truncation-and-queue-handoff

Fix streaming content truncation on unmount and queue handoff drop

Author: Mng-dev-ai

CLAUDE.md

@@ -119,6 +119,7 @@
 - Never use `useEffect` to derive state from other state or props — if state is always computable from existing values, use inline computation or `useMemo`; `useEffect(() => setX(f(y)), [y])` causes an unnecessary extra render cycle
 - When state must reset on a prop/ID change, use a ref-based render check instead of `useEffect` — pattern: `const prevRef = useRef(prop); if (prevRef.current !== prop) { prevRef.current = prop; setState(initial); }` — this runs synchronously during render and avoids the effect-induced double render
 - Distinguish "derived state" from "form state initialization" — if local state is a copy of server data that the user then edits independently (e.g., secrets form, settings form), syncing it via `useEffect` on query data change is the correct pattern; do not convert these to inline derivation, as the local state intentionally diverges from the source
+- `useEffect` cleanup closures must not rely on hook-scoped utilities that close over the current entity ID (e.g., `updateMessageInCache` scoped to the rendered `chatId`) when the cleanup may serve sessions from a different entity — use refs or the underlying API directly (e.g., `queryClient.setQueryData` with `session.chatId`) to target the correct entity
 
 ### Component Variants
 - Create explicit variant components instead of one component with many boolean modes (e.g., `ThreadComposer`, `EditComposer` instead of `<Composer isThread isEditing />`)
@@ -145,6 +146,7 @@
 - Do not wrap trivial expressions in `useMemo` (e.g., `useMemo(() => x || [], [x])`) — use direct expressions (`x ?? []`)
 - Hoist regex patterns to module-level constants — never create RegExp inside loops or frequently-called functions
 - Prefer single-pass iteration (`.reduce()`) over chained `.filter().map()` in render paths
+- When reordering a function call to run earlier in a per-event hot path (e.g., stream envelope processing), gate the call with the cheapest possible condition check at the call site — avoid paying function-call overhead for the 99% of events that will just early-return
 - Keep `useEffect` for external system subscriptions and DOM side effects — keyboard shortcuts, resize observers, WebSocket lifecycle, scroll-into-view, and focus management require post-render timing and cleanup; do not convert these to ref-based render checks
 
 ### Async Patterns

frontend/src/hooks/useStreamCallbacks.ts

@@ -22,6 +22,9 @@ import type { StreamOptions } from '@/services/streamService';
 import { useChatSettingsStore } from '@/store/chatSettingsStore';
 import type { PaginatedMessages } from '@/types/api.types';
 
+// Batching window for streaming content updates. Envelopes arrive at token-level
+// granularity (~10-50ms apart); flushing on every token would thrash React state
+// and the query cache. 130ms collects a visible chunk of text per paint cycle.
 const STREAM_FLUSH_INTERVAL_MS = 130;
 
 function findMessageInCache(
@@ -59,46 +62,6 @@ function buildProjectionUpdate(
   });
 }
 
-interface UseStreamCallbacksParams {
-  messages: Message[];
-  chatId: string | undefined;
-  currentChat: Chat | undefined;
-  queryClient: QueryClient;
-  refetchFilesMetadata: () => Promise<unknown>;
-  onContextUsageUpdate?: (data: ContextUsage, chatId?: string) => void;
-  onPermissionRequest?: (request: PermissionRequest) => void;
-  setMessages: Dispatch<SetStateAction<Message[]>>;
-  setStreamState: Dispatch<SetStateAction<StreamState>>;
-  setCurrentMessageId: Dispatch<SetStateAction<string | null>>;
-  setError: Dispatch<SetStateAction<Error | null>>;
-  pendingStopRef: React.MutableRefObject<Set<string>>;
-  onPendingUserMessageIdChange?: (id: string | null) => void;
-}
-
-interface UseStreamCallbacksResult {
-  onEnvelope: (envelope: StreamEnvelope) => void;
-  onComplete: (
-    messageId?: string,
-    streamId?: string,
-    terminalKind?: 'complete' | 'cancelled',
-  ) => void;
-  onError: (error: Error, messageId?: string, streamId?: string) => void;
-  onQueueProcess: (data: QueueProcessingData) => void;
-  startStream: (request: StreamOptions['request']) => Promise<string>;
-  replayStream: (messageId: string, afterSeq?: number) => Promise<string>;
-  stopStream: (messageId: string) => Promise<void>;
-  updateMessageInCache: ReturnType<typeof useMessageCache>['updateMessageInCache'];
-  addMessageToCache: ReturnType<typeof useMessageCache>['addMessageToCache'];
-  removeMessagesFromCache: ReturnType<typeof useMessageCache>['removeMessagesFromCache'];
-  setPendingUserMessageId: (id: string | null) => void;
-}
-
-interface StreamSessionState {
-  messageId: string;
-  lastSeq: number;
-  chatId: string;
-}
-
 function envelopeToRenderEvent(envelope: StreamEnvelope): AssistantStreamEvent | null {
   const payload = envelope.payload as Record<string, unknown>;
 
@@ -148,6 +111,46 @@ function envelopeToRenderEvent(envelope: StreamEnvelope): AssistantStreamEvent |
   }
 }
 
+interface UseStreamCallbacksParams {
+  messages: Message[];
+  chatId: string | undefined;
+  currentChat: Chat | undefined;
+  queryClient: QueryClient;
+  refetchFilesMetadata: () => Promise<unknown>;
+  onContextUsageUpdate?: (data: ContextUsage, chatId?: string) => void;
+  onPermissionRequest?: (request: PermissionRequest) => void;
+  setMessages: Dispatch<SetStateAction<Message[]>>;
+  setStreamState: Dispatch<SetStateAction<StreamState>>;
+  setCurrentMessageId: Dispatch<SetStateAction<string | null>>;
+  setError: Dispatch<SetStateAction<Error | null>>;
+  pendingStopRef: React.MutableRefObject<Set<string>>;
+  onPendingUserMessageIdChange?: (id: string | null) => void;
+}
+
+interface UseStreamCallbacksResult {
+  onEnvelope: (envelope: StreamEnvelope) => void;
+  onComplete: (
+    messageId?: string,
+    streamId?: string,
+    terminalKind?: 'complete' | 'cancelled',
+  ) => void;
+  onError: (error: Error, messageId?: string, streamId?: string) => void;
+  onQueueProcess: (data: QueueProcessingData) => void;
+  startStream: (request: StreamOptions['request']) => Promise<string>;
+  replayStream: (messageId: string, afterSeq?: number) => Promise<string>;
+  stopStream: (messageId: string) => Promise<void>;
+  updateMessageInCache: ReturnType<typeof useMessageCache>['updateMessageInCache'];
+  addMessageToCache: ReturnType<typeof useMessageCache>['addMessageToCache'];
+  removeMessagesFromCache: ReturnType<typeof useMessageCache>['removeMessagesFromCache'];
+  setPendingUserMessageId: (id: string | null) => void;
+}
+
+interface StreamSessionState {
+  messageId: string;
+  lastSeq: number;
+  chatId: string;
+}
+
 export function useStreamCallbacks({
   messages,
   chatId,
@@ -215,6 +218,10 @@ export function useStreamCallbacks({
     return undefined;
   }, []);
 
+  // Snapshot the accumulator's collected content into both React state (for the
+  // visible chat) and the query cache (so navigating away and back shows progress).
+  // Only updates React state when the session's chat is the one on screen — off-screen
+  // streams still write to the cache so content isn't lost.
   const applyProjection = useCallback(
     (streamId: string, { writeToCache }: { writeToCache: boolean }) => {
       const accumulator = accumulatorsRef.current.get(streamId);
@@ -236,6 +243,8 @@ export function useStreamCallbacks({
     [setMessages, updateMessageInCache],
   );
 
+  // Coalescing timer: only one pending flush per stream. Multiple envelopes arriving
+  // within the same 130ms window are batched into a single applyProjection call.
   const scheduleProjection = useCallback(
     (streamId: string) => {
       if (flushTimersRef.current.has(streamId)) {
@@ -252,6 +261,9 @@ export function useStreamCallbacks({
     [applyProjection],
   );
 
+  // Returns (or creates) the accumulator for a stream. On first call for a given
+  // streamId, seeds the accumulator from the message's existing content so that
+  // reconnections append to prior content rather than starting from blank.
   const ensureAccumulator = useCallback(
     (
       streamId: string,
@@ -270,6 +282,7 @@ export function useStreamCallbacks({
         return existing;
       }
 
+      // Seed from the on-screen message (fast path) or the query cache (off-screen chat).
       let seedEvents: AssistantStreamEvent[] = [];
       let seedText = '';
       const existingMessage =
@@ -308,7 +321,34 @@ export function useStreamCallbacks({
       timerIdsRef.current.forEach(clearTimeout);
       timerIdsRef.current = [];
 
-      flushTimers.forEach((timer) => clearTimeout(timer));
+      // Flush any pending projections to the query cache before clearing,
+      // so content already cursor-acked in chatStorage is not lost on unmount.
+      // Write directly via queryClient using session.chatId — the hook-level
+      // updateMessageInCache is scoped to the current chatId, which may differ
+      // from the session's owning chat after cross-chat navigation.
+      for (const [streamId, timer] of flushTimers.entries()) {
+        clearTimeout(timer);
+        const accumulator = accumulators.get(streamId);
+        const session = streamSessions.get(streamId);
+        if (accumulator && session) {
+          const update = buildProjectionUpdate(streamId, accumulator, session);
+          queryClient.setQueryData(
+            queryKeys.messages(session.chatId),
+            (oldData: { pages: PaginatedMessages[]; pageParams: unknown[] } | undefined) => {
+              if (!oldData?.pages) return oldData;
+              return {
+                ...oldData,
+                pages: oldData.pages.map((page: PaginatedMessages) => ({
+                  ...page,
+                  items: page.items.map((msg: Message) =>
+                    msg.id === session.messageId ? update(msg) : msg,
+                  ),
+                })),
+              };
+            },
+          );
+        }
+      }
       flushTimers.clear();
 
       accumulators.clear();
@@ -324,16 +364,23 @@ export function useStreamCallbacks({
     [onPendingUserMessageIdChange],
   );
 
+  // Central dispatch for every envelope arriving from the EventSource. Handles
+  // side-effect-only events (permissions, system metadata, plan mode transitions)
+  // inline, then delegates renderable content (text, thinking, tools) to the
+  // accumulator → scheduleProjection pipeline for batched UI updates.
   const onEnvelope = useCallback(
     (envelope: StreamEnvelope) => {
       if (pendingStopRef.current.has(envelope.messageId)) {
         return;
       }
 
+      // First envelope for a message clears the optimistic "sending…" indicator.
       if (pendingUserMessageIdRef.current && chatId === chatIdRef.current) {
         setPendingUserMessageId(null);
       }
 
+      // Permission requests are dispatched directly to the modal system and
+      // never accumulate into the message content — early return skips the pipeline.
       if (envelope.kind === 'permission_request' && onPermissionRequest) {
         const payload = envelope.payload as Record<string, unknown>;
         const request_id = typeof payload.request_id === 'string' ? payload.request_id : undefined;
@@ -497,6 +544,8 @@ export function useStreamCallbacks({
       timerIdsRef.current.forEach(clearTimeout);
       timerIdsRef.current = [];
 
+      // Delay usage refetch — the backend finalizes token counts asynchronously
+      // after the stream ends, so an immediate query would return stale totals.
       timerIdsRef.current.push(
         setTimeout(() => {
           queryClient.invalidateQueries({ queryKey: [queryKeys.auth.usage] });
@@ -509,6 +558,8 @@ export function useStreamCallbacks({
           queryClient.invalidateQueries({ queryKey: [queryKeys.chats, 'infinite'] });
           queryClient.invalidateQueries({ queryKey: queryKeys.workspaces });
         }
+        // Context usage aggregation runs as a background task after the stream
+        // completes; 6s gives it time to finish before we refetch.
         timerIdsRef.current.push(
           setTimeout(() => {
             queryClient.invalidateQueries({ queryKey: queryKeys.contextUsage(chatId) });
@@ -648,6 +699,9 @@ export function useStreamCallbacks({
     ],
   );
 
+  // Stash the latest callbacks in a ref so startStream/replayStream — which are
+  // intentionally stable (empty deps) to avoid re-registering the EventSource —
+  // always dispatch through the freshest closures.
   useEffect(() => {
     optionsRef.current = chatId
       ? { chatId, onEnvelope, onComplete, onError, onQueueProcess }

frontend/src/services/streamService.ts

@@ -170,13 +170,18 @@ class StreamService {
       chatStorage.setEventId(chatId, String(seq));
     }
 
+    // Queue handoff events carry the old assistant message ID but must be processed
+    // before the messageId filter — they update the stream's messageId for subsequent events.
+    if (parsed.kind === 'queue_processing') {
+      this.maybeHandleQueueEvent(parsed, chatId);
+    }
+
     const isForActiveMessage = parsed.messageId === currentStream.messageId;
     if (!isForActiveMessage) {
       return;
     }
 
     currentStream.callbacks?.onEnvelope?.(parsed);
-    this.maybeHandleQueueEvent(parsed, chatId);
 
     if (parsed.kind === 'complete' || parsed.kind === 'cancelled') {
       const { callbacks } = currentStream;