DataDog
diff --git a/‎packages/datadog-instrumentations/src/aws-sdk.js‎
Lines changed: 3 additions & 2 deletions b/‎packages/datadog-instrumentations/src/aws-sdk.js‎
Lines changed: 3 additions & 2 deletions
diff --git a/‎packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/tracing.js‎
Lines changed: 1 addition & 1 deletion b/‎packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/tracing.js‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/utils.js‎
Lines changed: 218 additions & 4 deletions b/‎packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/utils.js‎
Lines changed: 218 additions & 4 deletions
diff --git a/‎packages/datadog-plugin-aws-sdk/test/bedrockruntime.spec.js‎
Lines changed: 41 additions & 1 deletion b/‎packages/datadog-plugin-aws-sdk/test/bedrockruntime.spec.js‎
Lines changed: 41 additions & 1 deletion
diff --git a/‎packages/datadog-plugin-aws-sdk/test/bedrockruntime.util.spec.js‎
Lines changed: 25 additions & 0 deletions b/‎packages/datadog-plugin-aws-sdk/test/bedrockruntime.util.spec.js‎
Lines changed: 25 additions & 0 deletions
@@ -233,13 +233,14 @@ function wrapSmithySend (send) {
 }
 
 function handleCompletion (result, ctx, channels) {
-  const iterator = result?.body?.[Symbol.asyncIterator]
+  const streamable = result?.body ?? result?.stream
+  const iterator = streamable?.[Symbol.asyncIterator]
   if (!iterator) {
     channels.complete.publish(ctx)
     return
   }
 
-  shimmer.wrap(result.body, Symbol.asyncIterator, function (asyncIterator) {
+  shimmer.wrap(streamable, Symbol.asyncIterator, function (asyncIterator) {
     return function (...args) {
       const iterator = asyncIterator.apply(this, args)
       shimmer.wrap(iterator, 'next', function (next) {
 
@@ -3,7 +3,7 @@
 const BaseAwsSdkPlugin = require('../../base')
 const { parseModelId } = require('./utils')
 
-const enabledOperations = new Set(['invokeModel', 'invokeModelWithResponseStream'])
+const enabledOperations = new Set(['invokeModel', 'invokeModelWithResponseStream', 'converse', 'converseStream'])
 
 class BedrockRuntime extends BaseAwsSdkPlugin {
   static id = 'bedrockruntime'
 
@@ -131,6 +131,7 @@ class Generation {
     outputTokens,
     cacheReadTokens,
     cacheWriteTokens,
+    messages,
   } = {}) {
     // stringify message as it could be a single generated message as well as a list of embeddings
     this.message = typeof message === 'string' ? message : JSON.stringify(message) || ''
@@ -143,6 +144,7 @@ class Generation {
       cacheReadTokens,
       cacheWriteTokens,
     }
+    this.messages = messages ?? [{ content: this.message, role: this.role }]
   }
 }
 
@@ -401,10 +403,7 @@ function extractTextAndResponseReason (response, provider, modelName) {
             message: output.message?.content[0]?.text ?? 'Unsupported content type',
             finishReason: body.stopReason,
             role: output.message?.role,
-            inputTokens: body.usage?.inputTokens,
-            outputTokens: body.usage?.outputTokens,
-            cacheReadInputTokenCount: body.usage?.cacheReadInputTokenCount,
-            cacheWriteInputTokenCount: body.usage?.cacheWriteInputTokenCount,
+            ...buildUsage(body.usage),
           })
         }
         break
@@ -476,12 +475,227 @@ function extractTextAndResponseReason (response, provider, modelName) {
   return new Generation()
 }
 
+/**
+ * Convert a Converse content-block array to an LLMObs message array.
+ *
+ * @param {string} role
+ * @param {Array<object>} contentBlocks
+ * @returns {{ content?: string, role: string, toolCalls?: Array, toolResults?: Array } | undefined}
+ */
+function extractMessagesFromConverseContent (role, contentBlocks) {
+  let content = ''
+  const toolCalls = []
+  const toolResults = []
+
+  for (const block of contentBlocks || []) {
+    if (block == null || typeof block !== 'object') continue
+    if (typeof block.text === 'string') {
+      content += block.text
+    } else if (block.toolUse) {
+      toolCalls.push(buildToolCall(block.toolUse))
+    } else if (block.toolResult) {
+      toolResults.push(buildToolResult(block.toolResult))
+    } else {
+      content += `[Unsupported content type: ${getContentBlockType(block)}]`
+    }
+  }
+
+  if (!content && toolCalls.length === 0 && toolResults.length === 0) return
+
+  const message = { role }
+  if (content) message.content = content
+  if (toolCalls.length > 0) message.toolCalls = toolCalls
+  if (toolResults.length > 0) message.toolResults = toolResults
+  return message
+}
+
+/**
+ * Resolve a Converse `ContentBlock`'s member type. The block is a key-presence
+ * tagged union (no `type` discriminator), so the active member is its sole own
+ * key. For forward-compat `$unknown` members the real type is the first element
+ * of the `[name, value]` tuple.
+ *
+ * @param {object} block
+ * @returns {string}
+ */
+function getContentBlockType (block) {
+  const key = Object.keys(block)[0]
+  if (key === '$unknown') return block.$unknown?.[0] ?? 'unknown'
+  return key ?? 'unknown'
+}
+
+// Always emit at least one output message so downstream tagging has a role to attach to.
+function toOutputMessages (role, contentBlocks) {
+  const message = extractMessagesFromConverseContent(role, contentBlocks)
+  return message ? [message] : [{ role, content: '' }]
+}
+
+function buildToolCall ({ name, input, toolUseId }) {
+  return { name: name ?? '', arguments: input ?? {}, toolId: toolUseId ?? '', type: 'toolUse' }
+}
+
+function parseToolInput (inputStr) {
+  try {
+    return JSON.parse(inputStr)
+  } catch {
+    log.warn('Failed to parse Converse stream toolUse.input JSON; emitting empty arguments')
+    return {}
+  }
+}
+
+function buildToolResult ({ toolUseId, content }) {
+  const result = (content || []).map(resolveToolResultItem).join('')
+  return { name: '', result, toolId: toolUseId ?? '', type: 'tool_result' }
+}
+
+function resolveToolResultItem (item) {
+  if (typeof item.text === 'string') return item.text
+  if (item.json != null) return JSON.stringify(item.json)
+  return `[Unsupported content type(s): ${getContentBlockType(item)}]`
+}
+
+function buildUsage (usage = {}) {
+  return {
+    inputTokens: usage.inputTokens,
+    outputTokens: usage.outputTokens,
+    cacheReadTokens: usage.cacheReadInputTokens ?? usage.cacheReadInputTokenCount,
+    cacheWriteTokens: usage.cacheWriteInputTokens ?? usage.cacheWriteInputTokenCount,
+  }
+}
+
+/**
+ * Extract tool definitions from a Converse request's `toolConfig.tools`,
+ * mapping Bedrock's `toolSpec` shape to LLMObs `ToolDefinition` shape.
+ *
+ * @param {object} params - Converse request params with optional `toolConfig.tools[].toolSpec`.
+ * @returns {Array<{ name: string, description: string, schema: object }>}
+ */
+function extractConverseToolDefinitions (params) {
+  const toolDefinitions = []
+  for (const tool of params.toolConfig?.tools || []) {
+    const toolSpec = tool?.toolSpec
+    if (!toolSpec?.name) continue
+    toolDefinitions.push({
+      name: toolSpec.name,
+      description: toolSpec.description ?? '',
+      schema: toolSpec.inputSchema ?? {},
+    })
+  }
+  return toolDefinitions
+}
+
+/**
+ * Extract request metadata + rendered input messages from a Converse /
+ * ConverseStream request.
+ *
+ * @param {{ modelId?: string, messages?: Array, system?: Array, inferenceConfig?: object, toolConfig?: object }} params
+ * @returns {RequestParams}
+ */
+function extractRequestParamsConverse (params) {
+  const prompt = []
+  for (const block of params.system || []) {
+    if (typeof block?.text === 'string') prompt.push({ content: block.text, role: 'system' })
+  }
+  for (const msg of params.messages || []) {
+    if (msg == null || typeof msg !== 'object') continue
+    const message = extractMessagesFromConverseContent(msg.role || 'user', msg.content)
+    if (message) prompt.push(message)
+  }
+
+  const { temperature, topP, maxTokens, stopSequences } = params.inferenceConfig || {}
+  return new RequestParams({ prompt, temperature, topP, maxTokens, stopSequences })
+}
+
+/**
+ * Extract output messages + usage from a non-stream Converse response.
+ *
+ * @param {{ output?: { message?: { role?: string, content?: Array } }, stopReason?: string, usage?: object }} response
+ * @returns {Generation}
+ */
+function extractTextAndResponseReasonConverse (response) {
+  const outputMessage = response?.output?.message
+  const role = outputMessage?.role || 'assistant'
+
+  return new Generation({
+    role,
+    finishReason: response?.stopReason || '',
+    ...buildUsage(response?.usage),
+    messages: toOutputMessages(role, outputMessage?.content),
+  })
+}
+
+/**
+ * Aggregate Converse stream events into a single output message + usage.
+ * One messageStart / messageStop pair per response, so one message out.
+ *
+ * Stream events describe the same content-block structure as the non-stream
+ * response, spread across start/delta chunks. We reassemble those chunks
+ * into a normalized content-block array and reuse the non-stream extractor.
+ *
+ * @param {Array<object>} chunks - Ordered ConverseStreamOutput events.
+ * @returns {Generation}
+ */
+function extractTextAndResponseReasonConverseFromStream (chunks) {
+  let role = 'assistant'
+  let stopReason = ''
+  let usage = {}
+  const blocksByIdx = new Map()
+
+  for (const chunk of chunks || []) {
+    if (chunk.messageStart?.role) {
+      role = chunk.messageStart.role
+    } else if (chunk.messageStop?.stopReason) {
+      stopReason = chunk.messageStop.stopReason
+    } else if (chunk.metadata?.usage) {
+      usage = chunk.metadata.usage
+    } else if (chunk.contentBlockStart?.start?.toolUse) {
+      const { contentBlockIndex, start: { toolUse } } = chunk.contentBlockStart
+      blocksByIdx.set(contentBlockIndex, {
+        toolUse: { toolUseId: toolUse.toolUseId, name: toolUse.name, inputStr: '' },
+      })
+    } else if (chunk.contentBlockDelta) {
+      const { contentBlockIndex, delta } = chunk.contentBlockDelta
+      if (typeof delta?.text === 'string') {
+        const block = blocksByIdx.get(contentBlockIndex) ?? {}
+        block.text = (block.text ?? '') + delta.text
+        blocksByIdx.set(contentBlockIndex, block)
+      } else if (typeof delta?.toolUse?.input === 'string') {
+        const block = blocksByIdx.get(contentBlockIndex) ?? { toolUse: { inputStr: '' } }
+        block.toolUse ??= { inputStr: '' }
+        block.toolUse.inputStr += delta.toolUse.input
+        blocksByIdx.set(contentBlockIndex, block)
+      }
+    }
+  }
+
+  const contentBlocks = [...blocksByIdx.keys()].sort((a, b) => a - b).map(i => {
+    const block = blocksByIdx.get(i)
+    if (block.toolUse) {
+      const { toolUseId, name, inputStr } = block.toolUse
+      block.toolUse = { toolUseId, name, input: parseToolInput(inputStr) }
+    }
+    return block
+  })
+
+  return new Generation({
+    role,
+    finishReason: stopReason,
+    ...buildUsage(usage),
+    messages: toOutputMessages(role, contentBlocks),
+  })
+}
+
 module.exports = {
   Generation,
   RequestParams,
   extractTextAndResponseReasonFromStream,
   parseModelId,
   extractRequestParams,
   extractTextAndResponseReason,
+  extractMessagesFromConverseContent,
+  extractConverseToolDefinitions,
+  extractRequestParamsConverse,
+  extractTextAndResponseReasonConverse,
+  extractTextAndResponseReasonConverseFromStream,
   PROVIDER,
 }
@@ -5,7 +5,7 @@ const { describe, it, before, after } = require('mocha')
 
 const agent = require('../../dd-trace/test/plugins/agent')
 const { setup, withAwsSdkVersions } = require('./spec_helpers')
-const { models } = require('./fixtures/bedrockruntime')
+const { models, converseRequest } = require('./fixtures/bedrockruntime')
 const serviceName = 'bedrock-service-name-test'
 
 describe('Plugin', () => {
@@ -97,6 +97,46 @@ describe('Plugin', () => {
             await tracesPromise
           })
         })
+
+        it('should converse', async function () {
+          if (typeof AWS.ConverseCommand !== 'function') return this.skip()
+          const command = new AWS.ConverseCommand({ modelId: converseRequest.modelId, ...converseRequest.request })
+
+          const tracesPromise = agent.assertFirstTraceSpan({
+            meta: {
+              'aws.operation': 'converse',
+              'aws.bedrock.request.model': converseRequest.modelId.split('.')[1],
+              'aws.bedrock.request.model_provider': converseRequest.provider.toLowerCase(),
+            },
+          })
+          tracesPromise.catch(() => {}) // silence unhandled rejection if `send` throws first
+
+          await bedrockRuntimeClient.send(command)
+          await tracesPromise
+        })
+
+        it('should converse-stream', async function () {
+          if (typeof AWS.ConverseStreamCommand !== 'function') return this.skip()
+          const command = new AWS.ConverseStreamCommand({
+            modelId: converseRequest.modelId,
+            ...converseRequest.request,
+          })
+
+          const tracesPromise = agent.assertFirstTraceSpan({
+            meta: {
+              'aws.operation': 'converseStream',
+              'aws.bedrock.request.model': converseRequest.modelId.split('.')[1],
+              'aws.bedrock.request.model_provider': converseRequest.provider.toLowerCase(),
+            },
+          })
+          tracesPromise.catch(() => {}) // silence unhandled rejection if stream iteration throws first
+
+          const result = await bedrockRuntimeClient.send(command)
+          for await (const _event of result.stream) { // eslint-disable-line no-unused-vars
+            // drain
+          }
+          await tracesPromise
+        })
       })
     })
   })
 
@@ -0,0 +1,25 @@
+'use strict'
+
+const assert = require('node:assert/strict')
+
+const { describe, it } = require('mocha')
+
+const {
+  extractTextAndResponseReasonConverseFromStream,
+} = require('../src/services/bedrockruntime/utils')
+
+describe('bedrockruntime converse stream extractor', () => {
+  it('emits empty tool-call arguments when the streamed tool-use input is malformed JSON', () => {
+    const generation = extractTextAndResponseReasonConverseFromStream([
+      { messageStart: { role: 'assistant' } },
+      { contentBlockStart: { contentBlockIndex: 0, start: { toolUse: { toolUseId: 't-1', name: 'get_weather' } } } },
+      { contentBlockDelta: { contentBlockIndex: 0, delta: { toolUse: { input: 'not valid json{' } } } },
+      { messageStop: { stopReason: 'tool_use' } },
+    ])
+
+    assert.deepStrictEqual(generation.messages, [{
+      role: 'assistant',
+      toolCalls: [{ name: 'get_weather', arguments: {}, toolId: 't-1', type: 'toolUse' }],
+    }])
+  })
+})