From 97cd55f7792aa912a8becf0411227f41f5f592fa Mon Sep 17 00:00:00 2001 From: Sergiy Dybskiy Date: Tue, 27 Oct 2026 11:39:16 -0700 Subject: [PATCH] fix(server-utils): Trace streamObject calls on AI SDK v4-v6 Co-Authored-By: Claude Opus 5.5 --- .../v4-static/scenario-stream-object.mjs | 38 +++++++++ .../suites/tracing/vercelai/v4-static/test.ts | 30 +++++++ .../v6_v7/scenario-stream-object-invalid.mjs | 40 ++++++++++ .../vercelai/v6_v7/scenario-stream-object.mjs | 42 ++++++++++ .../suites/tracing/vercelai/v6_v7/test.ts | 79 +++++++++++++++++++ .../vercel-ai/vercel-ai-dc-subscriber.ts | 16 +++- .../vercel-ai-orchestrion-subscriber.ts | 47 ++++++----- .../src/orchestrion/config/vercel-ai.ts | 8 +- 8 files changed, 274 insertions(+), 26 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/vercelai/v4-static/scenario-stream-object.mjs create mode 100644 dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-object-invalid.mjs create mode 100644 dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-object.mjs diff --git a/dev-packages/node-integration-tests/suites/tracing/vercelai/v4-static/scenario-stream-object.mjs b/dev-packages/node-integration-tests/suites/tracing/vercelai/v4-static/scenario-stream-object.mjs new file mode 100644 index 000000000000..a263f5896eac --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/vercelai/v4-static/scenario-stream-object.mjs @@ -0,0 +1,38 @@ +import * as Sentry from '@sentry/node'; +import { streamObject } from 'ai'; +import { convertArrayToReadableStream, MockLanguageModelV1 } from 'ai/test'; +import { z } from 'zod'; + +async function run() { + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const result = streamObject({ + model: new MockLanguageModelV1({ + defaultObjectGenerationMode: 'json', + doStream: async () => ({ + stream: convertArrayToReadableStream([ + { type: 'response-metadata', id: 'resp-1', modelId: 'mock-model-id' }, + { type: 'text-delta', textDelta: '{ "name": "John Doe",' }, + { type: 'text-delta', textDelta: ' "age": 30 }' }, + { + type: 'finish', + finishReason: 'stop', + usage: { promptTokens: 15, completionTokens: 25 }, + }, + ]), + rawCall: { rawPrompt: null, rawSettings: {} }, + }), + }), + schema: z.object({ name: z.string(), age: z.number() }), + prompt: 'Generate a person object', + }); + + // streamObject returns synchronously; drive the lazy stream to completion so the spans finish. + for await (const _part of result.partialObjectStream) { + void _part; + } + }); + + await Sentry.flush(2000); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/vercelai/v4-static/test.ts b/dev-packages/node-integration-tests/suites/tracing/vercelai/v4-static/test.ts index 01d9aa1450e3..088eafc6b796 100644 --- a/dev-packages/node-integration-tests/suites/tracing/vercelai/v4-static/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/vercelai/v4-static/test.ts @@ -697,6 +697,36 @@ describe('Vercel AI integration (v4)', () => { }); }); + createEsmAndCjsTests(__dirname, 'scenario-stream-object.mjs', 'instrument.mjs', (createRunner, test) => { + test('creates ai spans for streamObject (doStream)', async () => { + await createRunner() + .expect({ transaction: { transaction: 'main' } }) + .expect({ + span: container => { + expect(container.items).toHaveLength(2); + + const invokeAgentSpan = container.items.find(span => span.name === 'invoke_agent'); + expect(invokeAgentSpan!.status).toBe('ok'); + expect(invokeAgentSpan!.attributes[SENTRY_OP].value).toBe('gen_ai.invoke_agent'); + expect(invokeAgentSpan!.attributes[SENTRY_ORIGIN].value).toBe(expectedOrigin); + expect(invokeAgentSpan!.attributes['vercel.ai.operationId'].value).toBe('ai.streamObject'); + expect(invokeAgentSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS].value).toBe(15); + expect(invokeAgentSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS].value).toBe(25); + expect(invokeAgentSpan!.attributes[GEN_AI_USAGE_TOTAL_TOKENS].value).toBe(40); + + const generateContentSpan = container.items.find(span => span.name === 'generate_content mock-model-id'); + expect(generateContentSpan!.status).toBe('ok'); + expect(generateContentSpan!.parent_span_id).toBe(invokeAgentSpan!.span_id); + expect(generateContentSpan!.attributes['vercel.ai.operationId'].value).toBe('ai.streamObject.doStream'); + expect(generateContentSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS].value).toBe(15); + expect(generateContentSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS].value).toBe(25); + }, + }) + .start() + .completed(); + }); + }); + createEsmAndCjsTests(__dirname, 'scenario-generate-object.mjs', 'instrument-with-pii.mjs', (createRunner, test) => { test('captures generateObject spans with schema attributes', async () => { await createRunner() diff --git a/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-object-invalid.mjs b/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-object-invalid.mjs new file mode 100644 index 000000000000..ca2973a96c64 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-object-invalid.mjs @@ -0,0 +1,40 @@ +import * as Sentry from '@sentry/node'; +import { streamObject } from 'ai'; +import { MockLanguageModelV3, simulateReadableStream } from 'ai/test'; +import { z } from 'zod'; + +async function run() { + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const result = streamObject({ + model: new MockLanguageModelV3({ + doStream: async () => ({ + stream: simulateReadableStream({ + chunks: [ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: '0' }, + { type: 'text-delta', id: '0', delta: 'not json' }, + { 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 }, + }, + }, + ], + }), + }), + }), + schema: z.object({ name: z.string(), age: z.number() }), + prompt: 'Generate a person object', + }); + + for await (const _part of result.partialObjectStream) { + void _part; + } + }); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-object.mjs b/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-object.mjs new file mode 100644 index 000000000000..ee86684a57bf --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/vercelai/v6_v7/scenario-stream-object.mjs @@ -0,0 +1,42 @@ +import * as Sentry from '@sentry/node'; +import { streamObject } from 'ai'; +import { MockLanguageModelV3, simulateReadableStream } from 'ai/test'; +import { z } from 'zod'; + +async function run() { + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const result = streamObject({ + model: new MockLanguageModelV3({ + doStream: async () => ({ + stream: simulateReadableStream({ + chunks: [ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: '0' }, + { type: 'text-delta', id: '0', delta: '{ "name": "John Doe",' }, + { type: 'text-delta', id: '0', delta: ' "age": 30 }' }, + { 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 }, + }, + }, + ], + }), + }), + }), + schema: z.object({ name: z.string(), age: z.number() }), + prompt: 'Generate a person object', + }); + + // streamObject returns synchronously; drive the lazy stream to completion so the spans finish. + for await (const _part of result.partialObjectStream) { + 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 ea89bf21ddaf..1f332143cd1d 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 @@ -1236,6 +1236,85 @@ describe.each(matrix)('Vercel AI integration (version %s)', (version, vercelAiVe }, ); + createEsmTests( + __dirname, + 'scenario-stream-object.mjs', + 'instrument.mjs', + (createRunner, test) => { + // Like `generateObject`, v7's native channel never publishes `streamObject`; v6 uses orchestrion. + test.skipIf(version === '7')( + 'creates streamObject spans with the model call parented to invoke_agent', + async () => { + await createRunner() + .expect({ + span: container => { + const segment = container.items.find(span => span.is_segment && span.name === 'main'); + expect(segment).toBeDefined(); + const spans = container.items.filter( + span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.vercelai', + ); + + const invokeAgent = spans.find(span => span.attributes[SENTRY_OP]?.value === 'gen_ai.invoke_agent')!; + expect(invokeAgent.status).toBe('ok'); + expect(invokeAgent.attributes['vercel.ai.operationId']?.value).toBe('ai.streamObject'); + expect(invokeAgent.attributes[GEN_AI_USAGE_INPUT_TOKENS]?.value).toBe(15); + expect(invokeAgent.attributes[GEN_AI_USAGE_OUTPUT_TOKENS]?.value).toBe(25); + expect(invokeAgent.attributes[GEN_AI_USAGE_TOTAL_TOKENS]?.value).toBe(40); + + const generateContent = spans.find( + span => span.attributes[SENTRY_OP]?.value === 'gen_ai.generate_content', + )!; + expect(generateContent.status).toBe('ok'); + expect(generateContent.parent_span_id).toBe(invokeAgent.span_id); + expect(generateContent.attributes['vercel.ai.operationId']?.value).toBe('ai.streamObject.doStream'); + expect(generateContent.attributes[GEN_AI_REQUEST_MODEL]?.value).toBe('mock-model-id'); + expect(generateContent.attributes[GEN_AI_USAGE_INPUT_TOKENS]?.value).toBe(15); + expect(generateContent.attributes[GEN_AI_USAGE_OUTPUT_TOKENS]?.value).toBe(25); + }, + }) + .unordered() + .start() + .completed(); + }, + ); + }, + { + additionalDependencies: { + ai: vercelAiVersion, + }, + }, + ); + + createEsmTests( + __dirname, + 'scenario-stream-object-invalid.mjs', + 'instrument.mjs', + (createRunner, test) => { + test.skipIf(version === '7')('marks streamObject as failed when the output fails the schema', async () => { + await createRunner() + .expect({ + span: container => { + const invokeAgent = container.items.find( + span => + span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.vercelai' && + span.attributes[SENTRY_OP]?.value === 'gen_ai.invoke_agent', + )!; + expect(invokeAgent.attributes['vercel.ai.operationId']?.value).toBe('ai.streamObject'); + expect(invokeAgent.status).toBe('error'); + }, + }) + .unordered() + .start() + .completed(); + }); + }, + { + additionalDependencies: { + ai: vercelAiVersion, + }, + }, + ); + createEsmTests( __dirname, 'scenario-aborted-stream-text.mjs', diff --git a/packages/server-utils/src/integrations/vercel-ai/vercel-ai-dc-subscriber.ts b/packages/server-utils/src/integrations/vercel-ai/vercel-ai-dc-subscriber.ts index 1079886c6762..6368fd27eade 100644 --- a/packages/server-utils/src/integrations/vercel-ai/vercel-ai-dc-subscriber.ts +++ b/packages/server-utils/src/integrations/vercel-ai/vercel-ai-dc-subscriber.ts @@ -114,6 +114,7 @@ const ROOT_OPERATION_TYPES = new Set([ 'generateText', 'streamText', 'generateObject', + 'streamObject', 'embed', 'embedMany', 'rerank', @@ -235,6 +236,7 @@ export type ChannelEventType = | 'generateText' | 'streamText' | 'generateObject' + | 'streamObject' | 'step' | 'languageModelCall' | 'executeTool' @@ -463,10 +465,16 @@ export function createSpanFromMessage( case 'generateText': case 'streamText': case 'generateObject': - // `generateObject` builds the same `invoke_agent` span as `generateText` (non-streaming); its - // distinct `ai.generateObject` operationId rides on `event.operationId`. The JSON-schema attribute - // the OTel path derives from the SDK's Zod schema is not reconstructed on the channel path. - return buildInvokeAgentSpan(event, baseAttributes, recordInputs, callId, type === 'streamText'); + case 'streamObject': + // The object variants' operationId rides on `event.operationId`. The JSON-schema attribute the OTel + // path derives from the SDK's Zod schema is not reconstructed on the channel path. + return buildInvokeAgentSpan( + event, + baseAttributes, + recordInputs, + callId, + type === 'streamText' || type === 'streamObject', + ); case 'languageModelCall': _INTERNAL_skipAiProviderWrapping([WORKERS_AI_INTEGRATION_NAME]); diff --git a/packages/server-utils/src/integrations/vercel-ai/vercel-ai-orchestrion-subscriber.ts b/packages/server-utils/src/integrations/vercel-ai/vercel-ai-orchestrion-subscriber.ts index 2a59de79b4f8..ba46ce47ae5c 100644 --- a/packages/server-utils/src/integrations/vercel-ai/vercel-ai-orchestrion-subscriber.ts +++ b/packages/server-utils/src/integrations/vercel-ai/vercel-ai-orchestrion-subscriber.ts @@ -150,6 +150,7 @@ export function subscribeVercelAiOrchestrionChannels( bindOperation(tracingChannel, CHANNELS.VERCEL_AI_GENERATE_TEXT, buildTextMessage('generateText'), options); bindOperation(tracingChannel, CHANNELS.VERCEL_AI_STREAM_TEXT, buildTextMessage('streamText'), options); bindOperation(tracingChannel, CHANNELS.VERCEL_AI_GENERATE_OBJECT, buildTextMessage('generateObject'), options); + bindOperation(tracingChannel, CHANNELS.VERCEL_AI_STREAM_OBJECT, buildTextMessage('streamObject'), options); bindOperation( tracingChannel, CHANNELS.VERCEL_AI_EMBED, @@ -305,33 +306,35 @@ function bindOperation( message.result = message.type === 'executeTool' ? { output: data.result } : data.result; enrichSpanOnEnd(span, message, options); } - // A `streamText` model call runs after this (synchronously-returning) operation's span has - // already ended, so its `callId` entry must outlive the operation — it's cleared once the model - // call settles (see `patchModelMethod`). Every other operation can clear here. - if (message.type !== 'streamText') { + // A `streamText`/`streamObject` model call runs after this (synchronously-returning) operation's + // span has already ended, so its `callId` entry must outlive the operation — it's cleared once the + // model call settles (see `patchModelMethod`). Every other operation can clear here. + if (!isStreamOperation(message.type)) { clearOperationId(message); } 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), + // `streamText`/`streamObject` return synchronously, so their 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 }) => deferStreamOperationEnd(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. + * Keep a streamed `streamText`/`streamObject` operation span open until the stream drains, then enrich it + * from the result's completion promises (usage/output/finish reason) and end it. A `StreamObjectResult` + * has no `text`/`toolCalls`/`totalUsage`, so those resolve to `undefined`. Returns `false` for anything + * that isn't a streamed result, so the helper ends the span as usual. */ -function deferStreamTextOperationEnd( +function deferStreamOperationEnd( data: TracingChannelPayloadWithSpan, end: (error?: unknown) => void, ): boolean { - if (messages.get(data)?.type !== 'streamText' || 'error' in data || !isStreamingResult(data.result)) { + const type = messages.get(data)?.type; + if (!type || !isStreamOperation(type) || 'error' in data || !isStreamingResult(data.result)) { return false; } @@ -344,6 +347,8 @@ function deferStreamTextOperationEnd( streamResult.toolCalls, streamResult.finishReason, streamResult.response, + // Only `object` rejects when the output fails the schema; awaiting it fails the span like `generateObject`. + type === 'streamObject' ? streamResult.object : undefined, ]); // Feed the resolved values back through the shared enrichment: `beforeSpanEnd` reads `data.result`. data.result = { usage, text, toolCalls, finishReason, response }; @@ -356,7 +361,11 @@ function deferStreamTextOperationEnd( return true; } -/** A `StreamTextResult` exposes its aggregates as promises (`totalUsage`/`usage`); a settled result does not. */ +function isStreamOperation(type: VercelAiChannelMessage['type']): boolean { + return type === 'streamText' || type === 'streamObject'; +} + +/** A streamed result exposes its aggregates as promises (`totalUsage`/`usage`); a settled result does not. */ function isStreamingResult(result: unknown): result is Record | undefined> { return isObjectLike(result) && (isThenable(result.totalUsage) || isThenable(result.usage)); } @@ -502,8 +511,8 @@ function patchModelMethod( return Promise.resolve(original.apply(this, args)); } - // `streamText` ends its operation span synchronously, so its `callId` entry was deliberately left in - // place for this later model call; drop it now that we've used it. + // `streamText`/`streamObject` end their operation span synchronously, so the `callId` entry was + // deliberately left in place for this later model call; drop it now that we've used it. const clearStreamCallId = (): void => { if (method === 'doStream' && callId) { clearOperationCallId(callId); @@ -525,7 +534,7 @@ function patchModelMethod( // 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. + // separately by `deferStreamOperationEnd`, which awaits the operation's own result promises. if (method === 'doStream' && isObjectLike(value) && isReadableStream(value.stream)) { value.stream = tapModelCallStream( value.stream, @@ -654,7 +663,7 @@ function patchToolExecute( }; } -function buildTextMessage(type: 'generateText' | 'streamText' | 'generateObject'): MessageBuilder { +function buildTextMessage(type: 'generateText' | 'streamText' | 'generateObject' | 'streamObject'): MessageBuilder { return (options, telemetry) => ({ type, event: { diff --git a/packages/server-utils/src/orchestrion/config/vercel-ai.ts b/packages/server-utils/src/orchestrion/config/vercel-ai.ts index 10f404b56a0b..e1b1c1cf408f 100644 --- a/packages/server-utils/src/orchestrion/config/vercel-ai.ts +++ b/packages/server-utils/src/orchestrion/config/vercel-ai.ts @@ -12,14 +12,15 @@ export const vercelAiConfig = [ // not to span it, but so the subscriber can monkey-patch `doGenerate`/ // `doStream` on the returned model (the only way to span the model call, // which is an inline call with no injectable definition in `ai`). - // `streamText` returns its result synchronously (streaming is lazy), so it's - // `Sync`; the subscriber binds the span via `bindTracingChannelToSpan`, which - // ends it when the (synchronous) call returns. + // `streamText` and `streamObject` return their result synchronously (streaming + // is lazy), so they're `Sync`; the subscriber binds the span via + // `bindTracingChannelToSpan`, which ends it when the (synchronous) call returns. // The majority of entrypoints are present in all versions we support ...vercelAiEntries('>=4.0.0 <7.0.0', 'generateText', 'generateText', 'Async'), ...vercelAiEntries('>=4.0.0 <7.0.0', 'streamText', 'streamText', 'Sync'), ...vercelAiEntries('>=4.0.0 <7.0.0', 'generateObject', 'generateObject', 'Async'), + ...vercelAiEntries('>=4.0.0 <7.0.0', 'streamObject', 'streamObject', 'Sync'), ...vercelAiEntries('>=4.0.0 <7.0.0', 'embed', 'embed', 'Async'), ...vercelAiEntries('>=4.0.0 <7.0.0', 'embedMany', 'embedMany', 'Async'), @@ -40,6 +41,7 @@ export const vercelAiChannels = { VERCEL_AI_GENERATE_TEXT: 'orchestrion:ai:generateText', VERCEL_AI_STREAM_TEXT: 'orchestrion:ai:streamText', VERCEL_AI_GENERATE_OBJECT: 'orchestrion:ai:generateObject', + VERCEL_AI_STREAM_OBJECT: 'orchestrion:ai:streamObject', VERCEL_AI_EMBED: 'orchestrion:ai:embed', VERCEL_AI_EMBED_MANY: 'orchestrion:ai:embedMany', VERCEL_AI_EXECUTE_TOOL_CALL: 'orchestrion:ai:executeToolCall',