fix(mothership): re-read the transcript until the server has saved a finished turn (#8676)

* fix(mothership): re-read the transcript until the server has saved a finished turn

A tab finalizes a turn on the stream's `complete` event and reads the saved
transcript. The server sends that event before it persists the turn and clears
the chat's stream marker, so the read can land in the gap and return the
in-flight copy: the finished stream still listed as active and the answer under
its live id. That copy matches the optimistic one, so the history effect saw
nothing new, and the chat's `completed` event only marks a cached live stream
stale, so the tab kept it until a reload. On the local QA stack this happened
in roughly one turn in four: the finished answer had no Fork action, and the
cached stream marker kept a queued message from draining on a later return.

Finalize now re-reads with a short backoff while the server still lists the
stream it saw end. The first pass joins finalize's own read, so a turn saved in
time is still read once.

* fix(mothership): keep waiting for a slow save, and wait with a follow-up queued

- The re-read gave up after six reads (about 8s), so a slow save left the
  finished turn on its in-flight copy. It now keeps reading on a capped
  backoff for up to two minutes, and still stops as soon as the chat moves on
  (another send or another chat).
- It skipped the wait when a follow-up was queued, but a queued follow-up that
  does not go out at once (held for an edit) left the finished stream listed
  as running in the cache. It now runs for every finished turn; a follow-up
  that does go out ends it, since its send cancels the read and replaces the
  stream.

* improvement(mothership): pace the saved-turn re-read with the shared jittered backoff

Author: waleedlatif1

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

@@ -2138,6 +2138,125 @@ describe('useChat remount send recovery', () => {
         ?.messages.map((message) => message.id)
     ).toEqual(['saved-user', 'saved-assistant'])
   })
+
+  /**
+   * The tab finalizes on the `complete` event, which reaches it before the server
+   * saves the turn. A transcript read in that gap is the server's in-flight copy;
+   * the tab must read again rather than keep it (live ids, the finished stream
+   * still listed as running) until something else happens to refetch. That holds
+   * when the save is slow, and when a follow-up is queued but not yet sent.
+   */
+  it.each([
+    { label: 'right after the first read', unsavedReads: 1, heldFollowUp: false },
+    { label: 'only after a slow save', unsavedReads: 9, heldFollowUp: false },
+    { label: 'with a follow-up queued but held', unsavedReads: 1, heldFollowUp: true },
+  ])(
+    're-reads a transcript fetched before the server saved the finished turn ($label)',
+    async ({ unsavedReads, heldFollowUp }) => {
+      vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] })
+      try {
+        const chatId = `chat-saved-after-complete-${unsavedReads}-${heldFollowUp}`
+        const history: MothershipChatHistory = {
+          id: chatId,
+          mode: 'agent',
+          title: 'Saved late',
+          messages: [],
+          activeStreamId: null,
+          resources: [],
+        }
+        let streamId: string | undefined
+        let completed = false
+        let detailReads = 0
+        mockRequestJson.mockImplementation((contract: AnyApiRouteContract) => {
+          if (contract.path !== '/api/mothership/chats/[chatId]') {
+            return Promise.resolve({ chats: [] })
+          }
+          if (!completed) return Promise.resolve({ chat: history })
+          detailReads++
+          if (detailReads <= unsavedReads && streamId) {
+            return Promise.resolve({
+              chat: {
+                ...history,
+                activeStreamId: streamId,
+                messages: [
+                  { id: streamId, role: 'user', content: 'Summarize the run' },
+                  { id: `live-assistant:${streamId}`, role: 'assistant', content: 'Done.' },
+                ],
+              },
+            })
+          }
+          return Promise.resolve({
+            chat: {
+              ...history,
+              messages: [
+                { id: streamId, role: 'user', content: 'Summarize the run' },
+                { id: 'saved-assistant', role: 'assistant', content: 'Done.' },
+              ],
+            },
+          })
+        })
+        let stream: ReadableStreamDefaultController<Uint8Array> | undefined
+        vi.stubGlobal('fetch', async (input: RequestInfo | URL, init?: RequestInit) => {
+          if (String(input) !== '/api/mothership/chat' || init?.method !== 'POST') {
+            return fetchStub(input, init)
+          }
+          streamId = JSON.parse(String(init.body)).userMessageId
+          return new Response(
+            new ReadableStream<Uint8Array>({
+              start(controller) {
+                stream = controller
+              },
+            }),
+            { headers: { 'Content-Type': 'text/event-stream', 'x-mothership-chat-id': chatId } }
+          )
+        })
+        const { getResult } = renderUseChatInChat(chatId, history)
+
+        await act(async () => {
+          void getResult().sendMessage('Summarize the run')
+          await vi.advanceTimersByTimeAsync(50)
+        })
+        expect(stream).toBeDefined()
+        if (heldFollowUp) {
+          useMothershipQueueStore
+            .getState()
+            .enqueue(chatId, { id: 'held-follow-up', content: 'And the next one' })
+          useMothershipQueueStore.getState().setEditing(chatId, 'held-follow-up')
+        }
+        const emit = (event: Omit<MothershipStreamV1EventEnvelope, 'v' | 'ts' | 'stream'>) =>
+          stream?.enqueue(
+            new TextEncoder().encode(
+              `data: ${JSON.stringify({ v: 1, ts: '', stream: { streamId }, ...event })}\n\n`
+            )
+          )
+        const saved = () =>
+          queryClient
+            .getQueryData<MothershipChatHistory>(mothershipChatKeys.detail(chatId))
+            ?.messages.some((message) => message.id === 'saved-assistant') === true
+        await act(async () => {
+          emit({ seq: 1, type: 'text', payload: { channel: 'assistant', text: 'Done.' } })
+          completed = true
+          emit({ seq: 2, type: 'complete', payload: { status: 'complete' } })
+          stream?.close()
+          await vi.advanceTimersByTimeAsync(50)
+        })
+        for (let second = 0; second < 90 && !saved(); second++) {
+          await act(async () => vi.advanceTimersByTimeAsync(1_000))
+        }
+
+        expect(saved()).toBe(true)
+        expect(
+          queryClient.getQueryData<MothershipChatHistory>(mothershipChatKeys.detail(chatId))
+            ?.activeStreamId
+        ).toBeNull()
+        expect(detailReads).toBe(unsavedReads + 1)
+        expect(getResult().isSending).toBe(false)
+      } finally {
+        vi.useRealTimers()
+      }
+    }
+  )
+
   describe.each([
     {
       kind: 'browser action',

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

@@ -299,6 +299,11 @@ const STREAM_BATCH_FETCH_TIMEOUT_MS = 10_000
 const STREAM_IDLE_TIMEOUT_MS = 45_000
 const STREAM_CHAT_ID_RESOLVE_TIMEOUT_MS = 10_000
 const CHAT_HISTORY_RECOVERY_TIMEOUT_MS = 10_000
+/** Backoff for re-reading a transcript the server has not yet saved a finished turn into. */
+const PERSISTED_TURN_REFETCH_BASE_MS = 250
+const PERSISTED_TURN_REFETCH_MAX_DELAY_MS = 5_000
+/** How long a finished turn's save is waited for; a slow save still lands well inside it. */
+const PERSISTED_TURN_WAIT_MS = 120_000
 const STOP_REQUEST_TIMEOUT_MS = 15_000
 const DETACHED_CHAT_RETRY_BASE_MS = 1000
 const DETACHED_CHAT_RETRY_MAX_MS = 30_000
@@ -992,6 +997,8 @@ export function useChat(
   // the copy-request-ID button functional after refetch).
   const streamRequestIdRef = useRef<string | undefined>(undefined)
   const locallyTerminalStreamIdRef = useRef<string | undefined>(undefined)
+  /** The finished stream whose saved transcript is being waited for, if any. */
+  const persistedTurnWaitRef = useRef<string | null>(null)
   const lastCursorRef = useRef('0')
   const logResyncedStreamIdRef = useRef<string | null>(null)
   const activeStreamReturnRecoveryRef = useRef<ActiveStreamRecovery | null>(null)
@@ -1915,6 +1922,47 @@ export function useChat(
     resetHomeChatState()
   }, [isHomePage, resetHomeChatState])
 
+  /**
+   * This tab finalizes on the stream's `complete` event, which the server sends
+   * before it saves the turn, so the transcript read right after can still be the
+   * in-flight copy: the stream listed as active and the answer under its live id.
+   * That copy matches the optimistic one, so nothing would read it again; re-read
+   * until the saved turn is there. The first pass joins finalize's own read. The
+   * wait ends as soon as this chat moves on: another send, or another chat.
+   */
+  const awaitPersistedTurn = useCallback(
+    async (chatId: string, streamId: string) => {
+      if (persistedTurnWaitRef.current === streamId) return
+      persistedTurnWaitRef.current = streamId
+      const deadline = Date.now() + PERSISTED_TURN_WAIT_MS
+      try {
+        for (let attempt = 0; Date.now() < deadline; attempt++) {
+          if (attempt > 0) {
+            await sleep(
+              backoffWithJitter(attempt, null, {
+                baseMs: PERSISTED_TURN_REFETCH_BASE_MS,
+                maxMs: PERSISTED_TURN_REFETCH_MAX_DELAY_MS,
+              })
+            )
+          }
+          if (locallyTerminalStreamIdRef.current !== streamId || chatIdRef.current !== chatId)
+            return
+          await queryClient.refetchQueries(
+            { queryKey: mothershipChatKeys.detail(chatId), exact: true },
+            { cancelRefetch: false }
+          )
+          const history = queryClient.getQueryData<MothershipChatHistory>(
+            mothershipChatKeys.detail(chatId)
+          )
+          if (history?.activeStreamId !== streamId) return
+        }
+      } finally {
+        if (persistedTurnWaitRef.current === streamId) persistedTurnWaitRef.current = null
+      }
+    },
+    [queryClient]
+  )
+
   useEffect(() => {
     if (!chatHistory) return
 
@@ -3347,9 +3395,12 @@ export function useChat(
       if (completedActivityTracker?.generation === streamGenRef.current) {
         clearResourceActivity(completedActivityTracker, true)
       }
+      const terminalStreamId =
+        options?.streamTerminal !== false
+          ? (streamIdRef.current ?? activeTurnRef.current?.userMessageId ?? undefined)
+          : undefined
       if (options?.streamTerminal !== false) {
-        locallyTerminalStreamIdRef.current =
-          streamIdRef.current ?? activeTurnRef.current?.userMessageId ?? undefined
+        locallyTerminalStreamIdRef.current = terminalStreamId
       }
       clearActiveTurn()
       setTransportIdle()
@@ -3358,9 +3409,13 @@ export function useChat(
         includeDetail: !hasQueuedFollowUp,
         ...(options?.targetChatId ? { targetChatId: options.targetChatId } : {}),
       })
+      if (terminalStreamId && completedChatId) {
+        void awaitPersistedTurn(completedChatId, terminalStreamId)
+      }
       notifyTurnEnded({ error: isError })
     },
     [
+      awaitPersistedTurn,
       clearResourceActivity,
       clearActiveTurn,
       invalidateChatQueries,