chore: import upstream snapshot with attribution

apps/sim/lib/copilot/request/handlers/tool.ts

@@ -0,0 +1,560 @@
+import { createLogger } from '@sim/logger'
+import { getErrorMessage, toError } from '@sim/utils/errors'
+import { ASYNC_TOOL_CONFIRMATION_STATUS } from '@/lib/copilot/async-runs/lifecycle'
+import { markAsyncToolDelivered, upsertAsyncToolCall } from '@/lib/copilot/async-runs/repository'
+import { STREAM_TIMEOUT_MS } from '@/lib/copilot/constants'
+import {
+  MothershipStreamV1AsyncToolRecordStatus,
+  type MothershipStreamV1ToolCallDescriptor,
+  MothershipStreamV1ToolOutcome,
+  type MothershipStreamV1ToolResultPayload,
+} from '@/lib/copilot/generated/mothership-stream-v1'
+import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
+import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
+import { withCopilotSpan } from '@/lib/copilot/request/otel'
+import {
+  isToolArgsDeltaStreamEvent,
+  isToolCallStreamEvent,
+  isToolResultStreamEvent,
+  TOOL_CALL_STATUS,
+} from '@/lib/copilot/request/session'
+import { markToolResultSeen, wasToolResultSeen } from '@/lib/copilot/request/sse-utils'
+import { setTerminalToolCallState } from '@/lib/copilot/request/tool-call-state'
+import { executeToolAndReport, waitForToolCompletion } from '@/lib/copilot/request/tools/executor'
+import type {
+  ExecutionContext,
+  OrchestratorOptions,
+  StreamEvent,
+  StreamingContext,
+  ToolCallState,
+} from '@/lib/copilot/request/types'
+import { getToolEntry, isSimExecuted } from '@/lib/copilot/tool-executor'
+import { isToolHiddenInUi } from '@/lib/copilot/tools/client/hidden-tools'
+import { getToolDisplayTitle } from '@/lib/copilot/tools/tool-display'
+import { isWorkflowToolName } from '@/lib/copilot/tools/workflow-tools'
+import type { ToolScope } from './types'
+import {
+  abortPendingToolIfStreamDead,
+  addContentBlock,
+  emitSyntheticToolResult,
+  ensureTerminalToolCallState,
+  flushSubagentThinkingBlock,
+  flushThinkingBlock,
+  getScopedParentToolCallId,
+  getScopedSpanIdentity,
+  getToolCallUI,
+  getToolResultErrorMessage,
+  handleClientCompletion,
+  inferToolSuccess,
+  registerPendingToolPromise,
+} from './types'
+
+const logger = createLogger('CopilotToolHandler')
+
+function applyToolDisplay(toolCall: ToolCallState | undefined): void {
+  if (!toolCall?.name) return
+  toolCall.displayTitle = getToolDisplayTitle(
+    toolCall.name,
+    toolCall.params as Record<string, unknown> | undefined
+  )
+}
+
+/**
+ * Upsert the durable `async_tool_calls` row before the authoritative tool-call
+ * SSE frame is forwarded to the client, so `/api/copilot/confirm` can never
+ * race ahead of the row that identifies the call. This is the sole
+ * persistence point for client-executable tools; gating mirrors the
+ * client-wait branch in `dispatchToolExecution`.
+ */
+export async function prePersistClientExecutableToolCall(
+  event: StreamEvent,
+  context: StreamingContext
+): Promise<void> {
+  if (event.type !== 'tool') return
+  if (!isToolCallStreamEvent(event)) return
+
+  const data = event.payload
+  const isGenerating = data.status === TOOL_CALL_STATUS.generating
+  const isPartial = data.partial === true || isGenerating
+  if (isPartial) return
+
+  const ui = getToolCallUI(data)
+  if (!ui.clientExecutable) return
+
+  const catalogEntry = getToolEntry(data.toolName)
+  const isInternal = ui.internal === true || catalogEntry?.internal === true
+  if (isInternal) return
+
+  const delegateWorkflowRunToClient = isWorkflowToolName(data.toolName)
+  if (isSimExecuted(data.toolName) && !delegateWorkflowRunToClient) return
+
+  if (!context.runId) return
+
+  await upsertAsyncToolCall({
+    runId: context.runId,
+    toolCallId: data.toolCallId,
+    toolName: data.toolName,
+    args: data.arguments,
+    status: MothershipStreamV1AsyncToolRecordStatus.running,
+  }).catch((err) => {
+    logger.warn('Failed to pre-persist async tool row before forwarding call frame', {
+      toolCallId: data.toolCallId,
+      toolName: data.toolName,
+      error: getErrorMessage(err),
+    })
+  })
+}
+
+/**
+ * Unified tool event handler for both main and subagent scopes.
+ *
+ * The main vs subagent differences are:
+ * - Subagent requires a parentToolCallId and tracks tool calls in subAgentToolCalls
+ * - Subagent result phase also updates the subAgentToolCalls record
+ * - Subagent call phase stores in both subAgentToolCalls and context.toolCalls
+ * - Main call phase only stores in context.toolCalls
+ */
+export async function handleToolEvent(
+  event: StreamEvent,
+  context: StreamingContext,
+  execContext: ExecutionContext,
+  options: OrchestratorOptions,
+  scope: ToolScope
+): Promise<void> {
+  const isSubagent = scope === 'subagent'
+  const parentToolCallId = isSubagent ? getScopedParentToolCallId(event, context) : undefined
+
+  if (isSubagent && !parentToolCallId) return
+
+  if (event.type !== 'tool') {
+    return
+  }
+
+  if (isToolArgsDeltaStreamEvent(event)) {
+    return
+  }
+
+  // A tool event breaks the thinking stream. Flush any open thinking
+  // block into contentBlocks BEFORE we add the tool_call block, or
+  // contentBlocks will end up with tool_call before thinking — which
+  // re-renders on reload in the wrong order (Mothership group above
+  // the Thinking block, even though thinking happened first). A subagent
+  // tool event flushes only its OWN lane so a concurrent sibling's thinking
+  // is left intact; a main tool event flushes all subagent lanes.
+  if (isSubagent && parentToolCallId) {
+    flushSubagentThinkingBlock(context, parentToolCallId)
+  } else {
+    flushSubagentThinkingBlock(context)
+  }
+  flushThinkingBlock(context)
+
+  if (isToolResultStreamEvent(event)) {
+    handleResultPhase(event.payload, context, parentToolCallId)
+    return
+  }
+
+  if (!isToolCallStreamEvent(event)) {
+    return
+  }
+
+  if (!parentToolCallId) {
+    context.sawMainToolCall = true
+    context.finalAssistantContent = ''
+  }
+
+  await handleCallPhase(
+    event.payload,
+    context,
+    execContext,
+    options,
+    parentToolCallId,
+    scope,
+    getScopedSpanIdentity(event)
+  )
+}
+
+function handleResultPhase(
+  data: MothershipStreamV1ToolResultPayload,
+  context: StreamingContext,
+  parentToolCallId: string | undefined
+): void {
+  const { toolCallId, toolName } = data
+  const mainToolCall = ensureTerminalToolCallState(context, toolCallId, toolName)
+  const { success, hasResultData } = inferToolSuccess(data)
+  let status: MothershipStreamV1ToolOutcome
+  if (data.status === MothershipStreamV1ToolOutcome.cancelled) {
+    status = MothershipStreamV1ToolOutcome.cancelled
+  } else if (data.status === MothershipStreamV1ToolOutcome.skipped) {
+    status = MothershipStreamV1ToolOutcome.skipped
+  } else if (data.status === MothershipStreamV1ToolOutcome.rejected) {
+    status = MothershipStreamV1ToolOutcome.rejected
+  } else {
+    status = success ? MothershipStreamV1ToolOutcome.success : MothershipStreamV1ToolOutcome.error
+  }
+  const endTime = Date.now()
+  const errorMessage =
+    !success && status !== MothershipStreamV1ToolOutcome.skipped
+      ? getToolResultErrorMessage(data) ||
+        (status === MothershipStreamV1ToolOutcome.cancelled
+          ? 'Tool cancelled'
+          : status === MothershipStreamV1ToolOutcome.rejected
+            ? 'Tool rejected'
+            : 'Tool failed')
+      : undefined
+
+  if (parentToolCallId) {
+    const toolCalls = context.subAgentToolCalls[parentToolCallId] || []
+    const subAgentToolCall = toolCalls.find((tc) => tc.id === toolCallId)
+    if (subAgentToolCall) {
+      setTerminalToolCallState(subAgentToolCall, {
+        status,
+        ...(hasResultData ? { output: data.output } : {}),
+        ...(errorMessage ? { error: errorMessage } : {}),
+        endTime,
+      })
+    }
+  }
+
+  setTerminalToolCallState(mainToolCall, {
+    status,
+    ...(hasResultData ? { output: data.output } : {}),
+    ...(errorMessage ? { error: errorMessage } : {}),
+    endTime,
+  })
+  stampToolCallBlockEnd(context, toolCallId, endTime)
+  markToolResultSeen(toolCallId)
+}
+
+function stampToolCallBlockEnd(
+  context: StreamingContext,
+  toolCallId: string,
+  endTime: number
+): void {
+  for (let i = context.contentBlocks.length - 1; i >= 0; i--) {
+    const block = context.contentBlocks[i]
+    if (block.type === 'tool_call' && block.toolCall?.id === toolCallId) {
+      if (block.endedAt === undefined) block.endedAt = endTime
+      return
+    }
+  }
+}
+
+async function handleCallPhase(
+  data: MothershipStreamV1ToolCallDescriptor,
+  context: StreamingContext,
+  execContext: ExecutionContext,
+  options: OrchestratorOptions,
+  parentToolCallId: string | undefined,
+  scope: ToolScope,
+  spanIdentity: { spanId?: string; parentSpanId?: string }
+): Promise<void> {
+  const { toolCallId, toolName } = data
+  const args = data.arguments
+  const isGenerating = data.status === TOOL_CALL_STATUS.generating
+  const isPartial = data.partial === true || isGenerating
+  const existing = context.toolCalls.get(toolCallId)
+  const isSubagent = scope === 'subagent'
+  const ui = getToolCallUI(data)
+
+  if (isPartial && shouldDelayVfsPlaceholder(toolName, args)) return
+
+  if (isSubagent) {
+    if (wasToolResultSeen(toolCallId) || existing?.endTime) {
+      if (existing && !existing.name && toolName) existing.name = toolName
+      if (existing && !existing.params && args) existing.params = args
+      applyToolDisplay(existing)
+      return
+    }
+  } else {
+    if (
+      existing?.endTime ||
+      (existing && existing.status !== 'pending' && existing.status !== 'executing')
+    ) {
+      if (!existing.name && toolName) existing.name = toolName
+      if (!existing.params && args) existing.params = args
+      applyToolDisplay(existing)
+      return
+    }
+  }
+
+  if (isSubagent) {
+    registerSubagentToolCall(
+      context,
+      toolCallId,
+      toolName,
+      args,
+      parentToolCallId!,
+      ui,
+      spanIdentity
+    )
+  } else {
+    registerMainToolCall(context, toolCallId, toolName, args, existing, ui)
+  }
+
+  if (isPartial) return
+  if (!isSubagent && wasToolResultSeen(toolCallId)) return
+  if (context.pendingToolPromises.has(toolCallId) || existing?.status === 'executing') {
+    return
+  }
+
+  const toolCall = context.toolCalls.get(toolCallId)
+  if (!toolCall) return
+
+  // Capture the invoking subagent's channel id so the executor can thread it
+  // into the server tool context — this is what scopes the workspace_file ->
+  // edit_content intent handoff to one file subagent under concurrency.
+  if (parentToolCallId) toolCall.parentToolCallId = parentToolCallId
+
+  const readPath = typeof args?.path === 'string' ? args.path : undefined
+  if (toolName === 'read' && readPath?.startsWith('internal/')) return
+
+  const { clientExecutable, simExecutable, internal } = ui
+  const catalogEntry = getToolEntry(toolName)
+  const isInternal = internal || catalogEntry?.internal === true
+  const staticSimExecuted = isSimExecuted(toolName)
+  const willDispatch = !isInternal && (staticSimExecuted || simExecutable || clientExecutable)
+  logger.info('Tool call routing decision', {
+    toolCallId,
+    toolName,
+    scope,
+    isSubagent,
+    parentToolCallId,
+    executor: data.executor,
+    clientExecutable,
+    simExecutable,
+    staticSimExecuted,
+    internal: isInternal,
+    hasPendingPromise: context.pendingToolPromises.has(toolCallId),
+    existingStatus: existing?.status,
+    willDispatch,
+  })
+  if (isInternal) return
+  if (!willDispatch) return
+
+  await dispatchToolExecution(
+    toolCall,
+    toolCallId,
+    toolName,
+    args,
+    context,
+    execContext,
+    options,
+    clientExecutable,
+    scope
+  )
+}
+
+function shouldDelayVfsPlaceholder(
+  toolName: string,
+  args: Record<string, unknown> | undefined
+): boolean {
+  return (toolName === 'read' || toolName === 'glob') && !args
+}
+
+function removeToolCallContentBlock(context: StreamingContext, toolCallId: string): void {
+  for (let i = context.contentBlocks.length - 1; i >= 0; i--) {
+    const block = context.contentBlocks[i]
+    if (block.type === 'tool_call' && block.toolCall?.id === toolCallId) {
+      context.contentBlocks.splice(i, 1)
+    }
+  }
+}
+
+function registerSubagentToolCall(
+  context: StreamingContext,
+  toolCallId: string,
+  toolName: string,
+  args: Record<string, unknown> | undefined,
+  parentToolCallId: string,
+  ui: { title?: string; phaseLabel?: string; hidden?: boolean },
+  spanIdentity: { spanId?: string; parentSpanId?: string }
+): void {
+  if (!context.subAgentToolCalls[parentToolCallId]) {
+    context.subAgentToolCalls[parentToolCallId] = []
+  }
+  const hideFromUi = isToolHiddenInUi(toolName) || ui.hidden === true
+  let toolCall = context.toolCalls.get(toolCallId)
+  if (toolCall) {
+    if (!toolCall.name && toolName) toolCall.name = toolName
+    if (args && !toolCall.params) toolCall.params = args
+    applyToolDisplay(toolCall)
+    if (hideFromUi) removeToolCallContentBlock(context, toolCallId)
+  } else {
+    toolCall = {
+      id: toolCallId,
+      name: toolName,
+      status: 'pending',
+      params: args,
+      startTime: Date.now(),
+    }
+    applyToolDisplay(toolCall)
+    context.toolCalls.set(toolCallId, toolCall)
+    const parentToolCall = context.toolCalls.get(parentToolCallId)
+    if (!hideFromUi) {
+      addContentBlock(context, {
+        type: 'tool_call',
+        toolCall,
+        calledBy: parentToolCall?.name,
+        parentToolCallId,
+        ...spanIdentity,
+      })
+    }
+  }
+
+  const subagentToolCalls = context.subAgentToolCalls[parentToolCallId]
+  const existingSubagentToolCall = subagentToolCalls.find((tc) => tc.id === toolCallId)
+  if (existingSubagentToolCall) {
+    if (!existingSubagentToolCall.name && toolName) existingSubagentToolCall.name = toolName
+    if (args && !existingSubagentToolCall.params) existingSubagentToolCall.params = args
+    applyToolDisplay(existingSubagentToolCall)
+  } else {
+    subagentToolCalls.push(toolCall)
+  }
+}
+
+function registerMainToolCall(
+  context: StreamingContext,
+  toolCallId: string,
+  toolName: string,
+  args: Record<string, unknown> | undefined,
+  existing: ToolCallState | undefined,
+  ui: { title?: string; phaseLabel?: string; hidden?: boolean }
+): void {
+  const hideFromUi = isToolHiddenInUi(toolName) || ui.hidden === true
+  if (existing) {
+    if (args && !existing.params) existing.params = args
+    applyToolDisplay(existing)
+    if (hideFromUi) {
+      removeToolCallContentBlock(context, toolCallId)
+      return
+    }
+    if (
+      !hideFromUi &&
+      !context.contentBlocks.some((b) => b.type === 'tool_call' && b.toolCall?.id === toolCallId)
+    ) {
+      addContentBlock(context, { type: 'tool_call', toolCall: existing })
+    }
+  } else {
+    const created: ToolCallState = {
+      id: toolCallId,
+      name: toolName,
+      status: 'pending',
+      params: args,
+      startTime: Date.now(),
+    }
+    applyToolDisplay(created)
+    context.toolCalls.set(toolCallId, created)
+    if (!hideFromUi) {
+      addContentBlock(context, { type: 'tool_call', toolCall: created })
+    }
+  }
+}
+
+async function dispatchToolExecution(
+  toolCall: ToolCallState,
+  toolCallId: string,
+  toolName: string,
+  args: Record<string, unknown> | undefined,
+  context: StreamingContext,
+  execContext: ExecutionContext,
+  options: OrchestratorOptions,
+  clientExecutable: boolean,
+  scope: ToolScope
+): Promise<void> {
+  const scopeLabel = scope === 'subagent' ? 'subagent ' : ''
+
+  const fireToolExecution = () => {
+    const pendingPromise = (async () => {
+      return executeToolAndReport(toolCallId, context, execContext, options)
+    })().catch((err) => {
+      logger.error(`Parallel ${scopeLabel}tool execution failed`, {
+        toolCallId,
+        toolName,
+        error: toError(err).message,
+      })
+      return {
+        status: MothershipStreamV1ToolOutcome.error,
+        message: 'Tool execution failed',
+        data: { error: 'Tool execution failed' },
+      }
+    })
+    registerPendingToolPromise(context, toolCallId, pendingPromise)
+  }
+
+  if (options.interactive === false) {
+    if (options.autoExecuteTools !== false) {
+      if (!abortPendingToolIfStreamDead(toolCall, toolCallId, options, context)) {
+        fireToolExecution()
+      }
+    }
+    return
+  }
+
+  if (clientExecutable) {
+    const delegateWorkflowRunToClient = isWorkflowToolName(toolName)
+    if (isSimExecuted(toolName) && !delegateWorkflowRunToClient) {
+      if (!abortPendingToolIfStreamDead(toolCall, toolCallId, options, context)) {
+        fireToolExecution()
+      }
+    } else {
+      toolCall.status = 'executing'
+      const pendingPromise = withCopilotSpan(
+        TraceSpan.CopilotToolWaitForClientResult,
+        {
+          [TraceAttr.ToolName]: toolName,
+          [TraceAttr.ToolCallId]: toolCallId,
+          [TraceAttr.ToolTimeoutMs]: options.timeout || STREAM_TIMEOUT_MS,
+          ...(context.runId ? { [TraceAttr.RunId]: context.runId } : {}),
+        },
+        async (span) => {
+          const completion = await waitForToolCompletion(
+            toolCallId,
+            options.timeout || STREAM_TIMEOUT_MS,
+            options.abortSignal
+          )
+          span.setAttribute(TraceAttr.ToolCompletionReceived, completion !== undefined)
+          if (completion) {
+            span.setAttribute(TraceAttr.ToolOutcome, completion.status)
+          }
+          handleClientCompletion(toolCall, toolCallId, completion)
+          if (completion?.status === ASYNC_TOOL_CONFIRMATION_STATUS.background) {
+            await markAsyncToolDelivered(toolCallId).catch((err) => {
+              logger.warn(`Failed to mark background ${scopeLabel}tool delivered`, {
+                toolCallId,
+                toolName,
+                error: toError(err).message,
+              })
+            })
+          }
+          await emitSyntheticToolResult(toolCallId, toolCall.name, completion, options)
+          return (
+            completion ?? {
+              status: MothershipStreamV1ToolOutcome.error,
+              message: 'Tool completion missing',
+              data: { error: 'Tool completion missing' },
+            }
+          )
+        }
+      ).catch((err) => {
+        logger.error(`Client-executable ${scopeLabel}tool wait failed`, {
+          toolCallId,
+          toolName,
+          error: toError(err).message,
+        })
+        return {
+          status: MothershipStreamV1ToolOutcome.error,
+          message: 'Tool wait failed',
+          data: { error: 'Tool wait failed' },
+        }
+      })
+      registerPendingToolPromise(context, toolCallId, pendingPromise)
+    }
+    return
+  }
+
+  if (options.autoExecuteTools !== false) {
+    if (!abortPendingToolIfStreamDead(toolCall, toolCallId, options, context)) {
+      fireToolExecution()
+    }
+  }
+}