feat(server-utils): Record Bedrock responses as output messages

dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts

@@ -1,6 +1,7 @@
 import {
   GEN_AI_INPUT_MESSAGES,
   GEN_AI_OPERATION_NAME,
+  GEN_AI_OUTPUT_MESSAGES,
   GEN_AI_PROVIDER_NAME,
   GEN_AI_REQUEST_MAX_TOKENS,
   GEN_AI_REQUEST_MODEL,
@@ -59,7 +60,11 @@ function assertBedrockSpans(container: SerializedStreamedSpanContainer): void {
 
 const SYSTEM_INSTRUCTIONS = JSON.stringify([{ type: 'text', content: 'You are a weather bot.' }]);
 
-function assertBedrockInputs(container: SerializedStreamedSpanContainer): void {
+const OUTPUT_MESSAGES = JSON.stringify([
+  { role: 'assistant', parts: [{ type: 'text', content: 'Hello from Bedrock' }], finish_reason: 'end_turn' },
+]);
+
+function assertBedrockContent(container: SerializedStreamedSpanContainer): void {
   const converseSpan = container.items.find(span => span.name === `chat ${MODEL_ID}`);
   expect(converseSpan!.attributes[GEN_AI_SYSTEM_INSTRUCTIONS].value).toBe(SYSTEM_INSTRUCTIONS);
   expect(converseSpan!.attributes[GEN_AI_INPUT_MESSAGES].value).toBe(
@@ -68,12 +73,14 @@ function assertBedrockInputs(container: SerializedStreamedSpanContainer): void {
   expect(converseSpan!.attributes[GEN_AI_TOOL_DEFINITIONS].value).toBe(
     JSON.stringify([{ name: 'get_weather', inputSchema: { json: { type: 'object' } } }]),
   );
+  expect(converseSpan!.attributes[GEN_AI_OUTPUT_MESSAGES].value).toBe(OUTPUT_MESSAGES);
 
   const invokeModelSpan = container.items.find(span => span.name === `generate_content ${MODEL_ID}`);
   expect(invokeModelSpan!.attributes[GEN_AI_SYSTEM_INSTRUCTIONS].value).toBe(SYSTEM_INSTRUCTIONS);
   expect(invokeModelSpan!.attributes[GEN_AI_INPUT_MESSAGES].value).toBe(
     JSON.stringify([{ role: 'user', content: 'Hello' }]),
   );
+  expect(invokeModelSpan!.attributes[GEN_AI_OUTPUT_MESSAGES].value).toBe(OUTPUT_MESSAGES);
 }
 
 function assertStreamSpans(container: SerializedStreamedSpanContainer): void {
@@ -113,8 +120,8 @@ describe('awsIntegration - Bedrock', () => {
   });
 
   createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument-gen-ai.mjs', (createTestRunner, test) => {
-    test('records Bedrock inputs by default', { timeout: 90_000 }, async () => {
-      await createTestRunner().ignore('event').expect({ span: assertBedrockInputs }).start().completed();
+    test('records Bedrock inputs and outputs by default', { timeout: 90_000 }, async () => {
+      await createTestRunner().ignore('event').expect({ span: assertBedrockContent }).start().completed();
     });
   });
 

packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime-content.ts

@@ -1,6 +1,8 @@
+import type { Span } from '@sentry/core';
 import { isObjectLike, stringify } from '@sentry/core';
 import {
   GEN_AI_INPUT_MESSAGES,
+  GEN_AI_OUTPUT_MESSAGES,
   GEN_AI_SYSTEM_INSTRUCTIONS,
   GEN_AI_TOOL_DEFINITIONS,
 } from '@sentry/conventions/attributes';
@@ -11,11 +13,11 @@ import type { GenAiMessagePart } from '../../../ai/types';
  * Translate Converse `ContentBlock`s (also used by Amazon Nova `InvokeModel` bodies) into message parts.
  * The conventions require a `type` on each part, and Bedrock's blocks (`{ text }`, `{ toolUse }`, ...) have none.
  *
- * Media is reported by modality only: inline media is raw bytes, which the conventions require to be
- * dropped rather than recorded. Blocks without content worth recording (`cachePoint`, redacted
+ * Media is reported by modality only: inline media is raw bytes, which would serialize as a large array of
+ * numbers. Blocks without content worth recording (`cachePoint`, redacted
  * reasoning, unknown members) are left out rather than serialized, so no binary payload slips through.
  */
-function contentBlocksToParts(content: unknown): GenAiMessagePart[] {
+export function contentBlocksToParts(content: unknown): GenAiMessagePart[] {
   return Array.isArray(content) ? content.flatMap(contentBlockToParts) : [];
 }
 
@@ -156,3 +158,69 @@ export function getPromptInputAttributes(prompt: unknown): Record<string, unknow
   }
   return { [GEN_AI_INPUT_MESSAGES]: JSON.stringify([{ role: 'user', parts: [{ type: 'text', content: prompt }] }]) };
 }
+
+/** Output parts of an Anthropic Messages response body (Claude `InvokeModel`). */
+export function getClaudeOutputParts(content: unknown): GenAiMessagePart[] {
+  if (!Array.isArray(content)) {
+    return [];
+  }
+
+  return content.flatMap((block): GenAiMessagePart[] => {
+    if (!isObjectLike(block)) {
+      return [];
+    }
+    if (block.type === 'text') {
+      return textParts(block.text);
+    }
+    if (block.type === 'thinking') {
+      return typeof block.thinking === 'string' ? [{ type: 'reasoning', content: block.thinking }] : [];
+    }
+    if (block.type === 'tool_use') {
+      return [{ type: 'tool_call', id: block.id, name: block.name, arguments: stringify(block.input ?? {}, String) }];
+    }
+    return [];
+  });
+}
+
+export function textParts(text: unknown): GenAiMessagePart[] {
+  return typeof text === 'string' ? [{ type: 'text', content: text }] : [];
+}
+
+/** Record the parts as one assistant message. Parts of a stream are kept at their block index, so skip the gaps. */
+export function setOutputMessages(span: Span, parts: GenAiMessagePart[], finishReason: unknown): void {
+  const filledParts = parts.filter(Boolean);
+  if (filledParts.length) {
+    // An `undefined` finish reason is left out of the JSON.
+    span.setAttribute(
+      GEN_AI_OUTPUT_MESSAGES,
+      JSON.stringify([{ role: 'assistant', parts: filledParts, finish_reason: finishReason }]),
+    );
+  }
+}
+
+/** Append a streamed fragment to the part at the block `index`, and open a `type` part there if there is none. */
+export function appendStreamPart(
+  parts: GenAiMessagePart[] | undefined,
+  index: unknown,
+  type: 'text' | 'reasoning' | 'tool_call',
+  fragment: unknown,
+): void {
+  if (!parts || typeof fragment !== 'string') {
+    return;
+  }
+  // Tool call arguments stream in as fragments of their JSON.
+  const field = type === 'tool_call' ? 'arguments' : 'content';
+  const part = (parts[typeof index === 'number' ? index : 0] ??= { type, [field]: '' });
+  part[field] = `${part[field]}${fragment}`;
+}
+
+export function startStreamToolCall(
+  parts: GenAiMessagePart[] | undefined,
+  index: unknown,
+  id: unknown,
+  name: unknown,
+): void {
+  if (parts) {
+    parts[typeof index === 'number' ? index : 0] = { type: 'tool_call', id, name, arguments: '' };
+  }
+}