ref(server-utils): Add orchestrion support for Vercel AI v5 #21979
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,9 +1,11 @@ | ||
| /* 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'; | ||
| import { CHANNELS } from '../orchestrion/channels'; | ||
| import { bindTracingChannelToSpan, type TracingChannelPayloadWithSpan } from '../tracing-channel'; | ||
| import { | ||
| captureToolError, | ||
| clearOperationCallId, | ||
| clearOperationId, | ||
| createSpanFromMessage, | ||
|
|
@@ -14,22 +16,31 @@ import { | |
| } from './vercel-ai-dc-subscriber'; | ||
|
|
||
| /** | ||
| * v5 & v6 channel adapter for the Vercel AI (`ai`) SDK. | ||
| * | ||
| * `ai` >= 7 publishes a normalized `ai:telemetry` tracing channel natively | ||
| * (consumed by `subscribeVercelAiTracingChannel`). v5/v6 have no such channel, so | ||
| * orchestrion injects `orchestrion:ai:*` channels around the top-level | ||
| * functions (see `orchestrion/config/index.ts`). The injected channels carry only the | ||
| * wrapped call's `{ arguments, result, error }` — NOT v7's normalized `event` | ||
| * object — so this adapter reconstructs an equivalent {@link VercelAiChannelMessage} | ||
| * from v5/v6's argument/result shapes and delegates to the SAME span-building core | ||
| * (`createSpanFromMessage` / `enrichSpanOnEnd`) the v7 subscriber uses, so the | ||
| * emitted spans are identical between v5, v6 and v7. | ||
| * | ||
| * The model call (`languageModelCall` / `generate_content` span) has no | ||
| * injectable definition in `ai`, so we instead wrap `resolveLanguageModel` (the | ||
| * single chokepoint every model call flows through, present in both v5 and v6) | ||
| * and monkey-patch `doGenerate`/`doStream` on the returned model. | ||
| * | ||
| * Tool-call spans differ by version: v6 exposes a per-call `executeToolCall` | ||
| * function orchestrion wraps into its own channel. v5 has no such export (only a | ||
| * batch `executeTools`), so instead we monkey-patch each tool's `execute` from | ||
| * the operation's `tools`. On v6 that patch is inert: `executeToolCall` runs | ||
| * `execute` inside its own tool-call span (the active async context), so the | ||
| * patch sees that span as its parent and skips — only v5 (where the parent is | ||
| * the enclosing `invoke_agent` span) emits the patched span, so v6 never | ||
| * double-counts tool spans. | ||
| */ | ||
|
|
||
| /** Shape orchestrion's transform attaches to the tracing-channel context. */ | ||
|
|
@@ -51,10 +62,14 @@ interface ResolvedModel { | |
| } | ||
|
|
||
| const PATCHED = Symbol('SentryVercelAiModelPatched'); | ||
| const TOOL_PATCHED = Symbol('SentryVercelAiToolPatched'); | ||
|
|
||
| /** A resolved model with our patch bookkeeping (idempotency flag). */ | ||
| type PatchableModel = ResolvedModel & { [PATCHED]?: boolean }; | ||
|
|
||
| /** A tool definition off an operation's `tools`, with our patch bookkeeping. */ | ||
| type PatchableTool = { execute?: (...args: unknown[]) => unknown; [TOOL_PATCHED]?: boolean }; | ||
|
|
||
| // Per-operation correlation id. No Date/random (unavailable / non-deterministic) — a counter is enough. | ||
| let callIdCounter = 0; | ||
| function nextCallId(): string { | ||
|
|
@@ -68,6 +83,10 @@ const messages = new WeakMap<object, VercelAiChannelMessage>(); | |
| // parent against this set (so it never mis-attributes to the enclosing `main`/user span) and reads the | ||
| // parent's `callId` so its span can be named after the operation (e.g. `ai.streamText.doStream`). | ||
| const operationSpans = new WeakSet<Span>(); | ||
| // Tool-call spans — only v6's `executeToolCall` channel opens these. The v5 tool-`execute` patch | ||
| // consults this to avoid double-counting: `executeToolCall` runs `execute` inside its span (which is | ||
| // therefore the active parent), so a resolved parent in this set means v6 already spanned the call. | ||
| const toolCallSpans = new WeakSet<Span>(); | ||
| const callIdBySpan = new WeakMap<Span, string>(); | ||
| // The operation's per-call recording flags (`experimental_telemetry.recordInputs/recordOutputs`), keyed by | ||
| // its span. A model call carries no telemetry of its own, so it inherits the enclosing operation's flags — | ||
|
|
@@ -198,11 +217,21 @@ function bindOperation( | |
| if (span) { | ||
| messages.set(data, message); | ||
| operationSpans.add(span); | ||
| // v6's `executeToolCall` span: tracked so the v5 tool-`execute` patch can recognize (and skip) | ||
| // tool calls it runs, since that patch sees this span as its active parent (see `patchToolExecute`). | ||
| if (message.type === 'executeTool') { | ||
| toolCallSpans.add(span); | ||
| } | ||
| const callId = asString(message.event.callId); | ||
| if (callId) { | ||
| callIdBySpan.set(span, callId); | ||
| } | ||
| recordingBySpan.set(span, recording(telemetry)); | ||
| // v5 has no `executeToolCall` channel, so patch each tool's `execute` to emit the tool-call span. | ||
| // Inert on v6 (guarded inside `patchToolExecute` when the parent is `executeToolCall`'s own span). | ||
| if (isRecord(callOptions.tools)) { | ||
| patchOperationTools(callOptions.tools, options); | ||
| } | ||
| } | ||
| return span; | ||
| }; | ||
|
|
@@ -394,6 +423,92 @@ function patchModelMethod( | |
| }; | ||
| } | ||
|
|
||
| /** | ||
| * Patch every executable tool on an operation's `tools` (a record keyed by tool name) so each | ||
| * invocation emits a `gen_ai.execute_tool` span. v5 routes tool execution through the batch | ||
| * `executeTools`, which binds `tool.execute` at call time — so replacing the `execute` property here | ||
| * (before the call runs) is picked up. Tools without an `execute` (client-side tools) are skipped, | ||
| * matching `executeTools`. | ||
| */ | ||
| function patchOperationTools(tools: Record<string, unknown>, options: VercelAiChannelOptions): void { | ||
| // This runs inside the channel `start` transform (no upstream try/catch), so a throw here would break | ||
| // the user's `ai` call. Instrumentation must never do that — degrade to no tool span instead. | ||
| try { | ||
| for (const [toolName, tool] of Object.entries(tools)) { | ||
| if (isRecord(tool)) { | ||
| patchToolExecute(toolName, tool as PatchableTool, tools, options); | ||
| } | ||
| } | ||
| } catch { | ||
| DEBUG_BUILD && debug.log('Vercel AI orchestrion tool patching failed.'); | ||
| } | ||
|
Comment on lines
+434
to
+444
Member
There was a problem hiding this comment. m: Maybe we could move the try/catch inside the loop so we can continue patching tools even if earlier tools fail to patch for whatever reason, wdyt? Or does one failure indicate we can never patch the other tools anyway? |
||
| } | ||
|
|
||
| function patchToolExecute( | ||
| toolName: string, | ||
| tool: PatchableTool, | ||
| tools: Record<string, unknown>, | ||
| options: VercelAiChannelOptions, | ||
| ): void { | ||
| const original = tool.execute; | ||
| if (typeof original !== 'function' || tool[TOOL_PATCHED]) { | ||
| return; | ||
| } | ||
| tool[TOOL_PATCHED] = true; | ||
| tool.execute = function (this: unknown, input: unknown, ...rest: unknown[]): unknown { | ||
| // Skip if there's no enclosing operation span (telemetry disabled for the call), or if that parent | ||
| // is itself a tool-call span — the latter means v6's `executeToolCall` already opened a span for | ||
| // this call and is running `execute` inside it, so spanning again here would double-count. On v5 | ||
| // the parent is the enclosing `invoke_agent` operation span, so we proceed. | ||
| const parent = resolveModelCallParent(); | ||
| if (!parent || toolCallSpans.has(parent)) { | ||
| return original.apply(this, [input, ...rest]); | ||
| } | ||
|
There was a problem hiding this comment. streamText v5 tools miss spansMedium Severity For Vercel AI v5 under orchestrion,   Reviewed by Cursor Bugbot for commit d2d3306. Configure here.
Member
Author
There was a problem hiding this comment. this behaves the same as in v5, things may be emitted weirdly but it's fine to keep consistency with v5 in otel I'd say here. |
||
|
|
||
| // v5 passes `{ toolCallId, messages, abortSignal, ... }` as the second argument to `execute`. | ||
| const callOptions = isRecord(rest[0]) ? rest[0] : {}; | ||
| const message: VercelAiChannelMessage = { | ||
| type: 'executeTool', | ||
| event: { | ||
| callId: callIdBySpan.get(parent), | ||
| toolCall: { toolName, toolCallId: asString(callOptions.toolCallId), input }, | ||
| // The `tools` record (keyed by name) lets the shared core backfill the tool's `description`. | ||
| tools, | ||
| // Inherit the enclosing operation's per-call recording flags so tool inputs/outputs are recorded | ||
| // whenever they are on the parent `invoke_agent` span. | ||
| ...recordingBySpan.get(parent), | ||
| }, | ||
| }; | ||
| const span = withActiveSpan(parent, () => createSpanFromMessage(message, options)); | ||
| // `executeTool` always opens a span; the guard just keeps the wrapper safe if that changes. | ||
| if (!span) { | ||
| return original.apply(this, [input, ...rest]); | ||
| } | ||
|
|
||
| // v5's `executeTools` catches a thrown tool error and turns it into `tool-error` content rather | ||
| // than rejecting, so the user never sees a rejection — we must capture it here. Rethrow so | ||
| // `executeTools` still produces its `tool-error` result and the operation continues normally. | ||
| const failSpan = (error: unknown): never => { | ||
| captureToolError(span, message, error); | ||
| span.end(); | ||
| throw error; | ||
| }; | ||
|
|
||
| try { | ||
| const result = Promise.resolve(original.apply(this, [input, ...rest])); | ||
| return result.then(value => { | ||
| // The shared core (matching v6/v7) expects the tool result nested under `output`. | ||
| message.result = { output: value }; | ||
| enrichSpanOnEnd(span, message, options); | ||
| span.end(); | ||
| return value; | ||
| }, failSpan); | ||
| } catch (error) { | ||
| return failSpan(error); | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| function buildTextMessage(type: 'generateText' | 'streamText'): MessageBuilder { | ||
| return (options, telemetry) => ({ | ||
| type, | ||
|
|
@@ -422,7 +537,15 @@ function normalizePromptMessages(options: Record<string, unknown>): unknown { | |
| } | ||
|
|
||
| function recording(telemetry: Record<string, unknown>): { recordInputs: unknown; recordOutputs: unknown } { | ||
| // Match the OTel integration's per-call default: an explicit `recordInputs`/`recordOutputs` wins, but a | ||
| // call that merely enables telemetry (`isEnabled: true`, no explicit flags) defaults both to `true`. | ||
| // Unlike v7's native `ai:telemetry` channel, the orchestrion channels expose `isEnabled`, so — as noted | ||
| // in `resolveRecording` — we CAN reproduce that default here (the native-channel subscriber cannot). | ||
| const enabledDefault = telemetry.isEnabled === true ? true : undefined; | ||
| return { | ||
| recordInputs: telemetry.recordInputs ?? enabledDefault, | ||
| recordOutputs: telemetry.recordOutputs ?? enabledDefault, | ||
| }; | ||
| } | ||
|
|
||
| function modelFields(model: unknown): { provider?: string; modelId?: string } { | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
q: Why don't we just adjust the ranges for these 4 apis?