Skip to content

Commit 55bd48a

Browse files
JPeer264claude
andauthored
feat(core): Instrument workers-ai-provider (#22119)
closes #20839 <br>closes #20839 This is instrumenting `workers-ai`, but it needs to be wrapped, as in Cloudflare there is no monkey patching for bindings. In a follow-up PR this will be included in the `instrumentEnv` within Cloudflare. Tbh I would have added this into `packages/cloudflare`, but all AI integrations do live in core, and there are also some utils which are only there and would have needed to export, such as `endStreamSpan` and all the gen-ai-attributes. The downside of this is that the workers-ai types needed to be vendored in (in the `packages/cloudflare` we could have taken them from the `@cloudflare/workers-types` package. --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent c47b449 commit 55bd48a

9 files changed

Lines changed: 1223 additions & 0 deletions

File tree

‎packages/core/src/shared-exports.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,8 @@ export {
230230
export { wrapToolsWithSpans, extractLLMFromParams, extractAgentNameFromParams } from './tracing/langgraph/utils';
231231
export { LANGGRAPH_INTEGRATION_NAME } from './tracing/langgraph/constants';
232232
export type { LangGraphOptions, LangGraphIntegration, CompiledGraph } from './tracing/langgraph/types';
233+
export { instrumentWorkersAiClient } from './tracing/workers-ai';
234+
export type { WorkersAiClient, WorkersAiOptions } from './tracing/workers-ai/types';
233235
// eslint-disable-next-line typescript/no-deprecated
234236
export type { OpenAiClient, OpenAiOptions, InstrumentedMethod } from './tracing/openai/types';
235237
export type {
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
/**
2+
* The provider value for the `gen_ai.provider.name` attribute.
3+
* @see https://developers.cloudflare.com/workers-ai/
4+
*/
5+
export const WORKERS_AI_PROVIDER_NAME = 'cloudflare.workers_ai';
6+
7+
/**
8+
* The Sentry origin for spans created by the Workers AI instrumentation.
9+
*/
10+
export const WORKERS_AI_ORIGIN = 'auto.ai.cloudflare.workers_ai';
Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
import { SPAN_STATUS_ERROR } from '../../tracing';
2+
import { startSpan, startSpanManual } from '../../tracing/trace';
3+
import type { Span } from '../../types/span';
4+
import { isObjectLike } from '../../utils/is';
5+
import { resolveAIRecordingOptions, shouldEnableTruncation } from '../ai/utils';
6+
import { instrumentWorkersAiStream } from './streaming';
7+
import type { WorkersAiOptions } from './types';
8+
import { addRequestAttributes, addResponseAttributes, extractRequestAttributes, getOperationName } from './utils';
9+
10+
// Adapted from /server-utils/src/vercel-ai/util.ts
11+
// TODO(v11): Reuse this function once this gets moved to @sentry/server-utils
12+
// Workers AI streaming responses are SSE byte streams, so we narrow to `Uint8Array`.
13+
function isReadableStream(value: unknown): value is ReadableStream<Uint8Array> {
14+
return (
15+
isObjectLike(value) &&
16+
typeof (value as { pipeThrough?: unknown }).pipeThrough === 'function' &&
17+
typeof (value as { getReader?: unknown }).getReader === 'function'
18+
);
19+
}
20+
21+
/**
22+
* Wrap the `run` method of the Workers AI binding with Sentry tracing.
23+
*/
24+
function instrumentRun(
25+
originalRun: (...args: unknown[]) => Promise<unknown>,
26+
context: unknown,
27+
options: WorkersAiOptions & Required<Pick<WorkersAiOptions, 'recordInputs' | 'recordOutputs'>>,
28+
): (...args: unknown[]) => Promise<unknown> {
29+
return function instrumentedRun(...args: unknown[]): Promise<unknown> {
30+
const [model, inputs, runOptions] = args as [unknown, unknown, Record<string, unknown> | undefined];
31+
32+
const operationName = getOperationName(inputs);
33+
const requestAttributes = extractRequestAttributes(model, inputs, operationName);
34+
const modelName = typeof model === 'string' ? model : 'unknown';
35+
36+
const isStreamRequested =
37+
!!inputs && typeof inputs === 'object' && (inputs as { stream?: unknown }).stream === true;
38+
const returnsRawResponse =
39+
!!runOptions &&
40+
typeof runOptions === 'object' &&
41+
(runOptions.returnRawResponse === true || runOptions.websocket === true);
42+
43+
const spanConfig = {
44+
name: `${operationName} ${modelName}`,
45+
op: `gen_ai.${operationName}`,
46+
attributes: requestAttributes,
47+
};
48+
49+
if (isStreamRequested && !returnsRawResponse) {
50+
return startSpanManual(spanConfig, (span: Span) => {
51+
// `startSpanManual` does not auto-end the span, so we must end it on every exit path,
52+
// including a synchronous throw from `run`.
53+
const handleError = (error: unknown): never => {
54+
span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
55+
span.end();
56+
throw error;
57+
};
58+
59+
let originalResult: Promise<unknown>;
60+
61+
try {
62+
originalResult = originalRun.apply(context, args) as Promise<unknown>;
63+
} catch (error) {
64+
return handleError(error);
65+
}
66+
67+
if (options.recordInputs) {
68+
addRequestAttributes(span, inputs, operationName, shouldEnableTruncation(options.enableTruncation));
69+
}
70+
71+
return originalResult.then(result => {
72+
if (isReadableStream(result)) {
73+
return instrumentWorkersAiStream(result, span, options.recordOutputs);
74+
}
75+
76+
// The model did not actually return a stream — finalize the span eagerly.
77+
addResponseAttributes(span, result, options.recordOutputs);
78+
span.end();
79+
return result;
80+
}, handleError);
81+
});
82+
}
83+
84+
return startSpan(spanConfig, (span: Span) => {
85+
const originalResult = originalRun.apply(context, args) as Promise<unknown>;
86+
87+
if (options.recordInputs) {
88+
addRequestAttributes(span, inputs, operationName, shouldEnableTruncation(options.enableTruncation));
89+
}
90+
91+
return originalResult.then(result => {
92+
if (!returnsRawResponse) {
93+
addResponseAttributes(span, result, options.recordOutputs);
94+
}
95+
return result;
96+
});
97+
});
98+
};
99+
}
100+
101+
/**
102+
* Instrument a Cloudflare Workers AI binding (`env.AI`) with Sentry tracing.
103+
*
104+
* This wraps the binding's `run` method to create `gen_ai` spans following the
105+
* Sentry AI Agents conventions. All other methods are passed through untouched.
106+
*
107+
* In `@sentry/cloudflare`, the `env.AI` binding is instrumented automatically —
108+
* wrapping manually is only needed to pass custom options.
109+
*
110+
* @example
111+
* ```javascript
112+
* const ai = Sentry.instrumentWorkersAiClient(env.AI, { recordInputs: true, recordOutputs: true });
113+
* const result = await ai.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' });
114+
* ```
115+
*/
116+
export function instrumentWorkersAiClient<T extends object>(client: T, options?: WorkersAiOptions): T {
117+
const resolvedOptions = resolveAIRecordingOptions(options);
118+
119+
const instrumented = new Proxy(client, {
120+
get(target: object, prop: string | symbol, receiver: unknown): unknown {
121+
const value = Reflect.get(target, prop, receiver);
122+
123+
if (prop === 'run' && typeof value === 'function') {
124+
return instrumentRun(value as (...args: unknown[]) => Promise<unknown>, target, resolvedOptions);
125+
}
126+
127+
// Bind passed-through functions to the original target to preserve `this` (e.g. private fields).
128+
return typeof value === 'function' ? (value as (...args: unknown[]) => unknown).bind(target) : value;
129+
},
130+
}) as T;
131+
132+
return instrumented;
133+
}
Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,229 @@
1+
import { SPAN_STATUS_ERROR } from '../../tracing';
2+
import type { Span } from '../../types/span';
3+
import { endStreamSpan, type StreamResponseState } from '../ai/utils';
4+
import type { WorkersAiUsage } from './types';
5+
import { setOutputMessagesAttribute } from './utils';
6+
7+
interface WorkersAiStreamingToolCall {
8+
index?: number;
9+
id?: string;
10+
type?: string;
11+
function?: { name?: string; arguments?: string };
12+
// Some Workers AI models stream tool calls with the name/arguments at the top
13+
// level of the tool-call object instead of nested under `function`.
14+
name?: string;
15+
arguments?: string;
16+
}
17+
18+
interface WorkersAiStreamChunk {
19+
// Native Workers AI streaming shape (`env.AI.run` with `stream: true`).
20+
response?: unknown;
21+
tool_calls?: unknown[];
22+
// OpenAI-compatible streaming shape emitted for models routed through the
23+
// OpenAI-compatible endpoint (e.g. via `workers-ai-provider`).
24+
choices?: Array<{
25+
delta?: { content?: unknown; tool_calls?: WorkersAiStreamingToolCall[] };
26+
finish_reason?: unknown;
27+
}>;
28+
usage?: WorkersAiUsage & { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number };
29+
}
30+
31+
/**
32+
* Accumulate a fragmented OpenAI-compatible tool call (delivered across multiple
33+
* `choices[].delta.tool_calls` chunks) into the index-keyed accumulator.
34+
*/
35+
function accumulateStreamingToolCalls(
36+
toolCalls: WorkersAiStreamingToolCall[],
37+
accumulator: Record<number, WorkersAiStreamingToolCall>,
38+
): void {
39+
for (const toolCall of toolCalls) {
40+
// Normalize both shapes: name/arguments nested under `function`, or at the top level.
41+
const name = toolCall.function?.name ?? toolCall.name;
42+
const args = toolCall.function?.arguments ?? toolCall.arguments;
43+
44+
// A tool call must carry at least a name or argument fragment to be meaningful.
45+
if (name == null && args == null) {
46+
continue;
47+
}
48+
49+
const index = toolCall.index ?? 0;
50+
const existing = accumulator[index];
51+
52+
if (!existing) {
53+
accumulator[index] = {
54+
index,
55+
id: toolCall.id,
56+
type: toolCall.type,
57+
function: {
58+
name,
59+
arguments: args ?? '',
60+
},
61+
};
62+
} else if (existing.function) {
63+
if (name && !existing.function.name) {
64+
existing.function.name = name;
65+
}
66+
if (args) {
67+
existing.function.arguments = `${existing.function.arguments ?? ''}${args}`;
68+
}
69+
}
70+
}
71+
}
72+
73+
/**
74+
* Parse a single SSE line (`data: {...}`) and accumulate its data into the streaming state.
75+
*
76+
* Handles both the native Workers AI shape (top-level `response`/`tool_calls`) and the
77+
* OpenAI-compatible shape (`choices[].delta.content`/`choices[].delta.tool_calls`), because
78+
* the same `run()` call transparently yields either format depending on the model.
79+
*/
80+
function processLine(
81+
line: string,
82+
state: StreamResponseState,
83+
recordOutputs: boolean,
84+
toolCallAccumulator: Record<number, WorkersAiStreamingToolCall>,
85+
): void {
86+
const trimmed = line.trim();
87+
if (!trimmed.startsWith('data:')) {
88+
return;
89+
}
90+
91+
const data = trimmed.slice('data:'.length).trim();
92+
if (!data || data === '[DONE]') {
93+
return;
94+
}
95+
96+
let parsed: WorkersAiStreamChunk;
97+
try {
98+
parsed = JSON.parse(data) as WorkersAiStreamChunk;
99+
} catch {
100+
return;
101+
}
102+
103+
if (parsed.usage) {
104+
if (typeof parsed.usage.prompt_tokens === 'number') {
105+
state.promptTokens = parsed.usage.prompt_tokens;
106+
}
107+
if (typeof parsed.usage.completion_tokens === 'number') {
108+
state.completionTokens = parsed.usage.completion_tokens;
109+
}
110+
if (typeof parsed.usage.total_tokens === 'number') {
111+
state.totalTokens = parsed.usage.total_tokens;
112+
}
113+
}
114+
115+
if (recordOutputs && typeof parsed.response === 'string') {
116+
state.responseTexts.push(parsed.response);
117+
}
118+
119+
if (recordOutputs && Array.isArray(parsed.tool_calls) && parsed.tool_calls.length > 0) {
120+
state.toolCalls.push(...parsed.tool_calls);
121+
}
122+
123+
if (Array.isArray(parsed.choices)) {
124+
for (const choice of parsed.choices) {
125+
if (recordOutputs && typeof choice.delta?.content === 'string' && choice.delta.content) {
126+
state.responseTexts.push(choice.delta.content);
127+
}
128+
if (recordOutputs && Array.isArray(choice.delta?.tool_calls)) {
129+
accumulateStreamingToolCalls(choice.delta.tool_calls, toolCallAccumulator);
130+
}
131+
if (typeof choice.finish_reason === 'string') {
132+
state.finishReasons.push(choice.finish_reason);
133+
}
134+
}
135+
}
136+
}
137+
138+
/**
139+
* Wrap a Workers AI streaming response (a server-sent-events `ReadableStream`) so we can
140+
* accumulate the response text and token usage while passing the original bytes through untouched.
141+
*
142+
* The span is ended once the consumer finishes reading (or cancels) the stream.
143+
*/
144+
export function instrumentWorkersAiStream(
145+
stream: ReadableStream<Uint8Array>,
146+
span: Span,
147+
recordOutputs: boolean,
148+
): ReadableStream<Uint8Array> {
149+
const reader = stream.getReader();
150+
const decoder = new TextDecoder();
151+
152+
const state: StreamResponseState = {
153+
responseId: '',
154+
responseModel: '',
155+
finishReasons: [],
156+
responseTexts: [],
157+
toolCalls: [],
158+
promptTokens: undefined,
159+
completionTokens: undefined,
160+
totalTokens: undefined,
161+
};
162+
163+
// OpenAI-compatible tool calls arrive fragmented across chunks and are keyed by index;
164+
// accumulate them here and flatten into `state.toolCalls` once the stream ends.
165+
const toolCallAccumulator: Record<number, WorkersAiStreamingToolCall> = {};
166+
167+
let buffer = '';
168+
let spanEnded = false;
169+
170+
const finish = (): void => {
171+
if (spanEnded) {
172+
return;
173+
}
174+
spanEnded = true;
175+
176+
if (recordOutputs) {
177+
const accumulatedToolCalls = Object.values(toolCallAccumulator);
178+
if (accumulatedToolCalls.length > 0) {
179+
state.toolCalls.push(...accumulatedToolCalls);
180+
}
181+
182+
// Set the authoritative `gen_ai.output.messages` alongside the deprecated response
183+
// attributes `endStreamSpan` writes, so tool calls survive Relay's lossy migration.
184+
setOutputMessagesAttribute(span, {
185+
responseText: state.responseTexts.join(''),
186+
toolCalls: state.toolCalls,
187+
});
188+
}
189+
190+
endStreamSpan(span, state, recordOutputs);
191+
};
192+
193+
const flushBuffer = (isDone: boolean): void => {
194+
const lines = buffer.split('\n');
195+
// Keep the last (potentially incomplete) line in the buffer unless the stream is done.
196+
buffer = isDone ? '' : (lines.pop() ?? '');
197+
for (const line of lines) {
198+
processLine(line, state, recordOutputs, toolCallAccumulator);
199+
}
200+
};
201+
202+
return new ReadableStream<Uint8Array>({
203+
async pull(controller) {
204+
try {
205+
const { done, value } = await reader.read();
206+
207+
if (done) {
208+
buffer += decoder.decode();
209+
flushBuffer(true);
210+
finish();
211+
controller.close();
212+
return;
213+
}
214+
215+
buffer += decoder.decode(value, { stream: true });
216+
flushBuffer(false);
217+
controller.enqueue(value);
218+
} catch (error) {
219+
span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
220+
finish();
221+
controller.error(error);
222+
}
223+
},
224+
async cancel(reason) {
225+
finish();
226+
await reader.cancel(reason);
227+
},
228+
});
229+
}

0 commit comments

Comments
 (0)