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

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

@@ -610,6 +610,144 @@ describe.each([
                 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) => {
+      // See the Node-18 note on `scenario-stream-text.mjs` above.
+      test.skipIf(version === '7' && nodeVersion === 18)(
+        '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) => {
+      // See the Node-18 note on `scenario-stream-text.mjs` above.
+      test.skipIf(version === '7' && nodeVersion === 18)(
+        '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()