fix(chat): preserve live turns and recover interrupted streams

Author: Sg312

apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx

@@ -1153,9 +1153,13 @@ describe('useChat remount send recovery', () => {
     }
   )
 
-  it.each([false, true])(
-    'keeps the owning optimistic user through lagging resource history (model output: %s)',
-    async (hasOutput) => {
+  it.each(
+    [false, true].flatMap((hasOutput) =>
+      (['assistant', 'user', 'neither'] as const).map((retained) => ({ hasOutput, retained }))
+    )
+  )(
+    'keeps the live turn through lagging resource history (output: $hasOutput, retained: $retained)',
+    async ({ hasOutput, retained }) => {
       state.postBehavior = 'task'
       const history: MothershipChatHistory = {
         id: 'chat-lagging-user',
@@ -1208,19 +1212,22 @@ describe('useChat remount send recovery', () => {
         .messages.find((message) => message.role === 'assistant')!
       await act(async () => getResult().addResource(search))
       await waitFor(() => historyReads > 0)
+      const beforeHistoryRefresh = getResult()
       await act(async () =>
         staleHistory.resolve({
           chat: {
             ...history,
             title: 'Hydrated during run',
             activeStreamId: sentUser.id,
-            messages: [liveAssistant],
+            messages:
+              retained === 'assistant' ? [liveAssistant] : retained === 'user' ? [sentUser] : [],
             resources: [search],
           },
         })
       )
       await waitFor(
         () =>
+          getResult() !== beforeHistoryRefresh &&
           queryClient.getQueryData<MothershipChatHistory>(mothershipChatKeys.detail(history.id))
             ?.title === 'Hydrated during run'
       )
@@ -1229,6 +1236,9 @@ describe('useChat remount send recovery', () => {
         liveAssistant.id,
       ])
       expect(getResult().messages[0].content).toBe('Orion acceptance policy')
+      expect(getResult().isSending).toBe(true)
+      expect(getResult().messages[1].content).toBe(liveAssistant.content)
+      if (hasOutput) expect(getResult().messages[1].contentBlocks).toHaveLength(1)
       await act(async () =>
         queryClient.setQueryData(mothershipChatKeys.detail(history.id), {
           ...history,
@@ -1411,6 +1421,62 @@ describe('useChat remount send recovery', () => {
     expect(state.postBodies[0]).toHaveProperty('effort')
   })
 
+  it('recovers a running turn after reconnect exhaustion without reloading or resending', async () => {
+    vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] })
+    try {
+      let online = false
+      let recoveredTail = false
+      let failedReconnects = 0
+      const history: MothershipChatHistory = {
+        id: 'chat-reconnect-exhausted',
+        mode: 'agent',
+        title: 'Reconnect',
+        messages: [],
+        activeStreamId: null,
+        resources: [],
+      }
+      mockRequestJson.mockImplementation(() =>
+        Promise.resolve({
+          chat: {
+            ...history,
+            activeStreamId: state.postBodies[0]?.userMessageId ?? null,
+          },
+        })
+      )
+      state.postBehavior = 'accept'
+      vi.stubGlobal('fetch', async (input: RequestInfo | URL, init?: RequestInit) => {
+        const url = String(input)
+        if (!url.includes('/api/mothership/chat/stream')) return fetchStub(input, init)
+        if (!online) {
+          failedReconnects++
+          throw new TypeError('Failed to fetch')
+        }
+        if (url.includes('batch=true')) {
+          return Response.json({ success: true, events: [], status: 'streaming' })
+        }
+        recoveredTail = true
+        return new Response(new ReadableStream<Uint8Array>(), {
+          headers: { 'Content-Type': 'text/event-stream' },
+        })
+      })
+      const { getResult } = renderUseChatInChat(history.id, history)
+      await act(async () => {
+        void getResult().sendMessage('Continue working')
+      })
+      for (let second = 0; second < 240 && failedReconnects < 12; second++) {
+        await act(async () => vi.advanceTimersByTimeAsync(1_000))
+      }
+      expect(failedReconnects).toBeGreaterThanOrEqual(12)
+      online = true
+      await act(async () => vi.advanceTimersByTimeAsync(30_000))
+      expect(recoveredTail).toBe(true)
+      expect(getResult().isSending).toBe(true)
+      expect(state.postBodies).toHaveLength(1)
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
   it('preserves a visible workflow watch when Stop persists the partial response', async () => {
     state.postBehavior = 'task'
     const { getResult } = renderUseChatInChat('chat-a')

apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts

@@ -187,6 +187,13 @@ export interface SendMessageOptions {
   assistantSearchLevel?: AssistantSearchLevel
 }
 
+interface FinalizeOptions {
+  error?: boolean
+  targetChatId?: string
+  /** A lost transport must remain recoverable while the server is still running. */
+  streamTerminal?: boolean
+}
+
 /**
  * `true` when the send owns the transcript (rendered, or handed to reconnect),
  * `false` when the caller should restore the queue entry, and the object form
@@ -883,9 +890,7 @@ export function useChat(
   const resolveDetachedChatForStreamRef = useRef<
     (streamId: string, signal?: AbortSignal) => Promise<DetachedChatResolution>
   >(async () => ({ terminal: false }))
-  const finalizeRef = useRef<(options?: { error?: boolean; targetChatId?: string }) => void>(
-    () => {}
-  )
+  const finalizeRef = useRef<(options?: FinalizeOptions) => void>(() => {})
   const recoveringQueuedSendHandoffRef = useRef<ActiveQueuedSendHandoffRecovery | null>(null)
   const recoverActiveStreamRef = useRef<
     (reason: 'pageshow' | 'visible' | 'online' | 'exhausted_recheck') => Promise<void>
@@ -1238,21 +1243,34 @@ export function useChat(
     options?.requestMode ?? chatHistory?.mode ?? (organizationId ? 'assistant' : 'agent')
   const pendingTurn =
     chatHistory?.id === (initialChatId ?? chatIdRef.current) ? activeTurnRef.current : null
+  const liveContent = streamingContentRef.current
+  const liveBlocks = streamingBlocksRef.current
   const messages = useMemo(() => {
     const source = chatHistory?.messages.map(toDisplayMessage) ?? [...pendingMessages]
-    /** A resource-history read can lag admission; keep this chat's own optimistic user visible. */
-    if (pendingTurn && !source.some((message) => message.id === pendingTurn.userMessageId)) {
-      const assistantIndex = source.findIndex(
-        (message) => message.id === pendingTurn.assistantMessageId
-      )
-      source.splice(
-        assistantIndex < 0 ? source.length : assistantIndex,
-        0,
-        pendingTurn.optimisticUserMessage
-      )
+    /** History reads can lag the live turn; retain both its owner and current response. */
+    if (pendingTurn) {
+      let ownerIndex = source.findIndex((message) => message.id === pendingTurn.userMessageId)
+      if (ownerIndex < 0) {
+        const assistantIndex = source.findIndex(
+          (message) => message.id === pendingTurn.assistantMessageId
+        )
+        ownerIndex = assistantIndex < 0 ? source.length : assistantIndex
+        source.splice(ownerIndex, 0, pendingTurn.optimisticUserMessage)
+      }
+      const assistant = source[ownerIndex + 1]
+      const liveAssistant = {
+        ...pendingTurn.optimisticAssistantMessage,
+        content: liveContent,
+        contentBlocks: liveBlocks,
+      }
+      if (assistant?.id === pendingTurn.assistantMessageId) {
+        source[ownerIndex + 1] = { ...assistant, ...liveAssistant }
+      } else if (assistant?.role !== 'assistant') {
+        source.splice(ownerIndex + 1, 0, liveAssistant)
+      }
     }
     return source.map((m) => restoreRevealedSimKeysForMessage(m, revealedSimKeysRef.current))
-  }, [chatHistory, pendingMessages, pendingTurn])
+  }, [chatHistory, pendingMessages, pendingTurn, liveContent, liveBlocks])
   const addResource = useCallback(
     (resourceUpdate: MothershipResourceUpdate): boolean => {
       // The single fan-in for tab creation, so the invariant lives here.
@@ -2006,7 +2024,11 @@ export function useChat(
         }
         if (!succeeded && streamGenRef.current === gen) {
           try {
-            finalizeRef.current({ error: true, targetChatId: chatHistory.id })
+            finalizeRef.current({
+              error: true,
+              targetChatId: chatHistory.id,
+              streamTerminal: false,
+            })
           } catch {
             setTransportIdle()
             abortControllerRef.current = null
@@ -2941,7 +2963,7 @@ export function useChat(
             shouldContinue: isSameRecoverySubject,
           })
           if (!succeeded && streamGenRef.current === recoveryGen && isSameRecoverySubject()) {
-            finalizeRef.current({ error: true, targetChatId: chatId })
+            finalizeRef.current({ error: true, targetChatId: chatId, streamTerminal: false })
           }
         }
       })()
@@ -3170,7 +3192,7 @@ export function useChat(
   )
 
   const finalize = useCallback(
-    (options?: { error?: boolean; targetChatId?: string }) => {
+    (options?: FinalizeOptions) => {
       const isError = !!options?.error
       if (isError) {
         const blocks = streamingBlocksRef.current
@@ -3219,8 +3241,10 @@ export function useChat(
       if (completedActivityTracker?.generation === streamGenRef.current) {
         clearResourceActivity(completedActivityTracker, true)
       }
-      locallyTerminalStreamIdRef.current =
-        streamIdRef.current ?? activeTurnRef.current?.userMessageId ?? undefined
+      if (options?.streamTerminal !== false) {
+        locallyTerminalStreamIdRef.current =
+          streamIdRef.current ?? activeTurnRef.current?.userMessageId ?? undefined
+      }
       clearActiveTurn()
       setTransportIdle()
       abortControllerRef.current = null
@@ -3804,6 +3828,7 @@ export function useChat(
         if (gen !== undefined && streamGenRef.current === gen) {
           finalize({
             error: true,
+            streamTerminal: false,
             ...(streamTargetChatId ? { targetChatId: streamTargetChatId } : {}),
           })
         }