From bcb96ab76a5e4d44c7ba3b13df1a3b6aaeb35d87 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Mon, 29 Jun 2026 13:39:39 +0200 Subject: [PATCH 01/10] feat(core): Instrument workers-ai-provider --- packages/core/src/shared-exports.ts | 2 + .../core/src/tracing/workers-ai/constants.ts | 10 + packages/core/src/tracing/workers-ai/index.ts | 140 ++++++++++ .../core/src/tracing/workers-ai/streaming.ts | 131 ++++++++++ packages/core/src/tracing/workers-ai/types.ts | 58 +++++ packages/core/src/tracing/workers-ai/utils.ts | 173 +++++++++++++ .../lib/tracing/workers-ai-streaming.test.ts | 114 +++++++++ .../core/test/lib/tracing/workers-ai.test.ts | 41 +++ .../test/lib/utils/workers-ai-utils.test.ts | 242 ++++++++++++++++++ 9 files changed, 911 insertions(+) create mode 100644 packages/core/src/tracing/workers-ai/constants.ts create mode 100644 packages/core/src/tracing/workers-ai/index.ts create mode 100644 packages/core/src/tracing/workers-ai/streaming.ts create mode 100644 packages/core/src/tracing/workers-ai/types.ts create mode 100644 packages/core/src/tracing/workers-ai/utils.ts create mode 100644 packages/core/test/lib/tracing/workers-ai-streaming.test.ts create mode 100644 packages/core/test/lib/tracing/workers-ai.test.ts create mode 100644 packages/core/test/lib/utils/workers-ai-utils.test.ts diff --git a/packages/core/src/shared-exports.ts b/packages/core/src/shared-exports.ts index a7b3f070c2c4..983367c882f9 100644 --- a/packages/core/src/shared-exports.ts +++ b/packages/core/src/shared-exports.ts @@ -230,6 +230,8 @@ export { export { wrapToolsWithSpans, extractLLMFromParams, extractAgentNameFromParams } from './tracing/langgraph/utils'; export { LANGGRAPH_INTEGRATION_NAME } from './tracing/langgraph/constants'; export type { LangGraphOptions, LangGraphIntegration, CompiledGraph } from './tracing/langgraph/types'; +export { instrumentWorkersAiClient } from './tracing/workers-ai'; +export type { WorkersAiClient, WorkersAiOptions } from './tracing/workers-ai/types'; // eslint-disable-next-line typescript/no-deprecated export type { OpenAiClient, OpenAiOptions, InstrumentedMethod } from './tracing/openai/types'; export type { diff --git a/packages/core/src/tracing/workers-ai/constants.ts b/packages/core/src/tracing/workers-ai/constants.ts new file mode 100644 index 000000000000..89847abbae1c --- /dev/null +++ b/packages/core/src/tracing/workers-ai/constants.ts @@ -0,0 +1,10 @@ +/** + * The provider value for the `gen_ai.system` attribute. + * @see https://developers.cloudflare.com/workers-ai/ + */ +export const WORKERS_AI_SYSTEM_NAME = 'cloudflare.workers_ai'; + +/** + * The Sentry origin for spans created by the Workers AI instrumentation. + */ +export const WORKERS_AI_ORIGIN = 'auto.ai.cloudflare.workers_ai'; diff --git a/packages/core/src/tracing/workers-ai/index.ts b/packages/core/src/tracing/workers-ai/index.ts new file mode 100644 index 000000000000..e38d4bbd53ce --- /dev/null +++ b/packages/core/src/tracing/workers-ai/index.ts @@ -0,0 +1,140 @@ +import { captureException } from '../../exports'; +import { SPAN_STATUS_ERROR } from '../../tracing'; +import { startSpan, startSpanManual } from '../../tracing/trace'; +import type { Span } from '../../types/span'; +import { resolveAIRecordingOptions, shouldEnableTruncation } from '../ai/utils'; +import { WORKERS_AI_ORIGIN } from './constants'; +import { instrumentWorkersAiStream } from './streaming'; +import type { WorkersAiOptions } from './types'; +import { addRequestAttributes, addResponseAttributes, extractRequestAttributes, getOperationName } from './utils'; + +function isReadableStream(value: unknown): value is ReadableStream { + return typeof ReadableStream !== 'undefined' && value instanceof ReadableStream; +} + +/** + * Wrap the `run` method of the Workers AI binding with Sentry tracing. + */ +function instrumentRun( + originalRun: (...args: unknown[]) => Promise, + context: unknown, + options: WorkersAiOptions & Required>, +): (...args: unknown[]) => Promise { + return function instrumentedRun(...args: unknown[]): Promise { + const [model, inputs, runOptions] = args as [unknown, unknown, Record | undefined]; + + const operationName = getOperationName(inputs); + const requestAttributes = extractRequestAttributes(model, inputs, operationName); + const modelName = typeof model === 'string' ? model : 'unknown'; + + const isStreamRequested = + !!inputs && typeof inputs === 'object' && (inputs as { stream?: unknown }).stream === true; + const returnsRawResponse = + !!runOptions && + typeof runOptions === 'object' && + (runOptions.returnRawResponse === true || runOptions.websocket === true); + + const spanConfig = { + name: `${operationName} ${modelName}`, + op: `gen_ai.${operationName}`, + attributes: requestAttributes, + }; + + if (isStreamRequested && !returnsRawResponse) { + return startSpanManual(spanConfig, (span: Span) => { + const originalResult = originalRun.apply(context, args) as Promise; + + if (options.recordInputs) { + addRequestAttributes(span, inputs, operationName, shouldEnableTruncation(options.enableTruncation)); + } + + return originalResult.then( + result => { + if (isReadableStream(result)) { + return instrumentWorkersAiStream(result, span, options.recordOutputs); + } + + // The model did not actually return a stream — finalize the span eagerly. + addResponseAttributes(span, result, options.recordOutputs); + span.end(); + return result; + }, + error => { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + captureException(error, { + mechanism: { handled: false, type: `${WORKERS_AI_ORIGIN}.stream`, data: { function: 'run' } }, + }); + span.end(); + throw error; + }, + ); + }); + } + + return startSpan(spanConfig, (span: Span) => { + const originalResult = originalRun.apply(context, args) as Promise; + + if (options.recordInputs) { + addRequestAttributes(span, inputs, operationName, shouldEnableTruncation(options.enableTruncation)); + } + + return originalResult.then( + result => { + if (!returnsRawResponse) { + addResponseAttributes(span, result, options.recordOutputs); + } + return result; + }, + error => { + captureException(error, { + mechanism: { handled: false, type: WORKERS_AI_ORIGIN, data: { function: 'run' } }, + }); + throw error; + }, + ); + }); + }; +} + +// Guards against double-wrapping when a binding instrumented via the automatic +// env instrumentation in `@sentry/cloudflare` is additionally wrapped manually. +const instrumentedClients = new WeakSet(); + +/** + * Instrument a Cloudflare Workers AI binding (`env.AI`) with Sentry tracing. + * + * This wraps the binding's `run` method to create `gen_ai` spans following the + * Sentry AI Agents conventions. All other methods are passed through untouched. + * + * In `@sentry/cloudflare`, the `env.AI` binding is instrumented automatically — + * wrapping manually is only needed to pass custom options. + * + * @example + * ```javascript + * const ai = Sentry.instrumentWorkersAiClient(env.AI, { recordInputs: true, recordOutputs: true }); + * const result = await ai.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' }); + * ``` + */ +export function instrumentWorkersAiClient(client: T, options?: WorkersAiOptions): T { + if (instrumentedClients.has(client)) { + return client; + } + + const resolvedOptions = resolveAIRecordingOptions(options); + + const instrumented = new Proxy(client, { + get(target: object, prop: string | symbol, receiver: unknown): unknown { + const value = Reflect.get(target, prop, receiver); + + if (prop === 'run' && typeof value === 'function') { + return instrumentRun(value as (...args: unknown[]) => Promise, target, resolvedOptions); + } + + // Bind passed-through functions to the original target to preserve `this` (e.g. private fields). + return typeof value === 'function' ? (value as (...args: unknown[]) => unknown).bind(target) : value; + }, + }) as T; + + instrumentedClients.add(instrumented); + return instrumented; +} diff --git a/packages/core/src/tracing/workers-ai/streaming.ts b/packages/core/src/tracing/workers-ai/streaming.ts new file mode 100644 index 000000000000..c2a1055a8be0 --- /dev/null +++ b/packages/core/src/tracing/workers-ai/streaming.ts @@ -0,0 +1,131 @@ +import { captureException } from '../../exports'; +import { SPAN_STATUS_ERROR } from '../../tracing'; +import type { Span } from '../../types/span'; +import { endStreamSpan, type StreamResponseState } from '../ai/utils'; +import { WORKERS_AI_ORIGIN } from './constants'; +import type { WorkersAiUsage } from './types'; + +interface WorkersAiStreamChunk { + response?: unknown; + usage?: WorkersAiUsage; + tool_calls?: unknown[]; +} + +/** + * Parse a single SSE line (`data: {...}`) and accumulate its data into the streaming state. + */ +function processLine(line: string, state: StreamResponseState, recordOutputs: boolean): void { + const trimmed = line.trim(); + if (!trimmed.startsWith('data:')) { + return; + } + + const data = trimmed.slice('data:'.length).trim(); + if (!data || data === '[DONE]') { + return; + } + + let parsed: WorkersAiStreamChunk; + try { + parsed = JSON.parse(data) as WorkersAiStreamChunk; + } catch { + return; + } + + if (parsed.usage) { + if (typeof parsed.usage.prompt_tokens === 'number') { + state.promptTokens = parsed.usage.prompt_tokens; + } + if (typeof parsed.usage.completion_tokens === 'number') { + state.completionTokens = parsed.usage.completion_tokens; + } + if (typeof parsed.usage.total_tokens === 'number') { + state.totalTokens = parsed.usage.total_tokens; + } + } + + if (recordOutputs && typeof parsed.response === 'string') { + state.responseTexts.push(parsed.response); + } + + if (recordOutputs && Array.isArray(parsed.tool_calls) && parsed.tool_calls.length > 0) { + state.toolCalls.push(...parsed.tool_calls); + } +} + +/** + * Wrap a Workers AI streaming response (a server-sent-events `ReadableStream`) so we can + * accumulate the response text and token usage while passing the original bytes through untouched. + * + * The span is ended once the consumer finishes reading (or cancels) the stream. + */ +export function instrumentWorkersAiStream( + stream: ReadableStream, + span: Span, + recordOutputs: boolean, +): ReadableStream { + const reader = stream.getReader(); + const decoder = new TextDecoder(); + + const state: StreamResponseState = { + responseId: '', + responseModel: '', + finishReasons: [], + responseTexts: [], + toolCalls: [], + promptTokens: undefined, + completionTokens: undefined, + totalTokens: undefined, + }; + + let buffer = ''; + let spanEnded = false; + + const finish = (): void => { + if (spanEnded) { + return; + } + spanEnded = true; + endStreamSpan(span, state, recordOutputs); + }; + + const flushBuffer = (isDone: boolean): void => { + const lines = buffer.split('\n'); + // Keep the last (potentially incomplete) line in the buffer unless the stream is done. + buffer = isDone ? '' : (lines.pop() ?? ''); + for (const line of lines) { + processLine(line, state, recordOutputs); + } + }; + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + + if (done) { + buffer += decoder.decode(); + flushBuffer(true); + finish(); + controller.close(); + return; + } + + buffer += decoder.decode(value, { stream: true }); + flushBuffer(false); + controller.enqueue(value); + } catch (error) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + captureException(error, { + mechanism: { handled: false, type: `${WORKERS_AI_ORIGIN}.stream` }, + }); + finish(); + controller.error(error); + } + }, + async cancel(reason) { + finish(); + await reader.cancel(reason); + }, + }); +} diff --git a/packages/core/src/tracing/workers-ai/types.ts b/packages/core/src/tracing/workers-ai/types.ts new file mode 100644 index 000000000000..33334a6ed1bf --- /dev/null +++ b/packages/core/src/tracing/workers-ai/types.ts @@ -0,0 +1,58 @@ +import type { AIRecordingOptions } from '../ai/utils'; + +export interface WorkersAiOptions extends AIRecordingOptions { + /** + * Enable or disable truncation of recorded input messages. + * Defaults to `true`. + */ + enableTruncation?: boolean; +} + +/** + * Minimal shape of the Cloudflare Workers AI binding (`env.AI`). + * We only rely on the `run` method, everything else is passed through untouched. + * @see https://developers.cloudflare.com/workers-ai/configuration/bindings/ + */ +export interface WorkersAiClient { + run: (model: string, inputs: Record, options?: Record) => Promise; + [key: string]: unknown; +} + +/** + * The token usage reported by Workers AI text generation models. + */ +export interface WorkersAiUsage { + prompt_tokens?: number; + completion_tokens?: number; + total_tokens?: number; +} + +/** + * The (subset of) inputs accepted by Workers AI `run` calls that we read from. + * @see https://developers.cloudflare.com/workers-ai/models/ + */ +export interface WorkersAiInput { + // Text generation + prompt?: string; + messages?: Array<{ role?: string; content?: unknown }>; + stream?: boolean; + max_tokens?: number; + temperature?: number; + top_p?: number; + top_k?: number; + frequency_penalty?: number; + presence_penalty?: number; + tools?: unknown; + functions?: unknown; + // Text embeddings + text?: string | string[]; +} + +/** + * The (subset of) outputs returned by Workers AI text generation models that we read from. + */ +export interface WorkersAiOutput { + response?: unknown; + tool_calls?: unknown[]; + usage?: WorkersAiUsage; +} diff --git a/packages/core/src/tracing/workers-ai/utils.ts b/packages/core/src/tracing/workers-ai/utils.ts new file mode 100644 index 000000000000..94611c90ce19 --- /dev/null +++ b/packages/core/src/tracing/workers-ai/utils.ts @@ -0,0 +1,173 @@ +import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes'; +import type { Span, SpanAttributeValue } from '../../types/span'; +import { + GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, + GEN_AI_INPUT_MESSAGES_ATTRIBUTE, + GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, + GEN_AI_OPERATION_NAME_ATTRIBUTE, + GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE, + GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE, + GEN_AI_REQUEST_MODEL_ATTRIBUTE, + GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE, + GEN_AI_REQUEST_STREAM_ATTRIBUTE, + GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE, + GEN_AI_REQUEST_TOP_K_ATTRIBUTE, + GEN_AI_REQUEST_TOP_P_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, + GEN_AI_SYSTEM_ATTRIBUTE, + GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, +} from '../ai/gen-ai-attributes'; +import { extractSystemInstructions, getJsonString, getTruncatedJsonString, setTokenUsageAttributes } from '../ai/utils'; +import { WORKERS_AI_ORIGIN, WORKERS_AI_SYSTEM_NAME } from './constants'; +import type { WorkersAiInput, WorkersAiOutput } from './types'; + +/** + * Determine the gen_ai operation name from the inputs passed to `AI.run`. + * Workers AI exposes a single `run` method, so we infer the operation from the input shape. + */ +export function getOperationName(inputs: unknown): string { + if (inputs && typeof inputs === 'object') { + if ('messages' in inputs || 'prompt' in inputs) { + return 'chat'; + } + if ('text' in inputs) { + return 'embeddings'; + } + } + return 'chat'; +} + +/** + * Extract the request attributes (model, request parameters, system, origin) from a `run` call. + */ +export function extractRequestAttributes( + model: unknown, + inputs: unknown, + operationName: string, +): Record { + const attributes: Record = { + [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, + [GEN_AI_OPERATION_NAME_ATTRIBUTE]: operationName, + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, + [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: typeof model === 'string' ? model : 'unknown', + }; + + if (inputs && typeof inputs === 'object') { + const params = inputs as WorkersAiInput; + + if (typeof params.temperature === 'number') { + attributes[GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE] = params.temperature; + } + if (typeof params.max_tokens === 'number') { + attributes[GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE] = params.max_tokens; + } + if (typeof params.top_p === 'number') { + attributes[GEN_AI_REQUEST_TOP_P_ATTRIBUTE] = params.top_p; + } + if (typeof params.top_k === 'number') { + attributes[GEN_AI_REQUEST_TOP_K_ATTRIBUTE] = params.top_k; + } + if (typeof params.frequency_penalty === 'number') { + attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE] = params.frequency_penalty; + } + if (typeof params.presence_penalty === 'number') { + attributes[GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE] = params.presence_penalty; + } + if (params.stream === true) { + attributes[GEN_AI_REQUEST_STREAM_ATTRIBUTE] = true; + } + } + + return attributes; +} + +/** + * Record the request inputs (messages/prompt/embeddings input) on the span. + * Only called when `recordInputs` is enabled. + */ +export function addRequestAttributes( + span: Span, + inputs: unknown, + operationName: string, + enableTruncation: boolean, +): void { + if (!inputs || typeof inputs !== 'object') { + return; + } + const params = inputs as WorkersAiInput; + + // Store embeddings input on a separate attribute and do not truncate it + if (operationName === 'embeddings') { + const text = params.text; + + if (text == null) { + return; + } + if (typeof text === 'string' && text.length === 0) { + return; + } + if (Array.isArray(text) && text.length === 0) { + return; + } + + span.setAttribute(GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, typeof text === 'string' ? text : JSON.stringify(text)); + return; + } + + const src = params.messages ?? params.prompt; + if (src == null) { + return; + } + if (Array.isArray(src) && src.length === 0) { + return; + } + + const { systemInstructions, filteredMessages } = extractSystemInstructions(src); + + if (systemInstructions) { + span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, systemInstructions); + } + + span.setAttribute( + GEN_AI_INPUT_MESSAGES_ATTRIBUTE, + enableTruncation ? getTruncatedJsonString(filteredMessages) : getJsonString(filteredMessages), + ); + + span.setAttribute( + GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, + Array.isArray(filteredMessages) ? filteredMessages.length : 1, + ); +} + +/** + * Record the response attributes (token usage, response text, tool calls) on the span. + */ +export function addResponseAttributes(span: Span, result: unknown, recordOutputs: boolean): void { + if (!result || typeof result !== 'object') { + return; + } + + // Raw `Response` objects (from `returnRawResponse`/`websocket`) cannot be introspected without consuming them. + if (typeof Response !== 'undefined' && result instanceof Response) { + return; + } + + const response = result as WorkersAiOutput; + + if (response.usage) { + setTokenUsageAttributes(span, response.usage.prompt_tokens, response.usage.completion_tokens); + } + + if (recordOutputs) { + if (typeof response.response === 'string') { + span.setAttribute(GEN_AI_RESPONSE_TEXT_ATTRIBUTE, response.response); + } else if (response.response != null) { + span.setAttribute(GEN_AI_RESPONSE_TEXT_ATTRIBUTE, JSON.stringify(response.response)); + } + + if (Array.isArray(response.tool_calls) && response.tool_calls.length > 0) { + span.setAttribute(GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, JSON.stringify(response.tool_calls)); + } + } +} diff --git a/packages/core/test/lib/tracing/workers-ai-streaming.test.ts b/packages/core/test/lib/tracing/workers-ai-streaming.test.ts new file mode 100644 index 000000000000..16f7983b0d4e --- /dev/null +++ b/packages/core/test/lib/tracing/workers-ai-streaming.test.ts @@ -0,0 +1,114 @@ +import { describe, expect, it } from 'vitest'; +import type { Span } from '../../../src'; +import { + GEN_AI_RESPONSE_STREAMING_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, +} from '../../../src/tracing/ai/gen-ai-attributes'; +import { instrumentWorkersAiStream } from '../../../src/tracing/workers-ai/streaming'; + +function createMockSpan(): { span: Span; attributes: Record; ended: () => boolean } { + const attributes: Record = {}; + let isEnded = false; + const span = { + isRecording: () => !isEnded, + setAttribute: (key: string, value: unknown) => { + attributes[key] = value; + }, + setAttributes: (attrs: Record) => { + Object.assign(attributes, attrs); + }, + setStatus: () => {}, + end: () => { + isEnded = true; + }, + } as unknown as Span; + return { span, attributes, ended: () => isEnded }; +} + +function streamFromChunks(chunks: string[]): ReadableStream { + const encoder = new TextEncoder(); + const queued = [...chunks]; + return new ReadableStream({ + pull(controller) { + const next = queued.shift(); + if (next === undefined) { + controller.close(); + return; + } + controller.enqueue(encoder.encode(next)); + }, + }); +} + +describe('instrumentWorkersAiStream', () => { + it('passes the original bytes through untouched', async () => { + const { span } = createMockSpan(); + const raw = 'data: {"response":"Hello"}\n\ndata: [DONE]\n\n'; + + const instrumented = instrumentWorkersAiStream(streamFromChunks([raw]), span, true); + const passedThrough = await new Response(instrumented).text(); + + expect(passedThrough).toBe(raw); + }); + + it('accumulates response text and usage across SSE events split across read boundaries', async () => { + const { span, attributes, ended } = createMockSpan(); + // The second event is deliberately split mid-line across two reads. + const chunks = [ + 'data: {"response":"The capital "}\n\ndata: {"respo', + 'nse":"of France "}\n\ndata: {"response":"is Paris."', + ',"usage":{"prompt_tokens":12,"completion_tokens":7,"total_tokens":19}}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + expect(attributes[GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]).toBe(true); + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('The capital of France is Paris.'); + expect(attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]).toBe(12); + expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]).toBe(7); + expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]).toBe(19); + expect(ended()).toBe(true); + }); + + it('records usage but not response text when recordOutputs is false', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"response":"secret"}\n\n', + 'data: {"response":"","usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, false); + await new Response(instrumented).text(); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBeUndefined(); + expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]).toBe(1); + }); + + it('ignores malformed SSE payloads without throwing', async () => { + const { span, attributes, ended } = createMockSpan(); + const chunks = ['data: not-json\n\n', 'data: {"response":"ok"}\n\ndata: [DONE]\n\n']; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('ok'); + expect(ended()).toBe(true); + }); + + it('ends the span when the consumer cancels the stream', async () => { + const { span, attributes, ended } = createMockSpan(); + + const instrumented = instrumentWorkersAiStream(streamFromChunks(['data: {"response":"partial"}\n\n']), span, true); + + const reader = instrumented.getReader(); + await reader.read(); + await reader.cancel('no longer needed'); + + expect(ended()).toBe(true); + expect(attributes[GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]).toBe(true); + }); +}); diff --git a/packages/core/test/lib/tracing/workers-ai.test.ts b/packages/core/test/lib/tracing/workers-ai.test.ts new file mode 100644 index 000000000000..9c15e7c0b5a5 --- /dev/null +++ b/packages/core/test/lib/tracing/workers-ai.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, it, vi } from 'vitest'; +import { instrumentWorkersAiClient } from '../../../src/tracing/workers-ai'; + +describe('instrumentWorkersAiClient', () => { + it('returns an already-instrumented client as-is instead of double-wrapping', () => { + const client = { run: vi.fn().mockResolvedValue({ response: 'ok' }) }; + + const instrumented = instrumentWorkersAiClient(client); + const instrumentedTwice = instrumentWorkersAiClient(instrumented); + + expect(instrumented).not.toBe(client); + expect(instrumentedTwice).toBe(instrumented); + }); + + it('passes through non-run methods bound to the original client', () => { + class FakeAi { + #internal = 'secret'; + + public run = vi.fn().mockResolvedValue({ response: 'ok' }); + + public readInternal(): string { + return this.#internal; + } + } + + const client = new FakeAi(); + const instrumented = instrumentWorkersAiClient(client); + + expect(instrumented.readInternal()).toBe('secret'); + }); + + it('creates a span and forwards arguments for run calls', async () => { + const client = { run: vi.fn().mockResolvedValue({ response: 'Paris' }) }; + const instrumented = instrumentWorkersAiClient(client); + + const result = await instrumented.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' }); + + expect(client.run).toHaveBeenCalledWith('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' }); + expect(result).toEqual({ response: 'Paris' }); + }); +}); diff --git a/packages/core/test/lib/utils/workers-ai-utils.test.ts b/packages/core/test/lib/utils/workers-ai-utils.test.ts new file mode 100644 index 000000000000..eed6cf170ed3 --- /dev/null +++ b/packages/core/test/lib/utils/workers-ai-utils.test.ts @@ -0,0 +1,242 @@ +import { describe, expect, it } from 'vitest'; +import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../../src'; +import type { Span } from '../../../src'; +import { + GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, + GEN_AI_INPUT_MESSAGES_ATTRIBUTE, + GEN_AI_OPERATION_NAME_ATTRIBUTE, + GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE, + GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE, + GEN_AI_REQUEST_MODEL_ATTRIBUTE, + GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE, + GEN_AI_REQUEST_STREAM_ATTRIBUTE, + GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE, + GEN_AI_REQUEST_TOP_K_ATTRIBUTE, + GEN_AI_REQUEST_TOP_P_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, + GEN_AI_SYSTEM_ATTRIBUTE, + GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, + GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, +} from '../../../src/tracing/ai/gen-ai-attributes'; +import { WORKERS_AI_ORIGIN, WORKERS_AI_SYSTEM_NAME } from '../../../src/tracing/workers-ai/constants'; +import { + addRequestAttributes, + addResponseAttributes, + extractRequestAttributes, + getOperationName, +} from '../../../src/tracing/workers-ai/utils'; + +const MODEL = '@cf/meta/llama-3.1-8b-instruct'; + +function createMockSpan(): { span: Span; attributes: Record } { + const attributes: Record = {}; + const span = { + setAttribute: (key: string, value: unknown) => { + attributes[key] = value; + }, + setAttributes: (attrs: Record) => { + Object.assign(attributes, attrs); + }, + } as unknown as Span; + return { span, attributes }; +} + +describe('workers-ai utils', () => { + describe('getOperationName', () => { + it('returns "chat" for prompt inputs', () => { + expect(getOperationName({ prompt: 'Hello' })).toBe('chat'); + }); + + it('returns "chat" for messages inputs', () => { + expect(getOperationName({ messages: [{ role: 'user', content: 'Hi' }] })).toBe('chat'); + }); + + it('returns "embeddings" for text inputs', () => { + expect(getOperationName({ text: 'embed me' })).toBe('embeddings'); + }); + + it('prefers "chat" when both messages and text are present', () => { + expect(getOperationName({ messages: [{ role: 'user', content: 'Hi' }], text: 'embed me' })).toBe('chat'); + }); + + it('falls back to "chat" for null, undefined and empty inputs', () => { + expect(getOperationName(null)).toBe('chat'); + expect(getOperationName(undefined)).toBe('chat'); + expect(getOperationName({})).toBe('chat'); + }); + }); + + describe('extractRequestAttributes', () => { + it('sets exactly the base attributes for a minimal request', () => { + expect(extractRequestAttributes(MODEL, { prompt: 'Hello' }, 'chat')).toEqual({ + [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, + [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'chat', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, + [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: MODEL, + }); + }); + + it('maps all supported request parameters', () => { + expect( + extractRequestAttributes( + MODEL, + { + prompt: 'Hello', + temperature: 0.5, + max_tokens: 100, + top_p: 0.9, + top_k: 40, + frequency_penalty: 0.1, + presence_penalty: 0.2, + stream: true, + }, + 'chat', + ), + ).toEqual({ + [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, + [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'chat', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, + [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: MODEL, + [GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE]: 0.5, + [GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE]: 100, + [GEN_AI_REQUEST_TOP_P_ATTRIBUTE]: 0.9, + [GEN_AI_REQUEST_TOP_K_ATTRIBUTE]: 40, + [GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE]: 0.1, + [GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE]: 0.2, + [GEN_AI_REQUEST_STREAM_ATTRIBUTE]: true, + }); + }); + + it('does not set the stream attribute when stream is false', () => { + const attrs = extractRequestAttributes(MODEL, { prompt: 'Hello', stream: false }, 'chat'); + expect(attrs[GEN_AI_REQUEST_STREAM_ATTRIBUTE]).toBeUndefined(); + }); + + it('falls back to "unknown" model when model is not a string', () => { + const attrs = extractRequestAttributes(undefined, { prompt: 'Hello' }, 'chat'); + expect(attrs[GEN_AI_REQUEST_MODEL_ATTRIBUTE]).toBe('unknown'); + }); + }); + + describe('addRequestAttributes', () => { + it('records messages and extracts system instructions', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes( + span, + { + messages: [ + { role: 'system', content: 'You are helpful.' }, + { role: 'user', content: 'Hi' }, + ], + }, + 'chat', + false, + ); + + expect(attributes[GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE]).toBe( + JSON.stringify([{ type: 'text', content: 'You are helpful.' }]), + ); + expect(attributes[GEN_AI_INPUT_MESSAGES_ATTRIBUTE]).toBe(JSON.stringify([{ role: 'user', content: 'Hi' }])); + }); + + it('records the prompt string directly', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, { prompt: 'Hello world' }, 'chat', false); + + expect(attributes[GEN_AI_INPUT_MESSAGES_ATTRIBUTE]).toBe('Hello world'); + }); + + it('records embeddings input on a dedicated attribute', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, { text: ['embed a', 'embed b'] }, 'embeddings', false); + + expect(attributes).toEqual({ [GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: JSON.stringify(['embed a', 'embed b']) }); + }); + + it('records nothing for an empty messages array', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, { messages: [] }, 'chat', false); + + expect(attributes).toEqual({}); + }); + + it('records nothing for empty embeddings input', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, { text: '' }, 'embeddings', false); + + expect(attributes).toEqual({}); + }); + + it('records nothing when inputs are missing', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, undefined, 'chat', false); + + expect(attributes).toEqual({}); + }); + }); + + describe('addResponseAttributes', () => { + it('sets token usage and computes the total from input and output tokens', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { response: 'Paris', usage: { prompt_tokens: 12, completion_tokens: 7 } }, false); + + expect(attributes).toEqual({ + [GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 12, + [GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]: 7, + [GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 19, + }); + }); + + it('does not record response text when recordOutputs is false', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { response: 'Paris' }, false); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBeUndefined(); + }); + + it('records response text and tool calls when recordOutputs is true', () => { + const { span, attributes } = createMockSpan(); + const toolCalls = [{ name: 'lookup', arguments: { city: 'Paris' } }]; + + addResponseAttributes(span, { response: 'Paris', tool_calls: toolCalls }, true); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('Paris'); + expect(attributes[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]).toBe(JSON.stringify(toolCalls)); + }); + + it('serializes non-string response payloads as JSON', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { response: { translated_text: 'Bonjour' } }, true); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe(JSON.stringify({ translated_text: 'Bonjour' })); + }); + + it('ignores raw Response objects', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, new Response('raw'), true); + + expect(attributes).toEqual({}); + }); + + it('ignores non-object results', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, null, true); + + expect(attributes).toEqual({}); + }); + }); +}); From aecd6d6ad632bb43db4f804d35186c5d673928d6 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Mon, 13 Jul 2026 17:02:58 +0200 Subject: [PATCH 02/10] fixup! feat(core): Instrument workers-ai-provider --- packages/core/src/tracing/workers-ai/index.ts | 9 ------- packages/core/src/tracing/workers-ai/utils.ts | 26 +++++++------------ .../core/test/lib/tracing/workers-ai.test.ts | 10 ------- 3 files changed, 9 insertions(+), 36 deletions(-) diff --git a/packages/core/src/tracing/workers-ai/index.ts b/packages/core/src/tracing/workers-ai/index.ts index e38d4bbd53ce..9351578d4823 100644 --- a/packages/core/src/tracing/workers-ai/index.ts +++ b/packages/core/src/tracing/workers-ai/index.ts @@ -96,10 +96,6 @@ function instrumentRun( }; } -// Guards against double-wrapping when a binding instrumented via the automatic -// env instrumentation in `@sentry/cloudflare` is additionally wrapped manually. -const instrumentedClients = new WeakSet(); - /** * Instrument a Cloudflare Workers AI binding (`env.AI`) with Sentry tracing. * @@ -116,10 +112,6 @@ const instrumentedClients = new WeakSet(); * ``` */ export function instrumentWorkersAiClient(client: T, options?: WorkersAiOptions): T { - if (instrumentedClients.has(client)) { - return client; - } - const resolvedOptions = resolveAIRecordingOptions(options); const instrumented = new Proxy(client, { @@ -135,6 +127,5 @@ export function instrumentWorkersAiClient(client: T, options?: }, }) as T; - instrumentedClients.add(instrumented); return instrumented; } diff --git a/packages/core/src/tracing/workers-ai/utils.ts b/packages/core/src/tracing/workers-ai/utils.ts index 94611c90ce19..aee94eddd911 100644 --- a/packages/core/src/tracing/workers-ai/utils.ts +++ b/packages/core/src/tracing/workers-ai/utils.ts @@ -101,13 +101,7 @@ export function addRequestAttributes( if (operationName === 'embeddings') { const text = params.text; - if (text == null) { - return; - } - if (typeof text === 'string' && text.length === 0) { - return; - } - if (Array.isArray(text) && text.length === 0) { + if (text == null || (typeof text === 'string' && text.length === 0) || (Array.isArray(text) && text.length === 0)) { return; } @@ -116,10 +110,8 @@ export function addRequestAttributes( } const src = params.messages ?? params.prompt; - if (src == null) { - return; - } - if (Array.isArray(src) && src.length === 0) { + + if (src == null || (Array.isArray(src) && src.length === 0)) { return; } @@ -144,12 +136,12 @@ export function addRequestAttributes( * Record the response attributes (token usage, response text, tool calls) on the span. */ export function addResponseAttributes(span: Span, result: unknown, recordOutputs: boolean): void { - if (!result || typeof result !== 'object') { - return; - } - - // Raw `Response` objects (from `returnRawResponse`/`websocket`) cannot be introspected without consuming them. - if (typeof Response !== 'undefined' && result instanceof Response) { + if ( + !result || + typeof result !== 'object' || + // Raw `Response` objects (from `returnRawResponse`/`websocket`) cannot be introspected without consuming them. + (typeof Response !== 'undefined' && result instanceof Response) + ) { return; } diff --git a/packages/core/test/lib/tracing/workers-ai.test.ts b/packages/core/test/lib/tracing/workers-ai.test.ts index 9c15e7c0b5a5..d9c563b59962 100644 --- a/packages/core/test/lib/tracing/workers-ai.test.ts +++ b/packages/core/test/lib/tracing/workers-ai.test.ts @@ -2,16 +2,6 @@ import { describe, expect, it, vi } from 'vitest'; import { instrumentWorkersAiClient } from '../../../src/tracing/workers-ai'; describe('instrumentWorkersAiClient', () => { - it('returns an already-instrumented client as-is instead of double-wrapping', () => { - const client = { run: vi.fn().mockResolvedValue({ response: 'ok' }) }; - - const instrumented = instrumentWorkersAiClient(client); - const instrumentedTwice = instrumentWorkersAiClient(instrumented); - - expect(instrumented).not.toBe(client); - expect(instrumentedTwice).toBe(instrumented); - }); - it('passes through non-run methods bound to the original client', () => { class FakeAi { #internal = 'secret'; From 9c7da1139ab23b93415099391a9996d4b53c9871 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 17 Jul 2026 11:14:35 +0200 Subject: [PATCH 03/10] ref: Use conventions directly --- .../core/src/tracing/workers-ai/constants.ts | 4 +- packages/core/src/tracing/workers-ai/utils.ts | 61 ++++++------- .../test/lib/utils/workers-ai-utils.test.ts | 88 +++++++++---------- 3 files changed, 76 insertions(+), 77 deletions(-) diff --git a/packages/core/src/tracing/workers-ai/constants.ts b/packages/core/src/tracing/workers-ai/constants.ts index 89847abbae1c..8f532a0a0372 100644 --- a/packages/core/src/tracing/workers-ai/constants.ts +++ b/packages/core/src/tracing/workers-ai/constants.ts @@ -1,8 +1,8 @@ /** - * The provider value for the `gen_ai.system` attribute. + * The provider value for the `gen_ai.provider.name` attribute. * @see https://developers.cloudflare.com/workers-ai/ */ -export const WORKERS_AI_SYSTEM_NAME = 'cloudflare.workers_ai'; +export const WORKERS_AI_PROVIDER_NAME = 'cloudflare.workers_ai'; /** * The Sentry origin for spans created by the Workers AI instrumentation. diff --git a/packages/core/src/tracing/workers-ai/utils.ts b/packages/core/src/tracing/workers-ai/utils.ts index aee94eddd911..f9abbe1c3e5d 100644 --- a/packages/core/src/tracing/workers-ai/utils.ts +++ b/packages/core/src/tracing/workers-ai/utils.ts @@ -1,25 +1,26 @@ +import { + GEN_AI_EMBEDDINGS_INPUT, + GEN_AI_INPUT_MESSAGES, + GEN_AI_OPERATION_NAME, + GEN_AI_REQUEST_FREQUENCY_PENALTY, + GEN_AI_REQUEST_MAX_TOKENS, + GEN_AI_REQUEST_MODEL, + GEN_AI_REQUEST_PRESENCE_PENALTY, + GEN_AI_REQUEST_TEMPERATURE, + GEN_AI_REQUEST_TOP_K, + GEN_AI_REQUEST_TOP_P, + GEN_AI_PROVIDER_NAME, + GEN_AI_OUTPUT_MESSAGES, + GEN_AI_SYSTEM_INSTRUCTIONS, +} from '@sentry/conventions/attributes'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes'; import type { Span, SpanAttributeValue } from '../../types/span'; import { - GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, - GEN_AI_INPUT_MESSAGES_ATTRIBUTE, GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, - GEN_AI_OPERATION_NAME_ATTRIBUTE, - GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE, - GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE, - GEN_AI_REQUEST_MODEL_ATTRIBUTE, - GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE, GEN_AI_REQUEST_STREAM_ATTRIBUTE, - GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE, - GEN_AI_REQUEST_TOP_K_ATTRIBUTE, - GEN_AI_REQUEST_TOP_P_ATTRIBUTE, - GEN_AI_RESPONSE_TEXT_ATTRIBUTE, - GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, - GEN_AI_SYSTEM_ATTRIBUTE, - GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, } from '../ai/gen-ai-attributes'; import { extractSystemInstructions, getJsonString, getTruncatedJsonString, setTokenUsageAttributes } from '../ai/utils'; -import { WORKERS_AI_ORIGIN, WORKERS_AI_SYSTEM_NAME } from './constants'; +import { WORKERS_AI_ORIGIN, WORKERS_AI_PROVIDER_NAME } from './constants'; import type { WorkersAiInput, WorkersAiOutput } from './types'; /** @@ -47,32 +48,32 @@ export function extractRequestAttributes( operationName: string, ): Record { const attributes: Record = { - [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, - [GEN_AI_OPERATION_NAME_ATTRIBUTE]: operationName, + [GEN_AI_PROVIDER_NAME]: WORKERS_AI_PROVIDER_NAME, + [GEN_AI_OPERATION_NAME]: operationName, [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, - [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: typeof model === 'string' ? model : 'unknown', + [GEN_AI_REQUEST_MODEL]: typeof model === 'string' ? model : 'unknown', }; if (inputs && typeof inputs === 'object') { const params = inputs as WorkersAiInput; if (typeof params.temperature === 'number') { - attributes[GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE] = params.temperature; + attributes[GEN_AI_REQUEST_TEMPERATURE] = params.temperature; } if (typeof params.max_tokens === 'number') { - attributes[GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE] = params.max_tokens; + attributes[GEN_AI_REQUEST_MAX_TOKENS] = params.max_tokens; } if (typeof params.top_p === 'number') { - attributes[GEN_AI_REQUEST_TOP_P_ATTRIBUTE] = params.top_p; + attributes[GEN_AI_REQUEST_TOP_P] = params.top_p; } if (typeof params.top_k === 'number') { - attributes[GEN_AI_REQUEST_TOP_K_ATTRIBUTE] = params.top_k; + attributes[GEN_AI_REQUEST_TOP_K] = params.top_k; } if (typeof params.frequency_penalty === 'number') { - attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE] = params.frequency_penalty; + attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY] = params.frequency_penalty; } if (typeof params.presence_penalty === 'number') { - attributes[GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE] = params.presence_penalty; + attributes[GEN_AI_REQUEST_PRESENCE_PENALTY] = params.presence_penalty; } if (params.stream === true) { attributes[GEN_AI_REQUEST_STREAM_ATTRIBUTE] = true; @@ -105,7 +106,7 @@ export function addRequestAttributes( return; } - span.setAttribute(GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, typeof text === 'string' ? text : JSON.stringify(text)); + span.setAttribute(GEN_AI_EMBEDDINGS_INPUT, typeof text === 'string' ? text : JSON.stringify(text)); return; } @@ -118,11 +119,11 @@ export function addRequestAttributes( const { systemInstructions, filteredMessages } = extractSystemInstructions(src); if (systemInstructions) { - span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, systemInstructions); + span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, systemInstructions); } span.setAttribute( - GEN_AI_INPUT_MESSAGES_ATTRIBUTE, + GEN_AI_INPUT_MESSAGES, enableTruncation ? getTruncatedJsonString(filteredMessages) : getJsonString(filteredMessages), ); @@ -153,13 +154,13 @@ export function addResponseAttributes(span: Span, result: unknown, recordOutputs if (recordOutputs) { if (typeof response.response === 'string') { - span.setAttribute(GEN_AI_RESPONSE_TEXT_ATTRIBUTE, response.response); + span.setAttribute(GEN_AI_OUTPUT_MESSAGES, response.response); } else if (response.response != null) { - span.setAttribute(GEN_AI_RESPONSE_TEXT_ATTRIBUTE, JSON.stringify(response.response)); + span.setAttribute(GEN_AI_OUTPUT_MESSAGES, JSON.stringify(response.response)); } if (Array.isArray(response.tool_calls) && response.tool_calls.length > 0) { - span.setAttribute(GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, JSON.stringify(response.tool_calls)); + span.setAttribute(GEN_AI_OUTPUT_MESSAGES, JSON.stringify(response.tool_calls)); } } } diff --git a/packages/core/test/lib/utils/workers-ai-utils.test.ts b/packages/core/test/lib/utils/workers-ai-utils.test.ts index eed6cf170ed3..58a6c0b7224b 100644 --- a/packages/core/test/lib/utils/workers-ai-utils.test.ts +++ b/packages/core/test/lib/utils/workers-ai-utils.test.ts @@ -2,26 +2,25 @@ import { describe, expect, it } from 'vitest'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../../src'; import type { Span } from '../../../src'; import { - GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, - GEN_AI_INPUT_MESSAGES_ATTRIBUTE, - GEN_AI_OPERATION_NAME_ATTRIBUTE, - GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE, - GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE, - GEN_AI_REQUEST_MODEL_ATTRIBUTE, - GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE, - GEN_AI_REQUEST_STREAM_ATTRIBUTE, - GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE, - GEN_AI_REQUEST_TOP_K_ATTRIBUTE, - GEN_AI_REQUEST_TOP_P_ATTRIBUTE, - GEN_AI_RESPONSE_TEXT_ATTRIBUTE, - GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, - GEN_AI_SYSTEM_ATTRIBUTE, - GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, - GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, - GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, - GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, -} from '../../../src/tracing/ai/gen-ai-attributes'; -import { WORKERS_AI_ORIGIN, WORKERS_AI_SYSTEM_NAME } from '../../../src/tracing/workers-ai/constants'; + GEN_AI_PROVIDER_NAME, + GEN_AI_EMBEDDINGS_INPUT, + GEN_AI_INPUT_MESSAGES, + GEN_AI_OPERATION_NAME, + GEN_AI_REQUEST_FREQUENCY_PENALTY, + GEN_AI_REQUEST_MAX_TOKENS, + GEN_AI_REQUEST_MODEL, + GEN_AI_REQUEST_PRESENCE_PENALTY, + GEN_AI_REQUEST_TEMPERATURE, + GEN_AI_REQUEST_TOP_K, + GEN_AI_REQUEST_TOP_P, + GEN_AI_OUTPUT_MESSAGES, + GEN_AI_SYSTEM_INSTRUCTIONS, + GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_USAGE_TOTAL_TOKENS, +} from '@sentry/conventions/attributes'; +import { GEN_AI_REQUEST_STREAM_ATTRIBUTE } from '../../../src/tracing/ai/gen-ai-attributes'; +import { WORKERS_AI_ORIGIN, WORKERS_AI_PROVIDER_NAME } from '../../../src/tracing/workers-ai/constants'; import { addRequestAttributes, addResponseAttributes, @@ -72,10 +71,10 @@ describe('workers-ai utils', () => { describe('extractRequestAttributes', () => { it('sets exactly the base attributes for a minimal request', () => { expect(extractRequestAttributes(MODEL, { prompt: 'Hello' }, 'chat')).toEqual({ - [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, - [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'chat', + [GEN_AI_PROVIDER_NAME]: WORKERS_AI_PROVIDER_NAME, + [GEN_AI_OPERATION_NAME]: 'chat', [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, - [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: MODEL, + [GEN_AI_REQUEST_MODEL]: MODEL, }); }); @@ -96,16 +95,16 @@ describe('workers-ai utils', () => { 'chat', ), ).toEqual({ - [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, - [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'chat', + [GEN_AI_PROVIDER_NAME]: WORKERS_AI_PROVIDER_NAME, + [GEN_AI_OPERATION_NAME]: 'chat', [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, - [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: MODEL, - [GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE]: 0.5, - [GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE]: 100, - [GEN_AI_REQUEST_TOP_P_ATTRIBUTE]: 0.9, - [GEN_AI_REQUEST_TOP_K_ATTRIBUTE]: 40, - [GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE]: 0.1, - [GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE]: 0.2, + [GEN_AI_REQUEST_MODEL]: MODEL, + [GEN_AI_REQUEST_TEMPERATURE]: 0.5, + [GEN_AI_REQUEST_MAX_TOKENS]: 100, + [GEN_AI_REQUEST_TOP_P]: 0.9, + [GEN_AI_REQUEST_TOP_K]: 40, + [GEN_AI_REQUEST_FREQUENCY_PENALTY]: 0.1, + [GEN_AI_REQUEST_PRESENCE_PENALTY]: 0.2, [GEN_AI_REQUEST_STREAM_ATTRIBUTE]: true, }); }); @@ -117,7 +116,7 @@ describe('workers-ai utils', () => { it('falls back to "unknown" model when model is not a string', () => { const attrs = extractRequestAttributes(undefined, { prompt: 'Hello' }, 'chat'); - expect(attrs[GEN_AI_REQUEST_MODEL_ATTRIBUTE]).toBe('unknown'); + expect(attrs[GEN_AI_REQUEST_MODEL]).toBe('unknown'); }); }); @@ -137,10 +136,10 @@ describe('workers-ai utils', () => { false, ); - expect(attributes[GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE]).toBe( + expect(attributes[GEN_AI_SYSTEM_INSTRUCTIONS]).toBe( JSON.stringify([{ type: 'text', content: 'You are helpful.' }]), ); - expect(attributes[GEN_AI_INPUT_MESSAGES_ATTRIBUTE]).toBe(JSON.stringify([{ role: 'user', content: 'Hi' }])); + expect(attributes[GEN_AI_INPUT_MESSAGES]).toBe(JSON.stringify([{ role: 'user', content: 'Hi' }])); }); it('records the prompt string directly', () => { @@ -148,7 +147,7 @@ describe('workers-ai utils', () => { addRequestAttributes(span, { prompt: 'Hello world' }, 'chat', false); - expect(attributes[GEN_AI_INPUT_MESSAGES_ATTRIBUTE]).toBe('Hello world'); + expect(attributes[GEN_AI_INPUT_MESSAGES]).toBe('Hello world'); }); it('records embeddings input on a dedicated attribute', () => { @@ -156,7 +155,7 @@ describe('workers-ai utils', () => { addRequestAttributes(span, { text: ['embed a', 'embed b'] }, 'embeddings', false); - expect(attributes).toEqual({ [GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: JSON.stringify(['embed a', 'embed b']) }); + expect(attributes).toEqual({ [GEN_AI_EMBEDDINGS_INPUT]: JSON.stringify(['embed a', 'embed b']) }); }); it('records nothing for an empty messages array', () => { @@ -191,9 +190,9 @@ describe('workers-ai utils', () => { addResponseAttributes(span, { response: 'Paris', usage: { prompt_tokens: 12, completion_tokens: 7 } }, false); expect(attributes).toEqual({ - [GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 12, - [GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]: 7, - [GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 19, + [GEN_AI_USAGE_INPUT_TOKENS]: 12, + [GEN_AI_USAGE_OUTPUT_TOKENS]: 7, + [GEN_AI_USAGE_TOTAL_TOKENS]: 19, }); }); @@ -202,17 +201,16 @@ describe('workers-ai utils', () => { addResponseAttributes(span, { response: 'Paris' }, false); - expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBeUndefined(); + expect(attributes[GEN_AI_OUTPUT_MESSAGES]).toBeUndefined(); }); it('records response text and tool calls when recordOutputs is true', () => { const { span, attributes } = createMockSpan(); const toolCalls = [{ name: 'lookup', arguments: { city: 'Paris' } }]; - addResponseAttributes(span, { response: 'Paris', tool_calls: toolCalls }, true); + addResponseAttributes(span, { tool_calls: toolCalls }, true); - expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('Paris'); - expect(attributes[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]).toBe(JSON.stringify(toolCalls)); + expect(attributes[GEN_AI_OUTPUT_MESSAGES]).toBe(JSON.stringify(toolCalls)); }); it('serializes non-string response payloads as JSON', () => { @@ -220,7 +218,7 @@ describe('workers-ai utils', () => { addResponseAttributes(span, { response: { translated_text: 'Bonjour' } }, true); - expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe(JSON.stringify({ translated_text: 'Bonjour' })); + expect(attributes[GEN_AI_OUTPUT_MESSAGES]).toBe(JSON.stringify({ translated_text: 'Bonjour' })); }); it('ignores raw Response objects', () => { From 4b72e3de664bcc4be7bb7ceb083197959684e2aa Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 17 Jul 2026 11:30:04 +0200 Subject: [PATCH 04/10] fix: Duck-type isReadableStream --- packages/core/src/tracing/workers-ai/index.ts | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/packages/core/src/tracing/workers-ai/index.ts b/packages/core/src/tracing/workers-ai/index.ts index 9351578d4823..365acdbb9824 100644 --- a/packages/core/src/tracing/workers-ai/index.ts +++ b/packages/core/src/tracing/workers-ai/index.ts @@ -2,14 +2,21 @@ import { captureException } from '../../exports'; import { SPAN_STATUS_ERROR } from '../../tracing'; import { startSpan, startSpanManual } from '../../tracing/trace'; import type { Span } from '../../types/span'; +import { isObjectLike } from '../../utils/is'; import { resolveAIRecordingOptions, shouldEnableTruncation } from '../ai/utils'; import { WORKERS_AI_ORIGIN } from './constants'; import { instrumentWorkersAiStream } from './streaming'; import type { WorkersAiOptions } from './types'; import { addRequestAttributes, addResponseAttributes, extractRequestAttributes, getOperationName } from './utils'; -function isReadableStream(value: unknown): value is ReadableStream { - return typeof ReadableStream !== 'undefined' && value instanceof ReadableStream; +// Copied from /server-utils/src/vercel-ai/util.ts +// TODO(v11): Reuse this function once this gets moved to @sentry/server-utils +function isReadableStream(value: unknown): value is ReadableStream { + return ( + isObjectLike(value) && + typeof (value as { pipeThrough?: unknown }).pipeThrough === 'function' && + typeof (value as { getReader?: unknown }).getReader === 'function' + ); } /** From 08a7b5a3b34de3468cd36ba2a837501f92a834c0 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 17 Jul 2026 11:44:46 +0200 Subject: [PATCH 05/10] =?UTF-8?q?fix:=20Defensive=20calling=20=F0=9F=93=9E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- packages/core/src/tracing/workers-ai/index.ts | 47 +++++++++++-------- 1 file changed, 27 insertions(+), 20 deletions(-) diff --git a/packages/core/src/tracing/workers-ai/index.ts b/packages/core/src/tracing/workers-ai/index.ts index 365acdbb9824..5b1ab6734ff3 100644 --- a/packages/core/src/tracing/workers-ai/index.ts +++ b/packages/core/src/tracing/workers-ai/index.ts @@ -49,32 +49,39 @@ function instrumentRun( if (isStreamRequested && !returnsRawResponse) { return startSpanManual(spanConfig, (span: Span) => { - const originalResult = originalRun.apply(context, args) as Promise; + // `startSpanManual` does not auto-end the span, so we must end it on every exit path, + // including a synchronous throw from `run`. + const handleError = (error: unknown): never => { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + captureException(error, { + mechanism: { handled: false, type: `${WORKERS_AI_ORIGIN}.stream`, data: { function: 'run' } }, + }); + span.end(); + throw error; + }; + + let originalResult: Promise; + + try { + originalResult = originalRun.apply(context, args) as Promise; + } catch (error) { + return handleError(error); + } if (options.recordInputs) { addRequestAttributes(span, inputs, operationName, shouldEnableTruncation(options.enableTruncation)); } - return originalResult.then( - result => { - if (isReadableStream(result)) { - return instrumentWorkersAiStream(result, span, options.recordOutputs); - } + return originalResult.then(result => { + if (isReadableStream(result)) { + return instrumentWorkersAiStream(result, span, options.recordOutputs); + } - // The model did not actually return a stream — finalize the span eagerly. - addResponseAttributes(span, result, options.recordOutputs); - span.end(); - return result; - }, - error => { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - captureException(error, { - mechanism: { handled: false, type: `${WORKERS_AI_ORIGIN}.stream`, data: { function: 'run' } }, - }); - span.end(); - throw error; - }, - ); + // The model did not actually return a stream — finalize the span eagerly. + addResponseAttributes(span, result, options.recordOutputs); + span.end(); + return result; + }, handleError); }); } From 743733b0e7971beb6a25a02887281a44d1224062 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 17 Jul 2026 11:59:44 +0200 Subject: [PATCH 06/10] fixup! fix: Duck-type isReadableStream --- packages/core/src/tracing/workers-ai/index.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/core/src/tracing/workers-ai/index.ts b/packages/core/src/tracing/workers-ai/index.ts index 5b1ab6734ff3..51a3141ca815 100644 --- a/packages/core/src/tracing/workers-ai/index.ts +++ b/packages/core/src/tracing/workers-ai/index.ts @@ -9,9 +9,10 @@ import { instrumentWorkersAiStream } from './streaming'; import type { WorkersAiOptions } from './types'; import { addRequestAttributes, addResponseAttributes, extractRequestAttributes, getOperationName } from './utils'; -// Copied from /server-utils/src/vercel-ai/util.ts +// Adapted from /server-utils/src/vercel-ai/util.ts // TODO(v11): Reuse this function once this gets moved to @sentry/server-utils -function isReadableStream(value: unknown): value is ReadableStream { +// Workers AI streaming responses are SSE byte streams, so we narrow to `Uint8Array`. +function isReadableStream(value: unknown): value is ReadableStream { return ( isObjectLike(value) && typeof (value as { pipeThrough?: unknown }).pipeThrough === 'function' && From 9388819fda2e098310fa3f083424f7e87e681a1b Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 17 Jul 2026 12:10:00 +0200 Subject: [PATCH 07/10] fix: Rebase removal --- packages/core/src/tracing/workers-ai/utils.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/core/src/tracing/workers-ai/utils.ts b/packages/core/src/tracing/workers-ai/utils.ts index f9abbe1c3e5d..d640e0fa74f4 100644 --- a/packages/core/src/tracing/workers-ai/utils.ts +++ b/packages/core/src/tracing/workers-ai/utils.ts @@ -19,9 +19,10 @@ import { GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, GEN_AI_REQUEST_STREAM_ATTRIBUTE, } from '../ai/gen-ai-attributes'; -import { extractSystemInstructions, getJsonString, getTruncatedJsonString, setTokenUsageAttributes } from '../ai/utils'; +import { extractSystemInstructions, getTruncatedJsonString, setTokenUsageAttributes } from '../ai/utils'; import { WORKERS_AI_ORIGIN, WORKERS_AI_PROVIDER_NAME } from './constants'; import type { WorkersAiInput, WorkersAiOutput } from './types'; +import { stringify } from '../../shared-exports'; /** * Determine the gen_ai operation name from the inputs passed to `AI.run`. @@ -124,7 +125,7 @@ export function addRequestAttributes( span.setAttribute( GEN_AI_INPUT_MESSAGES, - enableTruncation ? getTruncatedJsonString(filteredMessages) : getJsonString(filteredMessages), + enableTruncation ? getTruncatedJsonString(filteredMessages) : stringify(filteredMessages), ); span.setAttribute( From a8526499d47917a5d317b6a5e9cc296f2bdafdf2 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 17 Jul 2026 12:28:24 +0200 Subject: [PATCH 08/10] fixup! fix: Rebase removal --- packages/core/src/tracing/workers-ai/utils.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/core/src/tracing/workers-ai/utils.ts b/packages/core/src/tracing/workers-ai/utils.ts index d640e0fa74f4..f9fa147788ec 100644 --- a/packages/core/src/tracing/workers-ai/utils.ts +++ b/packages/core/src/tracing/workers-ai/utils.ts @@ -20,9 +20,9 @@ import { GEN_AI_REQUEST_STREAM_ATTRIBUTE, } from '../ai/gen-ai-attributes'; import { extractSystemInstructions, getTruncatedJsonString, setTokenUsageAttributes } from '../ai/utils'; +import { stringify } from '../../utils/string'; import { WORKERS_AI_ORIGIN, WORKERS_AI_PROVIDER_NAME } from './constants'; import type { WorkersAiInput, WorkersAiOutput } from './types'; -import { stringify } from '../../shared-exports'; /** * Determine the gen_ai operation name from the inputs passed to `AI.run`. From 5ff111c59f9b822ee9178b84b444c037d136cd64 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 17 Jul 2026 17:09:03 +0200 Subject: [PATCH 09/10] fixup! feat(core): Instrument workers-ai-provider --- packages/core/src/tracing/workers-ai/index.ts | 25 +++++-------------- .../core/src/tracing/workers-ai/streaming.ts | 5 ---- 2 files changed, 6 insertions(+), 24 deletions(-) diff --git a/packages/core/src/tracing/workers-ai/index.ts b/packages/core/src/tracing/workers-ai/index.ts index 51a3141ca815..c6be4824305a 100644 --- a/packages/core/src/tracing/workers-ai/index.ts +++ b/packages/core/src/tracing/workers-ai/index.ts @@ -1,10 +1,8 @@ -import { captureException } from '../../exports'; import { SPAN_STATUS_ERROR } from '../../tracing'; import { startSpan, startSpanManual } from '../../tracing/trace'; import type { Span } from '../../types/span'; import { isObjectLike } from '../../utils/is'; import { resolveAIRecordingOptions, shouldEnableTruncation } from '../ai/utils'; -import { WORKERS_AI_ORIGIN } from './constants'; import { instrumentWorkersAiStream } from './streaming'; import type { WorkersAiOptions } from './types'; import { addRequestAttributes, addResponseAttributes, extractRequestAttributes, getOperationName } from './utils'; @@ -54,9 +52,6 @@ function instrumentRun( // including a synchronous throw from `run`. const handleError = (error: unknown): never => { span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - captureException(error, { - mechanism: { handled: false, type: `${WORKERS_AI_ORIGIN}.stream`, data: { function: 'run' } }, - }); span.end(); throw error; }; @@ -93,20 +88,12 @@ function instrumentRun( addRequestAttributes(span, inputs, operationName, shouldEnableTruncation(options.enableTruncation)); } - return originalResult.then( - result => { - if (!returnsRawResponse) { - addResponseAttributes(span, result, options.recordOutputs); - } - return result; - }, - error => { - captureException(error, { - mechanism: { handled: false, type: WORKERS_AI_ORIGIN, data: { function: 'run' } }, - }); - throw error; - }, - ); + return originalResult.then(result => { + if (!returnsRawResponse) { + addResponseAttributes(span, result, options.recordOutputs); + } + return result; + }); }); }; } diff --git a/packages/core/src/tracing/workers-ai/streaming.ts b/packages/core/src/tracing/workers-ai/streaming.ts index c2a1055a8be0..4551a27ad407 100644 --- a/packages/core/src/tracing/workers-ai/streaming.ts +++ b/packages/core/src/tracing/workers-ai/streaming.ts @@ -1,8 +1,6 @@ -import { captureException } from '../../exports'; import { SPAN_STATUS_ERROR } from '../../tracing'; import type { Span } from '../../types/span'; import { endStreamSpan, type StreamResponseState } from '../ai/utils'; -import { WORKERS_AI_ORIGIN } from './constants'; import type { WorkersAiUsage } from './types'; interface WorkersAiStreamChunk { @@ -116,9 +114,6 @@ export function instrumentWorkersAiStream( controller.enqueue(value); } catch (error) { span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - captureException(error, { - mechanism: { handled: false, type: `${WORKERS_AI_ORIGIN}.stream` }, - }); finish(); controller.error(error); } From 8608d9323c351383d6a53cc3a9e260205e577cde Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Sat, 18 Jul 2026 14:14:25 +0200 Subject: [PATCH 10/10] feat(core): Emit gen_ai.output.messages from Workers AI instrumentation The Sentry product reads model output from `gen_ai.output.messages`, treating `gen_ai.response.text` and `gen_ai.response.tool_calls` as deprecated. Relay migrates `response.text` into `output.messages` at ingestion, but the tool-calls half of that migration is lossy, so tool-call turns rendered an empty Output tab. Build `gen_ai.output.messages` directly in both the streaming and non-streaming paths (mirroring the Vercel AI integration), normalizing the OpenAI-compatible (`function.{name,arguments}`) and native (top-level `name`/`arguments`) tool-call shapes. The deprecated attributes are still written for backward compatibility. The streaming parser is also extended to read the OpenAI-compatible SSE shape (`choices[].delta.content` / `choices[].delta.tool_calls`) that models routed through the OpenAI-compatible endpoint emit, which previously dropped both text and tool calls while still capturing usage. Co-Authored-By: Claude Opus 4.8 --- .../core/src/tracing/workers-ai/streaming.ts | 109 ++++++++++++++++- packages/core/src/tracing/workers-ai/utils.ts | 68 ++++++++++- .../lib/tracing/workers-ai-streaming.test.ts | 110 ++++++++++++++++++ .../test/lib/utils/workers-ai-utils.test.ts | 85 ++++++++++++-- 4 files changed, 357 insertions(+), 15 deletions(-) diff --git a/packages/core/src/tracing/workers-ai/streaming.ts b/packages/core/src/tracing/workers-ai/streaming.ts index 4551a27ad407..5a36d5eff2dc 100644 --- a/packages/core/src/tracing/workers-ai/streaming.ts +++ b/packages/core/src/tracing/workers-ai/streaming.ts @@ -2,17 +2,87 @@ import { SPAN_STATUS_ERROR } from '../../tracing'; import type { Span } from '../../types/span'; import { endStreamSpan, type StreamResponseState } from '../ai/utils'; import type { WorkersAiUsage } from './types'; +import { setOutputMessagesAttribute } from './utils'; + +interface WorkersAiStreamingToolCall { + index?: number; + id?: string; + type?: string; + function?: { name?: string; arguments?: string }; + // Some Workers AI models stream tool calls with the name/arguments at the top + // level of the tool-call object instead of nested under `function`. + name?: string; + arguments?: string; +} interface WorkersAiStreamChunk { + // Native Workers AI streaming shape (`env.AI.run` with `stream: true`). response?: unknown; - usage?: WorkersAiUsage; tool_calls?: unknown[]; + // OpenAI-compatible streaming shape emitted for models routed through the + // OpenAI-compatible endpoint (e.g. via `workers-ai-provider`). + choices?: Array<{ + delta?: { content?: unknown; tool_calls?: WorkersAiStreamingToolCall[] }; + finish_reason?: unknown; + }>; + usage?: WorkersAiUsage & { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number }; +} + +/** + * Accumulate a fragmented OpenAI-compatible tool call (delivered across multiple + * `choices[].delta.tool_calls` chunks) into the index-keyed accumulator. + */ +function accumulateStreamingToolCalls( + toolCalls: WorkersAiStreamingToolCall[], + accumulator: Record, +): void { + for (const toolCall of toolCalls) { + // Normalize both shapes: name/arguments nested under `function`, or at the top level. + const name = toolCall.function?.name ?? toolCall.name; + const args = toolCall.function?.arguments ?? toolCall.arguments; + + // A tool call must carry at least a name or argument fragment to be meaningful. + if (name == null && args == null) { + continue; + } + + const index = toolCall.index ?? 0; + const existing = accumulator[index]; + + if (!existing) { + accumulator[index] = { + index, + id: toolCall.id, + type: toolCall.type, + function: { + name, + arguments: args ?? '', + }, + }; + } else if (existing.function) { + if (name && !existing.function.name) { + existing.function.name = name; + } + if (args) { + existing.function.arguments = `${existing.function.arguments ?? ''}${args}`; + } + } + } } /** * Parse a single SSE line (`data: {...}`) and accumulate its data into the streaming state. + * + * Handles both the native Workers AI shape (top-level `response`/`tool_calls`) and the + * OpenAI-compatible shape (`choices[].delta.content`/`choices[].delta.tool_calls`), because + * the same `run()` call transparently yields either format depending on the model. */ -function processLine(line: string, state: StreamResponseState, recordOutputs: boolean): void { +function processLine( + line: string, + state: StreamResponseState, + recordOutputs: boolean, + toolCallAccumulator: Record, +): void { const trimmed = line.trim(); if (!trimmed.startsWith('data:')) { return; @@ -49,6 +119,20 @@ function processLine(line: string, state: StreamResponseState, recordOutputs: bo if (recordOutputs && Array.isArray(parsed.tool_calls) && parsed.tool_calls.length > 0) { state.toolCalls.push(...parsed.tool_calls); } + + if (Array.isArray(parsed.choices)) { + for (const choice of parsed.choices) { + if (recordOutputs && typeof choice.delta?.content === 'string' && choice.delta.content) { + state.responseTexts.push(choice.delta.content); + } + if (recordOutputs && Array.isArray(choice.delta?.tool_calls)) { + accumulateStreamingToolCalls(choice.delta.tool_calls, toolCallAccumulator); + } + if (typeof choice.finish_reason === 'string') { + state.finishReasons.push(choice.finish_reason); + } + } + } } /** @@ -76,6 +160,10 @@ export function instrumentWorkersAiStream( totalTokens: undefined, }; + // OpenAI-compatible tool calls arrive fragmented across chunks and are keyed by index; + // accumulate them here and flatten into `state.toolCalls` once the stream ends. + const toolCallAccumulator: Record = {}; + let buffer = ''; let spanEnded = false; @@ -84,6 +172,21 @@ export function instrumentWorkersAiStream( return; } spanEnded = true; + + if (recordOutputs) { + const accumulatedToolCalls = Object.values(toolCallAccumulator); + if (accumulatedToolCalls.length > 0) { + state.toolCalls.push(...accumulatedToolCalls); + } + + // Set the authoritative `gen_ai.output.messages` alongside the deprecated response + // attributes `endStreamSpan` writes, so tool calls survive Relay's lossy migration. + setOutputMessagesAttribute(span, { + responseText: state.responseTexts.join(''), + toolCalls: state.toolCalls, + }); + } + endStreamSpan(span, state, recordOutputs); }; @@ -92,7 +195,7 @@ export function instrumentWorkersAiStream( // Keep the last (potentially incomplete) line in the buffer unless the stream is done. buffer = isDone ? '' : (lines.pop() ?? ''); for (const line of lines) { - processLine(line, state, recordOutputs); + processLine(line, state, recordOutputs, toolCallAccumulator); } }; diff --git a/packages/core/src/tracing/workers-ai/utils.ts b/packages/core/src/tracing/workers-ai/utils.ts index f9fa147788ec..23eac0eaa6fc 100644 --- a/packages/core/src/tracing/workers-ai/utils.ts +++ b/packages/core/src/tracing/workers-ai/utils.ts @@ -10,14 +10,16 @@ import { GEN_AI_REQUEST_TOP_K, GEN_AI_REQUEST_TOP_P, GEN_AI_PROVIDER_NAME, - GEN_AI_OUTPUT_MESSAGES, GEN_AI_SYSTEM_INSTRUCTIONS, } from '@sentry/conventions/attributes'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes'; import type { Span, SpanAttributeValue } from '../../types/span'; import { GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, + GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE, GEN_AI_REQUEST_STREAM_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, } from '../ai/gen-ai-attributes'; import { extractSystemInstructions, getTruncatedJsonString, setTokenUsageAttributes } from '../ai/utils'; import { stringify } from '../../utils/string'; @@ -134,6 +136,55 @@ export function addRequestAttributes( ); } +/** + * Build the `gen_ai.output.messages` value (a single assistant message with text and/or + * tool-call parts) from the response text and tool calls. + * + * We set this in addition to the deprecated `gen_ai.response.text` / `gen_ai.response.tool_calls` + * attributes because Sentry's product reads the model output from `gen_ai.output.messages` first. + * Relay migrates `gen_ai.response.text` into `gen_ai.output.messages`, but the tool-calls half of + * that migration is lossy — so tool-call turns would otherwise render an empty Output. Emitting the + * normalized message here (mirroring the Vercel AI integration) keeps tool calls visible. + */ +export function setOutputMessagesAttribute( + span: Span, + { responseText, toolCalls }: { responseText?: string; toolCalls?: unknown[] }, +): void { + const parts: Array> = []; + + if (typeof responseText === 'string' && responseText.length > 0) { + parts.push({ type: 'text', content: responseText }); + } + + if (Array.isArray(toolCalls)) { + for (const toolCall of toolCalls) { + if (!toolCall || typeof toolCall !== 'object') { + continue; + } + const call = toolCall as { + id?: unknown; + function?: { name?: unknown; arguments?: unknown }; + name?: unknown; + arguments?: unknown; + }; + // Normalize both the OpenAI-compatible shape (name/arguments nested under `function`) + // and the native Workers AI shape (name/arguments at the top level). + const name = call.function?.name ?? call.name; + const args = call.function?.arguments ?? call.arguments; + parts.push({ + type: 'tool_call', + id: call.id, + name, + arguments: typeof args === 'string' ? args : JSON.stringify(args ?? {}), + }); + } + } + + if (parts.length > 0) { + span.setAttribute(GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE, JSON.stringify([{ role: 'assistant', parts }])); + } +} + /** * Record the response attributes (token usage, response text, tool calls) on the span. */ @@ -154,14 +205,21 @@ export function addResponseAttributes(span: Span, result: unknown, recordOutputs } if (recordOutputs) { + let responseText: string | undefined; if (typeof response.response === 'string') { - span.setAttribute(GEN_AI_OUTPUT_MESSAGES, response.response); + responseText = response.response; + span.setAttribute(GEN_AI_RESPONSE_TEXT_ATTRIBUTE, response.response); } else if (response.response != null) { - span.setAttribute(GEN_AI_OUTPUT_MESSAGES, JSON.stringify(response.response)); + responseText = JSON.stringify(response.response); + span.setAttribute(GEN_AI_RESPONSE_TEXT_ATTRIBUTE, responseText); } - if (Array.isArray(response.tool_calls) && response.tool_calls.length > 0) { - span.setAttribute(GEN_AI_OUTPUT_MESSAGES, JSON.stringify(response.tool_calls)); + const toolCalls = + Array.isArray(response.tool_calls) && response.tool_calls.length > 0 ? response.tool_calls : undefined; + if (toolCalls) { + span.setAttribute(GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, JSON.stringify(toolCalls)); } + + setOutputMessagesAttribute(span, { responseText, toolCalls }); } } diff --git a/packages/core/test/lib/tracing/workers-ai-streaming.test.ts b/packages/core/test/lib/tracing/workers-ai-streaming.test.ts index 16f7983b0d4e..8de48172d6d8 100644 --- a/packages/core/test/lib/tracing/workers-ai-streaming.test.ts +++ b/packages/core/test/lib/tracing/workers-ai-streaming.test.ts @@ -1,8 +1,10 @@ import { describe, expect, it } from 'vitest'; import type { Span } from '../../../src'; import { + GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE, GEN_AI_RESPONSE_STREAMING_ATTRIBUTE, GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, @@ -99,6 +101,114 @@ describe('instrumentWorkersAiStream', () => { expect(ended()).toBe(true); }); + // Models routed through the OpenAI-compatible endpoint (e.g. via `workers-ai-provider`, + // which the Cloudflare Agents SDK uses) stream `choices[].delta.content` instead of the + // native top-level `response` field. Usage still arrives top-level, which is why the + // pre-fix parser captured tokens but dropped the response text entirely. + it('accumulates response text from OpenAI-compatible choices[].delta.content chunks', async () => { + const { span, attributes, ended } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"content":"The capital "},"finish_reason":null}]}\n\n', + 'data: {"choices":[{"delta":{"content":"of France "},"finish_reason":null}]}\n\n', + 'data: {"choices":[{"delta":{"content":"is Paris."},"finish_reason":"stop"}]}\n\n', + 'data: {"usage":{"prompt_tokens":12,"completion_tokens":7,"total_tokens":19}}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + expect(attributes[GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]).toBe(true); + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('The capital of France is Paris.'); + expect(attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]).toBe(12); + expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]).toBe(7); + expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]).toBe(19); + expect(ended()).toBe(true); + }); + + it('assembles fragmented OpenAI-compatible tool calls from choices[].delta.tool_calls', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"getRepoInfo","arguments":"{\\"owner\\":"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\\"cloudflare\\",\\"name\\":\\"agents\\"}"}}]}}]}\n\n', + 'data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + const toolCalls = JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE] as string); + expect(toolCalls).toEqual([ + { + index: 0, + id: 'call_1', + type: 'function', + function: { name: 'getRepoInfo', arguments: '{"owner":"cloudflare","name":"agents"}' }, + }, + ]); + + // The product reads model output from `gen_ai.output.messages`; tool calls must appear there too. + expect(JSON.parse(attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE] as string)).toEqual([ + { + role: 'assistant', + parts: [ + { + type: 'tool_call', + id: 'call_1', + name: 'getRepoInfo', + arguments: '{"owner":"cloudflare","name":"agents"}', + }, + ], + }, + ]); + }); + + // Some Workers AI models (e.g. `@cf/moonshotai/kimi-k2.6`) stream tool calls with the + // name/arguments at the top level of the tool-call object rather than nested under `function`. + it('assembles tool calls whose name/arguments are at the top level of the delta', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","name":"getRepoInfo","arguments":"{\\"owner\\":"}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"arguments":"\\"cloudflare\\"}"}]}}]}\n\n', + 'data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + const toolCalls = JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE] as string); + expect(toolCalls).toEqual([ + { + index: 0, + id: 'call_1', + type: undefined, + function: { name: 'getRepoInfo', arguments: '{"owner":"cloudflare"}' }, + }, + ]); + + expect(JSON.parse(attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE] as string)).toEqual([ + { + role: 'assistant', + parts: [{ type: 'tool_call', id: 'call_1', name: 'getRepoInfo', arguments: '{"owner":"cloudflare"}' }], + }, + ]); + }); + + it('does not record OpenAI-compatible output when recordOutputs is false', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"choices":[{"delta":{"content":"secret"}}]}\n\n', + 'data: {"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, false); + await new Response(instrumented).text(); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBeUndefined(); + expect(attributes[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]).toBeUndefined(); + expect(attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]).toBeUndefined(); + expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]).toBe(1); + }); + it('ends the span when the consumer cancels the stream', async () => { const { span, attributes, ended } = createMockSpan(); diff --git a/packages/core/test/lib/utils/workers-ai-utils.test.ts b/packages/core/test/lib/utils/workers-ai-utils.test.ts index 58a6c0b7224b..3eefe0bf79bc 100644 --- a/packages/core/test/lib/utils/workers-ai-utils.test.ts +++ b/packages/core/test/lib/utils/workers-ai-utils.test.ts @@ -13,13 +13,17 @@ import { GEN_AI_REQUEST_TEMPERATURE, GEN_AI_REQUEST_TOP_K, GEN_AI_REQUEST_TOP_P, - GEN_AI_OUTPUT_MESSAGES, GEN_AI_SYSTEM_INSTRUCTIONS, GEN_AI_USAGE_INPUT_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_TOTAL_TOKENS, } from '@sentry/conventions/attributes'; -import { GEN_AI_REQUEST_STREAM_ATTRIBUTE } from '../../../src/tracing/ai/gen-ai-attributes'; +import { + GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE, + GEN_AI_REQUEST_STREAM_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, +} from '../../../src/tracing/ai/gen-ai-attributes'; import { WORKERS_AI_ORIGIN, WORKERS_AI_PROVIDER_NAME } from '../../../src/tracing/workers-ai/constants'; import { addRequestAttributes, @@ -201,16 +205,83 @@ describe('workers-ai utils', () => { addResponseAttributes(span, { response: 'Paris' }, false); - expect(attributes[GEN_AI_OUTPUT_MESSAGES]).toBeUndefined(); + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBeUndefined(); + }); + + it('records response text when recordOutputs is true', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { response: 'Paris' }, true); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('Paris'); + expect(attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]).toBe( + JSON.stringify([{ role: 'assistant', parts: [{ type: 'text', content: 'Paris' }] }]), + ); + }); + + it('records tool calls when recordOutputs is true', () => { + const { span, attributes } = createMockSpan(); + const toolCalls = [{ id: 'call_1', name: 'lookup', arguments: { city: 'Paris' } }]; + + addResponseAttributes(span, { tool_calls: toolCalls }, true); + + expect(attributes[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]).toBe(JSON.stringify(toolCalls)); + // The product reads model output from `gen_ai.output.messages`; tool calls must appear there too. + expect(attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]).toBe( + JSON.stringify([ + { + role: 'assistant', + parts: [{ type: 'tool_call', id: 'call_1', name: 'lookup', arguments: JSON.stringify({ city: 'Paris' }) }], + }, + ]), + ); }); - it('records response text and tool calls when recordOutputs is true', () => { + it('normalizes OpenAI-compatible tool calls (name/arguments nested under function) into output messages', () => { const { span, attributes } = createMockSpan(); - const toolCalls = [{ name: 'lookup', arguments: { city: 'Paris' } }]; + const toolCalls = [ + { id: 'call_1', type: 'function', function: { name: 'lookup', arguments: '{"city":"Paris"}' } }, + ]; addResponseAttributes(span, { tool_calls: toolCalls }, true); - expect(attributes[GEN_AI_OUTPUT_MESSAGES]).toBe(JSON.stringify(toolCalls)); + expect(attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]).toBe( + JSON.stringify([ + { + role: 'assistant', + parts: [{ type: 'tool_call', id: 'call_1', name: 'lookup', arguments: '{"city":"Paris"}' }], + }, + ]), + ); + }); + + it('records both response text and tool calls without overwriting each other', () => { + const { span, attributes } = createMockSpan(); + const toolCalls = [{ id: 'call_1', name: 'lookup', arguments: { city: 'Paris' } }]; + + addResponseAttributes(span, { response: 'Looking that up', tool_calls: toolCalls }, true); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('Looking that up'); + expect(attributes[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]).toBe(JSON.stringify(toolCalls)); + expect(attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]).toBe( + JSON.stringify([ + { + role: 'assistant', + parts: [ + { type: 'text', content: 'Looking that up' }, + { type: 'tool_call', id: 'call_1', name: 'lookup', arguments: JSON.stringify({ city: 'Paris' }) }, + ], + }, + ]), + ); + }); + + it('does not set output messages when recordOutputs is false', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { response: 'Paris', tool_calls: [{ name: 'lookup' }] }, false); + + expect(attributes[GEN_AI_OUTPUT_MESSAGES_ATTRIBUTE]).toBeUndefined(); }); it('serializes non-string response payloads as JSON', () => { @@ -218,7 +289,7 @@ describe('workers-ai utils', () => { addResponseAttributes(span, { response: { translated_text: 'Bonjour' } }, true); - expect(attributes[GEN_AI_OUTPUT_MESSAGES]).toBe(JSON.stringify({ translated_text: 'Bonjour' })); + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe(JSON.stringify({ translated_text: 'Bonjour' })); }); it('ignores raw Response objects', () => {