improvement(mothership): streaming state transitions (#4439)

* improvement(mothership): improve streaming state transitions

* address comments

apps/sim/app/api/copilot/chat/stop/route.test.ts

@@ -29,6 +29,16 @@ const {
   mockSql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })),
 }))
 
+vi.mock('@sim/db/schema', () => ({
+  copilotChats: {
+    id: 'copilotChats.id',
+    userId: 'copilotChats.userId',
+    workspaceId: 'copilotChats.workspaceId',
+    messages: 'copilotChats.messages',
+    conversationId: 'copilotChats.conversationId',
+  },
+}))
+
 vi.mock('@sim/db', () => ({
   db: {
     select: mockSelect,
@@ -140,6 +150,7 @@ describe('copilot chat stop route', () => {
       workspaceId: 'ws-1',
       chatId: 'chat-1',
       type: 'completed',
+      streamId: 'stream-1',
     })
   })
 })

apps/sim/app/api/copilot/chat/stop/route.ts

@@ -111,6 +111,7 @@ export const POST = withRouteHandler((req: NextRequest) =>
           workspaceId: updated.workspaceId,
           chatId,
           type: 'completed',
+          streamId,
         })
       }
 

apps/sim/app/api/copilot/chat/stream/route.ts

@@ -248,6 +248,7 @@ async function handleResumeRequestBody({
       events: batchEvents,
       previewSessions,
       status: run.status,
+      ...(run.chatId ? { chatId: run.chatId } : {}),
     })
   }
 

apps/sim/app/api/mothership/events/route.ts

@@ -27,6 +27,7 @@ const mothershipEventsHandler = createWorkspaceSSE({
           send('task_status', {
             chatId: event.chatId,
             type: event.type,
+            ...(event.streamId ? { streamId: event.streamId } : {}),
             timestamp: Date.now(),
           })
         })

apps/sim/hooks/use-task-events.test.ts

@@ -9,14 +9,46 @@ import { handleTaskStatusEvent } from '@/hooks/use-task-events'
 
 describe('handleTaskStatusEvent', () => {
   const queryClient = {
+    getQueryData: vi.fn(),
     invalidateQueries: vi.fn().mockResolvedValue(undefined),
-  } satisfies Pick<QueryClient, 'invalidateQueries'>
+    removeQueries: vi.fn(),
+  } satisfies Pick<QueryClient, 'getQueryData' | 'invalidateQueries' | 'removeQueries'>
 
   beforeEach(() => {
     vi.clearAllMocks()
+    queryClient.getQueryData.mockReturnValue(undefined)
   })
 
-  it('invalidates only the task list for completed task events', () => {
+  it('invalidates the task list and detail for completed task events', () => {
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'completed',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(2)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.detail('chat-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('keeps completed task detail when an unkeyed completion races an active stream', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'new-stream' }, { id: 'live-assistant:new-stream' }],
+      activeStreamId: 'new-stream',
+      resources: [],
+    })
+
     handleTaskStatusEvent(
       queryClient,
       'ws-1',
@@ -31,9 +63,225 @@ describe('handleTaskStatusEvent', () => {
     expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
       queryKey: taskKeys.list('ws-1'),
     })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
   })
 
-  it('keeps list invalidation only for non-completed task events', () => {
+  it('keeps completed task detail when a newer optimistic stream is active', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'old-stream' }, { id: 'new-stream' }],
+      activeStreamId: 'new-stream',
+      resources: [],
+    })
+
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'completed',
+        streamId: 'old-stream',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('keeps completed task detail when only a newer optimistic stream is cached', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'new-stream' }, { id: 'live-assistant:new-stream' }],
+      activeStreamId: 'new-stream',
+      resources: [],
+    })
+
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'completed',
+        streamId: 'old-stream',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('invalidates completed task detail when the active stream disagreement is only stale cache', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'new-stream' }, { id: 'old-stream' }],
+      activeStreamId: 'new-stream',
+      resources: [],
+    })
+
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'completed',
+        streamId: 'old-stream',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(2)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.detail('chat-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('invalidates completed task detail when a missing stream may be newer server state', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'old-stream' }],
+      activeStreamId: 'old-stream',
+      resources: [],
+    })
+
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'completed',
+        streamId: 'new-stream',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(2)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.detail('chat-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('invalidates completed task detail when the completed stream is active', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [],
+      activeStreamId: 'stream-1',
+      resources: [],
+    })
+
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'completed',
+        streamId: 'stream-1',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(2)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.detail('chat-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('invalidates the task list and detail for metadata-changing task events', () => {
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'renamed',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(2)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.detail('chat-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('invalidates the task list and removes detail cache for deleted task events', () => {
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'deleted',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.removeQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.removeQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.detail('chat-1'),
+    })
+  })
+
+  it('invalidates the task list and detail for started task events', () => {
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'started',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(2)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.detail('chat-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('keeps started task detail when an unkeyed started event races an active stream', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'new-stream' }, { id: 'live-assistant:new-stream' }],
+      activeStreamId: 'new-stream',
+      resources: [],
+    })
+
     handleTaskStatusEvent(
       queryClient,
       'ws-1',
@@ -48,11 +296,115 @@ describe('handleTaskStatusEvent', () => {
     expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
       queryKey: taskKeys.list('ws-1'),
     })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('keeps started task detail when the started stream is already active', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'stream-1' }],
+      activeStreamId: 'stream-1',
+      resources: [],
+    })
+
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'started',
+        streamId: 'stream-1',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('keeps started task detail when a stale started stream is older than the active stream', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'old-stream' }, { id: 'new-stream' }],
+      activeStreamId: 'new-stream',
+      resources: [],
+    })
+
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'started',
+        streamId: 'old-stream',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('invalidates started task detail when a missing stream may be newer server state', () => {
+    queryClient.getQueryData.mockReturnValue({
+      id: 'chat-1',
+      title: null,
+      messages: [{ id: 'old-stream' }],
+      activeStreamId: 'old-stream',
+      resources: [],
+    })
+
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'started',
+        streamId: 'new-stream',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(2)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.detail('chat-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
+  })
+
+  it('keeps list invalidation only for unknown task event types', () => {
+    handleTaskStatusEvent(
+      queryClient,
+      'ws-1',
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'archived',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.list('ws-1'),
+    })
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
   })
 
   it('does not invalidate when task event payload is invalid', () => {
     handleTaskStatusEvent(queryClient, 'ws-1', '{')
 
     expect(queryClient.invalidateQueries).not.toHaveBeenCalled()
+    expect(queryClient.removeQueries).not.toHaveBeenCalled()
   })
 })

apps/sim/hooks/use-task-events.ts

@@ -2,13 +2,68 @@ import { useEffect } from 'react'
 import { createLogger } from '@sim/logger'
 import type { QueryClient } from '@tanstack/react-query'
 import { useQueryClient } from '@tanstack/react-query'
-import { taskKeys } from '@/hooks/queries/tasks'
+import { getLiveAssistantMessageId } from '@/lib/copilot/chat/effective-transcript'
+import { type TaskChatHistory, taskKeys } from '@/hooks/queries/tasks'
 
 const logger = createLogger('TaskEvents')
 
+const TASK_STATUS_TYPES = ['started', 'completed', 'created', 'deleted', 'renamed'] as const
+type TaskStatusEventType = (typeof TASK_STATUS_TYPES)[number]
+const TASK_STATUS_TYPE_SET = new Set<string>(TASK_STATUS_TYPES)
+
 interface TaskStatusEventPayload {
   chatId?: string
-  type?: 'started' | 'completed' | 'created' | 'deleted' | 'renamed'
+  type?: TaskStatusEventType
+  streamId?: string
+}
+
+const DETAIL_INVALIDATING_TASK_STATUS_TYPES = new Set<TaskStatusEventType>([
+  'started',
+  'completed',
+  'renamed',
+])
+
+function isTaskStatusEventType(value: unknown): value is TaskStatusEventType {
+  return typeof value === 'string' && TASK_STATUS_TYPE_SET.has(value)
+}
+
+function isLocalOptimisticActiveStream(current: TaskChatHistory | undefined) {
+  if (!current?.activeStreamId) return false
+  const liveAssistantId = getLiveAssistantMessageId(current.activeStreamId)
+  return current.messages.some((message) => message.id === liveAssistantId)
+}
+
+/**
+ * Returns true when the cached active stream is known to be later in the
+ * chronological transcript than the stream that emitted this status event.
+ * If either stream is absent from the transcript, callers should refetch
+ * instead of inferring order from incomplete cache state.
+ */
+function hasNewerKnownActiveStream(current: TaskChatHistory | undefined, streamId: string) {
+  if (!current?.activeStreamId || current.activeStreamId === streamId) return false
+
+  const activeIndex = current.messages.findIndex((message) => message.id === current.activeStreamId)
+  const eventStreamIndex = current.messages.findIndex((message) => message.id === streamId)
+  if (activeIndex === -1) return false
+  if (eventStreamIndex === -1) return false
+  return activeIndex > eventStreamIndex
+}
+
+function shouldSkipDetailInvalidationForStreamEvent(
+  current: TaskChatHistory | undefined,
+  payload: TaskStatusEventPayload
+) {
+  if (payload.type !== 'started' && payload.type !== 'completed') return false
+  if (!current?.activeStreamId) return false
+  if (!payload.streamId) return 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' &&
+    isLocalOptimisticActiveStream(current) &&
+    !current.messages.some((message) => message.id === payload.streamId)
+  )
 }
 
 function parseTaskStatusEventPayload(data: unknown): TaskStatusEventPayload | null {
@@ -30,14 +85,13 @@ function parseTaskStatusEventPayload(data: unknown): TaskStatusEventPayload | nu
 
   return {
     ...(typeof record.chatId === 'string' ? { chatId: record.chatId } : {}),
-    ...(typeof record.type === 'string'
-      ? { type: record.type as TaskStatusEventPayload['type'] }
-      : {}),
+    ...(isTaskStatusEventType(record.type) ? { type: record.type } : {}),
+    ...(typeof record.streamId === 'string' ? { streamId: record.streamId } : {}),
   }
 }
 
 export function handleTaskStatusEvent(
-  queryClient: Pick<QueryClient, 'invalidateQueries'>,
+  queryClient: Pick<QueryClient, 'getQueryData' | 'invalidateQueries' | 'removeQueries'>,
   workspaceId: string,
   data: unknown
 ): void {
@@ -48,6 +102,20 @@ export function handleTaskStatusEvent(
   }
 
   queryClient.invalidateQueries({ queryKey: taskKeys.list(workspaceId) })
+  if (!payload.chatId) return
+  if (payload.type === 'deleted') {
+    queryClient.removeQueries({ queryKey: taskKeys.detail(payload.chatId) })
+    return
+  }
+  if (payload.type === 'started' || payload.type === 'completed') {
+    const current = queryClient.getQueryData<TaskChatHistory>(taskKeys.detail(payload.chatId))
+    if (shouldSkipDetailInvalidationForStreamEvent(current, payload)) {
+      return
+    }
+  }
+  if (payload.type && DETAIL_INVALIDATING_TASK_STATUS_TYPES.has(payload.type)) {
+    queryClient.invalidateQueries({ queryKey: taskKeys.detail(payload.chatId) })
+  }
 }
 
 /**

apps/sim/lib/copilot/chat/post.ts

@@ -329,6 +329,7 @@ async function persistUserMessage(params: {
           workspaceId,
           chatId,
           type: 'started',
+          streamId: userMessageId,
         })
       }
 
@@ -430,6 +431,7 @@ function buildOnComplete(params: {
           workspaceId,
           chatId,
           type: 'completed',
+          streamId: userMessageId,
         })
       }
     } catch (error) {
@@ -461,6 +463,7 @@ function buildOnError(params: {
           workspaceId,
           chatId,
           type: 'completed',
+          streamId: userMessageId,
         })
       }
     } catch (error) {

apps/sim/lib/copilot/tasks.ts

@@ -13,6 +13,7 @@ interface TaskStatusEvent {
   workspaceId: string
   chatId: string
   type: 'started' | 'completed' | 'created' | 'deleted' | 'renamed'
+  streamId?: string
 }
 
 const channel =

apps/sim/lib/mothership/inbox/executor.ts

@@ -131,11 +131,14 @@ export async function executeInboxTask(taskId: string): Promise<void> {
       })
     }
 
+    const userMessageId = generateId()
+
     if (chatId) {
       taskPubSub?.publishStatusChanged({
         workspaceId: ws.id,
         chatId,
         type: 'started',
+        streamId: userMessageId,
       })
     }
 
@@ -178,7 +181,6 @@ export async function executeInboxTask(taskId: string): Promise<void> {
     }
     const messageContent = formatEmailAsMessage(truncatedTask, attachments)
 
-    const userMessageId = generateId()
     const requestPayload: Record<string, unknown> = {
       message: messageContent,
       userId,
@@ -244,6 +246,7 @@ export async function executeInboxTask(taskId: string): Promise<void> {
         workspaceId: ws.id,
         chatId,
         type: 'completed',
+        streamId: userMessageId,
       })
     }