dd-trace 6.18.0 → 6.19.0

package/packages/dd-trace/src/aiguard/integrations/stream.js

@@ -0,0 +1,60 @@
+'use strict'
+
+const log = require('../../log')
+
+/**
+ * Splits an SDK stream, consumes one branch, and returns the other after inspection.
+ *
+ * @param {object} stream
+ * @param {(chunks: Array<object>) => void|Promise<void>} inspect
+ * @returns {object|Promise<object>}
+ */
+function interceptStream (stream, inspect) {
+  if (typeof stream?.tee !== 'function') return stream
+
+  try {
+    const [inspectionStream, resultStream] = stream.tee()
+    return drainStream(inspectionStream).then(chunks => {
+      return Promise.resolve(inspect(chunks)).then(() => resultStream)
+    })
+  } catch {
+    return stream
+  }
+}
+
+/**
+ * Buffers the whole stream. A stream that ends early still yields what it delivered: truncated
+ * output can carry the very violation the evaluation is looking for, so it is judged rather than
+ * skipped, and the read failure is logged because nothing downstream reports it.
+ *
+ * @param {object} stream
+ * @returns {Promise<Array<object>>}
+ */
+function drainStream (stream) {
+  const chunks = []
+  const iterator = stream[Symbol.asyncIterator]()
+
+  function onReadError (error) {
+    log.error('AIGuard: the streamed response ended after %s chunks: %s', chunks.length, error)
+    return chunks
+  }
+
+  function readAll () {
+    let next
+    try {
+      next = iterator.next()
+    } catch (error) {
+      return Promise.resolve(onReadError(error))
+    }
+
+    return next.then(({ done, value }) => {
+      if (done) return chunks
+      chunks.push(value)
+      return readAll()
+    }, onReadError)
+  }
+
+  return readAll()
+}
+
+module.exports = { interceptStream }

package/packages/dd-trace/src/aiguard/messages/anthropic.js

@@ -431,10 +431,74 @@ function getMessagesOutputMessages (body) {
   return convertAnthropicMessage({ role, content: body.content })
 }
 
+/**
+ * Combines Anthropic message stream events into regular output messages.
+ *
+ * Accumulates exactly the way the SDK does: every `content_block_start` appends a block, and
+ * deltas address blocks by position. Keying blocks by `event.index` instead would let a repeated
+ * or out-of-range index hide output that the caller still receives.
+ *
+ * @param {Array<object>} events
+ * @returns {Array<object>}
+ */
+function getStreamedMessagesOutputMessages (events) {
+  let message
+  let contentBlocks
+  // Keyed by block, not by index: distinct indices can resolve to one block, and a second
+  // message must not inherit partial JSON accumulated for the first.
+  const inputJson = new Map()
+
+  for (const event of events) {
+    if (!event || typeof event !== 'object') continue
+
+    if (event.type === 'message_start' && event.message && typeof event.message === 'object') {
+      contentBlocks = Array.isArray(event.message.content)
+        ? event.message.content.map(block => ({ ...block }))
+        : []
+      message = { role: event.message.role || 'assistant' }
+      continue
+    }
+
+    if (event.type === 'content_block_start' && event.content_block && typeof event.content_block === 'object') {
+      message ??= { role: 'assistant' }
+      contentBlocks ??= []
+      contentBlocks.push({ ...event.content_block })
+      continue
+    }
+
+    if (event.type !== 'content_block_delta' || !event.delta || typeof event.delta !== 'object') continue
+
+    const block = contentBlocks?.at(event.index ?? 0)
+    if (!block) continue
+
+    if (event.delta.type === 'text_delta' && block.type === 'text' && typeof event.delta.text === 'string') {
+      block.text = (block.text || '') + event.delta.text
+    } else if (event.delta.type === 'input_json_delta' && typeof event.delta.partial_json === 'string') {
+      inputJson.set(block, (inputJson.get(block) || '') + event.delta.partial_json)
+    }
+  }
+
+  if (!message) return []
+
+  for (const [block, json] of inputJson) {
+    // An empty buffer is what a no-argument tool call accumulates; the SDK keeps `{}` there.
+    if (!json) continue
+    try {
+      block.input = JSON.parse(json)
+    } catch {
+      block.input = json
+    }
+  }
+
+  message.content = contentBlocks
+  return getMessagesOutputMessages(message)
+}
+
 module.exports = {
   convertAnthropicSystem,
   convertAnthropicBlocksToContent,
   convertAnthropicMessage,
   getMessagesInputMessages,
   getMessagesOutputMessages,
+  getStreamedMessagesOutputMessages,
 }

package/packages/dd-trace/src/config/generated-config-types.d.ts

@@ -574,8 +574,8 @@ export interface GeneratedConfig {
     DD_CODE_COVERAGE_FLAGS: string | undefined;
     DD_TEST_EARLY_FLAKE_DETECTION_RETRY_COUNT: number | undefined;
     DD_TEST_FAILED_TEST_REPLAY_ENABLED: boolean;
-    DD_TEST_FAILURE_SCREENSHOTS_ENABLED: boolean | undefined;
-    DD_TEST_FAILURE_VIDEOS_ENABLED: boolean | undefined;
+    DD_TEST_FAILURE_SCREENSHOTS_ENABLED: boolean;
+    DD_TEST_FAILURE_VIDEOS_ENABLED: boolean;
     DD_TEST_FLEET_CONFIG_PATH: string | undefined;
     DD_TEST_LOCAL_CONFIG_PATH: string | undefined;
     DD_TEST_MANAGEMENT_ATTEMPT_TO_FIX_RETRIES: number;
@@ -835,8 +835,8 @@ export interface GeneratedEnvVarConfig {
   DD_TELEMETRY_METRICS_ENABLED: boolean;
   DD_TEST_EARLY_FLAKE_DETECTION_RETRY_COUNT: number | undefined;
   DD_TEST_FAILED_TEST_REPLAY_ENABLED: boolean;
-  DD_TEST_FAILURE_SCREENSHOTS_ENABLED: boolean | undefined;
-  DD_TEST_FAILURE_VIDEOS_ENABLED: boolean | undefined;
+  DD_TEST_FAILURE_SCREENSHOTS_ENABLED: boolean;
+  DD_TEST_FAILURE_VIDEOS_ENABLED: boolean;
   DD_TEST_FLEET_CONFIG_PATH: string | undefined;
   DD_TEST_LOCAL_CONFIG_PATH: string | undefined;
   DD_TEST_MANAGEMENT_ATTEMPT_TO_FIX_RETRIES: number;

package/packages/dd-trace/src/config/supported-configurations.json

@@ -59,7 +59,7 @@
         "type": "boolean",
         "namespace": "aiguard",
         "default": "false",
-        "description": "Analyze streamed OpenAI and Vercel AI responses with AI Guard After Model evaluation."
+        "description": "Analyze streamed Anthropic, OpenAI, and Vercel AI responses with AI Guard After Model evaluation."
       }
     ],
     "DD_AI_GUARD_BLOCK": [
@@ -2055,17 +2055,17 @@
     ],
     "DD_TEST_FAILURE_SCREENSHOTS_ENABLED": [
       {
-        "implementation": "A",
+        "implementation": "B",
         "type": "boolean",
-        "default": null,
+        "default": "true",
         "namespace": "testOptimization"
       }
     ],
     "DD_TEST_FAILURE_VIDEOS_ENABLED": [
       {
-        "implementation": "A",
+        "implementation": "B",
         "type": "boolean",
-        "default": null,
+        "default": "true",
         "namespace": "testOptimization"
       }
     ],

package/packages/dd-trace/src/constants.js

@@ -22,6 +22,8 @@ module.exports = {
   SPAN_SAMPLING_RULE_RATE: '_dd.span_sampling.rule_rate',
   SPAN_SAMPLING_MAX_PER_SECOND: '_dd.span_sampling.max_per_second',
   SVC_SRC_KEY: '_dd.svc_src',
+  SDK_OTLP_EXPORT_KEY: '_dd.sdk.otlp_export',
+  SDK_SEMANTICS_KEY: 'datadog.sdk.semantics',
   DATADOG_LAMBDA_EXTENSION_PATH: '/opt/extensions/datadog-agent',
   DATADOG_MINI_AGENT_PATH: '/tmp/datadog/mini_agent_ready',
   DECISION_MAKER_KEY: '_dd.p.dm',

package/packages/dd-trace/src/llmobs/constants/tags.js

@@ -1,5 +1,14 @@
 'use strict'
 
+const INPUT_TOKENS_METRIC_KEY = 'input_tokens'
+const OUTPUT_TOKENS_METRIC_KEY = 'output_tokens'
+const TOTAL_TOKENS_METRIC_KEY = 'total_tokens'
+const CACHE_READ_INPUT_TOKENS_METRIC_KEY = 'cache_read_input_tokens'
+const CACHE_WRITE_INPUT_TOKENS_METRIC_KEY = 'cache_write_input_tokens'
+const CACHE_WRITE_5M_INPUT_TOKENS_METRIC_KEY = 'ephemeral_5m_input_tokens'
+const CACHE_WRITE_1H_INPUT_TOKENS_METRIC_KEY = 'ephemeral_1h_input_tokens'
+const REASONING_OUTPUT_TOKENS_METRIC_KEY = 'reasoning_output_tokens'
+
 module.exports = {
   SPAN_KINDS: ['llm', 'agent', 'workflow', 'task', 'tool', 'embedding', 'retrieval', 'experiment'],
   SPAN_KIND: '_ml_obs.meta.span.kind',
@@ -42,6 +51,7 @@ module.exports = {
   MODEL_NAME: '_ml_obs.meta.model_name',
   MODEL_PROVIDER: '_ml_obs.meta.model_provider',
   UNKNOWN_MODEL_PROVIDER: 'unknown',
+  DEFAULT_MODEL: 'custom',
 
   INPUT_DOCUMENTS: '_ml_obs.meta.input.documents',
   INPUT_MESSAGES: '_ml_obs.meta.input.messages',
@@ -52,14 +62,27 @@ module.exports = {
   OUTPUT_MESSAGES: '_ml_obs.meta.output.messages',
   OUTPUT_VALUE: '_ml_obs.meta.output.value',
 
-  INPUT_TOKENS_METRIC_KEY: 'input_tokens',
-  OUTPUT_TOKENS_METRIC_KEY: 'output_tokens',
-  TOTAL_TOKENS_METRIC_KEY: 'total_tokens',
-  CACHE_READ_INPUT_TOKENS_METRIC_KEY: 'cache_read_input_tokens',
-  CACHE_WRITE_INPUT_TOKENS_METRIC_KEY: 'cache_write_input_tokens',
-  CACHE_WRITE_5M_INPUT_TOKENS_METRIC_KEY: 'ephemeral_5m_input_tokens',
-  CACHE_WRITE_1H_INPUT_TOKENS_METRIC_KEY: 'ephemeral_1h_input_tokens',
-  REASONING_OUTPUT_TOKENS_METRIC_KEY: 'reasoning_output_tokens',
+  INPUT_TOKENS_METRIC_KEY,
+  OUTPUT_TOKENS_METRIC_KEY,
+  TOTAL_TOKENS_METRIC_KEY,
+  CACHE_READ_INPUT_TOKENS_METRIC_KEY,
+  CACHE_WRITE_INPUT_TOKENS_METRIC_KEY,
+  CACHE_WRITE_5M_INPUT_TOKENS_METRIC_KEY,
+  CACHE_WRITE_1H_INPUT_TOKENS_METRIC_KEY,
+  REASONING_OUTPUT_TOKENS_METRIC_KEY,
+
+  // integrations build metric objects with these camelCase spellings. Null prototype: a metric
+  // named after an `Object.prototype` member must not resolve to an inherited property.
+  METRIC_KEY_ALIASES: Object.assign(Object.create(null), {
+    inputTokens: INPUT_TOKENS_METRIC_KEY,
+    outputTokens: OUTPUT_TOKENS_METRIC_KEY,
+    totalTokens: TOTAL_TOKENS_METRIC_KEY,
+    cacheReadTokens: CACHE_READ_INPUT_TOKENS_METRIC_KEY,
+    cacheWriteTokens: CACHE_WRITE_INPUT_TOKENS_METRIC_KEY,
+    cacheWrite5mTokens: CACHE_WRITE_5M_INPUT_TOKENS_METRIC_KEY,
+    cacheWrite1hTokens: CACHE_WRITE_1H_INPUT_TOKENS_METRIC_KEY,
+    reasoningOutputTokens: REASONING_OUTPUT_TOKENS_METRIC_KEY,
+  }),
 
   DROPPED_IO_COLLECTION_ERROR: 'dropped_io',
 

package/packages/dd-trace/src/llmobs/gen-ai-tags.js

@@ -0,0 +1,127 @@
+'use strict'
+
+const {
+  ARTIFICIAL_GEN_AI_TAGS,
+  CACHE_READ_INPUT_TOKENS_METRIC_KEY,
+  CACHE_WRITE_INPUT_TOKENS_METRIC_KEY,
+  DEFAULT_MODEL,
+  GEN_AI_APPLICATION_NAME,
+  GEN_AI_CONVERSATION_ID,
+  GEN_AI_OPERATION_NAME,
+  GEN_AI_PROVIDER_NAME,
+  GEN_AI_REQUEST_MODEL,
+  GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS_METRIC_KEY,
+  GEN_AI_USAGE_CACHE_WRITE_INPUT_TOKENS_METRIC_KEY,
+  GEN_AI_USAGE_INPUT_TOKENS_METRIC_KEY,
+  GEN_AI_USAGE_OUTPUT_TOKENS_METRIC_KEY,
+  GEN_AI_USAGE_REASONING_OUTPUT_TOKENS_METRIC_KEY,
+  GEN_AI_USAGE_TOTAL_TOKENS_METRIC_KEY,
+  INPUT_TOKENS_METRIC_KEY,
+  METRIC_KEY_ALIASES,
+  OUTPUT_TOKENS_METRIC_KEY,
+  REASONING_OUTPUT_TOKENS_METRIC_KEY,
+  TOTAL_TOKENS_METRIC_KEY,
+} = require('./constants/tags')
+
+/** @type {Set<string | undefined>} */
+const MODEL_BACKED_SPAN_KINDS = new Set(['llm', 'embedding'])
+
+// null prototype: a metric named after an `Object.prototype` member must not resolve to an
+// inherited property
+const GEN_AI_USAGE_METRIC_KEYS = Object.assign(Object.create(null), {
+  [INPUT_TOKENS_METRIC_KEY]: GEN_AI_USAGE_INPUT_TOKENS_METRIC_KEY,
+  [OUTPUT_TOKENS_METRIC_KEY]: GEN_AI_USAGE_OUTPUT_TOKENS_METRIC_KEY,
+  [TOTAL_TOKENS_METRIC_KEY]: GEN_AI_USAGE_TOTAL_TOKENS_METRIC_KEY,
+  [CACHE_READ_INPUT_TOKENS_METRIC_KEY]: GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS_METRIC_KEY,
+  [CACHE_WRITE_INPUT_TOKENS_METRIC_KEY]: GEN_AI_USAGE_CACHE_WRITE_INPUT_TOKENS_METRIC_KEY,
+  [REASONING_OUTPUT_TOKENS_METRIC_KEY]: GEN_AI_USAGE_REASONING_OUTPUT_TOKENS_METRIC_KEY,
+})
+
+/**
+ * @typedef {object} GenAiApmTags
+ * @property {string} [spanKind] LLMObs span kind
+ * @property {string} [modelName]
+ * @property {string} [modelProvider]
+ * @property {string} [mlApp]
+ * @property {string} [sessionId]
+ * @property {Record<string, unknown>} [metrics] LLMObs metrics, in either spelling
+ */
+
+/**
+ * Writes the scalar `gen_ai.*` attributes onto the APM span, so model, provider, application,
+ * conversation and token usage are searchable in APM. Message bodies stay off the APM span.
+ *
+ * @param {import('../opentracing/span')} span
+ * @param {GenAiApmTags} tags
+ */
+function setGenAiApmTags (span, tags) {
+  // mirrors the LLMObs span event: a model-backed span always reports a model and provider
+  updateGenAiApmTags(span, MODEL_BACKED_SPAN_KINDS.has(tags.spanKind)
+    ? { ...tags, modelName: tags.modelName || DEFAULT_MODEL, modelProvider: tags.modelProvider || DEFAULT_MODEL }
+    : tags)
+}
+
+/**
+ * Writes only the `gen_ai.*` attributes present in `tags`, for values an integration resolves
+ * after the span started. Absent fields keep whatever the span already carries.
+ *
+ * @param {import('../opentracing/span')} span
+ * @param {GenAiApmTags} tags
+ */
+function updateGenAiApmTags (span, { spanKind, modelName, modelProvider, mlApp, sessionId, metrics }) {
+  const spanContext = span.context()
+
+  // written ahead of the attributes it covers: a throw partway through would otherwise leave
+  // `gen_ai.*` tags behind with nothing marking them as tracer-emitted
+  if (spanKind || modelName || modelProvider || mlApp || sessionId) {
+    markArtificialGenAiTags(spanContext)
+  }
+
+  if (spanKind) spanContext.setTag(GEN_AI_OPERATION_NAME, spanKind)
+  if (modelName) spanContext.setTag(GEN_AI_REQUEST_MODEL, modelName)
+  if (modelProvider) spanContext.setTag(GEN_AI_PROVIDER_NAME, modelProvider.toLowerCase())
+  if (mlApp) spanContext.setTag(GEN_AI_APPLICATION_NAME, mlApp)
+  if (sessionId) spanContext.setTag(GEN_AI_CONVERSATION_ID, sessionId)
+  if (metrics) setGenAiApmUsageMetrics(span, spanKind, metrics)
+}
+
+/**
+ * @param {import('../opentracing/span_context')} spanContext
+ */
+function markArtificialGenAiTags (spanContext) {
+  spanContext.setTag(ARTIFICIAL_GEN_AI_TAGS, 'true')
+}
+
+/**
+ * Writes the `gen_ai.usage.*` metrics onto the APM span. Accepts both the LLMObs metric keys and
+ * the camelCase spellings integrations extract before the tagger normalizes them.
+ *
+ * @param {import('../opentracing/span')} span
+ * @param {string | undefined} spanKind LLMObs span kind
+ * @param {Record<string, unknown>} metrics
+ */
+function setGenAiApmUsageMetrics (span, spanKind, metrics) {
+  // Other kinds carry unrelated metrics that would be misleading under a `gen_ai.usage.*` key.
+  if (!MODEL_BACKED_SPAN_KINDS.has(spanKind)) return
+
+  const spanContext = span.context()
+
+  // ahead of the metrics, for the same reason `updateGenAiApmTags` marks before its scalars
+  markArtificialGenAiTags(spanContext)
+
+  for (const [key, value] of Object.entries(metrics)) {
+    if (typeof value !== 'number') continue
+
+    const genAiKey = GEN_AI_USAGE_METRIC_KEYS[METRIC_KEY_ALIASES[key] ?? key]
+    if (!genAiKey) continue
+
+    spanContext.setTag(genAiKey, value)
+  }
+}
+
+module.exports = {
+  MODEL_BACKED_SPAN_KINDS,
+  setGenAiApmTags,
+  setGenAiApmUsageMetrics,
+  updateGenAiApmTags,
+}

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

@@ -89,6 +89,24 @@ class DdTelemetryPlugin extends BaseLLMObsPlugin {
     return { kind, name: getLlmObsSpanName(operation, ctx.attributes['ai.telemetry.functionId']) }
   }
 
+  /**
+   * @override
+   */
+  getGenAiApmEndTags (ctx, spanKind) {
+    if (spanKind !== 'llm' && spanKind !== 'embedding') return
+
+    const tags = /** @type {Record<string, string>} */ (getSpanTags(ctx))
+    const embeddingUsage = tags['ai.usage.tokens']
+
+    return {
+      modelName: tags['ai.model.id'],
+      modelProvider: getModelProvider(tags),
+      metrics: spanKind === 'embedding'
+        ? { inputTokens: embeddingUsage, totalTokens: embeddingUsage }
+        : getUsage(tags),
+    }
+  }
+
   /**
    * @override
    */

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

@@ -140,6 +140,20 @@ function formatLanguageModelOutputMessages (content) {
   return outputMessages
 }
 
+/**
+ * @param {object} [usage] AI SDK usage, from the result or the stream's `finish` chunk
+ * @returns {Record<string, number | undefined>}
+ */
+function extractUsageMetrics (usage) {
+  return {
+    inputTokens: usage?.inputTokens?.total,
+    cacheWriteTokens: usage?.inputTokens?.cacheWrite ?? 0,
+    cacheReadTokens: usage?.inputTokens?.cacheRead ?? 0,
+    outputTokens: usage?.outputTokens?.total,
+    reasoningOutputTokens: usage?.outputTokens?.reasoning ?? 0,
+  }
+}
+
 class VercelAiTelemetryPlugin extends BaseLLMObsPlugin {
   static id = 'ai_llmobs_vercel_telemetry'
   static integration = 'ai'
@@ -152,6 +166,13 @@ class VercelAiTelemetryPlugin extends BaseLLMObsPlugin {
     super(...arguments)
 
     this.addSub('dd-trace:vercel-ai:chunk', ({ ctx, chunk, done }) => {
+      if (!this._llmobsEnabledFor(ctx)) {
+        // only the token usage is needed, for the `gen_ai.usage.*` metrics; the message bodies and
+        // `ctx.result` are left to the LLMObs path
+        if (chunk?.type === 'finish') ctx.streamedUsage = chunk.usage
+        return
+      }
+
       ctx.chunks ??= []
       const chunks = ctx.chunks
       if (chunk) chunks.push(chunk)
@@ -202,6 +223,17 @@ class VercelAiTelemetryPlugin extends BaseLLMObsPlugin {
     super.asyncEnd(ctx)
   }
 
+  /**
+   * @override
+   */
+  getGenAiApmEndTags (ctx, spanKind) {
+    const usage = ctx.result?.usage ?? ctx.streamedUsage
+    if (!usage) return {}
+
+    // `embed` reports a single token count, the generation operations a structured breakdown
+    return { metrics: spanKind === 'embedding' ? { inputTokens: usage.tokens } : extractUsageMetrics(usage) }
+  }
+
   /**
    * @override
    */
@@ -367,14 +399,7 @@ class VercelAiTelemetryPlugin extends BaseLLMObsPlugin {
     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,
-    })
+    this._tagger.tagMetrics(span, extractUsageMetrics(result.usage))
   }
 
   setToolTags (span, ctx) {

package/packages/dd-trace/src/llmobs/plugins/anthropic/index.js

@@ -22,6 +22,13 @@ class AnthropicLLMObsPlugin extends LLMObsPlugin {
     super(...arguments)
 
     this.addSub('apm:anthropic:request:chunk', ({ ctx, chunk, done }) => {
+      if (!this._llmobsEnabledFor(ctx)) {
+        // only the token usage is needed, for the `gen_ai.usage.*` metrics; the message bodies
+        // and the aggregated response are left to the LLMObs path
+        if (chunk) ctx.streamedUsage = mergeChunkUsage(ctx.streamedUsage, chunk)
+        return
+      }
+
       ctx.chunks ??= []
       const chunks = ctx.chunks
       if (chunk) chunks.push(chunk)
@@ -36,9 +43,8 @@ class AnthropicLLMObsPlugin extends LLMObsPlugin {
             const { message } = chunk
             if (!message) continue
 
-            const { role, usage } = message
-            if (role) response.role = role
-            if (usage) response.usage = usage
+            if (message.role) response.role = message.role
+            response.usage = mergeChunkUsage(response.usage, chunk)
             break
           }
           case 'content_block_start': {
@@ -86,22 +92,10 @@ class AnthropicLLMObsPlugin extends LLMObsPlugin {
             break
           }
           case 'message_delta': {
-            const { delta } = chunk
-
-            const finishReason = delta?.stop_reason
+            const finishReason = chunk.delta?.stop_reason
             if (finishReason) response.finish_reason = finishReason
 
-            const { usage } = chunk
-            if (usage) {
-              const responseUsage = (response.usage ??= { input_tokens: 0, output_tokens: 0 })
-              responseUsage.output_tokens = usage.output_tokens
-
-              const cacheCreationTokens = usage.cache_creation_input_tokens
-              const cacheReadTokens = usage.cache_read_input_tokens
-              if (cacheCreationTokens) responseUsage.cache_creation_input_tokens = cacheCreationTokens
-              if (cacheReadTokens) responseUsage.cache_read_input_tokens = cacheReadTokens
-            }
-
+            response.usage = mergeChunkUsage(response.usage, chunk)
             break
           }
           case 'error': {
@@ -140,6 +134,13 @@ class AnthropicLLMObsPlugin extends LLMObsPlugin {
     return UNKNOWN_MODEL_PROVIDER
   }
 
+  /**
+   * @override
+   */
+  getGenAiApmEndTags (ctx) {
+    return { metrics: extractUsage(ctx.result ?? { usage: ctx.streamedUsage }) }
+  }
+
   setLLMObsTags (ctx) {
     const span = ctx.currentStore?.span
     if (!span) return
@@ -213,38 +214,73 @@ class AnthropicLLMObsPlugin extends LLMObsPlugin {
   }
 
   #tagAnthropicUsage (span, result) {
-    if (!result) return
+    const metrics = extractUsage(result)
+    if (metrics) this._tagger.tagMetrics(span, metrics)
+  }
+}
 
-    const { usage } = result
-    if (!usage) return
+/**
+ * Merges the token usage a streamed chunk carries into the usage accumulated so far.
+ *
+ * @param {Record<string, number> | undefined} usage
+ * @param {object} chunk
+ * @returns {Record<string, number> | undefined}
+ */
+function mergeChunkUsage (usage, chunk) {
+  if (chunk.type === 'message_start') {
+    // copied: the `message_delta` counts are merged into this object, and the chunk is handed back
+    // to the application
+    const startUsage = chunk.message?.usage
+    return startUsage ? { ...startUsage } : usage
+  }
 
-    const inputTokens = usage.input_tokens
-    const outputTokens = usage.output_tokens
-    const cacheWriteTokens = usage.cache_creation_input_tokens
-    const cacheReadTokens = usage.cache_read_input_tokens
+  if (chunk.type !== 'message_delta' || !chunk.usage) return usage
 
-    const metrics = {
-      inputTokens: (inputTokens ?? 0) + (cacheWriteTokens ?? 0) + (cacheReadTokens ?? 0),
-    }
+  const mergedUsage = usage ?? { input_tokens: 0, output_tokens: 0 }
+  mergedUsage.output_tokens = chunk.usage.output_tokens
+
+  const cacheCreationTokens = chunk.usage.cache_creation_input_tokens
+  const cacheReadTokens = chunk.usage.cache_read_input_tokens
+  if (cacheCreationTokens) mergedUsage.cache_creation_input_tokens = cacheCreationTokens
+  if (cacheReadTokens) mergedUsage.cache_read_input_tokens = cacheReadTokens
 
-    if (outputTokens) metrics.outputTokens = outputTokens
-    const totalTokens = metrics.inputTokens + (outputTokens ?? 0)
-    if (totalTokens) metrics.totalTokens = totalTokens
+  return mergedUsage
+}
 
-    if (cacheWriteTokens != null) metrics.cacheWriteTokens = cacheWriteTokens
-    if (cacheReadTokens != null) metrics.cacheReadTokens = cacheReadTokens
+/**
+ * @param {object} [result]
+ * @returns {Record<string, number> | undefined}
+ */
+function extractUsage (result) {
+  const usage = result?.usage
+  if (!usage) return
+
+  const inputTokens = usage.input_tokens
+  const outputTokens = usage.output_tokens
+  const cacheWriteTokens = usage.cache_creation_input_tokens
+  const cacheReadTokens = usage.cache_read_input_tokens
+
+  const metrics = {
+    inputTokens: (inputTokens ?? 0) + (cacheWriteTokens ?? 0) + (cacheReadTokens ?? 0),
+  }
 
-    const cacheCreation = usage.cache_creation
-    if (cacheCreation) {
-      metrics.cacheWrite5mTokens = cacheCreation.ephemeral_5m_input_tokens ?? 0
-      metrics.cacheWrite1hTokens = cacheCreation.ephemeral_1h_input_tokens ?? 0
-    } else if (cacheWriteTokens != null) {
-      metrics.cacheWrite5mTokens = cacheWriteTokens
-      metrics.cacheWrite1hTokens = 0
-    }
+  if (outputTokens) metrics.outputTokens = outputTokens
+  const totalTokens = metrics.inputTokens + (outputTokens ?? 0)
+  if (totalTokens) metrics.totalTokens = totalTokens
 
-    this._tagger.tagMetrics(span, metrics)
+  if (cacheWriteTokens != null) metrics.cacheWriteTokens = cacheWriteTokens
+  if (cacheReadTokens != null) metrics.cacheReadTokens = cacheReadTokens
+
+  const cacheCreation = usage.cache_creation
+  if (cacheCreation) {
+    metrics.cacheWrite5mTokens = cacheCreation.ephemeral_5m_input_tokens ?? 0
+    metrics.cacheWrite1hTokens = cacheCreation.ephemeral_1h_input_tokens ?? 0
+  } else if (cacheWriteTokens != null) {
+    metrics.cacheWrite5mTokens = cacheWriteTokens
+    metrics.cacheWrite1hTokens = 0
   }
+
+  return metrics
 }
 
 module.exports = AnthropicLLMObsPlugin