chore: import upstream snapshot with attribution

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

@@ -0,0 +1,203 @@
+/**
+ * @vitest-environment node
+ */
+
+import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing'
+import { NextRequest } from 'next/server'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import {
+  MothershipStreamV1CompletionStatus,
+  MothershipStreamV1EventType,
+} from '@/lib/copilot/generated/mothership-stream-v1'
+
+const { getLatestRunForStream, readEvents, readFilePreviewSessions, checkForReplayGap } =
+  vi.hoisted(() => ({
+    getLatestRunForStream: vi.fn(),
+    readEvents: vi.fn(),
+    readFilePreviewSessions: vi.fn(),
+    checkForReplayGap: vi.fn(),
+  }))
+
+vi.mock('@/lib/copilot/async-runs/repository', () => ({
+  getLatestRunForStream,
+}))
+
+vi.mock('@/lib/copilot/request/session', () => ({
+  readEvents,
+  readFilePreviewSessions,
+  checkForReplayGap,
+  createEvent: (event: Record<string, unknown>) => ({
+    stream: {
+      streamId: event.streamId,
+      cursor: event.cursor,
+    },
+    seq: event.seq,
+    trace: { requestId: event.requestId ?? '' },
+    type: event.type,
+    payload: event.payload,
+  }),
+  encodeSSEEnvelope: (event: Record<string, unknown>) =>
+    new TextEncoder().encode(`data: ${JSON.stringify(event)}\n\n`),
+  encodeSSEComment: (comment: string) => new TextEncoder().encode(`: ${comment}\n\n`),
+  SSE_RESPONSE_HEADERS: {
+    'Content-Type': 'text/event-stream',
+  },
+}))
+
+vi.mock('@/lib/copilot/request/http', () => copilotHttpMock)
+
+import { GET } from './route'
+
+async function readAllChunks(response: Response): Promise<string[]> {
+  const reader = response.body?.getReader()
+  expect(reader).toBeTruthy()
+
+  const chunks: string[] = []
+  while (true) {
+    const { done, value } = await reader!.read()
+    if (done) {
+      break
+    }
+    chunks.push(new TextDecoder().decode(value))
+  }
+  return chunks
+}
+
+describe('copilot chat stream replay route', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValue({
+      userId: 'user-1',
+      isAuthenticated: true,
+    })
+    readEvents.mockResolvedValue([])
+    readFilePreviewSessions.mockResolvedValue([])
+    checkForReplayGap.mockResolvedValue(null)
+  })
+
+  it('returns preview sessions in batch mode', async () => {
+    getLatestRunForStream.mockResolvedValue({
+      status: 'active',
+      executionId: 'exec-1',
+      id: 'run-1',
+    })
+    readFilePreviewSessions.mockResolvedValue([
+      {
+        schemaVersion: 1,
+        id: 'preview-1',
+        streamId: 'stream-1',
+        toolCallId: 'preview-1',
+        status: 'streaming',
+        fileName: 'draft.md',
+        previewText: 'hello',
+        previewVersion: 2,
+        updatedAt: '2026-04-10T00:00:00.000Z',
+      },
+    ])
+
+    const response = await GET(
+      new NextRequest(
+        'http://localhost:3000/api/copilot/chat/stream?streamId=stream-1&after=0&batch=true'
+      )
+    )
+
+    expect(response.status).toBe(200)
+    await expect(response.json()).resolves.toMatchObject({
+      success: true,
+      previewSessions: [
+        expect.objectContaining({
+          id: 'preview-1',
+          previewText: 'hello',
+          previewVersion: 2,
+        }),
+      ],
+      status: 'active',
+    })
+  })
+
+  it('stops replay polling when run becomes cancelled', async () => {
+    getLatestRunForStream
+      .mockResolvedValueOnce({
+        status: 'active',
+        executionId: 'exec-1',
+        id: 'run-1',
+      })
+      .mockResolvedValueOnce({
+        status: 'cancelled',
+        executionId: 'exec-1',
+        id: 'run-1',
+      })
+
+    const response = await GET(
+      new NextRequest('http://localhost:3000/api/copilot/chat/stream?streamId=stream-1&after=0')
+    )
+
+    const chunks = await readAllChunks(response)
+    expect(chunks[0]).toBe(': accepted\n\n')
+    expect(chunks.join('')).toContain(
+      JSON.stringify({
+        status: MothershipStreamV1CompletionStatus.cancelled,
+        reason: 'terminal_status',
+      })
+    )
+    expect(getLatestRunForStream).toHaveBeenCalledTimes(2)
+  })
+
+  it('emits structured terminal replay error when run metadata disappears', async () => {
+    getLatestRunForStream
+      .mockResolvedValueOnce({
+        status: 'active',
+        executionId: 'exec-1',
+        id: 'run-1',
+      })
+      .mockResolvedValueOnce(null)
+
+    const response = await GET(
+      new NextRequest('http://localhost:3000/api/copilot/chat/stream?streamId=stream-1&after=0')
+    )
+
+    const chunks = await readAllChunks(response)
+    const body = chunks.join('')
+    expect(body).toContain(`"type":"${MothershipStreamV1EventType.error}"`)
+    expect(body).toContain('"code":"resume_run_unavailable"')
+    expect(body).toContain(`"type":"${MothershipStreamV1EventType.complete}"`)
+  })
+
+  it('uses the latest live request id for synthetic terminal replay events', async () => {
+    getLatestRunForStream
+      .mockResolvedValueOnce({
+        status: 'active',
+        executionId: 'exec-1',
+        id: 'run-1',
+      })
+      .mockResolvedValueOnce({
+        status: 'cancelled',
+        executionId: 'exec-1',
+        id: 'run-1',
+      })
+    readEvents
+      .mockResolvedValueOnce([
+        {
+          stream: { streamId: 'stream-1', cursor: '1' },
+          seq: 1,
+          trace: { requestId: 'req-live-123' },
+          type: MothershipStreamV1EventType.text,
+          payload: {
+            channel: 'assistant',
+            text: 'hello',
+          },
+        },
+      ])
+      .mockResolvedValueOnce([])
+
+    const response = await GET(
+      new NextRequest('http://localhost:3000/api/copilot/chat/stream?streamId=stream-1&after=0')
+    )
+
+    const chunks = await readAllChunks(response)
+    const terminalChunk = chunks[chunks.length - 1] ?? ''
+    expect(terminalChunk).toContain(`"type":"${MothershipStreamV1EventType.complete}"`)
+    expect(terminalChunk).toContain('"requestId":"req-live-123"')
+    expect(terminalChunk).toContain('"status":"cancelled"')
+  })
+})

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

@@ -0,0 +1,483 @@
+import { type Context, context as otelContext, type Span, trace } from '@opentelemetry/api'
+import { createLogger } from '@sim/logger'
+import { getErrorMessage } from '@sim/utils/errors'
+import { sleep } from '@sim/utils/helpers'
+import { type NextRequest, NextResponse } from 'next/server'
+import { copilotChatStreamContract } from '@/lib/api/contracts/copilot'
+import { parseRequest } from '@/lib/api/server'
+import { getLatestRunForStream } from '@/lib/copilot/async-runs/repository'
+import {
+  MothershipStreamV1CompletionStatus,
+  MothershipStreamV1EventType,
+} from '@/lib/copilot/generated/mothership-stream-v1'
+import {
+  CopilotResumeOutcome,
+  CopilotTransport,
+} from '@/lib/copilot/generated/trace-attribute-values-v1'
+import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
+import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
+import { contextFromRequestHeaders } from '@/lib/copilot/request/go/propagation'
+import { authenticateCopilotRequestSessionOnly } from '@/lib/copilot/request/http'
+import { getCopilotTracer, markSpanForError } from '@/lib/copilot/request/otel'
+import {
+  checkForReplayGap,
+  createEvent,
+  encodeSSEComment,
+  encodeSSEEnvelope,
+  readEvents,
+  readFilePreviewSessions,
+  SSE_RESPONSE_HEADERS,
+} from '@/lib/copilot/request/session'
+import { toStreamBatchEvent } from '@/lib/copilot/request/session/types'
+import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
+
+export const maxDuration = 3600
+
+const logger = createLogger('CopilotChatStreamAPI')
+const POLL_INTERVAL_MS = 250
+const REPLAY_KEEPALIVE_INTERVAL_MS = 15_000
+const MAX_STREAM_MS = 60 * 60 * 1000
+
+function extractCanonicalRequestId(value: unknown): string {
+  return typeof value === 'string' && value.length > 0 ? value : ''
+}
+
+function extractRunRequestId(run: { requestContext?: unknown } | null | undefined): string {
+  if (!run || typeof run.requestContext !== 'object' || run.requestContext === null) {
+    return ''
+  }
+  const requestContext = run.requestContext as Record<string, unknown>
+  return (
+    extractCanonicalRequestId(requestContext.requestId) ||
+    extractCanonicalRequestId(requestContext.simRequestId)
+  )
+}
+
+function extractEnvelopeRequestId(envelope: { trace?: { requestId?: unknown } }): string {
+  return extractCanonicalRequestId(envelope.trace?.requestId)
+}
+
+function isTerminalStatus(
+  status: string | null | undefined
+): status is MothershipStreamV1CompletionStatus {
+  return (
+    status === MothershipStreamV1CompletionStatus.complete ||
+    status === MothershipStreamV1CompletionStatus.error ||
+    status === MothershipStreamV1CompletionStatus.cancelled
+  )
+}
+
+function buildResumeTerminalEnvelopes(options: {
+  streamId: string
+  afterCursor: string
+  status: MothershipStreamV1CompletionStatus
+  message?: string
+  code: string
+  reason?: string
+  requestId?: string
+}) {
+  const baseSeq = Number(options.afterCursor || '0')
+  const seq = Number.isFinite(baseSeq) ? baseSeq : 0
+  const envelopes: ReturnType<typeof createEvent>[] = []
+  const rid = options.requestId ?? ''
+
+  if (options.status === MothershipStreamV1CompletionStatus.error) {
+    envelopes.push(
+      createEvent({
+        streamId: options.streamId,
+        cursor: String(seq + 1),
+        seq: seq + 1,
+        requestId: rid,
+        type: MothershipStreamV1EventType.error,
+        payload: {
+          message: options.message || 'Stream recovery failed before completion.',
+          code: options.code,
+        },
+      })
+    )
+  }
+
+  envelopes.push(
+    createEvent({
+      streamId: options.streamId,
+      cursor: String(seq + envelopes.length + 1),
+      seq: seq + envelopes.length + 1,
+      requestId: rid,
+      type: MothershipStreamV1EventType.complete,
+      payload: {
+        status: options.status,
+        ...(options.reason ? { reason: options.reason } : {}),
+      },
+    })
+  )
+
+  return envelopes
+}
+
+export const GET = withRouteHandler(async (request: NextRequest) => {
+  const { userId: authenticatedUserId, isAuthenticated } =
+    await authenticateCopilotRequestSessionOnly()
+
+  if (!isAuthenticated || !authenticatedUserId) {
+    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
+  }
+
+  const parsed = await parseRequest(copilotChatStreamContract, request, {})
+  if (!parsed.success) return parsed.response
+  const { streamId, after: afterCursor, batch: batchMode } = parsed.data.query
+
+  if (!streamId) {
+    return NextResponse.json({ error: 'streamId is required' }, { status: 400 })
+  }
+
+  // Root span for the whole resume/reconnect request. In stream mode the
+  // work happens inside `ReadableStream.start`, which the Node runtime
+  // invokes after this function returns and OUTSIDE the AsyncLocalStorage
+  // scope installed by `startActiveSpan`. We therefore start the span
+  // manually, capture its context, and re-enter that context inside the
+  // stream callback so every nested `withCopilotSpan` / `withDbSpan` call
+  // attaches to this root.
+  //
+  // `contextFromRequestHeaders` extracts the W3C `traceparent` the
+  // client echoed (set via `streamTraceparentRef` on Sim's chat POST
+  // response), so the resume span becomes a child of the original
+  // chat's `gen_ai.agent.execute` trace instead of a disconnected
+  // new root. On reconnects after page reload (client ref was wiped)
+  // the header is absent and extraction leaves the ambient context
+  // alone → the resume span becomes its own root. Same as pre-
+  // linking behavior; no regression.
+  const incomingContext = contextFromRequestHeaders(request.headers)
+  const rootSpan = getCopilotTracer().startSpan(
+    TraceSpan.CopilotResumeRequest,
+    {
+      attributes: {
+        [TraceAttr.CopilotTransport]: batchMode ? CopilotTransport.Batch : CopilotTransport.Stream,
+        [TraceAttr.StreamId]: streamId,
+        [TraceAttr.UserId]: authenticatedUserId,
+        [TraceAttr.CopilotResumeAfterCursor]: afterCursor || '0',
+      },
+    },
+    incomingContext
+  )
+  const rootContext = trace.setSpan(incomingContext, rootSpan)
+
+  try {
+    return await otelContext.with(rootContext, () =>
+      handleResumeRequestBody({
+        request,
+        streamId,
+        afterCursor,
+        batchMode,
+        authenticatedUserId,
+        rootSpan,
+        rootContext,
+      })
+    )
+  } catch (err) {
+    markSpanForError(rootSpan, err)
+    rootSpan.end()
+    throw err
+  }
+})
+
+async function handleResumeRequestBody({
+  request,
+  streamId,
+  afterCursor,
+  batchMode,
+  authenticatedUserId,
+  rootSpan,
+  rootContext,
+}: {
+  request: NextRequest
+  streamId: string
+  afterCursor: string
+  batchMode: boolean
+  authenticatedUserId: string
+  rootSpan: Span
+  rootContext: Context
+}) {
+  const run = await getLatestRunForStream(streamId, authenticatedUserId).catch((err) => {
+    logger.warn('Failed to fetch latest run for stream', {
+      streamId,
+      error: getErrorMessage(err),
+    })
+    return null
+  })
+  logger.info('[Resume] Stream lookup', {
+    streamId,
+    afterCursor,
+    batchMode,
+    hasRun: !!run,
+    runStatus: run?.status,
+  })
+  if (!run) {
+    rootSpan.setAttribute(TraceAttr.CopilotResumeOutcome, CopilotResumeOutcome.StreamNotFound)
+    rootSpan.end()
+    return NextResponse.json({ error: 'Stream not found' }, { status: 404 })
+  }
+  rootSpan.setAttribute(TraceAttr.CopilotRunStatus, run.status)
+
+  if (batchMode) {
+    const afterSeq = afterCursor || '0'
+    const [events, previewSessions] = await Promise.all([
+      readEvents(streamId, afterSeq),
+      readFilePreviewSessions(streamId).catch((error) => {
+        logger.warn('Failed to read preview sessions for stream batch', {
+          streamId,
+          error: getErrorMessage(error),
+        })
+        return []
+      }),
+    ])
+    const batchEvents = events.map(toStreamBatchEvent)
+    logger.info('[Resume] Batch response', {
+      streamId,
+      afterCursor: afterSeq,
+      eventCount: batchEvents.length,
+      previewSessionCount: previewSessions.length,
+      runStatus: run.status,
+    })
+    rootSpan.setAttributes({
+      [TraceAttr.CopilotResumeOutcome]: CopilotResumeOutcome.BatchDelivered,
+      [TraceAttr.CopilotResumeEventCount]: batchEvents.length,
+      [TraceAttr.CopilotResumePreviewSessionCount]: previewSessions.length,
+    })
+    rootSpan.end()
+    return NextResponse.json({
+      success: true,
+      events: batchEvents,
+      previewSessions,
+      status: run.status,
+      ...(run.chatId ? { chatId: run.chatId } : {}),
+    })
+  }
+
+  const startTime = Date.now()
+  let totalEventsFlushed = 0
+  let pollIterations = 0
+
+  const stream = new ReadableStream({
+    async start(controller) {
+      // Re-enter the root OTel context so any `withCopilotSpan` call below
+      // (inside flushEvents/checkForReplayGap/etc.) parents under
+      // copilot.resume.request instead of becoming an orphan.
+      return otelContext.with(rootContext, () => startInner(controller))
+    },
+  })
+
+  async function startInner(controller: ReadableStreamDefaultController) {
+    let cursor = afterCursor || '0'
+    let controllerClosed = false
+    let sawTerminalEvent = false
+    let currentRequestId = extractRunRequestId(run)
+    let lastWriteTime = Date.now()
+    // Stamp the logical request id + chat id on the resume root as soon
+    // as we resolve them from the run row, so TraceQL joins work on
+    // resume legs the same way they do on the original POST.
+    if (currentRequestId) {
+      rootSpan.setAttribute(TraceAttr.RequestId, currentRequestId)
+      rootSpan.setAttribute(TraceAttr.SimRequestId, currentRequestId)
+    }
+    if (run?.chatId) {
+      rootSpan.setAttribute(TraceAttr.ChatId, run.chatId)
+    }
+
+    const closeController = () => {
+      if (controllerClosed) return
+      controllerClosed = true
+      try {
+        controller.close()
+      } catch {
+        // Controller already closed by runtime/client
+      }
+    }
+
+    const enqueueEvent = (payload: unknown) => {
+      if (controllerClosed) return false
+      try {
+        controller.enqueue(encodeSSEEnvelope(payload))
+        lastWriteTime = Date.now()
+        return true
+      } catch {
+        controllerClosed = true
+        return false
+      }
+    }
+
+    const enqueueComment = (comment: string) => {
+      if (controllerClosed) return false
+      try {
+        controller.enqueue(encodeSSEComment(comment))
+        lastWriteTime = Date.now()
+        return true
+      } catch {
+        controllerClosed = true
+        return false
+      }
+    }
+
+    const abortListener = () => {
+      controllerClosed = true
+    }
+    request.signal.addEventListener('abort', abortListener, { once: true })
+
+    const flushEvents = async () => {
+      const events = await readEvents(streamId, cursor)
+      if (events.length > 0) {
+        logger.debug('[Resume] Flushing events', {
+          streamId,
+          afterCursor: cursor,
+          eventCount: events.length,
+        })
+      }
+      for (const envelope of events) {
+        if (!enqueueEvent(envelope)) {
+          break
+        }
+        totalEventsFlushed += 1
+        cursor = envelope.stream.cursor ?? String(envelope.seq)
+        currentRequestId = extractEnvelopeRequestId(envelope) || currentRequestId
+        if (envelope.type === MothershipStreamV1EventType.complete) {
+          sawTerminalEvent = true
+        }
+      }
+    }
+
+    const emitTerminalIfMissing = (
+      status: MothershipStreamV1CompletionStatus,
+      options?: { message?: string; code: string; reason?: string }
+    ) => {
+      if (controllerClosed || sawTerminalEvent) {
+        return
+      }
+      for (const envelope of buildResumeTerminalEnvelopes({
+        streamId,
+        afterCursor: cursor,
+        status,
+        message: options?.message,
+        code: options?.code ?? 'resume_terminal',
+        reason: options?.reason,
+        requestId: currentRequestId,
+      })) {
+        if (!enqueueEvent(envelope)) {
+          break
+        }
+        cursor = envelope.stream.cursor ?? String(envelope.seq)
+        if (envelope.type === MothershipStreamV1EventType.complete) {
+          sawTerminalEvent = true
+        }
+      }
+    }
+
+    try {
+      enqueueComment('accepted')
+
+      const gap = await checkForReplayGap(streamId, afterCursor, currentRequestId)
+      if (gap) {
+        for (const envelope of gap.envelopes) {
+          if (!enqueueEvent(envelope)) {
+            break
+          }
+          cursor = envelope.stream.cursor ?? String(envelope.seq)
+          currentRequestId = extractEnvelopeRequestId(envelope) || currentRequestId
+          if (envelope.type === MothershipStreamV1EventType.complete) {
+            sawTerminalEvent = true
+          }
+        }
+        return
+      }
+
+      await flushEvents()
+
+      while (!controllerClosed && Date.now() - startTime < MAX_STREAM_MS) {
+        pollIterations += 1
+        const currentRun = await getLatestRunForStream(streamId, authenticatedUserId).catch(
+          (err) => {
+            logger.warn('Failed to poll latest run for stream', {
+              streamId,
+              error: getErrorMessage(err),
+            })
+            return null
+          }
+        )
+        if (!currentRun) {
+          emitTerminalIfMissing(MothershipStreamV1CompletionStatus.error, {
+            message: 'The stream could not be recovered because its run metadata is unavailable.',
+            code: 'resume_run_unavailable',
+            reason: 'run_unavailable',
+          })
+          break
+        }
+
+        currentRequestId = extractRunRequestId(currentRun) || currentRequestId
+
+        await flushEvents()
+
+        if (controllerClosed) {
+          break
+        }
+        if (isTerminalStatus(currentRun.status)) {
+          emitTerminalIfMissing(currentRun.status, {
+            message:
+              currentRun.status === MothershipStreamV1CompletionStatus.error
+                ? typeof currentRun.error === 'string'
+                  ? currentRun.error
+                  : 'The recovered stream ended with an error.'
+                : undefined,
+            code: 'resume_terminal_status',
+            reason: 'terminal_status',
+          })
+          break
+        }
+
+        if (request.signal.aborted) {
+          controllerClosed = true
+          break
+        }
+
+        if (Date.now() - lastWriteTime >= REPLAY_KEEPALIVE_INTERVAL_MS) {
+          enqueueComment('keepalive')
+        }
+
+        await sleep(POLL_INTERVAL_MS)
+      }
+      if (!controllerClosed && Date.now() - startTime >= MAX_STREAM_MS) {
+        emitTerminalIfMissing(MothershipStreamV1CompletionStatus.error, {
+          message: 'The stream recovery timed out before completion.',
+          code: 'resume_timeout',
+          reason: 'timeout',
+        })
+      }
+    } catch (error) {
+      if (!controllerClosed && !request.signal.aborted) {
+        logger.warn('Stream replay failed', {
+          streamId,
+          error: getErrorMessage(error),
+        })
+        emitTerminalIfMissing(MothershipStreamV1CompletionStatus.error, {
+          message: 'The stream replay failed before completion.',
+          code: 'resume_internal',
+          reason: 'stream_replay_failed',
+        })
+      }
+      markSpanForError(rootSpan, error)
+    } finally {
+      request.signal.removeEventListener('abort', abortListener)
+      closeController()
+      rootSpan.setAttributes({
+        [TraceAttr.CopilotResumeOutcome]: sawTerminalEvent
+          ? CopilotResumeOutcome.TerminalDelivered
+          : controllerClosed
+            ? CopilotResumeOutcome.ClientDisconnected
+            : CopilotResumeOutcome.EndedWithoutTerminal,
+        [TraceAttr.CopilotResumeEventCount]: totalEventsFlushed,
+        [TraceAttr.CopilotResumePollIterations]: pollIterations,
+        [TraceAttr.CopilotResumeDurationMs]: Date.now() - startTime,
+      })
+      rootSpan.end()
+    }
+  }
+
+  return new Response(stream, { headers: SSE_RESPONSE_HEADERS })
+}