dd-trace 6.0.0 → 6.1.0

package/packages/datadog-instrumentations/src/claude-agent-sdk.js

@@ -0,0 +1,587 @@
+'use strict'
+
+const { tracingChannel } = require('dc-polyfill')
+const shimmer = require('../../datadog-shimmer')
+const { addHook, getHooks } = require('./helpers/instrument')
+const queryChannel = tracingChannel('orchestrion:@anthropic-ai/claude-agent-sdk:query')
+
+const stepCh = tracingChannel('apm:claude-agent-sdk:step')
+const llmCh = tracingChannel('apm:claude-agent-sdk:llm')
+const toolCh = tracingChannel('apm:claude-agent-sdk:tool')
+
+const LOCAL_LIFECYCLE_LOOKAHEAD = 4
+
+const chunkEmitTimes = new WeakMap()
+
+function hasDownstreamSubscribers () {
+  return queryChannel.asyncEnd.hasSubscribers ||
+    queryChannel.error.hasSubscribers ||
+    stepCh.start.hasSubscribers ||
+    llmCh.start.hasSubscribers ||
+    toolCh.start.hasSubscribers
+}
+
+function mergeHooks (userHooks, tracerHooks) {
+  const merged = {}
+
+  for (const event of Object.keys(tracerHooks)) {
+    const userMatchers = userHooks?.[event] || []
+    merged[event] = [...userMatchers, ...tracerHooks[event]]
+  }
+
+  if (userHooks) {
+    for (const event of Object.keys(userHooks)) {
+      if (!merged[event]) merged[event] = userHooks[event]
+    }
+  }
+
+  return merged
+}
+
+function getTool (sessionCtx, id) {
+  let tool = sessionCtx.tools.get(id)
+  if (!tool) {
+    tool = { id }
+    sessionCtx.tools.set(id, tool)
+  }
+  return tool
+}
+
+function buildTracerHooks (sessionCtx) {
+  function onSessionStart (input) {
+    sessionCtx.sessionId = input.session_id
+    sessionCtx.source = input.source
+    sessionCtx.cwd = input.cwd
+    sessionCtx.transcriptPath = input.transcript_path
+    sessionCtx.agentType = input.agent_type
+    sessionCtx.permissionMode = sessionCtx.permissionMode || input.permission_mode
+    return {}
+  }
+
+  function onSessionEnd (input) {
+    sessionCtx.endReason = input.reason
+    return {}
+  }
+
+  function onUserPromptSubmit (input) {
+    sessionCtx.sessionId = sessionCtx.sessionId || input.session_id
+    sessionCtx.prompt = sessionCtx.prompt || input.prompt
+    return {}
+  }
+
+  function onStop (input) {
+    sessionCtx.stopReason = input.stop_reason
+    sessionCtx.lastAssistantMessage = input.last_assistant_message
+    return {}
+  }
+
+  function onPreToolUse (input, toolUseId) {
+    const id = toolUseId || input.tool_use_id
+    if (!id) return {}
+
+    Object.assign(getTool(sessionCtx, id), {
+      id,
+      name: input.tool_name,
+      input: input.tool_input,
+      sessionId: input.session_id,
+      hookStartTime: Date.now(),
+    })
+    return {}
+  }
+
+  function onPostToolUse (input, toolUseId) {
+    const id = toolUseId || input.tool_use_id
+    if (!id) return {}
+
+    const tool = getTool(sessionCtx, id)
+    Object.assign(tool, {
+      id,
+      name: input.tool_name || tool.name,
+      input: input.tool_input === undefined ? tool.input : input.tool_input,
+      output: input.tool_response,
+      sessionId: input.session_id,
+      hookFinishTime: Date.now(),
+    })
+    return {}
+  }
+
+  function onPostToolUseFailure (input, toolUseId) {
+    const id = toolUseId || input.tool_use_id
+    if (!id) return {}
+
+    const tool = getTool(sessionCtx, id)
+    Object.assign(tool, {
+      id,
+      name: input.tool_name || tool.name,
+      input: input.tool_input === undefined ? tool.input : input.tool_input,
+      error: input.error,
+      isInterrupt: input.is_interrupt,
+      sessionId: input.session_id,
+      hookFinishTime: Date.now(),
+    })
+    return {}
+  }
+
+  function onSubagentStart (input) {
+    const id = input.agent_id
+    if (!id) return {}
+
+    sessionCtx.subagents.set(id, {
+      id,
+      sessionId: input.session_id,
+      agentType: input.agent_type,
+      hookStartTime: Date.now(),
+    })
+    return {}
+  }
+
+  function onSubagentStop (input) {
+    const id = input.agent_id
+    if (!id) return {}
+
+    const subagent = sessionCtx.subagents.get(id) || { id }
+    Object.assign(subagent, {
+      sessionId: input.session_id,
+      agentType: input.agent_type || subagent.agentType,
+      transcriptPath: input.agent_transcript_path,
+      output: input.last_assistant_message,
+      hookFinishTime: Date.now(),
+    })
+    sessionCtx.subagents.set(id, subagent)
+    return {}
+  }
+
+  return {
+    SessionStart: [{ hooks: [onSessionStart] }],
+    SessionEnd: [{ hooks: [onSessionEnd] }],
+    UserPromptSubmit: [{ hooks: [onUserPromptSubmit] }],
+    Stop: [{ hooks: [onStop] }],
+    PreToolUse: [{ hooks: [onPreToolUse] }],
+    PostToolUse: [{ hooks: [onPostToolUse] }],
+    PostToolUseFailure: [{ hooks: [onPostToolUseFailure] }],
+    SubagentStart: [{ hooks: [onSubagentStart] }],
+    SubagentStop: [{ hooks: [onSubagentStop] }],
+  }
+}
+
+function onQueryStart (ctx) {
+  if (!hasDownstreamSubscribers()) return
+
+  const { arguments: args } = ctx
+  const queryArg = args?.[0]
+  if (!queryArg) return
+
+  const options = queryArg.options || {}
+  const prompt = queryArg.prompt
+  const sessionCtx = {
+    prompt: typeof prompt === 'string' ? prompt : undefined,
+    model: options.model,
+    resume: options.resume,
+    maxTurns: options.maxTurns,
+    permissionMode: options.permissionMode,
+    tools: new Map(),
+    subagents: new Map(),
+  }
+
+  args[0] = {
+    ...queryArg,
+    options: {
+      ...options,
+      hooks: mergeHooks(options.hooks, buildTracerHooks(sessionCtx)),
+    },
+  }
+  ctx.sessionCtx = sessionCtx
+  ctx.claudeAgentSdkTracing = true
+}
+
+function buildStreamIndex (chunks) {
+  const lifecycleByToolId = new Map()
+
+  for (let idx = 0; idx < chunks.length; idx++) {
+    const chunk = chunks[idx]
+    if (chunk.type === 'system' && chunk.subtype === 'task_started' && chunk.tool_use_id) {
+      const lifecycle = lifecycleByToolId.get(chunk.tool_use_id) || {}
+      lifecycle.taskStartedChunk = chunk
+      lifecycleByToolId.set(chunk.tool_use_id, lifecycle)
+    } else if (chunk.type === 'user') {
+      const content = chunk.message?.content
+      if (!Array.isArray(content)) continue
+      for (const block of content) {
+        if (block.type === 'tool_result' && block.tool_use_id) {
+          const lifecycle = lifecycleByToolId.get(block.tool_use_id) || {}
+          if (lifecycle.toolResultIndex === undefined) lifecycle.toolResultIndex = idx
+          lifecycleByToolId.set(block.tool_use_id, lifecycle)
+        }
+      }
+    }
+  }
+
+  return lifecycleByToolId
+}
+
+function scanLocalLifecycle (chunks, startIndex, toolUseId, lifecycle) {
+  lifecycle.taskStartedChunk = undefined
+  lifecycle.toolResultIndex = undefined
+
+  const scanEnd = Math.min(chunks.length, startIndex + LOCAL_LIFECYCLE_LOOKAHEAD)
+
+  for (let idx = startIndex; idx < scanEnd; idx++) {
+    const chunk = chunks[idx]
+    if (chunk.type === 'system' && chunk.subtype === 'task_started' && chunk.tool_use_id === toolUseId) {
+      lifecycle.taskStartedChunk = chunk
+    } else if (chunk.type === 'user') {
+      const content = chunk.message?.content
+      if (!Array.isArray(content)) continue
+      for (const block of content) {
+        if (block.type === 'tool_result' && block.tool_use_id === toolUseId) {
+          lifecycle.toolResultIndex = idx
+          return lifecycle
+        }
+      }
+    }
+  }
+
+  return lifecycle
+}
+
+function createStreamLookup (chunks) {
+  let streamIndex
+  const localLifecycle = {}
+
+  return function getLifecycle (startIndex, toolUseId) {
+    if (streamIndex) return streamIndex.get(toolUseId) || {}
+
+    scanLocalLifecycle(chunks, startIndex, toolUseId, localLifecycle)
+    if (localLifecycle.toolResultIndex !== undefined) return localLifecycle
+
+    streamIndex = streamIndex || buildStreamIndex(chunks)
+    const indexedLifecycle = streamIndex.get(toolUseId)
+    return {
+      taskStartedChunk: localLifecycle.taskStartedChunk || indexedLifecycle?.taskStartedChunk,
+      toolResultIndex: indexedLifecycle?.toolResultIndex,
+    }
+  }
+}
+
+function getToolScanEnd (chunks, startIndex, toolUseId, toolResultIndex) {
+  if (toolResultIndex !== undefined) return toolResultIndex + 1
+
+  for (let idx = startIndex; idx < chunks.length; idx++) {
+    const chunk = chunks[idx]
+    if (chunk.type === 'result') return idx
+    if (chunk.type === 'assistant' && chunk.parent_tool_use_id !== toolUseId) return idx
+  }
+
+  return chunks.length
+}
+
+function getToolResultContent (chunk, toolUseId) {
+  if (chunk?.type !== 'user') return
+
+  const content = chunk.message?.content
+  if (!Array.isArray(content)) return
+
+  for (const block of content) {
+    if (block.type === 'tool_result' && block.tool_use_id === toolUseId) return [block]
+  }
+}
+
+/**
+ *
+ * @param {Array<Record<string, unknown>>} chunks
+ * @param {number} startIndex
+ * @param {string} toolUseId
+ * @returns {number} index in chunks where the next step should start iteration
+ */
+function processTool (chunks, startIndex, toolUseId, sessionCtx, getLifecycle, lifecycle) {
+  let chunkIndex = startIndex
+  let stepIndex = 0
+  const tool = sessionCtx?.tools.get(toolUseId)
+  const toolResultIndex = lifecycle.toolResultIndex
+  const scanEnd = getToolScanEnd(chunks, startIndex, toolUseId, toolResultIndex)
+
+  while (chunkIndex < scanEnd) {
+    const chunk = chunks[chunkIndex]
+
+    if (chunkIndex === toolResultIndex) return chunkIndex + 1
+
+    // only process steps for assistant chunks that belong to this subagent invocation
+    if (chunk.type === 'assistant' && chunk.parent_tool_use_id === toolUseId) {
+      let prevIdx = chunkIndex - 1
+      while (prevIdx >= 0 && chunks[prevIdx].type === 'system') prevIdx--
+      const prevChunk = prevIdx >= 0 ? chunks[prevIdx] : chunk
+      const stepCtx = {
+        stepIndex,
+        startTime: chunkEmitTimes.get(prevChunk),
+        parentToolUseId: toolUseId,
+        sessionId: chunks[0]?.session_id,
+      }
+      chunkIndex = stepCh.traceSync(() => {
+        const nextIdx = processStep(
+          chunks,
+          chunkIndex,
+          stepCtx.startTime,
+          toolUseId,
+          undefined,
+          stepCtx,
+          sessionCtx,
+          getLifecycle
+        )
+        stepCtx.finishTime = chunkEmitTimes.get(chunks[Math.min(nextIdx - 1, chunks.length - 1)])
+        return nextIdx
+      }, stepCtx)
+      stepIndex++
+    } else {
+      chunkIndex++
+    }
+  }
+
+  if (tool?.hookFinishTime) return scanEnd
+
+  return chunks.length + 1
+}
+
+/**
+ *
+ * @param {Array<Record<string, unknown>>} chunks
+ * @param {number} startIndex
+ * @returns {number} index in chunks where the next step should start iteration
+ */
+function processStep (
+  chunks,
+  startIndex,
+  stepStartTime,
+  parentToolUseId = null,
+  initialPrompt,
+  stepCtx = null,
+  sessionCtx = null,
+  getLifecycle
+) {
+  for (let idx = startIndex; idx < chunks.length; idx++) {
+    const chunk = chunks[idx]
+
+    if (chunk.type !== 'assistant') continue
+
+    const { id: messageId, model, usage } = chunk.message
+
+    // collect all chunks belonging to the same message (parallel tool_uses arrive as separate chunks)
+    const toolUses = []
+    let messageEndIdx = idx
+    for (let j = idx; j < chunks.length; j++) {
+      const c = chunks[j]
+      if (c.type !== 'assistant' || c.message.id !== messageId) break
+      if (c.message.content[0]?.type === 'tool_use') {
+        toolUses.push({ toolUse: c.message.content[0], startTime: chunkEmitTimes.get(c) })
+      }
+      messageEndIdx = j + 1
+    }
+
+    if (stepCtx) {
+      stepCtx.chunks = chunks
+      stepCtx.llmStartIdx = idx
+      stepCtx.llmEndIdx = messageEndIdx
+      stepCtx.toolOutputs = []
+    }
+
+    llmCh.traceSync(() => {}, {
+      model,
+      usage,
+      startTime: stepStartTime,
+      finishTime: chunkEmitTimes.get(chunks[messageEndIdx - 1]),
+      chunks,
+      llmStartIdx: idx,
+      llmEndIdx: messageEndIdx,
+      parentToolUseId,
+      initialPrompt,
+      sessionId: chunks[0]?.session_id,
+    })
+
+    if (toolUses.length === 0) {
+      return messageEndIdx
+    }
+
+    // for parallel tool calls, each processTool starts from the same messageEndIdx
+    // and independently scans for its own tool_result; take the max to advance past all results
+    let nextIdx = messageEndIdx
+    for (const { toolUse: { id, name, input }, startTime: toolUseStartTime } of toolUses) {
+      const hookTool = sessionCtx?.tools.get(id)
+      const lifecycle = getLifecycle(messageEndIdx, id)
+      // Agent tools emit a system:task_started chunk with a more precise start time
+      const taskStartedChunk = lifecycle.taskStartedChunk
+      const toolCtx = {
+        id,
+        name: hookTool?.name || name,
+        input: hookTool?.input || input,
+        startTime: hookTool?.hookStartTime ||
+          (taskStartedChunk ? chunkEmitTimes.get(taskStartedChunk) : toolUseStartTime),
+        sessionId: chunks[0]?.session_id,
+      }
+      const toolEndIdx = toolCh.traceSync(() => {
+        const endIdx = processTool(chunks, messageEndIdx, id, sessionCtx, getLifecycle, lifecycle)
+        const resultChunk = chunks[endIdx - 1]
+        if (hookTool?.isInterrupt) toolCtx.isInterrupt = hookTool.isInterrupt
+        if (hookTool?.error) {
+          toolCtx.error = hookTool.error
+          toolCh.error.publish(toolCtx)
+        }
+        const resultOutput = getToolResultContent(resultChunk, id)
+        if (resultOutput) {
+          toolCtx.output = resultOutput
+        } else if (hookTool?.output) {
+          toolCtx.output = hookTool.output
+        }
+        toolCtx.finishTime = hookTool?.hookFinishTime ||
+          chunkEmitTimes.get(chunks[Math.min(endIdx - 1, chunks.length - 1)])
+        return endIdx
+      }, toolCtx)
+
+      if (stepCtx && toolCtx.output) stepCtx.toolOutputs.push(toolCtx.output)
+
+      nextIdx = Math.max(nextIdx, toolEndIdx)
+    }
+
+    return nextIdx
+  }
+
+  return chunks.length + 1
+}
+
+/**
+ * @param {Array<Record<string, unknown>>} chunks
+ */
+function processChunks (chunks, agentCtx) {
+  let chunkIndex = 0
+  let stepIndex = 0
+  const sessionCtx = agentCtx.sessionCtx
+  const getLifecycle = createStreamLookup(chunks)
+
+  const { type, subtype, ...rest } = chunks[0]
+  Object.assign(agentCtx, rest)
+  if (sessionCtx) {
+    if (sessionCtx.sessionId) agentCtx.session_id = sessionCtx.sessionId
+    if (sessionCtx.cwd) agentCtx.cwd = sessionCtx.cwd
+    if (sessionCtx.permissionMode) agentCtx.permissionMode = sessionCtx.permissionMode
+    if (sessionCtx.lastAssistantMessage && !agentCtx.output) agentCtx.output = sessionCtx.lastAssistantMessage
+  }
+
+  while (chunkIndex < chunks.length) {
+    if (chunks[chunkIndex].type === 'result') break
+
+    const prevChunk = chunkIndex > 0 ? chunks[chunkIndex - 1] : chunks[chunkIndex]
+    const stepCtx = { stepIndex, startTime: chunkEmitTimes.get(prevChunk), sessionId: agentCtx.session_id }
+
+    const run = agentCtx.runInContext ?? (fn => fn())
+    chunkIndex = run(() => stepCh.traceSync(() => {
+      const initialPrompt = sessionCtx?.prompt || agentCtx.arguments?.[0]?.prompt
+      const nextIdx = processStep(
+        chunks,
+        chunkIndex,
+        stepCtx.startTime,
+        null,
+        initialPrompt,
+        stepCtx,
+        sessionCtx,
+        getLifecycle
+      )
+      stepCtx.finishTime = chunkEmitTimes.get(chunks[Math.min(nextIdx - 1, chunks.length - 1)])
+      return nextIdx
+    }, stepCtx))
+
+    stepIndex++
+  }
+}
+
+function finishStream (chunks, ctx, error) {
+  if (ctx.streamResolved) return
+
+  let processError
+  if (chunks.length > 0) {
+    const lastChunk = chunks[chunks.length - 1]
+    if (lastChunk?.type === 'result') ctx.output = lastChunk.result
+
+    try {
+      processChunks(chunks, ctx)
+    } catch (e) {
+      processError = e
+    }
+  }
+
+  const finalError = error || processError
+  ctx.finishTime = Date.now()
+
+  if (finalError) {
+    ctx.error = finalError
+    queryChannel.error.publish(ctx)
+  }
+
+  ctx.streamResolved = true
+  queryChannel.asyncEnd.publish(ctx)
+
+  if (processError && !error) throw processError
+}
+
+function wrapQueryAsyncIterator (asyncIterator, ctx) {
+  const chunks = []
+
+  return function () {
+    const iterator = asyncIterator.apply(this, arguments)
+    iterator.next = shimmer.wrapCallback(iterator.next, next => function () {
+      return next.apply(this, arguments).then(result => {
+        const { done, value } = result
+
+        if (!done && value) {
+          const chunkEmitTime = Date.now()
+          chunks.push(value)
+          chunkEmitTimes.set(value, chunkEmitTime)
+        } else {
+          finishStream(chunks, ctx)
+        }
+
+        return result
+      }).catch(error => {
+        finishStream(chunks, ctx, error)
+        throw error
+      })
+    })
+
+    if (typeof iterator.return === 'function') {
+      iterator.return = shimmer.wrapCallback(iterator.return, iteratorReturn => function () {
+        return iteratorReturn.apply(this, arguments).then(result => {
+          finishStream(chunks, ctx)
+          return result
+        }).catch(error => {
+          finishStream(chunks, ctx, error)
+          throw error
+        })
+      })
+    }
+
+    return iterator
+  }
+}
+
+let querySubscribed = false
+
+for (const hook of getHooks('@anthropic-ai/claude-agent-sdk')) {
+  hook.file = null
+
+  addHook(hook, exports => {
+    if (!querySubscribed) {
+      querySubscribed = true
+      queryChannel.subscribe({
+        start: onQueryStart,
+        end (ctx) {
+          if (!ctx.claudeAgentSdkTracing) return
+
+          const { result } = ctx
+
+          ctx.streamResolved = false
+
+          shimmer.wrap(result, Symbol.asyncIterator, asyncIterator => wrapQueryAsyncIterator(asyncIterator, ctx))
+        },
+      })
+    }
+
+    return exports
+  })
+}