feat(server-utils): Capture usage/output on streamed Vercel AI spans (v7 + orchestrion)

Author: logaretm

dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-structured-output.mjs

@@ -0,0 +1,48 @@
+import * as Sentry from '@sentry/node';
+import { Output, streamText } from 'ai';
+import { MockLanguageModelV3, simulateReadableStream } from 'ai/test';
+import { z } from 'zod';
+
+// `streamObject` is deprecated in `ai` v7 (it publishes no telemetry channel events); structured
+// output now flows through the `streamText` primitive with an `experimental_output`, which does. So a
+// streamed structured-output call must still produce the usual streamText/doStream spans, with the
+// streamed JSON captured as the output.
+async function run() {
+  await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
+    const result = streamText({
+      experimental_telemetry: { isEnabled: true, recordInputs: true, recordOutputs: true },
+      experimental_output: Output.object({
+        schema: z.object({ city: z.string(), weather: z.string() }),
+      }),
+      model: new MockLanguageModelV3({
+        doStream: async () => ({
+          stream: simulateReadableStream({
+            chunks: [
+              { type: 'stream-start', warnings: [] },
+              { type: 'text-start', id: '0' },
+              { type: 'text-delta', id: '0', delta: '{"city":"San Francisco",' },
+              { type: 'text-delta', id: '0', delta: '"weather":"sunny"}' },
+              { type: 'text-end', id: '0' },
+              {
+                type: 'finish',
+                finishReason: { unified: 'stop', raw: 'stop' },
+                usage: {
+                  inputTokens: { total: 12, noCache: 12, cached: 0 },
+                  outputTokens: { total: 18, noCache: 18, cached: 0 },
+                  totalTokens: { total: 30, noCache: 30, cached: 0 },
+                },
+              },
+            ],
+          }),
+        }),
+      }),
+      prompt: 'What is the weather in San Francisco?',
+    });
+
+    for await (const _part of result.fullStream) {
+      void _part;
+    }
+  });
+}
+
+run();

dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-tools.mjs

@@ -0,0 +1,78 @@
+import * as Sentry from '@sentry/node';
+import { stepCountIs, streamText, tool } from 'ai';
+import { MockLanguageModelV3, simulateReadableStream } from 'ai/test';
+import { z } from 'zod';
+
+async function run() {
+  await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
+    let callCount = 0;
+
+    const result = streamText({
+      experimental_telemetry: { isEnabled: true, recordInputs: true, recordOutputs: true },
+      model: new MockLanguageModelV3({
+        doStream: async () => {
+          // First step streams a tool call; after the tool runs, the second step streams the answer.
+          if (callCount++ === 0) {
+            return {
+              stream: simulateReadableStream({
+                chunks: [
+                  { type: 'stream-start', warnings: [] },
+                  {
+                    type: 'tool-call',
+                    toolCallId: 'call-1',
+                    toolName: 'getWeather',
+                    input: JSON.stringify({ location: 'San Francisco' }),
+                  },
+                  {
+                    type: 'finish',
+                    finishReason: { unified: 'tool-calls', raw: 'tool_calls' },
+                    usage: {
+                      inputTokens: { total: 10, noCache: 10, cached: 0 },
+                      outputTokens: { total: 20, noCache: 20, cached: 0 },
+                      totalTokens: { total: 30, noCache: 30, cached: 0 },
+                    },
+                  },
+                ],
+              }),
+            };
+          }
+          return {
+            stream: simulateReadableStream({
+              chunks: [
+                { type: 'stream-start', warnings: [] },
+                { type: 'text-start', id: '0' },
+                { type: 'text-delta', id: '0', delta: 'Sunny, ' },
+                { type: 'text-delta', id: '0', delta: '72°F.' },
+                { type: 'text-end', id: '0' },
+                {
+                  type: 'finish',
+                  finishReason: { unified: 'stop', raw: 'stop' },
+                  usage: {
+                    inputTokens: { total: 15, noCache: 15, cached: 0 },
+                    outputTokens: { total: 25, noCache: 25, cached: 0 },
+                    totalTokens: { total: 40, noCache: 40, cached: 0 },
+                  },
+                },
+              ],
+            }),
+          };
+        },
+      }),
+      tools: {
+        getWeather: tool({
+          description: 'Get the current weather for a location',
+          inputSchema: z.object({ location: z.string() }),
+          execute: async ({ location }) => `Weather in ${location}: Sunny, 72°F`,
+        }),
+      },
+      stopWhen: stepCountIs(5),
+      prompt: 'What is the weather in San Francisco?',
+    });
+
+    for await (const _part of result.fullStream) {
+      void _part;
+    }
+  });
+}
+
+run();

dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/test.ts

@@ -595,6 +595,136 @@ describe.each(matrix)('Vercel AI integration (version %s)', (version, vercelAiVe
               expect(generateContent).toBeDefined();
               expect(generateContent.parent_span_id).toBe(invokeAgent.span_id);
               expect(generateContent.attributes['vercel.ai.operationId']?.value).toBe('ai.streamText.doStream');
+
+              // The stream's final usage/finish/output arrive only as the stream drains, after the
+              // channel already resolved the model call. Tapping the stream recovers them onto the
+              // model-call span on every path (v7 channel, v6 OTel, v6 orchestrion).
+              expect(generateContent.attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]?.value).toBe(10);
+              expect(generateContent.attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]?.value).toBe(20);
+              expect(generateContent.attributes[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]?.value).toBe(30);
+              expect(generateContent.attributes[GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE]?.value).toBe('["stop"]');
+              expect(generateContent.attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]?.value).toBe(
+                '[{"role":"assistant","parts":[{"type":"text","content":"Stream response!"}],"finish_reason":"stop"}]',
+              );
+
+              // The summed usage and output also land on the parent invoke_agent span, whose own
+              // channel result is otherwise undefined for a stream.
+              expect(invokeAgent.attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]?.value).toBe(10);
+              expect(invokeAgent.attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]?.value).toBe(20);
+              expect(invokeAgent.attributes[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]?.value).toBe(30);
+              expect(invokeAgent.attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]?.value).toBe(
+                '[{"role":"assistant","parts":[{"type":"text","content":"Stream response!"}],"finish_reason":"stop"}]',
+              );
+            },
+          })
+          .start()
+          .completed();
+      });
+    },
+    {
+      additionalDependencies: {
+        ai: vercelAiVersion,
+      },
+    },
+  );
+
+  createEsmTests(
+    __dirname,
+    'scenario-stream-tools.mjs',
+    'instrument.mjs',
+    (createRunner, test) => {
+      test('captures usage, tool calls and output across a multi-step streamText', async () => {
+        await createRunner()
+          .expect({ transaction: { transaction: 'main' } })
+          .expect({
+            span: container => {
+              const invokeAgent = container.items.find(
+                span => span.attributes['sentry.op']?.value === 'gen_ai.invoke_agent',
+              )!;
+              expect(invokeAgent).toBeDefined();
+              expect(invokeAgent.status).toBe('ok');
+              expect(invokeAgent.attributes['vercel.ai.operationId']?.value).toBe('ai.streamText');
+              // Usage is summed across the two streamed model calls (10+15, 20+25, 30+40).
+              expect(invokeAgent.attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]?.value).toBe(25);
+              expect(invokeAgent.attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]?.value).toBe(45);
+              expect(invokeAgent.attributes[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]?.value).toBe(70);
+
+              const generateContents = container.items.filter(
+                span => span.attributes['sentry.op']?.value === 'gen_ai.generate_content',
+              );
+              expect(generateContents).toHaveLength(2);
+              generateContents.forEach(span => expect(span.parent_span_id).toBe(invokeAgent.span_id));
+
+              // The step that streamed a tool call: tool-call output part + tool-calls finish reason.
+              const toolStep = generateContents.find(
+                span => span.attributes[GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE]?.value === '["tool-calls"]',
+              )!;
+              expect(toolStep).toBeDefined();
+              expect(toolStep.attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]?.value).toBe(10);
+              expect(toolStep.attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]?.value).toBe(20);
+              const toolStepOutput = toolStep.attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]?.value as string;
+              expect(toolStepOutput).toContain('"type":"tool_call"');
+              expect(toolStepOutput).toContain('getWeather');
+
+              // The step that streamed the final answer text.
+              const textStep = generateContents.find(
+                span => span.attributes[GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE]?.value === '["stop"]',
+              )!;
+              expect(textStep).toBeDefined();
+              expect(textStep.attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]?.value).toBe(15);
+              expect(textStep.attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]?.value).toContain('Sunny, 72°F.');
+
+              // A tool span is emitted for the streamed tool call. Its parent and recorded input/output
+              // vary by path during stream consumption (tool i/o is covered by the non-stream scenario
+              // and the v7 path), so here we just assert the span exists with the right name/status.
+              const executeTool = container.items.find(span => span.name === 'execute_tool getWeather')!;
+              expect(executeTool).toBeDefined();
+              expect(executeTool.status).toBe('ok');
+              expect(executeTool.attributes[GEN_AI_TOOL_NAME_ATTRIBUTE]?.value).toBe('getWeather');
+            },
+          })
+          .start()
+          .completed();
+      });
+    },
+    {
+      additionalDependencies: {
+        ai: vercelAiVersion,
+      },
+    },
+  );
+
+  createEsmTests(
+    __dirname,
+    'scenario-stream-structured-output.mjs',
+    'instrument.mjs',
+    (createRunner, test) => {
+      test('captures streamed structured output (streamText with experimental_output)', async () => {
+        await createRunner()
+          .expect({ transaction: { transaction: 'main' } })
+          .expect({
+            span: container => {
+              const invokeAgent = container.items.find(
+                span => span.attributes['sentry.op']?.value === 'gen_ai.invoke_agent',
+              )!;
+              expect(invokeAgent).toBeDefined();
+              expect(invokeAgent.status).toBe('ok');
+              expect(invokeAgent.attributes['vercel.ai.operationId']?.value).toBe('ai.streamText');
+              expect(invokeAgent.attributes[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]?.value).toBe(30);
+
+              const generateContent = container.items.find(
+                span => span.attributes['sentry.op']?.value === 'gen_ai.generate_content',
+              )!;
+              expect(generateContent).toBeDefined();
+              expect(generateContent.parent_span_id).toBe(invokeAgent.span_id);
+              expect(generateContent.attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]?.value).toBe(12);
+              expect(generateContent.attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]?.value).toBe(18);
+              expect(generateContent.attributes[GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE]?.value).toBe('["stop"]');
+              // The streamed JSON object is accumulated from the text deltas and captured as the
+              // model's output text (embedded as an escaped JSON string in the output message).
+              const output = generateContent.attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]?.value as string;
+              expect(output).toContain('San Francisco');
+              expect(output).toContain('sunny');
             },
           })
           .start()