From f34c2fe6696959844005d92de7bbe980f06dfd17 Mon Sep 17 00:00:00 2001 From: Abdelrahman Awad Date: Wed, 8 Jul 2026 11:12:30 -0400 Subject: [PATCH] feat(server-utils): Capture usage/output on streamed Vercel AI spans (v7 + orchestrion) --- .../scenario-stream-structured-output.mjs | 48 +++++ .../vercelai/v6_v7/scenario-stream-tools.mjs | 78 +++++++ .../suites/tracing/vercelai/v6_v7/test.ts | 130 ++++++++++++ packages/server-utils/src/vercel-ai/util.ts | 190 ++++++++++++++++++ .../src/vercel-ai/vercel-ai-dc-subscriber.ts | 169 ++++++++++++---- .../vercel-ai-orchestrion-subscriber.ts | 88 +++++++- 6 files changed, 657 insertions(+), 46 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-structured-output.mjs create mode 100644 dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-tools.mjs create mode 100644 packages/server-utils/src/vercel-ai/util.ts diff --git a/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-structured-output.mjs b/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-structured-output.mjs new file mode 100644 index 000000000000..d5057d323502 --- /dev/null +++ b/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(); diff --git a/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-tools.mjs b/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-tools.mjs new file mode 100644 index 000000000000..54efdac423da --- /dev/null +++ b/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(); diff --git a/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/test.ts b/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/test.ts index 9a13d02a3b65..b49783eebc2e 100644 --- a/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/test.ts +++ b/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() diff --git a/packages/server-utils/src/vercel-ai/util.ts b/packages/server-utils/src/vercel-ai/util.ts new file mode 100644 index 000000000000..c64dc76a1901 --- /dev/null +++ b/packages/server-utils/src/vercel-ai/util.ts @@ -0,0 +1,190 @@ +/** Shared, state-free helpers for the Vercel AI (`ai`) channel subscribers, plus the streamed model-call tap. */ + +/** Narrow to a string, or `undefined` for anything else. */ +export function asString(value: unknown): string | undefined { + return typeof value === 'string' ? value : undefined; +} + +/** Narrow to a finite number, or `undefined` for anything else (including `NaN`). */ +export function asNumber(value: unknown): number | undefined { + return typeof value === 'number' && !isNaN(value) ? value : undefined; +} + +/** Add two optional numbers, treating a missing operand as `0` but returning `undefined` when both are absent. */ +export function sum(a: number | undefined, b: number | undefined): number | undefined { + return a === undefined && b === undefined ? undefined : (a ?? 0) + (b ?? 0); +} + +/** Narrow to a non-null object (records and arrays alike). */ +export function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +/** Stringify a value, passing strings through and falling back to a placeholder on circular/unserializable input. */ +export function safeStringify(value: unknown): string { + if (typeof value === 'string') { + return value; + } + try { + return JSON.stringify(value); + } catch { + return '[unserializable]'; + } +} + +/* + * Streaming support for the `ai:telemetry` tracing channel. + * + * For a streamed model call (`doStream`) the SDK resolves the `languageModelCall` channel promise as + * soon as it hands back the *unconsumed* stream — so `result.stream` is a `ReadableStream` and the + * final usage / finish reason / output only arrive later, as the stream is drained (they ride the + * stream's own `finish`/`text-delta`/`tool-call` chunks, never the channel context). The SDK + * deliberately leaves `result` undefined on the `streamText`/`step` contexts and exposes the stream + * only on the model call, expecting consumers to tap it themselves. + * @see https://github.com/vercel/ai/pull/15660 (discussion: "tee off and aggregate ourselves") + * + * We replace `result.stream` with a passthrough that forwards every chunk untouched (so the SDK's own + * consumption is unaffected) while accumulating the data we need, then hand the aggregate back once the + * stream settles so the model-call span can be enriched and ended out-of-band. + */ + +/** The subset of a streamed provider chunk we read. Unknown chunk types are forwarded and ignored. */ +interface StreamChunk { + type?: unknown; + delta?: unknown; + id?: unknown; + modelId?: unknown; + toolCallId?: unknown; + toolName?: unknown; + input?: unknown; + args?: unknown; + finishReason?: unknown; + usage?: unknown; + providerMetadata?: unknown; + error?: unknown; +} + +/** The aggregate handed back once a streamed model call finishes, in the shape `enrichSpanOnEnd` expects. */ +export interface StreamedModelCallResult { + text?: string; + toolCalls: Array>; + usage?: unknown; + finishReason?: unknown; + responseId?: string; + responseModel?: string; + providerMetadata?: unknown; +} + +/** A minimal structural check — the streamed model call exposes a web `ReadableStream` on `result.stream`. */ +export function isReadableStream(value: unknown): value is ReadableStream { + return ( + isRecord(value) && + typeof (value as { pipeThrough?: unknown }).pipeThrough === 'function' && + typeof (value as { getReader?: unknown }).getReader === 'function' + ); +} + +/** + * Wrap a streamed model call's `ReadableStream` so its chunks are observed as the SDK consumes them, + * without altering what the SDK sees. Returns a replacement stream to swap onto `result.stream`. + * + * `onFinal` runs exactly once when the stream drains cleanly; `onError` runs exactly once if it errors + * or is cancelled. Reading one source chunk per `pull` preserves the SDK's backpressure. The + * `try/catch` around every read guarantees the owning span is always ended — a leaked open span on a + * mid-stream failure would be worse than a slightly-less-enriched one. + */ +export function tapModelCallStream( + stream: ReadableStream, + onFinal: (result: StreamedModelCallResult) => void, + onError: (error: unknown) => void, +): ReadableStream { + const reader = stream.getReader(); + const state: StreamedModelCallResult = { toolCalls: [] }; + let text = ''; + let settled = false; + + const finalize = (): void => { + if (settled) { + return; + } + settled = true; + if (text) { + state.text = text; + } + onFinal(state); + }; + + const fail = (error: unknown): void => { + if (settled) { + return; + } + settled = true; + onError(error); + }; + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + finalize(); + controller.close(); + + return; + } + text += accumulateChunk(state, value) ?? ''; + controller.enqueue(value); + } catch (error) { + fail(error); + controller.error(error); + } + }, + cancel(reason) { + // Consumer stopped reading early (e.g. a `break` out of `for await`); close out the span with + // whatever we have rather than leave it open. A non-error reason is a deliberate stop, not a failure. + finalize(); + + return reader.cancel(reason); + }, + }); +} + +/** + * Fold a single streamed chunk into the running aggregate. Returns any text delta so the caller can + * accumulate it (kept out of `state` until the end to avoid re-joining on every chunk). + */ +function accumulateChunk(state: StreamedModelCallResult, chunk: unknown): string | undefined { + if (!isRecord(chunk)) { + return undefined; + } + const { type, delta, id, modelId, toolCallId, toolName, input, args, finishReason, usage, providerMetadata } = + chunk as StreamChunk; + + switch (type) { + case 'text-delta': + return typeof delta === 'string' ? delta : undefined; + case 'tool-call': + state.toolCalls.push({ toolCallId, toolName, input: input ?? args }); + + return undefined; + case 'response-metadata': + if (typeof id === 'string') { + state.responseId = id; + } + if (typeof modelId === 'string') { + state.responseModel = modelId; + } + + return undefined; + case 'finish': + state.finishReason = finishReason; + state.usage = usage; + if (providerMetadata !== undefined) { + state.providerMetadata = providerMetadata; + } + + return undefined; + default: + return undefined; + } +} diff --git a/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts b/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts index da71fd98ea4e..c6493486150b 100644 --- a/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts +++ b/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts @@ -45,6 +45,16 @@ import { } from '@sentry/core'; import type { TracingChannel } from 'node:diagnostics_channel'; import { bindTracingChannelToSpan } from '../tracing-channel'; +import { + asNumber, + asString, + isReadableStream, + isRecord, + safeStringify, + type StreamedModelCallResult, + sum, + tapModelCallStream, +} from './util'; /** * The single tracing channel the `ai` package (>= 7) publishes all telemetry lifecycle events to @@ -84,6 +94,12 @@ const operationIdByCallId = new Map>(); +// A streamed `streamText` operation's own channel result is always `undefined` (the SDK exposes the +// stream only on the model call), so its `invoke_agent` span can't be enriched from the channel. We +// key the span by `callId` and enrich it as each streamed model call drains — the model calls flush +// before the operation's own span ends. The running token sum lives on the span's own attributes. +const invokeAgentSpanByCallId = new Map(); + // Only top-level operations own the `callId` → operationId mapping; `step`/`languageModelCall`/ // `executeTool` share the parent's `callId`, so they must not clear it. const ROOT_OPERATION_TYPES = new Set(['generateText', 'streamText', 'embed', 'embedMany', 'rerank']); @@ -108,6 +124,7 @@ export function clearOperationId(data: VercelAiChannelMessage): void { export function clearOperationCallId(callId: string): void { operationIdByCallId.delete(callId); toolDescriptionsByCallId.delete(callId); + invokeAgentSpanByCallId.delete(callId); } /** Record tool name → description from an event's `tools`, so tool spans can backfill the description. */ @@ -192,8 +209,6 @@ export interface VercelAiChannelOptions { enableTruncation?: boolean; } -let subscribed = false; - /** * Subscribe Sentry span handlers to the `ai` SDK's native telemetry tracing channel (`ai:telemetry`, * available in `ai` >= 7) and emit fully-formed `gen_ai.*` spans directly — no OpenTelemetry span @@ -204,17 +219,12 @@ let subscribed = false; * * Safe to always call: on `ai` versions that don't publish to the channel (e.g. < 7) nothing is * ever emitted and this is inert, so there is no double-instrumentation against the OTel-based - * patcher. Idempotent. + * patcher. The integration's `setupOnce` guarantees this runs a single time. */ export function subscribeVercelAiTracingChannel( tracingChannel: VercelAiTracingChannelFactory, options: VercelAiChannelOptions = {}, ): void { - if (subscribed) { - return; - } - subscribed = true; - bindTracingChannelToSpan( tracingChannel(AI_SDK_TELEMETRY_TRACING_CHANNEL), data => createSpanFromMessage(data, options), @@ -225,10 +235,119 @@ export function subscribeVercelAiTracingChannel( enrichSpanOnEnd(span, data, options); clearOperationId(data); }, + // A streamed model call resolves before its stream is drained, so we tap the stream, keep the + // span open, and end it (via `end`) once the final usage/finish/output chunks arrive. + deferSpanEnd: ({ data, end }) => deferStreamedModelCallEnd(data, options, end), }, ); } +/** + * When a `languageModelCall` resolves to a live `ReadableStream`, defer ending its span: swap in a + * passthrough that forwards chunks to the SDK untouched while aggregating usage/finish/output, then + * enrich and end the span once the stream settles. Returns `false` for anything that isn't a streamed + * model call so the helper ends the span as usual. + */ +function deferStreamedModelCallEnd( + data: VercelAiChannelMessage, + options: VercelAiChannelOptions, + end: (error?: unknown) => void, +): boolean { + if (data.type !== 'languageModelCall' || !isRecord(data.result)) { + return false; + } + const result = data.result; + const stream = result.stream; + if (!isReadableStream(stream)) { + return false; + } + + const callId = asString(data.event.callId); + const { recordOutputs } = getRecordingOptions(data.event, options); + result.stream = tapModelCallStream( + stream, + final => { + // Reshape the aggregate into the result `enrichSpanOnEnd` expects, then let `end` run it (it calls + // `beforeSpanEnd`). Enriching the model-call span and the parent from the same aggregate keeps + // streamed and non-streamed spans identical. + data.result = { ...result, ...streamedResultToChannelResult(final) }; + end(); + enrichInvokeAgentFromStream(callId, final, recordOutputs); + }, + error => end(error), + ); + + return true; +} + +/** Map the tapped stream aggregate onto the `languageModelCall` result shape `enrichSpanOnEnd` reads. */ +export function streamedResultToChannelResult(final: StreamedModelCallResult): Record { + const content: Array> = []; + if (final.text) { + content.push({ type: 'text', text: final.text }); + } + for (const toolCall of final.toolCalls) { + content.push({ type: 'tool-call', ...toolCall }); + } + + return { + content, + ...(final.usage !== undefined ? { usage: final.usage } : {}), + ...(final.finishReason !== undefined ? { finishReason: final.finishReason } : {}), + ...(final.providerMetadata !== undefined ? { providerMetadata: final.providerMetadata } : {}), + ...(final.responseId || final.responseModel + ? { + response: { + ...(final.responseId ? { id: final.responseId } : {}), + ...(final.responseModel ? { modelId: final.responseModel } : {}), + }, + } + : {}), + }; +} + +/** + * Propagate a streamed model call's usage (summed across the operation's model calls) and output onto + * the enclosing `invoke_agent` span, which has no channel result of its own. Runs before that span ends + * because the operation's completion promise settles only after every model-call stream has drained. + */ +function enrichInvokeAgentFromStream( + callId: string | undefined, + final: StreamedModelCallResult, + recordOutputs: boolean, +): void { + const span = callId ? invokeAgentSpanByCallId.get(callId) : undefined; + if (!span) { + return; + } + + const usage = isRecord(final.usage) ? final.usage : undefined; + if (usage) { + const input = tokenCount(usage.inputTokens) ?? tokenCount(usage.tokens); + const output = tokenCount(usage.outputTokens); + addTokensToSpan(span, GEN_AI_USAGE_INPUT_TOKENS, input); + addTokensToSpan(span, GEN_AI_USAGE_OUTPUT_TOKENS, output); + addTokensToSpan(span, GEN_AI_USAGE_TOTAL_TOKENS, tokenCount(usage.totalTokens) ?? sum(input, output)); + } + + if (recordOutputs) { + const parts = partsFromTextAndToolCalls(final.text, final.toolCalls); + const outputMessages = buildOutputMessages(parts, getFinishReason({ finishReason: final.finishReason })); + if (outputMessages) { + span.setAttribute(GEN_AI_OUTPUT_MESSAGES, outputMessages); + } + } +} + +/** Add `value` into a span's numeric token attribute, using the span itself as the running sum. */ +function addTokensToSpan(span: Span, attribute: string, value: number | undefined): void { + if (value === undefined) { + return; + } + const current = spanToJSON(span).data[attribute]; + span.setAttribute(attribute, (typeof current === 'number' ? current : 0) + value); +} + /** * Transform a channel `start` payload into the span that should be active for the operation. For * `step` we deliberately don't open a span (model calls and tool calls are siblings under the @@ -315,13 +434,18 @@ function buildInvokeAgentSpan( if (callId) { operationIdByCallId.set(callId, { operationId, isStream }); } - return startGenAiSpan(GEN_AI_INVOKE_AGENT_SPAN_OP, functionId, { + const span = startGenAiSpan(GEN_AI_INVOKE_AGENT_SPAN_OP, functionId, { ...baseAttributes, [VERCEL_AI_OPERATION_ID_ATTRIBUTE]: operationId, [GEN_AI_RESPONSE_STREAMING]: isStream, ...(functionId ? { [GEN_AI_FUNCTION_ID]: functionId } : {}), ...(recordInputs ? buildInputMessageAttributes(event, enableTruncation) : {}), }); + if (isStream && callId) { + invokeAgentSpanByCallId.set(callId, span); + } + + return span; } function buildModelCallSpan( @@ -620,30 +744,3 @@ function buildInputMessageAttributes( return attributes; } - -function asString(value: unknown): string | undefined { - return typeof value === 'string' ? value : undefined; -} - -function asNumber(value: unknown): number | undefined { - return typeof value === 'number' && !isNaN(value) ? value : undefined; -} - -function sum(a: number | undefined, b: number | undefined): number | undefined { - return a === undefined && b === undefined ? undefined : (a ?? 0) + (b ?? 0); -} - -function isRecord(value: unknown): value is Record { - return typeof value === 'object' && value !== null; -} - -function safeStringify(value: unknown): string { - if (typeof value === 'string') { - return value; - } - try { - return JSON.stringify(value); - } catch { - return '[unserializable]'; - } -} diff --git a/packages/server-utils/src/vercel-ai/vercel-ai-orchestrion-subscriber.ts b/packages/server-utils/src/vercel-ai/vercel-ai-orchestrion-subscriber.ts index 5227b1dc5072..5ff19853335e 100644 --- a/packages/server-utils/src/vercel-ai/vercel-ai-orchestrion-subscriber.ts +++ b/packages/server-utils/src/vercel-ai/vercel-ai-orchestrion-subscriber.ts @@ -10,10 +10,12 @@ import { clearOperationId, createSpanFromMessage, enrichSpanOnEnd, + streamedResultToChannelResult, type VercelAiChannelMessage, type VercelAiChannelOptions, type VercelAiTracingChannelFactory, } from './vercel-ai-dc-subscriber'; +import { asString, isReadableStream, isRecord, tapModelCallStream } from './util'; /** * v5 & v6 channel adapter for the Vercel AI (`ai`) SDK. @@ -260,10 +262,58 @@ function bindOperation( } messages.delete(data); }, + // `streamText` returns synchronously, so its operation span would otherwise end before the stream + // drains — losing the aggregate usage/output. Defer the end and await the result's completion + // promises (`totalUsage`/`text`/…, which resolve on drain), mirroring how v7's channel defers the + // operation span on the SDK's total-usage promise. + deferSpanEnd: ({ data, end }) => deferStreamTextOperationEnd(data, end), }, ); } +/** + * Keep a streamed `streamText` operation span open until the stream drains, then enrich it from the + * `StreamTextResult`'s completion promises (usage/output/finish reason) and end it. Returns `false` for + * anything that isn't a streamed `streamText` result, so the helper ends the span as usual. + */ +function deferStreamTextOperationEnd( + data: TracingChannelPayloadWithSpan, + end: (error?: unknown) => void, +): boolean { + if (messages.get(data)?.type !== 'streamText' || 'error' in data || !isStreamingResult(data.result)) { + return false; + } + + const streamResult = data.result; + void (async () => { + try { + const [usage, text, toolCalls, finishReason, response] = await Promise.all([ + streamResult.totalUsage ?? streamResult.usage, + streamResult.text, + streamResult.toolCalls, + streamResult.finishReason, + streamResult.response, + ]); + // Feed the resolved values back through the shared enrichment: `beforeSpanEnd` reads `data.result`. + data.result = { usage, text, toolCalls, finishReason, response }; + end(); + } catch (error) { + end(error); + } + })(); + + return true; +} + +/** A `StreamTextResult` exposes its aggregates as promises (`totalUsage`/`usage`); a settled result does not. */ +function isStreamingResult(result: unknown): result is Record | undefined> { + return isRecord(result) && (isThenable(result.totalUsage) || isThenable(result.usage)); +} + +function isThenable(value: unknown): boolean { + return isRecord(value) && typeof value.then === 'function'; +} + /** * Neutralize `ai`'s native OpenTelemetry instrumentation for this call by pointing * `experimental_telemetry` at a copy with `isEnabled: false`. `ai`'s `getTracer` then returns its @@ -408,13 +458,39 @@ function patchModelMethod( try { const result = Promise.resolve(original.apply(this, args)); - // `doStream` resolves to `{ stream, ... }` before the stream is consumed; we end here (start/end - // bracket the call) to match the channel timing. return result.then(value => { + // A streamed model call resolves to `{ stream, ... }` before the stream is consumed, so its + // usage/finish/output only arrive as the stream drains. Tap it (same helper as the v7 path) and + // defer ending this model-call span until then. The parent `invoke_agent` span is enriched + // separately by `deferStreamTextOperationEnd`, which awaits the operation's own result promises. + if (method === 'doStream' && isRecord(value) && isReadableStream(value.stream)) { + value.stream = tapModelCallStream( + value.stream, + final => { + message.result = { ...value, ...streamedResultToChannelResult(final) }; + enrichSpanOnEnd(span, message, options); + span.end(); + clearStreamCallId(); + }, + error => { + span.setStatus({ + code: SPAN_STATUS_ERROR, + message: error instanceof Error ? error.message : 'unknown_error', + }); + span.end(); + clearStreamCallId(); + }, + ); + + return value; + } + // `doGenerate` (and any non-stream result) settles with the full result; end here, start/end + // bracket the call to match the channel timing. message.result = value; enrichSpanOnEnd(span, message, options); span.end(); clearStreamCallId(); + return value; }, failSpan); } catch (error) { @@ -555,11 +631,3 @@ function modelFields(model: unknown): { provider?: string; modelId?: string } { function modelField(model: unknown, field: 'modelId' | 'provider'): string | undefined { return isRecord(model) ? asString(model[field]) : undefined; } - -function asString(value: unknown): string | undefined { - return typeof value === 'string' ? value : undefined; -} - -function isRecord(value: unknown): value is Record { - return typeof value === 'object' && value !== null; -}