feat(server-utils): Enrich invoke_agent on streamed v6 orchestrion calls

Defer the orchestrion streamText operation span (via the deferSpanEnd hook the
binding already supports) and await the StreamTextResult's completion promises
(totalUsage/text/finishReason/toolCalls, resolved on drain) to enrich the
invoke_agent span. This closes the earlier gap where that span ended
synchronously before the stream drained, so v6 orchestrion now matches v7 and
the v6 OTel path for streamed usage and output.

Author: logaretm

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

@@ -44,12 +44,6 @@ describe.each([
   // else, we use the OTel processor
   const expectedOrigin = usesChannels ? 'auto.vercelai.channel' : 'auto.vercelai.otel';
 
-  // Streamed model-call spans get usage/finish/output on every path. The parent invoke_agent span only
-  // gets the streamed aggregate where the operation span outlives the stream: the OTel path (v6, native
-  // spans) and the v7 channel (asyncEnd deferred on total usage). In v6 orchestrion the operation span
-  // ends synchronously when `streamText` returns, before the stream drains, so it can't be enriched.
-  const enrichesInvokeAgentForStream = !(version === '6' && isOrchestrionEnabled());
-
   // We only run this in ESM and CJS to verify full support
   // Other suites we only run in ESM to simplify the test setup
   createEsmAndCjsTests(
@@ -628,16 +622,14 @@ describe.each([
                   '[{"role":"assistant","parts":[{"type":"text","content":"Stream response!"}],"finish_reason":"stop"}]',
                 );
 
-                // The summed usage and output also land on the parent invoke_agent span, except where its
-                // span ends before the stream drains (see `enrichesInvokeAgentForStream`).
-                if (enrichesInvokeAgentForStream) {
-                  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"}]',
-                  );
-                }
+                // 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()
@@ -672,11 +664,9 @@ describe.each([
                 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).
-                if (enrichesInvokeAgentForStream) {
-                  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);
-                }
+                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',
@@ -743,9 +733,7 @@ describe.each([
                 expect(invokeAgent).toBeDefined();
                 expect(invokeAgent.status).toBe('ok');
                 expect(invokeAgent.attributes?.['vercel.ai.operationId']?.value).toBe('ai.streamText');
-                if (enrichesInvokeAgentForStream) {
-                  expect(invokeAgent.attributes?.[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]?.value).toBe(30);
-                }
+                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',

packages/server-utils/src/vercel-ai/vercel-ai-orchestrion-v6-subscriber.ts

@@ -1,3 +1,4 @@
+/* eslint-disable max-lines */
 import type { Span } from '@sentry/core';
 import { debug, getActiveSpan, SPAN_STATUS_ERROR, withActiveSpan } from '@sentry/core';
 import { DEBUG_BUILD } from '../debug-build';
@@ -233,10 +234,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<OrchestrionContext>,
+  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<string, PromiseLike<unknown> | 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