feat(core): Instrument workers-ai-provider

packages/core/src/shared-exports.ts

@@ -221,6 +221,8 @@ export type { LangChainOptions, LangChainIntegration } from './tracing/langchain
 export { instrumentStateGraphCompile, instrumentCreateReactAgent, instrumentLangGraph } from './tracing/langgraph';
 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 {

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';

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<Uint8Array> {
+  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<unknown>,
+  context: unknown,
+  options: WorkersAiOptions & Required<Pick<WorkersAiOptions, 'recordInputs' | 'recordOutputs'>>,
+): (...args: unknown[]) => Promise<unknown> {
+  return function instrumentedRun(...args: unknown[]): Promise<unknown> {
+    const [model, inputs, runOptions] = args as [unknown, unknown, Record<string, unknown> | 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<unknown>;
+
+        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<unknown>;
+
+      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<object>();
+
+/**
+ * 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<T extends object>(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<unknown>, 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;
+}

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<Uint8Array>,
+  span: Span,
+  recordOutputs: boolean,
+): ReadableStream<Uint8Array> {
+  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<Uint8Array>({
+    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);
+    },
+  });
+}

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<string, unknown>, options?: Record<string, unknown>) => Promise<unknown>;
+  [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;
+}