fix use chat race + headless trace issues

icecrasher321 · icecrasher321 · commit 3e21899911db · 2026-04-13T10:42:00.000-07:00
diff --git a/apps/sim/app/api/mcp/copilot/route.ts b/apps/sim/app/api/mcp/copilot/route.ts
@@ -18,7 +18,7 @@ import { type NextRequest, NextResponse } from 'next/server'
 import { validateOAuthAccessToken } from '@/lib/auth/oauth-token'
 import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
 import { ORCHESTRATION_TIMEOUT_MS, SIM_AGENT_API_URL } from '@/lib/copilot/constants'
-import { runCopilotLifecycle } from '@/lib/copilot/request/lifecycle/run'
+import { runHeadlessCopilotLifecycle } from '@/lib/copilot/request/lifecycle/headless'
 import { orchestrateSubagentStream } from '@/lib/copilot/request/subagent'
 import { ensureHandlersRegistered, executeTool } from '@/lib/copilot/tool-executor'
 import { prepareExecutionContext } from '@/lib/copilot/tools/handlers/context'
@@ -724,7 +724,7 @@ async function handleBuildToolCall(
       chatId,
     }
 
-    const result = await runCopilotLifecycle(requestPayload, {
+    const result = await runHeadlessCopilotLifecycle(requestPayload, {
       userId,
       workflowId: resolved.workflowId,
       chatId,
diff --git a/apps/sim/app/api/mothership/execute/route.ts b/apps/sim/app/api/mothership/execute/route.ts
@@ -4,7 +4,7 @@ import { z } from 'zod'
 import { checkInternalAuth } from '@/lib/auth/hybrid'
 import { buildIntegrationToolSchemas } from '@/lib/copilot/chat/payload'
 import { generateWorkspaceContext } from '@/lib/copilot/chat/workspace-context'
-import { runCopilotLifecycle } from '@/lib/copilot/request/lifecycle/run'
+import { runHeadlessCopilotLifecycle } from '@/lib/copilot/request/lifecycle/headless'
 import { generateId } from '@/lib/core/utils/uuid'
 import {
   assertActiveWorkspaceAccess,
@@ -72,7 +72,7 @@ export async function POST(req: NextRequest) {
       ...(userPermission ? { userPermission } : {}),
     }
 
-    const result = await runCopilotLifecycle(requestPayload, {
+    const result = await runHeadlessCopilotLifecycle(requestPayload, {
       userId,
       workspaceId,
       chatId: effectiveChatId,
diff --git a/apps/sim/app/api/v1/copilot/chat/route.ts b/apps/sim/app/api/v1/copilot/chat/route.ts
@@ -2,7 +2,7 @@ import { createLogger } from '@sim/logger'
 import { type NextRequest, NextResponse } from 'next/server'
 import { z } from 'zod'
 import { COPILOT_REQUEST_MODES } from '@/lib/copilot/constants'
-import { runCopilotLifecycle } from '@/lib/copilot/request/lifecycle/run'
+import { runHeadlessCopilotLifecycle } from '@/lib/copilot/request/lifecycle/headless'
 import { generateId } from '@/lib/core/utils/uuid'
 import { getWorkflowById, resolveWorkflowIdForUser } from '@/lib/workflows/utils'
 import { authenticateV1Request } from '@/app/api/v1/auth'
@@ -82,7 +82,7 @@ export async function POST(req: NextRequest) {
     // Always generate a chatId - required for artifacts system to work with subagents
     const chatId = parsed.chatId || generateId()
 
-    messageId = crypto.randomUUID()
+    messageId = generateId()
     logger.info(
       messageId
         ? `Received headless copilot chat start request [messageId:${messageId}]`
@@ -106,7 +106,7 @@ export async function POST(req: NextRequest) {
       chatId,
     }
 
-    const result = await runCopilotLifecycle(requestPayload, {
+    const result = await runHeadlessCopilotLifecycle(requestPayload, {
       userId: auth.userId,
       workflowId: resolved.workflowId,
       chatId,
diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts
@@ -1,7 +1,7 @@
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import { createLogger } from '@sim/logger'
 import { useQueryClient } from '@tanstack/react-query'
-import { usePathname } from 'next/navigation'
+import { usePathname, useRouter } from 'next/navigation'
 import { toDisplayMessage } from '@/lib/copilot/chat/display-message'
 import type {
   PersistedFileAttachment,
@@ -709,6 +709,7 @@ export function useChat(
   options?: UseChatOptions
 ): UseChatReturn {
   const pathname = usePathname()
+  const router = useRouter()
   const queryClient = useQueryClient()
   const [messages, setMessages] = useState<ChatMessage[]>([])
   const [isSending, setIsSending] = useState(false)
@@ -962,6 +963,8 @@ export function useChat(
   const recoveringClientWorkflowToolIdsRef = useRef<Set<string>>(new Set())
   const executionStream = useExecutionStream()
   const isHomePage = pathname.endsWith('/home')
+  const wasHomePageRef = useRef(isHomePage)
+  const pendingHomeResetRef = useRef(false)
 
   const setTransportIdle = useCallback(() => {
     sendingRef.current = false
@@ -994,6 +997,25 @@ export function useChat(
     resetStreamingBuffers()
   }, [resetStreamingBuffers])
 
+  const resetHomeChatState = useCallback(() => {
+    streamGenRef.current++
+    chatIdRef.current = undefined
+    lastCursorRef.current = '0'
+    locallyTerminalStreamIdRef.current = undefined
+    clearActiveTurn()
+    setResolvedChatId(undefined)
+    appliedChatHistoryKeyRef.current = undefined
+    abortControllerRef.current = null
+    setMessages([])
+    setError(null)
+    setTransportIdle()
+    setResources([])
+    setActiveResourceId(null)
+    resetEphemeralPreviewState()
+    setMessageQueue([])
+    clearQueueDispatchState()
+  }, [clearActiveTurn, clearQueueDispatchState, resetEphemeralPreviewState, setTransportIdle])
+
   const mergeServerMessagesWithActiveTurn = useCallback(
     (serverMessages: ChatMessage[], previousMessages: ChatMessage[]) => {
       const activeTurn = activeTurnRef.current
@@ -1196,31 +1218,33 @@ export function useChat(
   ])
 
   useEffect(() => {
-    if (workflowIdRef.current) return
-    if (!isHomePage || !chatIdRef.current) return
-    streamGenRef.current++
-    chatIdRef.current = undefined
-    lastCursorRef.current = '0'
-    locallyTerminalStreamIdRef.current = undefined
-    clearActiveTurn()
-    setResolvedChatId(undefined)
-    appliedChatHistoryKeyRef.current = undefined
-    abortControllerRef.current = null
-    setMessages([])
-    setError(null)
-    setTransportIdle()
-    setResources([])
-    setActiveResourceId(null)
-    resetEphemeralPreviewState()
-    setMessageQueue([])
-    clearQueueDispatchState()
-  }, [
-    isHomePage,
-    resetEphemeralPreviewState,
-    clearQueueDispatchState,
-    clearActiveTurn,
-    setTransportIdle,
-  ])
+    const wasHomePage = wasHomePageRef.current
+    wasHomePageRef.current = isHomePage
+
+    if (!isHomePage) {
+      pendingHomeResetRef.current = false
+      return
+    }
+    if (workflowIdRef.current || !chatIdRef.current) return
+
+    const shouldHandleHomeReset = pendingHomeResetRef.current || !wasHomePage
+    if (!shouldHandleHomeReset) return
+
+    const hasActiveTransport =
+      isSending ||
+      sendingRef.current ||
+      isReconnecting ||
+      abortControllerRef.current !== null ||
+      streamReaderRef.current !== null
+
+    if (hasActiveTransport) {
+      pendingHomeResetRef.current = true
+      return
+    }
+
+    pendingHomeResetRef.current = false
+    resetHomeChatState()
+  }, [isHomePage, isReconnecting, isSending, resetHomeChatState])
 
   useEffect(() => {
     if (!chatHistory) return
@@ -1583,11 +1607,9 @@ export function useChat(
                     })
                   }
                   if (!workflowIdRef.current) {
-                    window.history.replaceState(
-                      null,
-                      '',
-                      `/workspace/${workspaceId}/task/${payloadChatId}`
-                    )
+                    router.replace(`/workspace/${workspaceId}/task/${payloadChatId}`, {
+                      scroll: false,
+                    })
                   }
                 }
               }
@@ -2153,6 +2175,7 @@ export function useChat(
     },
     [
       workspaceId,
+      router,
       queryClient,
       addResource,
       removeResource,
diff --git a/apps/sim/lib/copilot/request/go/stream.test.ts b/apps/sim/lib/copilot/request/go/stream.test.ts
@@ -3,6 +3,7 @@
  */
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import {
+  MothershipStreamV1CompletionStatus,
   MothershipStreamV1EventType,
   MothershipStreamV1ToolExecutor,
   MothershipStreamV1ToolMode,
@@ -145,8 +146,18 @@ describe('copilot go stream helpers', () => {
         output: { value: 'ok' },
       },
     })
+    const complete = createEvent({
+      streamId: 'stream-1',
+      cursor: '2',
+      seq: 2,
+      requestId: 'req-1',
+      type: MothershipStreamV1EventType.complete,
+      payload: {
+        status: MothershipStreamV1CompletionStatus.complete,
+      },
+    })
 
-    vi.mocked(fetch).mockResolvedValueOnce(createSseResponse([toolResult, toolResult]))
+    vi.mocked(fetch).mockResolvedValueOnce(createSseResponse([toolResult, toolResult, complete]))
 
     const onEvent = vi.fn()
     const context = createStreamingContext()
@@ -160,7 +171,10 @@ describe('copilot go stream helpers', () => {
       timeout: 1000,
     })
 
-    expect(onEvent).toHaveBeenCalledTimes(1)
+    expect(onEvent.mock.calls.map(([event]) => event.type)).toEqual([
+      MothershipStreamV1EventType.tool,
+      MothershipStreamV1EventType.complete,
+    ])
     expect(onEvent).toHaveBeenCalledWith(
       expect.objectContaining({
         type: MothershipStreamV1EventType.tool,
@@ -180,6 +194,39 @@ describe('copilot go stream helpers', () => {
     )
   })
 
+  it('fails closed when the shared stream ends before a terminal event', async () => {
+    const textEvent = createEvent({
+      streamId: 'stream-1',
+      cursor: '1',
+      seq: 1,
+      requestId: 'req-1',
+      type: MothershipStreamV1EventType.text,
+      payload: {
+        channel: 'assistant',
+        text: 'partial response',
+      },
+    })
+
+    vi.mocked(fetch).mockResolvedValueOnce(createSseResponse([textEvent]))
+
+    const context = createStreamingContext()
+    const execContext: ExecutionContext = {
+      userId: 'user-1',
+      workflowId: 'workflow-1',
+    }
+
+    await expect(
+      runStreamLoop('https://example.com/mothership/stream', {}, context, execContext, {
+        timeout: 1000,
+      })
+    ).rejects.toThrow('Copilot backend stream ended before a terminal event')
+    expect(
+      context.errors.some((message) =>
+        message.includes('Copilot backend stream ended before a terminal event')
+      )
+    ).toBe(true)
+  })
+
   it('fails closed when the shared stream receives an invalid event', async () => {
     vi.mocked(fetch).mockResolvedValueOnce(
       createSseResponse([
diff --git a/apps/sim/lib/copilot/request/go/stream.ts b/apps/sim/lib/copilot/request/go/stream.ts
@@ -277,6 +277,18 @@ export async function runStreamLoop(
       }
       return context.streamComplete || undefined
     })
+
+    if (!context.streamComplete && !abortSignal?.aborted && !context.wasAborted) {
+      const streamPath = new URL(fetchUrl).pathname
+      const message = `Copilot backend stream ended before a terminal event on ${streamPath}`
+      context.errors.push(message)
+      logger.error('Copilot backend stream ended before a terminal event', {
+        path: streamPath,
+        requestId: context.requestId,
+        messageId: context.messageId,
+      })
+      throw new CopilotBackendError(message, { status: 503 })
+    }
   } catch (error) {
     if (error instanceof FatalSseEventError && !context.errors.includes(error.message)) {
       context.errors.push(error.message)
diff --git a/apps/sim/lib/copilot/request/lifecycle/headless.test.ts b/apps/sim/lib/copilot/request/lifecycle/headless.test.ts
diff --git a/apps/sim/lib/copilot/request/lifecycle/headless.ts b/apps/sim/lib/copilot/request/lifecycle/headless.ts
diff --git a/apps/sim/lib/mothership/inbox/executor.ts b/apps/sim/lib/mothership/inbox/executor.ts
