fix(bedrock): Log text and usage for streamed InvokeModel calls to OpenAI models

.changeset/bright-otters-stream.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix(bedrock): Log text and token usage for streamed InvokeModel calls to OpenAI models such as gpt-oss

js/src/instrumentation/plugins/bedrock-runtime-plugin.test.ts

@@ -594,6 +594,209 @@ describe("aggregateInvokeModelResponseStreamChunks", () => {
     });
   });
 
+  it("aggregates OpenAI (gpt-oss) stream text and invocation metrics from InvokeModel response stream chunks", () => {
+    // gpt-oss sends Chat Completions chunks with no `usage`; the token counts
+    // only arrive in the last chunk's amazon-bedrock-invocationMetrics.
+    expect(
+      aggregateInvokeModelResponseStreamChunks(
+        [
+          {
+            choices: [
+              {
+                delta: { content: "", role: "assistant" },
+                finish_reason: null,
+                index: 0,
+              },
+            ],
+            object: "chat.completion.chunk",
+          },
+          {
+            choices: [
+              { delta: { content: "Hello" }, finish_reason: null, index: 0 },
+            ],
+            object: "chat.completion.chunk",
+          },
+          {
+            choices: [
+              { delta: { content: " world" }, finish_reason: null, index: 0 },
+            ],
+            object: "chat.completion.chunk",
+          },
+          {
+            choices: [{ delta: {}, finish_reason: "stop", index: 0 }],
+            object: "chat.completion.chunk",
+            "amazon-bedrock-invocationMetrics": {
+              inputTokenCount: 126,
+              outputTokenCount: 125,
+              invocationLatency: 2006,
+              firstByteLatency: 685,
+            },
+          },
+        ].map((event) => ({
+          chunk: { bytes: new TextEncoder().encode(JSON.stringify(event)) },
+        })),
+      ),
+    ).toEqual({
+      output: [
+        {
+          index: 0,
+          message: { role: "assistant", content: "Hello world" },
+          finish_reason: "stop",
+          logprobs: null,
+        },
+      ],
+      metrics: {
+        completion_tokens: 125,
+        prompt_tokens: 126,
+        tokens: 251,
+      },
+    });
+  });
+
+  it.each([undefined, "Checking the weather."])(
+    "preserves OpenAI tool calls alongside content %j",
+    (content) => {
+      const events = [
+        {
+          choices: [
+            {
+              index: 0,
+              delta: {
+                role: "assistant",
+                content,
+                reasoning_content: "Look up the weather.",
+              },
+            },
+          ],
+        },
+        {
+          choices: [
+            {
+              index: 0,
+              delta: {
+                tool_calls: [
+                  {
+                    index: 0,
+                    id: "call_weather",
+                    type: "function",
+                    function: {
+                      name: "get_weather",
+                      arguments: '{"city":',
+                    },
+                  },
+                ],
+              },
+            },
+          ],
+        },
+        {
+          choices: [
+            {
+              index: 0,
+              delta: {
+                tool_calls: [
+                  {
+                    index: 0,
+                    function: { arguments: '"Berlin"}' },
+                  },
+                ],
+              },
+            },
+          ],
+        },
+        {
+          choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }],
+          "amazon-bedrock-invocationMetrics": {
+            inputTokenCount: 12,
+            outputTokenCount: 20,
+          },
+        },
+      ];
+
+      expect(
+        aggregateInvokeModelResponseStreamChunks(
+          events.map((event) => ({
+            chunk: { bytes: new TextEncoder().encode(JSON.stringify(event)) },
+          })),
+        ),
+      ).toEqual({
+        output: [
+          {
+            index: 0,
+            message: {
+              role: "assistant",
+              content,
+              reasoning_content: "Look up the weather.",
+              tool_calls: [
+                {
+                  id: "call_weather",
+                  type: "function",
+                  function: {
+                    name: "get_weather",
+                    arguments: '{"city":"Berlin"}',
+                  },
+                },
+              ],
+            },
+            finish_reason: "tool_calls",
+            logprobs: null,
+          },
+        ],
+        metrics: {
+          prompt_tokens: 12,
+          completion_tokens: 20,
+          tokens: 32,
+        },
+      });
+    },
+  );
+
+  it("prefers Nova metadata usage over its zeroed invocation metrics in InvokeModel response stream chunks", () => {
+    // Nova reports inputTokenCount/outputTokenCount as 0 in
+    // amazon-bedrock-invocationMetrics; the real counts are in metadata.usage.
+    expect(
+      aggregateInvokeModelResponseStreamChunks(
+        [
+          { messageStart: { role: "assistant" } },
+          {
+            contentBlockDelta: {
+              delta: { text: "Hello" },
+              contentBlockIndex: 0,
+            },
+          },
+          { contentBlockStop: { contentBlockIndex: 0 } },
+          { messageStop: { stopReason: "end_turn" } },
+          {
+            metadata: {
+              usage: {
+                inputTokens: 6,
+                outputTokens: 5,
+                cacheReadInputTokenCount: 0,
+                cacheWriteInputTokenCount: 0,
+              },
+              metrics: {},
+              trace: {},
+            },
+            "amazon-bedrock-invocationMetrics": {
+              inputTokenCount: 0,
+              outputTokenCount: 0,
+              invocationLatency: 428,
+              firstByteLatency: 365,
+            },
+          },
+        ].map((event) => ({
+          chunk: { bytes: new TextEncoder().encode(JSON.stringify(event)) },
+        })),
+      ).metrics,
+    ).toEqual({
+      completion_tokens: 5,
+      prompt_cache_creation_tokens: 0,
+      prompt_cached_tokens: 0,
+      prompt_tokens: 6,
+      tokens: 11,
+    });
+  });
+
   it("aggregates Titan stream outputText from InvokeModel response stream chunks", () => {
     expect(
       aggregateInvokeModelResponseStreamChunks([

js/src/instrumentation/plugins/bedrock-runtime-plugin.ts

@@ -2,6 +2,7 @@ import type { ReturnOf } from "../core/interceptor-definitions";
 import type { StartOf } from "../core/tracing-types";
 import { SpanTypeAttribute, isObject } from "../../../util/index";
 import type { Span } from "../../logger";
+import type { OpenAIChatCompletionChunk } from "../../vendor-sdk-types/openai";
 import { INSTRUMENTATION_NAMES } from "../../span-origin";
 import { getCurrentUnixTimestamp } from "../../util";
 import { BasePlugin } from "../core";
@@ -27,6 +28,7 @@ import {
   getBedrockRuntimeCommandName,
   getBedrockRuntimeOperation,
 } from "./bedrock-runtime-common";
+import { aggregateChatCompletionChunks } from "./openai-plugin";
 
 export class BedrockRuntimePlugin extends BasePlugin {
   protected onEnable(): void {
@@ -530,6 +532,9 @@ export function aggregateInvokeModelResponseStreamChunks(
     .map((chunk) => parseJsonBody(chunk.chunk?.bytes))
     .filter((chunk) => chunk !== undefined);
   const jsonLikeChunks = parsedChunks.filter(isObject);
+  const chatCompletionChunks = jsonLikeChunks.filter((chunk) =>
+    Array.isArray(chunk.choices),
+  ) as OpenAIChatCompletionChunk[];
   const text = parsedChunks.map(extractTextFromJsonLike).join("");
   const lastMetadataChunk = jsonLikeChunks
     .slice()
@@ -556,25 +561,42 @@ export function aggregateInvokeModelResponseStreamChunks(
 
   return {
     output:
-      text.length > 0
-        ? { text }
-        : {
-            chunk_count: chunks.length,
-            chunks: sanitizeBedrockValue(
-              (jsonLikeChunks.length > 0 ? jsonLikeChunks : chunks).slice(
-                0,
-                20,
+      chatCompletionChunks.length > 0
+        ? aggregateChatCompletionChunks(chatCompletionChunks).output
+        : text.length > 0
+          ? { text }
+          : {
+              chunk_count: chunks.length,
+              chunks: sanitizeBedrockValue(
+                (jsonLikeChunks.length > 0 ? jsonLikeChunks : chunks).slice(
+                  0,
+                  20,
+                ),
               ),
-            ),
-          },
+            },
     metrics: parseBedrockRuntimeMetrics(
-      usage ?? (isObject(metadata) ? metadata.usage : undefined),
+      usage ??
+        (isObject(metadata) ? metadata.usage : undefined) ??
+        lastInvocationMetrics(jsonLikeChunks),
       isObject(metadata) ? metadata.metrics : undefined,
     ),
     ...(metadata ? { metadata } : {}),
   };
 }
 
+// Models whose stream carries no usage of its own (gpt-oss, for example) only
+// report token counts in the amazon-bedrock-invocationMetrics of the last
+// chunk. Nova sends zeros there, so it is only a fallback after metadata.usage.
+function lastInvocationMetrics(
+  chunks: Record<string, unknown>[],
+): Record<string, unknown> | undefined {
+  const metrics = chunks
+    .map((chunk) => chunk["amazon-bedrock-invocationMetrics"])
+    .reverse()
+    .find(isObject);
+  return metrics ? sanitizeRecord(metrics) : undefined;
+}
+
 function logBedrockStreamAggregation(
   span: Span,
   aggregated: BedrockRuntimeStreamAggregation,
@@ -819,6 +841,5 @@ function extractTextFromJsonLike(value: unknown): string {
   if (Array.isArray(value.outputs)) {
     return value.outputs.map(extractTextFromJsonLike).join("");
   }
-
   return "";
 }