simstudioai--sim/apps/sim/lib/copilot/request/lifecycle/start.ts

import { type Context, context as otelContextApi } from '@opentelemetry/api'
import { db } from '@sim/db'
import { copilotChats } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { getErrorMessage } from '@sim/utils/errors'
import { eq } from 'drizzle-orm'
import { createRunSegment } from '@/lib/copilot/async-runs/repository'
import { chatPubSub } from '@/lib/copilot/chat-status'
import {
  MothershipStreamV1EventType,
  MothershipStreamV1SessionKind,
} from '@/lib/copilot/generated/mothership-stream-v1'
import {
  RequestTraceV1Outcome,
  RequestTraceV1SpanStatus,
} from '@/lib/copilot/generated/request-trace-v1'
import {
  CopilotRequestCancelReason,
  type CopilotRequestCancelReasonValue,
  CopilotTransport,
} from '@/lib/copilot/generated/trace-attribute-values-v1'
import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
import { TraceEvent } from '@/lib/copilot/generated/trace-events-v1'
import { finalizeStream } from '@/lib/copilot/request/lifecycle/finalize'
import type { CopilotLifecycleOptions } from '@/lib/copilot/request/lifecycle/run'
import { runCopilotLifecycle } from '@/lib/copilot/request/lifecycle/run'
import { type CopilotLifecycleOutcome, startCopilotOtelRoot } from '@/lib/copilot/request/otel'
import {
  cleanupAbortMarker,
  clearFilePreviewSessions,
  isExplicitStopReason,
  registerActiveStream,
  releasePendingChatStream,
  resetBuffer,
  StreamWriter,
  scheduleBufferCleanup,
  scheduleFilePreviewSessionCleanup,
  startAbortPoller,
  unregisterActiveStream,
} from '@/lib/copilot/request/session'
import { SSE_RESPONSE_HEADERS } from '@/lib/copilot/request/session/sse'
import { TraceCollector } from '@/lib/copilot/request/trace'
import { getMothershipBaseURL, getMothershipSourceEnvHeaders } from '@/lib/copilot/server/agent-url'
import { env } from '@/lib/core/config/env'

export { SSE_RESPONSE_HEADERS }

const logger = createLogger('CopilotChatStreaming')

type CurrentChatSummary = {
  title?: string | null
} | null

export interface StreamingOrchestrationParams {
  requestPayload: Record<string, unknown>
  userId: string
  streamId: string
  executionId: string
  runId: string
  chatId?: string
  currentChat: CurrentChatSummary
  isNewChat: boolean
  message: string
  titleModel: string
  titleProvider?: string
  requestId: string
  workspaceId?: string
  orchestrateOptions: Omit<CopilotLifecycleOptions, 'onEvent'>
  /**
   * Pre-started root; child spans bind to it and `finish()` fires on
   * termination. Omit to let the stream start its own root (headless).
   */
  otelRoot?: ReturnType<typeof startCopilotOtelRoot>
}

export function createSSEStream(params: StreamingOrchestrationParams): ReadableStream {
  const {
    requestPayload,
    userId,
    streamId,
    executionId,
    runId,
    chatId,
    currentChat,
    isNewChat,
    message,
    titleModel,
    titleProvider,
    requestId,
    workspaceId,
    orchestrateOptions,
    otelRoot,
  } = params

  // Reuse caller's root if provided; otherwise start our own.
  const activeOtelRoot =
    otelRoot ??
    startCopilotOtelRoot({
      requestId,
      route: orchestrateOptions.goRoute,
      chatId,
      workflowId: orchestrateOptions.workflowId,
      executionId,
      runId,
      streamId,
      transport: CopilotTransport.Stream,
    })

  const abortController = new AbortController()
  registerActiveStream(streamId, abortController)

  const publisher = new StreamWriter({ streamId, chatId, requestId })

  // Classify cancel: signal.reason (explicit-stop set) wins, then
  // clientDisconnected, else Unknown (latent contract bug — log it).
  const recordCancelled = (errorMessage?: string): CopilotRequestCancelReasonValue => {
    const rawReason = abortController.signal.reason
    let cancelReason: CopilotRequestCancelReasonValue
    if (isExplicitStopReason(rawReason)) {
      cancelReason = CopilotRequestCancelReason.ExplicitStop
    } else if (publisher.clientDisconnected) {
      cancelReason = CopilotRequestCancelReason.ClientDisconnect
    } else {
      cancelReason = CopilotRequestCancelReason.Unknown
      const serializedReason =
        rawReason === undefined
          ? 'undefined'
          : rawReason instanceof Error
            ? `${rawReason.name}: ${rawReason.message}`
            : typeof rawReason === 'string'
              ? rawReason
              : (() => {
                  try {
                    return JSON.stringify(rawReason)
                  } catch {
                    return String(rawReason)
                  }
                })()
      // Contract violation: add the new reason to AbortReason /
      // isExplicitStopReason or extend the classifier.
      logger.error(`[${requestId}] Stream cancelled with unknown abort reason`, {
        streamId,
        chatId,
        reason: serializedReason,
      })
      activeOtelRoot.span.setAttribute(TraceAttr.CopilotAbortUnknownReason, serializedReason)
    }
    activeOtelRoot.span.setAttribute(TraceAttr.CopilotRequestCancelReason, cancelReason)
    activeOtelRoot.span.addEvent(TraceEvent.RequestCancelled, {
      [TraceAttr.CopilotRequestCancelReason]: cancelReason,
      ...(errorMessage ? { [TraceAttr.ErrorMessage]: errorMessage } : {}),
    })
    return cancelReason
  }

  const collector = new TraceCollector()

  return new ReadableStream({
    async start(controller) {
      publisher.attach(controller)

      // Re-enter the root OTel context — ALS doesn't survive the
      // Next handler → ReadableStream.start boundary.
      await otelContextApi.with(activeOtelRoot.context, async () => {
        const otelContext = activeOtelRoot.context
        let rootOutcome: CopilotLifecycleOutcome = RequestTraceV1Outcome.error
        let rootError: unknown
        // `cancelReason` must be declared OUTSIDE the outer `try` so
        // it remains in scope for the outer `finally` that calls
        // `activeOtelRoot.finish(rootOutcome, rootError, cancelReason)`.
        // `let` bindings declared inside a `try` block are NOT visible
        // in the paired `finally`; referencing one there raises a
        // TDZ ReferenceError, skipping `finish()`, leaving the root
        // span never-ended, and making Tempo see every child as an
        // orphan under a phantom parent. (Regression landed 2026-04-21.)
        let cancelReason: CopilotRequestCancelReasonValue | undefined
        try {
          const requestSpan = collector.startSpan('Sim Agent Request', 'request', {
            streamId,
            chatId,
            runId,
          })
          let outcome: CopilotLifecycleOutcome = RequestTraceV1Outcome.error
          let lifecycleResult:
            | {
                usage?: { prompt: number; completion: number }
                cost?: { input: number; output: number; total: number }
              }
            | undefined

          await Promise.all([resetBuffer(streamId), clearFilePreviewSessions(streamId)])

          if (chatId) {
            createRunSegment({
              id: runId,
              executionId,
              chatId,
              userId,
              workflowId: (requestPayload.workflowId as string | undefined) || null,
              workspaceId,
              streamId,
              model: (requestPayload.model as string | undefined) || null,
              provider: (requestPayload.provider as string | undefined) || null,
              requestContext: { requestId },
            }).catch((error) => {
              logger.warn(`[${requestId}] Failed to create copilot run segment`, {
                error: getErrorMessage(error),
              })
            })
          }

          const abortPoller = startAbortPoller(streamId, abortController, {
            requestId,
            chatId,
          })
          publisher.startKeepalive()

          if (chatId) {
            publisher.publish({
              type: MothershipStreamV1EventType.session,
              payload: {
                kind: MothershipStreamV1SessionKind.chat,
                chatId,
              },
            })
          }

          fireTitleGeneration({
            chatId,
            currentChat,
            isNewChat,
            userId,
            message,
            titleModel,
            titleProvider,
            workspaceId,
            requestId,
            publisher,
            otelContext,
          })

          try {
            const result = await runCopilotLifecycle(requestPayload, {
              ...orchestrateOptions,
              executionId,
              runId,
              trace: collector,
              simRequestId: requestId,
              otelContext,
              abortSignal: abortController.signal,
              onEvent: async (event) => {
                await publisher.publish(event)
              },
              onAbortObserved: (reason) => {
                if (!abortController.signal.aborted) {
                  abortController.abort(reason)
                }
              },
            })

            lifecycleResult = result
            // Outcome classification (priority order):
            //   1. `result.success` → success. The orchestrator
            //      reporting "finished cleanly" wins over any later
            //      signal change. Matters for the narrow race where
            //      the user clicks Stop a beat after the stream
            //      completed.
            //   2. `signal.aborted` (from `abortActiveStream` or the
            //      Redis-marker poller) OR `clientDisconnected` with
            //      a non-success result → cancelled. `recordCancelled`
            //      further refines into explicit_stop / client_disconnect
            //      / unknown via `signal.reason`.
            //   3. Otherwise → error.
            outcome = result.success
              ? RequestTraceV1Outcome.success
              : result.cancelled || abortController.signal.aborted || publisher.clientDisconnected
                ? RequestTraceV1Outcome.cancelled
                : RequestTraceV1Outcome.error
            if (outcome === RequestTraceV1Outcome.cancelled) {
              cancelReason = recordCancelled()
            }
            // Pass the resolved outcome — not `signal.aborted` — so
            // `finalizeStream` classifies the same way we did above.
            // A client-disconnect-without-controller-abort still needs
            // to hit `handleAborted` (not `handleError`) so the chat
            // row gets `cancelled` terminal state instead of `error`.
            await finalizeStream(result, publisher, runId, outcome, requestId)
          } catch (error) {
            // Error-path classification: if the abort signal fired or
            // the client disconnected, treat the thrown error as a
            // cancel (same rationale as the try-path above).
            const wasCancelled = abortController.signal.aborted || publisher.clientDisconnected
            outcome = wasCancelled ? RequestTraceV1Outcome.cancelled : RequestTraceV1Outcome.error
            if (outcome === RequestTraceV1Outcome.cancelled) {
              cancelReason = recordCancelled(getErrorMessage(error))
            }
            if (publisher.clientDisconnected) {
              logger.info(`[${requestId}] Stream errored after client disconnect`, {
                error: getErrorMessage(error, 'Stream error'),
              })
            }
            // Demote to warn when the throw came from a user-initiated
            // cancel — it isn't an "unexpected" failure then, and the
            // error-level log pollutes alerting on normal Stop presses.
            const logFn = outcome === RequestTraceV1Outcome.cancelled ? logger.warn : logger.error
            logFn.call(logger, `[${requestId}] Orchestration ended with ${outcome}:`, error)

            const syntheticResult = {
              success: false as const,
              content: '',
              contentBlocks: [],
              toolCalls: [],
              error: 'An unexpected error occurred while processing the response.',
            }
            await finalizeStream(syntheticResult, publisher, runId, outcome, requestId)
          } finally {
            collector.endSpan(
              requestSpan,
              outcome === RequestTraceV1Outcome.success
                ? RequestTraceV1SpanStatus.ok
                : outcome === RequestTraceV1Outcome.cancelled
                  ? RequestTraceV1SpanStatus.cancelled
                  : RequestTraceV1SpanStatus.error
            )

            clearInterval(abortPoller)
            try {
              await publisher.close()
            } catch (error) {
              logger.warn(`[${requestId}] Failed to flush stream persistence during close`, {
                error: getErrorMessage(error),
              })
            }
            unregisterActiveStream(streamId)
            if (chatId) {
              await releasePendingChatStream(chatId, streamId)
            }
            await scheduleBufferCleanup(streamId)
            await scheduleFilePreviewSessionCleanup(streamId)
            await cleanupAbortMarker(streamId)

            rootOutcome = outcome
            if (lifecycleResult?.usage) {
              activeOtelRoot.span.setAttributes({
                [TraceAttr.GenAiUsageInputTokens]: lifecycleResult.usage.prompt ?? 0,
                [TraceAttr.GenAiUsageOutputTokens]: lifecycleResult.usage.completion ?? 0,
              })
            }
            if (lifecycleResult?.cost) {
              activeOtelRoot.span.setAttributes({
                [TraceAttr.BillingCostInputUsd]: lifecycleResult.cost.input ?? 0,
                [TraceAttr.BillingCostOutputUsd]: lifecycleResult.cost.output ?? 0,
                [TraceAttr.BillingCostTotalUsd]: lifecycleResult.cost.total ?? 0,
              })
            }
          }
        } catch (error) {
          rootOutcome = RequestTraceV1Outcome.error
          rootError = error
          throw error
        } finally {
          // `finish` is idempotent, so it's safe whether the POST
          // handler started the root (and may also call finish on an
          // error path before the stream ran) or we did. The cancel
          // reason (if any) determines whether `cancelled` is an
          // expected outcome (explicit_stop → status OK) or a real
          // error (client_disconnect / unknown → status ERROR).
          //
          // Belt-and-suspenders: if `finish()` itself throws (e.g. an
          // argument in the TDZ, a bad attribute, a regression in
          // status-setting), fall back to `span.end()` directly. A
          // root that never ends leaves every child orphaned in Tempo
          // under a phantom parent; force-ending it keeps the trace
          // shape intact even when the pretty-finalize path is
          // broken. The error is logged so Loki greps surface the
          // regression instead of it silently costing us trace
          // fidelity for hours.
          try {
            activeOtelRoot.finish(rootOutcome, rootError, cancelReason)
          } catch (finishError) {
            logger.error(`[${requestId}] activeOtelRoot.finish threw; force-ending root span`, {
              error: getErrorMessage(finishError),
            })
            try {
              activeOtelRoot.span.end()
            } catch {
              // Already ended or an OTel internal failure — nothing
              // more we can do. The export pipe has already had its
              // chance; swallow to avoid masking the original error
              // path.
            }
          }
        }
      })
    },
    cancel() {
      // The browser's SSE reader closed. Flip `clientDisconnected` so
      // in-flight `publisher.publish` calls silently no-op (prevents
      // enqueueing on a closed controller).
      //
      // Browser disconnect is NOT an abort — firing the controller
      // here retroactively reclassifies in-flight successful streams
      // as aborted and skips assistant persistence. Let the
      // orchestrator drain naturally; publish no-ops post-disconnect.
      // Explicit Stop still fires the controller via /chat/abort.
      publisher.markDisconnected()
    },
  })
}

// ---------------------------------------------------------------------------
// Title generation (fire-and-forget side effect)
// ---------------------------------------------------------------------------

function fireTitleGeneration(params: {
  chatId?: string
  currentChat: CurrentChatSummary
  isNewChat: boolean
  userId?: string
  message: string
  titleModel: string
  titleProvider?: string
  workspaceId?: string
  requestId: string
  publisher: StreamWriter
  otelContext?: Context
}): void {
  const {
    chatId,
    currentChat,
    isNewChat,
    userId,
    message,
    titleModel,
    titleProvider,
    workspaceId,
    requestId,
    publisher,
    otelContext,
  } = params
  if (!chatId || currentChat?.title || !isNewChat) return

  requestChatTitle({
    message,
    model: titleModel,
    provider: titleProvider,
    userId,
    workspaceId,
    otelContext,
  })
    .then(async (title) => {
      if (!title) return
      await db.update(copilotChats).set({ title }).where(eq(copilotChats.id, chatId))
      await publisher.publish({
        type: MothershipStreamV1EventType.session,
        payload: { kind: MothershipStreamV1SessionKind.title, title },
      })
      if (workspaceId) {
        chatPubSub?.publishStatusChanged({
          workspaceId,
          chatId,
          type: 'renamed',
        })
      }
    })
    .catch((error) => {
      logger.error(`[${requestId}] Title generation failed:`, error)
    })
}

// ---------------------------------------------------------------------------
// Chat title helper
// ---------------------------------------------------------------------------

export async function requestChatTitle(params: {
  message: string
  model: string
  provider?: string
  userId?: string
  workspaceId?: string
  otelContext?: Context
}): Promise<string | null> {
  const { message, model, provider, userId, workspaceId, otelContext } = params
  if (!message || !model) return null

  const headers: Record<string, string> = {
    'Content-Type': 'application/json',
  }
  if (env.COPILOT_API_KEY) {
    headers['x-api-key'] = env.COPILOT_API_KEY
  }
  Object.assign(headers, getMothershipSourceEnvHeaders())

  try {
    const { fetchGo } = await import('@/lib/copilot/request/go/fetch')
    const mothershipBaseURL = await getMothershipBaseURL({ userId })
    const response = await fetchGo(`${mothershipBaseURL}/api/generate-chat-title`, {
      method: 'POST',
      headers,
      body: JSON.stringify({
        message,
        model,
        ...(provider ? { provider } : {}),
        ...(workspaceId ? { workspaceId } : {}),
        ...(userId ? { userId } : {}),
      }),
      otelContext,
      spanName: 'sim → go /api/generate-chat-title',
      operation: 'generate_chat_title',
      attributes: {
        [TraceAttr.GenAiRequestModel]: model,
        ...(provider ? { [TraceAttr.GenAiSystem]: provider } : {}),
      },
    })

    const payload = await response.json().catch(() => ({}))
    if (!response.ok) {
      logger.warn('Failed to generate chat title via copilot backend', {
        status: response.status,
        error: payload,
      })
      return null
    }

    const title = typeof payload?.title === 'string' ? payload.title.trim() : ''
    return title || null
  } catch (error) {
    logger.error('Error generating chat title:', error)
    return null
  }
}