dd-trace 6.0.0 → 6.2.0

package/packages/dd-trace/src/llmobs/plugins/ai/vercelTelemetry.js

@@ -0,0 +1,363 @@
+'use strict'
+
+const { parseModelProvider } = require('../../../../../datadog-plugin-ai/src/utils')
+const BaseLLMObsPlugin = require('../base')
+const {
+  getLlmObsSpanName,
+  getGenerationMetadataFromEvent,
+  getJsonStringValue,
+  getToolCallResultContent,
+} = require('./util')
+
+// TODO: add in embedMany once it has tracingChannel support
+const SPAN_NAME_TO_KIND_MAPPING = {
+  // embeddings
+  embed: 'embedding',
+  // text generation
+  generateText: 'workflow',
+  streamText: 'workflow',
+  // llm operations
+  languageModelCall: 'llm',
+  // steps
+  step: 'step', // TODO: support step spans for manual instrumentation as well
+  // tools
+  executeTool: 'tool',
+}
+
+function nameFromOperation (operation, event) {
+  if (operation === 'executeTool') {
+    return event.toolCall?.toolName
+  }
+
+  return operation
+}
+
+function formatLanguageModelInputMessages (instructions, messages) {
+  if (!Array.isArray(messages)) return
+  const inputMessages = []
+
+  if (instructions) {
+    const systemPrompt = typeof instructions === 'string'
+      ? instructions
+      : Array.isArray(instructions)
+        ? instructions.map(instruction => instruction.content).join('')
+        : instructions.content
+
+    inputMessages.push({ role: 'system', content: systemPrompt })
+  }
+
+  for (const message of messages) {
+    const { role, content } = message
+
+    if (role === 'system') {
+      inputMessages.push({ role, content })
+    } else if (role === 'user') {
+      const userMessageContent =
+      typeof content === 'string'
+        ? content
+        : content
+          .filter(part => part.type === 'text')
+          .map(part => part.text)
+          .join('')
+
+      inputMessages.push({ role, content: userMessageContent })
+    } else if (role === 'assistant') {
+      if (typeof content === 'string') {
+        inputMessages.push({ role, content })
+      } else {
+        // re-use existing output message formatting
+        inputMessages.push(...formatLanguageModelOutputMessages(content))
+      }
+    } else if (role === 'tool') {
+      for (const part of content) {
+        if (part.type === 'tool-result') { // TODO: support tool approvals
+          const safeResult = getToolCallResultContent(part)
+
+          inputMessages.push({
+            role,
+            content: safeResult,
+            toolId: part.toolCallId,
+          })
+        }
+      }
+    }
+  }
+
+  return inputMessages
+}
+
+function formatLanguageModelOutputMessages (content) {
+  if (!Array.isArray(content)) return
+
+  const outputMessages = []
+  const toolCalls = []
+
+  let textContent = ''
+  let reasoningContent = ''
+
+  for (const part of content) {
+    const { type } = part
+
+    if (type === 'text') {
+      textContent += part.text
+    } else if (type === 'reasoning') {
+      reasoningContent += part.text
+    } else if (type === 'tool-call') {
+      const toolCallArguments = typeof part.input === 'string'
+        ? getJsonStringValue(part.input, {})
+        : part.input
+
+      toolCalls.push({
+        arguments: toolCallArguments,
+        name: part.toolName,
+        type: 'function',
+        toolId: part.toolCallId,
+      })
+    }
+  }
+
+  if (reasoningContent) {
+    outputMessages.push({ role: 'reasoning', content: reasoningContent })
+  }
+
+  const finalTextMessage = { role: 'assistant' }
+
+  if (textContent) {
+    finalTextMessage.content = textContent
+  }
+
+  if (toolCalls.length) {
+    finalTextMessage.toolCalls = toolCalls
+  }
+
+  outputMessages.push(finalTextMessage)
+
+  return outputMessages
+}
+
+class VercelAiTelemetryPlugin extends BaseLLMObsPlugin {
+  static id = 'ai_llmobs_vercel_telemetry'
+  static integration = 'ai'
+  static prefix = 'tracing:ai:telemetry'
+
+  /** @type {Map<string, Array<Record<string, unknown>>>} */
+  #lastOutputContentByCallId = new Map()
+
+  constructor () {
+    super(...arguments)
+
+    this.addSub('dd-trace:vercel-ai:chunk', ({ ctx, chunk, done }) => {
+      ctx.chunks ??= []
+      const chunks = ctx.chunks
+      if (chunk) chunks.push(chunk)
+
+      ctx.streamConsumed = done
+
+      if (!done) return
+
+      ctx.result ??= {}
+      const contentById = {}
+
+      for (const chunk of chunks) {
+        if (chunk.type === 'finish') {
+          ctx.result.usage = chunk.usage
+        } else if (chunk.type === 'reasoning-start') {
+          contentById[chunk.id] = { type: 'reasoning', text: '' }
+        } else if (chunk.type === 'text-start') {
+          contentById[chunk.id] = { type: 'text', text: '' }
+        } else if (chunk.type === 'reasoning-delta' || chunk.type === 'text-delta') {
+          if (!contentById[chunk.id]) {
+            // special-case handling when the start event did not come through as a chunk
+            contentById[chunk.id] = {
+              type: chunk.type === 'reasoning-delta' ? 'reasoning' : 'text',
+              text: '',
+            }
+          }
+          contentById[chunk.id].text += chunk.delta
+        } else if (chunk.type === 'tool-call') {
+          contentById[chunk.toolCallId] = chunk
+        }
+      }
+
+      const content = Object.values(contentById)
+      ctx.result.content = content
+
+      this.#lastOutputContentByCallId.set(ctx.event.callId, content)
+    })
+  }
+
+  /**
+   * @override
+   */
+  asyncEnd (ctx) {
+    // check if isStreamed and stream resolved
+    // this event will fire multiple times for the same channel
+    if (ctx.isStream && ctx.result?.stream && !ctx.streamConsumed) return
+
+    super.asyncEnd(ctx)
+  }
+
+  /**
+   * @override
+   */
+  getLLMObsSpanRegisterOptions (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+
+    const { type: operation, event } = ctx
+    const kind = SPAN_NAME_TO_KIND_MAPPING[operation]
+    if (!kind) return
+
+    const normalizedName = nameFromOperation(operation, event) || operation
+
+    const options = {
+      kind, name: getLlmObsSpanName(normalizedName, event.functionId),
+    }
+
+    if (kind === 'llm' || kind === 'embedding') {
+      const modelName = event.modelId
+
+      options.modelName = modelName
+      options.modelProvider = parseModelProvider(event.provider, modelName)
+    }
+
+    return options
+  }
+
+  /**
+   * @override
+   */
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+
+    const { type: operation } = ctx
+    const kind = SPAN_NAME_TO_KIND_MAPPING[operation]
+    if (!kind) return
+
+    switch (operation) {
+      case 'embed':
+        this.setEmbeddingTags(span, ctx)
+        break
+      case 'generateText':
+      case 'streamText':
+        this.setTextGenerationTags(span, ctx)
+        break
+      case 'languageModelCall':
+        this.setLanguageModelCallTags(span, ctx)
+        break
+      case 'step':
+        this.setStepTags(span, ctx)
+        break
+      case 'executeTool':
+        this.setToolTags(span, ctx)
+        break
+      default:
+        break
+    }
+  }
+
+  setEmbeddingTags (span, ctx) {
+    const { event, result } = ctx
+
+    const input = event.value
+    const embedding = result?.embedding
+    const embeddingTextResult = embedding ? `[1 embedding(s) returned with size ${embedding.length}]` : ''
+
+    this._tagger.tagEmbeddingIO(span, input, embeddingTextResult)
+
+    this._tagger.tagMetrics(span, {
+      inputTokens: result?.usage?.tokens,
+    })
+  }
+
+  setTextGenerationTags (span, ctx) {
+    const { event, result } = ctx
+
+    const lastUserPrompt = event.messages.findLast(message => message.role === 'user')?.content
+    const input = Array.isArray(lastUserPrompt) ? lastUserPrompt.map(part => part.text ?? '').join('') : lastUserPrompt
+
+    const output =
+      ctx.isStream
+        ? this.#lastOutputContentByCallId.get(event.callId)?.find(part => part.type === 'text')?.text
+        : result?._output
+
+    this._tagger.tagTextIO(span, input, output)
+
+    const metadata = getGenerationMetadataFromEvent(event)
+    this._tagger.tagMetadata(span, metadata)
+
+    if (ctx.isStream) {
+      this.#lastOutputContentByCallId.delete(event.callId)
+    }
+  }
+
+  setStepTags (span, ctx) {
+    const { result, event } = ctx
+
+    const content = ctx.isStream ? this.#lastOutputContentByCallId.get(event.callId) : result?.content
+
+    // capture reasoning if applicable
+    const reasoning = content?.find(part => part.type === 'reasoning')?.text
+    if (reasoning) {
+      this._tagger.tagTextIO(span, reasoning)
+    }
+  }
+
+  setLanguageModelCallTags (span, ctx) {
+    const { event, result } = ctx
+
+    // input messages
+    const { instructions, messages } = event
+    const inputMessages = formatLanguageModelInputMessages(instructions, messages)
+
+    // output messages
+    const outputMessages =
+      result ? formatLanguageModelOutputMessages(result.content) : [{ role: 'assistant', content: '' }]
+
+    this._tagger.tagLLMIO(span, inputMessages, outputMessages)
+
+    // tool definitions
+    const { tools } = event
+    if (Array.isArray(tools)) {
+      this._tagger.tagToolDefinitions(
+        span,
+        tools.map(({ inputSchema, ...rest }) => ({
+          ...rest,
+          schema: {
+            properties: inputSchema?.properties ?? {},
+            required: inputSchema?.required ?? [],
+            type: inputSchema?.type ?? '',
+          },
+        }))
+      )
+    }
+
+    // metadata
+    const metadata = getGenerationMetadataFromEvent(event)
+    this._tagger.tagMetadata(span, metadata)
+
+    if (!result) return
+
+    // metrics
+    const { usage } = result
+    this._tagger.tagMetrics(span, {
+      inputTokens: usage?.inputTokens?.total,
+      cacheWriteTokens: usage?.inputTokens?.cacheWrite ?? 0,
+      cacheReadTokens: usage?.inputTokens?.cacheRead ?? 0,
+      outputTokens: usage?.outputTokens?.total,
+      reasoningOutputTokens: usage?.outputTokens?.reasoning ?? 0,
+    })
+  }
+
+  setToolTags (span, ctx) {
+    const { event, result } = ctx
+
+    const { toolCall } = event
+    if (!toolCall) return
+
+    this._tagger.tagTextIO(span, toolCall.input, result?.output?.output)
+  }
+}
+
+module.exports = VercelAiTelemetryPlugin

package/packages/dd-trace/src/llmobs/plugins/claude-agent-sdk/index.js

@@ -0,0 +1,270 @@
+'use strict'
+
+const LLMObsPlugin = require('../base')
+const { storage: llmobsStorage } = require('../../storage')
+const { NAME, SESSION_ID } = require('../../constants/tags')
+const { splitModel } = require('../../../../../datadog-plugin-claude-agent-sdk/src/util')
+
+const subagentToolIds = new Set()
+
+function normalizeToolOutputString (raw) {
+  const footerIndex = raw.search(/\s*agentId: [^\s]+ \(use SendMessage\b/)
+  if (footerIndex === -1) return raw
+
+  return raw.slice(0, footerIndex).trimEnd()
+}
+
+function getToolOutputText (raw) {
+  if (raw == null) return
+
+  if (Array.isArray(raw)) {
+    const output = []
+    for (const block of raw) {
+      const text = getToolOutputText(block)
+      if (text) output.push(text)
+    }
+    return output.join('\n') || undefined
+  }
+
+  if (raw.type === 'tool_result') return getToolOutputText(raw.content)
+  if (raw.type === 'text') return normalizeToolOutputString(raw.text)
+  if (raw.type === 'tool_reference') return raw.tool_name
+  if (raw.content !== undefined) return getToolOutputText(raw.content)
+
+  if (typeof raw === 'string') return normalizeToolOutputString(raw)
+
+  return JSON.stringify(raw)
+}
+
+function buildOutputMessages (chunks, llmStartIdx, llmEndIdx) {
+  let thinking = ''
+  let text = ''
+  const toolCalls = []
+
+  for (let i = llmStartIdx; i < llmEndIdx; i++) {
+    const c = chunks[i]
+    if (c.type !== 'assistant') continue
+    const block = c.message?.content?.[0]
+    if (block?.type === 'thinking') thinking += block.thinking ?? ''
+    else if (block?.type === 'text') text += block.text ?? ''
+    else if (block?.type === 'tool_use') {
+      toolCalls.push({ name: block.name, arguments: block.input ?? {}, toolId: block.id, type: block.type })
+    }
+  }
+
+  const messages = []
+  if (thinking) messages.push({ role: 'thinking', content: thinking })
+  const msg = { role: 'assistant', content: text }
+  if (toolCalls.length) msg.toolCalls = toolCalls
+  messages.push(msg)
+  return messages
+}
+
+class QueryLLMObsPlugin extends LLMObsPlugin {
+  static integration = 'claude-agent-sdk'
+  static id = 'llmobs_claude_agent_sdk_query'
+  static prefix = 'tracing:orchestrion:@anthropic-ai/claude-agent-sdk:query'
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    return { kind: 'agent' }
+  }
+
+  start (ctx) {
+    super.start(ctx)
+    if (!this._tracerConfig.llmobs.DD_LLMOBS_ENABLED) return
+    const store = llmobsStorage.getStore()
+    const prev = ctx.runInContext ?? (fn => fn())
+    ctx.runInContext = fn => prev(() => llmobsStorage.run(store, fn))
+  }
+
+  asyncEnd (ctx) {
+    if (!ctx.streamResolved) return
+    super.asyncEnd(ctx)
+  }
+
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+
+    // post-populate session_id
+    if (ctx.session_id) this._tagger._setTag(span, SESSION_ID, ctx.session_id)
+    this._tagger.tagTextIO(span, ctx.arguments?.[0]?.prompt, ctx.output)
+
+    // metadata
+    const { cwd, permissionMode } = ctx
+    const metadata = {}
+
+    if (cwd) metadata.cwd = cwd
+    if (permissionMode) metadata.permissionMode = permissionMode
+
+    this._tagger.tagMetadata(span, metadata)
+  }
+}
+
+class StepLlmObsPlugin extends LLMObsPlugin {
+  static integration = 'claude-agent-sdk'
+  static id = 'claude_agent_sdk_step_llmobs'
+  static system = 'claude-agent-sdk'
+  static prefix = 'tracing:apm:claude-agent-sdk:step'
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    if (ctx.parentToolUseId) subagentToolIds.add(ctx.parentToolUseId)
+    return { kind: 'step', name: `step-${ctx.stepIndex}`, sessionId: ctx.sessionId }
+  }
+
+  end (ctx) {
+    super.end(ctx)
+    super.asyncEnd(ctx)
+  }
+
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+
+    const { chunks, llmStartIdx, llmEndIdx, toolOutputs } = ctx
+    if (!chunks) return
+
+    const outputMessages = buildOutputMessages(chunks, llmStartIdx, llmEndIdx)
+    const thinking = outputMessages.find(m => m.role === 'thinking')?.content ?? ''
+
+    const output = toolOutputs?.length
+      ? getToolOutputText(toolOutputs)
+      : outputMessages.find(m => m.role === 'assistant')?.content ?? ''
+
+    this._tagger.tagTextIO(span, thinking, output)
+  }
+}
+
+class LlmLlmObsPlugin extends LLMObsPlugin {
+  static integration = 'claude-agent-sdk'
+  static id = 'claude_agent_sdk_llm_llmobs'
+  static system = 'claude-agent-sdk'
+  static prefix = 'tracing:apm:claude-agent-sdk:llm'
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    const { modelName, modelProvider } = splitModel(ctx.model)
+    return { kind: 'llm', name: ctx.model, modelName, modelProvider, sessionId: ctx.sessionId }
+  }
+
+  end (ctx) {
+    super.end(ctx)
+    super.asyncEnd(ctx)
+  }
+
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+
+    const { chunks, llmStartIdx, llmEndIdx, parentToolUseId, initialPrompt, usage } = ctx
+
+    if (chunks) {
+      const inputMessages = this.#buildInputMessages(chunks, llmStartIdx, parentToolUseId, initialPrompt)
+      const outputMessages = buildOutputMessages(chunks, llmStartIdx, llmEndIdx)
+      this._tagger.tagLLMIO(span, inputMessages, outputMessages)
+    }
+
+    if (usage) {
+      const cacheWriteTokens = usage.cache_creation_input_tokens ?? 0
+      const cacheReadTokens = usage.cache_read_input_tokens ?? 0
+      const inputTokens = (usage.input_tokens ?? 0) + cacheWriteTokens + cacheReadTokens
+      const outputTokens = usage.output_tokens ?? 0
+      this._tagger.tagMetrics(span, {
+        input_tokens: inputTokens,
+        output_tokens: outputTokens,
+        cache_read_input_tokens: cacheReadTokens,
+        cache_write_input_tokens: cacheWriteTokens,
+        total_tokens: inputTokens + outputTokens,
+      })
+    }
+  }
+
+  #buildInputMessages (chunks, llmStartIdx, parentToolUseId, initialPrompt) {
+    const messages = []
+    if (initialPrompt) messages.push({ role: 'user', content: initialPrompt })
+    const seenIds = new Set()
+
+    for (let i = 0; i < llmStartIdx; i++) {
+      const c = chunks[i]
+      if (c.parent_tool_use_id !== parentToolUseId) continue
+
+      if (c.type === 'assistant') {
+        const msgId = c.message?.id
+        if (!msgId || seenIds.has(msgId)) continue
+        seenIds.add(msgId)
+
+        let thinking = ''
+        let text = ''
+        const toolCalls = []
+        for (let j = i; j < llmStartIdx; j++) {
+          const cc = chunks[j]
+          if (cc.type !== 'assistant' || cc.message?.id !== msgId) break
+          const block = cc.message?.content?.[0]
+          if (block?.type === 'thinking') thinking += block.thinking ?? ''
+          else if (block?.type === 'text') text += block.text ?? ''
+          else if (block?.type === 'tool_use') {
+            toolCalls.push({ name: block.name, arguments: block.input ?? {}, toolId: block.id, type: block.type })
+          }
+        }
+        if (thinking) messages.push({ role: 'thinking', content: thinking })
+        const msg = { role: 'assistant', content: text }
+        if (toolCalls.length) msg.toolCalls = toolCalls
+        messages.push(msg)
+      } else if (c.type === 'user') {
+        const content = c.message?.content
+        if (!content) continue
+        for (const block of content) {
+          if (block.type === 'text') {
+            messages.push({ role: 'user', content: block.text ?? '' })
+          } else if (block.type === 'tool_result') {
+            const text = getToolOutputText(block.content) ?? ''
+            messages.push({ role: 'tool', content: text })
+          }
+        }
+      }
+    }
+
+    return messages
+  }
+}
+
+class ToolLlmObsPlugin extends LLMObsPlugin {
+  static integration = 'claude-agent-sdk'
+  static id = 'claude_agent_sdk_tool_llmobs'
+  static system = 'claude-agent-sdk'
+  static prefix = 'tracing:apm:claude-agent-sdk:tool'
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    return { kind: 'tool', name: ctx.name, sessionId: ctx.sessionId }
+  }
+
+  end (ctx) {
+    super.end(ctx)
+    super.asyncEnd(ctx)
+  }
+
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+
+    if (subagentToolIds.has(ctx.id)) {
+      subagentToolIds.delete(ctx.id)
+      const description = ctx.input?.description
+      this._tagger.changeKind(span, 'agent')
+      if (description) this._tagger._setTag(span, NAME, `${ctx.name} (${description})`)
+      const output = getToolOutputText(ctx.output)
+      this._tagger.tagTextIO(span, ctx.input?.prompt, output)
+      return
+    }
+
+    const input = ctx.input ? JSON.stringify(ctx.input) : undefined
+    const output = getToolOutputText(ctx.output)
+    this._tagger.tagTextIO(span, input, output)
+  }
+}
+
+module.exports = [
+  QueryLLMObsPlugin,
+  StepLlmObsPlugin,
+  ToolLlmObsPlugin,
+  LlmLlmObsPlugin,
+]

package/packages/dd-trace/src/llmobs/plugins/claude-agent-sdk/utils.js

@@ -0,0 +1,10 @@
+'use strict'
+
+function splitModel (model) {
+  if (!model) return { modelName: undefined, modelProvider: 'anthropic' }
+  const idx = model.indexOf('/')
+  if (idx === -1) return { modelName: model, modelProvider: 'anthropic' }
+  return { modelName: model.slice(idx + 1), modelProvider: model.slice(0, idx) }
+}
+
+module.exports = { splitModel }

package/packages/dd-trace/src/msgpack/chunk.js

@@ -10,6 +10,33 @@ const SHRINK_AFTER_FLUSHES = 32
 // shape: after a halving step the post-shrink fill is the prior peak doubled,
 // still under 50 %.
 const SHRINK_USAGE_RATIO = 4
+// Hard cap on chunk growth. The agent's trace intake rejects payloads over
+// 50 MiB, so anything past that is dead on arrival anyway. A pathological
+// single trace (unsanitized meta tag the size of a media file, multi-MB
+// stack trace) used to grow the buffer without limit until either the
+// allocation failed or the process tripped its memory ceiling. `reserve`
+// now refuses the growth with a tagged `RangeError`; `AgentEncoder` catches
+// it, resets the in-flight payload, and logs.
+const MAX_SIZE = 50 * 1024 * 1024 // 50 MiB
+
+/**
+ * Thrown when a chunk — or an assembled payload stitched from several chunks —
+ * would cross `MAX_SIZE`. Shared so the `reserve` cap and the per-encoder
+ * assembled-size guards (`0.5`, agentless CI Visibility) raise the same `code`,
+ * which every writer's `flush` recognises to drop the payload instead of
+ * crashing the host.
+ */
+class OverflowError extends RangeError {
+  code = 'ERR_MSGPACK_CHUNK_OVERFLOW'
+
+  /**
+   * @param {number} needed Requested total byte size.
+   */
+  constructor (needed) {
+    super(`MsgpackChunk capped at ${MAX_SIZE} bytes; requested ${needed}`)
+    this.name = 'OverflowError'
+  }
+}
 
 /**
  * Resizable msgpack write buffer. Owns the byte-layout primitives the encoder
@@ -100,11 +127,14 @@ class MsgpackChunk {
     const needed = this.length + size
 
     if (needed > this.buffer.length) {
+      if (needed > MAX_SIZE) {
+        throw new OverflowError(needed)
+      }
       let newSize = this.buffer.length
       // `*= 2` instead of `<<= 1`: `1073741824 << 1` is negative as int32,
       // and msgpack values can legitimately reach the multi-GiB range.
       while (newSize < needed) newSize *= 2
-      this.#resize(newSize)
+      this.#resize(Math.min(newSize, MAX_SIZE))
     }
 
     this.length += size
@@ -454,3 +484,5 @@ class MsgpackChunk {
 }
 
 module.exports = MsgpackChunk
+module.exports.MAX_SIZE = MAX_SIZE
+module.exports.OverflowError = OverflowError

package/packages/dd-trace/src/msgpack/index.js

@@ -97,4 +97,9 @@ function writeMap (bytes, value) {
   }
 }
 
-module.exports = { MsgpackChunk, encode }
+module.exports = {
+  MsgpackChunk,
+  encode,
+  MAX_SIZE: MsgpackChunk.MAX_SIZE,
+  OverflowError: MsgpackChunk.OverflowError,
+}