fix(chat): mark a completed live-stream detail stale instead of skipping it

Author: waleedlatif1

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

@@ -82,6 +82,7 @@ import {
 } from '@/app/workspace/[workspaceId]/home/hooks/send-handoff'
 import { useChat } from '@/app/workspace/[workspaceId]/home/hooks/use-chat'
 import { type MothershipChatHistory, mothershipChatKeys } from '@/hooks/queries/mothership-chats'
+import { handleMothershipChatStatusEvent } from '@/hooks/use-mothership-chat-events'
 import { useExecutionStore } from '@/stores/execution/store'
 import { useMothershipQueueStore } from '@/stores/mothership-queue/store'
 
@@ -1821,4 +1822,81 @@ describe('useChat remount send recovery', () => {
     expect(getResult().messageQueue.map((entry) => entry.id)).toEqual(['unsent-entry'])
     expect(state.postBodies).toHaveLength(0)
   })
+
+  it('loads the saved transcript once when its own stream completes', async () => {
+    const chatId = 'chat-own-completion'
+    const history: MothershipChatHistory = {
+      id: chatId,
+      mode: 'agent',
+      title: 'Own stream',
+      messages: [],
+      activeStreamId: null,
+      resources: [],
+    }
+    const saved = [
+      { id: 'saved-user', role: 'user', content: 'Summarize the run' },
+      { id: 'saved-assistant', role: 'assistant', content: 'Done.' },
+    ]
+    const detailRequests: string[] = []
+    mockRequestJson.mockImplementation((contract: AnyApiRouteContract) => {
+      if (contract.path !== '/api/mothership/chats/[chatId]') {
+        return Promise.resolve({ chats: [] })
+      }
+      detailRequests.push(chatId)
+      return Promise.resolve({ chat: { ...history, messages: saved } })
+    })
+    let stream: ReadableStreamDefaultController<Uint8Array> | undefined
+    let streamId: string | undefined
+    const emit = (event: Omit<MothershipStreamV1EventEnvelope, 'v' | 'ts' | 'stream'>) =>
+      stream?.enqueue(
+        new TextEncoder().encode(
+          `data: ${JSON.stringify({ v: 1, ts: '', stream: { streamId }, ...event })}\n\n`
+        )
+      )
+    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 waitFor(() => stream !== undefined)
+    emit({ seq: 1, type: 'text', payload: { channel: 'assistant', text: 'Done.' } })
+    await waitFor(
+      () =>
+        queryClient.getQueryData<MothershipChatHistory>(mothershipChatKeys.detail(chatId))
+          ?.activeStreamId === streamId
+    )
+    /** The server publishes `completed` after persisting and before closing the stream. */
+    handleMothershipChatStatusEvent(queryClient, 'ws-1', {
+      chatId,
+      type: 'completed',
+      streamId,
+    })
+    emit({ seq: 2, type: 'complete', payload: { status: 'complete' } })
+    stream?.close()
+
+    await waitFor(() => !getResult().isSending && detailRequests.length > 0)
+    await act(async () => {
+      await sleep(50)
+    })
+    expect(detailRequests).toHaveLength(1)
+    expect(
+      queryClient
+        .getQueryData<MothershipChatHistory>(mothershipChatKeys.detail(chatId))
+        ?.messages.map((message) => message.id)
+    ).toEqual(['saved-user', 'saved-assistant'])
+  })
 })

apps/sim/hooks/use-mothership-chat-events.test.ts

@@ -223,17 +223,22 @@ describe('chat detail refetches driven by status events', () => {
     unsubscribe()
   })
 
-  it('reloads the transcript when a stream this viewer is not rendering completes', async () => {
-    const { queryClient, fetchTranscript, unsubscribe } = mountDetail({
-      ...liveStream,
-      messages: [{ id: 'stream-1' }] as MothershipChatHistory['messages'],
-    })
+  it('reloads the saved transcript when a cached mid-stream detail is opened after completion', async () => {
+    const queryClient = new QueryClient()
+    const fetchTranscript = vi.fn(async () => ({ ...liveStream, activeStreamId: null }))
+    queryClient.setQueryData(mothershipChatKeys.detail('chat-1'), liveStream)
 
     handleMothershipChatStatusEvent(queryClient, 'ws-1', {
       chatId: 'chat-1',
       type: 'completed',
       streamId: 'stream-1',
     })
+    const unsubscribe = new QueryObserver(queryClient, {
+      queryKey: mothershipChatKeys.detail('chat-1'),
+      queryFn: fetchTranscript,
+      staleTime: Number.POSITIVE_INFINITY,
+    }).subscribe(() => {})
+
     await vi.waitFor(() => expect(fetchTranscript).toHaveBeenCalledTimes(1))
     unsubscribe()
   })

apps/sim/hooks/use-mothership-chat-events.ts

@@ -59,22 +59,14 @@ function hasNewerKnownActiveStream(current: MothershipChatHistory | undefined, s
   return activeIndex > eventStreamIndex
 }
 
-/**
- * Returns true when refetching the chat detail for a stream event would only
- * reload the transcript this client already holds or is about to reload. A
- * completion of the viewer's own live stream is skipped because the server
- * persists the turn before closing that stream, and the client's own
- * finalization refetches the detail once it sees the close.
- */
 function shouldSkipDetailInvalidationForStreamEvent(
   current: MothershipChatHistory | undefined,
   payload: ChatStatusEventPayload
 ) {
   if (!current?.activeStreamId) return false
   if (!payload.streamId) return isLocalOptimisticActiveStream(current)
-  if (current.activeStreamId === payload.streamId) {
-    return payload.type === 'started' || isLocalOptimisticActiveStream(current)
-  }
+  if (payload.type === 'started' && current.activeStreamId === payload.streamId) return true
+  if (current.activeStreamId === payload.streamId) return false
   if (hasNewerKnownActiveStream(current, payload.streamId)) return true
   return (
     payload.type === 'completed' &&
@@ -145,7 +137,21 @@ export function handleMothershipChatStatusEvent(
     mothershipChatKeys.detail(payload.chatId)
   )
   if (shouldSkipDetailInvalidationForStreamEvent(current, payload)) return
-  queryClient.invalidateQueries({ queryKey: mothershipChatKeys.detail(payload.chatId) })
+  /**
+   * A completion of the cached live stream only marks the detail stale. The
+   * server persists the turn before closing the stream, so a surface rendering
+   * it refetches through its own finalization; the live message alone cannot
+   * tell this tab's stream from a server-loaded mid-stream snapshot, so any
+   * other cached copy reloads the saved transcript on its next mount.
+   */
+  const completesCachedLiveStream =
+    payload.type === 'completed' &&
+    current?.activeStreamId === payload.streamId &&
+    isLocalOptimisticActiveStream(current)
+  queryClient.invalidateQueries({
+    queryKey: mothershipChatKeys.detail(payload.chatId),
+    ...(completesCachedLiveStream ? { refetchType: 'none' as const } : {}),
+  })
 }
 
 /**