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

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()
    }
  }
}