feat(cloudflare): Auto-instrument Workers AI binding via env instrumentation

dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/index.ts

@@ -0,0 +1,44 @@
+import * as Sentry from '@sentry/cloudflare';
+import { MockAi } from './mocks';
+
+interface Env {
+  SENTRY_DSN: string;
+}
+
+const ai = Sentry.instrumentWorkersAiClient(new MockAi());
+
+export default Sentry.withSentry(
+  (env: Env) => ({
+    dsn: env.SENTRY_DSN,
+    tracesSampleRate: 1.0,
+    // Keep gen_ai spans embedded in the transaction (instead of streamed as a
+    // separate envelope container) so they can be asserted on `transaction.spans`.
+    streamGenAiSpans: false,
+  }),
+  {
+    async fetch(request) {
+      const url = new URL(request.url);
+
+      if (url.pathname === '/stream') {
+        const stream = (await ai.run('@cf/meta/llama-3.1-8b-instruct', {
+          messages: [{ role: 'user', content: 'What is the capital of France?' }],
+          stream: true,
+        })) as ReadableStream;
+
+        const text = await new Response(stream).text();
+        return new Response(text);
+      }
+
+      const result = await ai.run('@cf/meta/llama-3.1-8b-instruct', {
+        messages: [
+          { role: 'system', content: 'You are a helpful assistant.' },
+          { role: 'user', content: 'What is the capital of France?' },
+        ],
+        temperature: 0.7,
+        max_tokens: 100,
+      });
+
+      return new Response(JSON.stringify(result));
+    },
+  },
+);

dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/mocks.ts

@@ -0,0 +1,45 @@
+import { simulateReadableStream } from 'ai';
+
+function createSseStream(events: string[]): ReadableStream<Uint8Array> {
+  const encoder = new TextEncoder();
+  return simulateReadableStream({
+    initialDelayInMs: 0,
+    chunkDelayInMs: 0,
+    chunks: events.map(event => encoder.encode(`data: ${event}\n\n`)),
+  });
+}
+
+/**
+ * Minimal mock of the Cloudflare Workers AI binding (`env.AI`).
+ */
+export class MockAi {
+  public async run(model: string, inputs: Record<string, unknown>): Promise<unknown> {
+    // Simulate processing time
+    await new Promise(resolve => setTimeout(resolve, 10));
+
+    if (model === 'error-model') {
+      const error = new Error('Model not found');
+      (error as unknown as { status: number }).status = 404;
+      throw error;
+    }
+
+    if (inputs?.stream === true) {
+      return createSseStream([
+        '{"response":"The capital "}',
+        '{"response":"of France "}',
+        '{"response":"is Paris."}',
+        '{"response":"","usage":{"prompt_tokens":12,"completion_tokens":7,"total_tokens":19}}',
+        '[DONE]',
+      ]);
+    }
+
+    return {
+      response: 'The capital of France is Paris.',
+      usage: {
+        prompt_tokens: 12,
+        completion_tokens: 7,
+        total_tokens: 19,
+      },
+    };
+  }
+}

dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/test.ts

@@ -0,0 +1,110 @@
+import { expect, it } from 'vitest';
+import {
+  GEN_AI_OPERATION_NAME_ATTRIBUTE,
+  GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE,
+  GEN_AI_REQUEST_MODEL_ATTRIBUTE,
+  GEN_AI_REQUEST_STREAM_ATTRIBUTE,
+  GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE,
+  GEN_AI_RESPONSE_STREAMING_ATTRIBUTE,
+  GEN_AI_SYSTEM_ATTRIBUTE,
+  GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE,
+  GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
+  GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
+} from '../../../../../packages/core/src/tracing/ai/gen-ai-attributes';
+import { createRunner } from '../../../runner';
+
+// These tests are not exhaustive because the instrumentation is
+// already tested in the core unit tests and we merely want to test
+// that the instrumentation does not break in our cloudflare SDK.
+
+it('traces a basic Workers AI text generation request', async ({ signal }) => {
+  const runner = createRunner(__dirname)
+    .ignore('event')
+    .expect(envelope => {
+      const transactionEvent = envelope[1]?.[0]?.[1] as any;
+
+      // The transaction event is framework-generated and carries non-deterministic fields
+      // (random ports, ids, timestamps, sdk version), so we assert the stable subset.
+      expect(transactionEvent).toEqual(
+        expect.objectContaining({
+          type: 'transaction',
+          transaction: 'GET /',
+          transaction_info: { source: 'route' },
+          contexts: expect.objectContaining({
+            trace: expect.objectContaining({
+              op: 'http.server',
+              origin: 'auto.http.cloudflare',
+              status: 'ok',
+            }),
+          }),
+          spans: [
+            expect.objectContaining({
+              description: 'chat @cf/meta/llama-3.1-8b-instruct',
+              op: 'gen_ai.chat',
+              origin: 'auto.ai.cloudflare.workers_ai',
+              data: {
+                'sentry.origin': 'auto.ai.cloudflare.workers_ai',
+                'sentry.op': 'gen_ai.chat',
+                [GEN_AI_SYSTEM_ATTRIBUTE]: 'cloudflare.workers_ai',
+                [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'chat',
+                [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: '@cf/meta/llama-3.1-8b-instruct',
+                [GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE]: 0.7,
+                [GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE]: 100,
+                [GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 12,
+                [GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]: 7,
+                [GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 19,
+              },
+            }),
+          ],
+        }),
+      );
+    })
+    .start(signal);
+  await runner.makeRequest('get', '/');
+  await runner.completed();
+});
+
+it('traces a streaming Workers AI text generation request', async ({ signal }) => {
+  const runner = createRunner(__dirname)
+    .ignore('event')
+    .expect(envelope => {
+      const transactionEvent = envelope[1]?.[0]?.[1] as any;
+
+      expect(transactionEvent).toEqual(
+        expect.objectContaining({
+          type: 'transaction',
+          transaction: 'GET /stream',
+          transaction_info: { source: 'url' },
+          contexts: expect.objectContaining({
+            trace: expect.objectContaining({
+              op: 'http.server',
+              origin: 'auto.http.cloudflare',
+              status: 'ok',
+            }),
+          }),
+          spans: [
+            expect.objectContaining({
+              description: 'chat @cf/meta/llama-3.1-8b-instruct',
+              op: 'gen_ai.chat',
+              origin: 'auto.ai.cloudflare.workers_ai',
+              data: {
+                'sentry.origin': 'auto.ai.cloudflare.workers_ai',
+                'sentry.op': 'gen_ai.chat',
+                [GEN_AI_SYSTEM_ATTRIBUTE]: 'cloudflare.workers_ai',
+                [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'chat',
+                [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: '@cf/meta/llama-3.1-8b-instruct',
+                [GEN_AI_REQUEST_STREAM_ATTRIBUTE]: true,
+                [GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]: true,
+                [GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 12,
+                [GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]: 7,
+                [GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 19,
+              },
+            }),
+          ],
+        }),
+      );
+    })
+    .start(signal);
+  await runner.makeRequest('get', '/stream');
+  await runner.completed();
+});

dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/wrangler.jsonc

@@ -0,0 +1,6 @@
+{
+  "name": "worker-name",
+  "compatibility_date": "2025-06-17",
+  "main": "index.ts",
+  "compatibility_flags": ["nodejs_als"],
+}

packages/cloudflare/src/index.ts

@@ -79,6 +79,7 @@ export {
   instrumentOpenAiClient,
   instrumentGoogleGenAIClient,
   instrumentAnthropicAiClient,
+  instrumentWorkersAiClient,
   eventFiltersIntegration,
   linkedErrorsIntegration,
   requestDataIntegration,

packages/cloudflare/src/instrumentations/worker/instrumentEnv.ts

@@ -1,5 +1,13 @@
+import { instrumentWorkersAiClient } from '@sentry/core';
 import type { CloudflareOptions } from '../../client';
-import { isD1Database, isDurableObjectNamespace, isJSRPC, isQueue, isR2Bucket } from '../../utils/isBinding';
+import {
+  isAiBinding,
+  isD1Database,
+  isDurableObjectNamespace,
+  isJSRPC,
+  isQueue,
+  isR2Bucket,
+} from '../../utils/isBinding';
 import { instrumentD1 } from './instrumentD1';
 import { appendRpcMeta } from '../../utils/rpcMeta';
 import { getEffectiveRpcPropagation } from '../../utils/rpcOptions';
@@ -23,6 +31,7 @@ const instrumentedBindings = new WeakMap<object, unknown>();
  * - Service bindings / JSRPC proxies
  * - Queue producers (via `send` + `sendBatch` duck-typing)
  * - R2 Buckets (via `head` + `put` + `createMultipartUpload` duck-typing)
+ * - Workers AI (via `run` + `gateway` + `toMarkdown` duck-typing)
  *
  * @param env - The Cloudflare env object to instrument
  * @param options - Optional CloudflareOptions to control RPC trace propagation
@@ -68,6 +77,12 @@ export function instrumentEnv<Env extends Record<string, unknown>>(env: Env, opt
         return instrumented;
       }
 
+      if (isAiBinding(item)) {
+        const instrumented = instrumentWorkersAiClient(item);
+        instrumentedBindings.set(item, instrumented);
+        return instrumented;
+      }
+
       if (!rpcPropagation) {
         return item;
       }

packages/cloudflare/src/utils/isBinding.ts

@@ -31,7 +31,7 @@
  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
-import type { D1Database, DurableObjectNamespace, Queue, R2Bucket } from '@cloudflare/workers-types';
+import type { Ai, D1Database, DurableObjectNamespace, Queue, R2Bucket } from '@cloudflare/workers-types';
 
 /**
  * Checks if a value is a JSRPC proxy (service binding).
@@ -82,6 +82,20 @@ export function isD1Database(item: unknown): item is D1Database {
   );
 }
 
+/**
+ * Duck-type check for Workers AI bindings.
+ * The Ai binding has `run`, `gateway`, and `toMarkdown` methods.
+ */
+export function isAiBinding(item: unknown): item is Ai {
+  return (
+    item != null &&
+    isNotJSRPC(item) &&
+    typeof item.run === 'function' &&
+    typeof item.gateway === 'function' &&
+    typeof item.toMarkdown === 'function'
+  );
+}
+
 /**
  * Duck-type check for R2 Bucket bindings.
  * R2Bucket has `head`, `put`, and `createMultipartUpload` methods.

packages/cloudflare/test/instrumentations/instrumentEnv.test.ts

@@ -256,6 +256,49 @@ describe('instrumentEnv', () => {
     expect(instrumentDurableObjectNamespace).toHaveBeenCalledWith(doNamespace);
   });
 
+  describe('Workers AI bindings', () => {
+    function createMockAiBinding() {
+      return {
+        run: vi.fn().mockResolvedValue({ response: 'Paris', usage: { prompt_tokens: 1, completion_tokens: 2 } }),
+        gateway: vi.fn(),
+        toMarkdown: vi.fn(),
+        models: vi.fn(),
+        autorag: vi.fn(),
+      };
+    }
+
+    it('detects and wraps AI bindings, forwarding run calls unchanged', async () => {
+      const ai = createMockAiBinding();
+      const env = { AI: ai };
+      const instrumented = instrumentEnv(env);
+
+      const wrapped = instrumented.AI as typeof ai;
+      expect(wrapped).not.toBe(ai);
+
+      const result = await wrapped.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' });
+
+      expect(ai.run).toHaveBeenCalledTimes(1);
+      expect(ai.run).toHaveBeenCalledWith('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' });
+      expect(result).toEqual({ response: 'Paris', usage: { prompt_tokens: 1, completion_tokens: 2 } });
+    });
+
+    it('caches the wrapped AI binding across repeated access', () => {
+      const ai = createMockAiBinding();
+      const env = { AI: ai };
+      const instrumented = instrumentEnv(env);
+
+      expect(instrumented.AI).toBe(instrumented.AI);
+    });
+
+    it('does not treat bindings with only a run method as AI bindings', () => {
+      const notAi = { run: vi.fn() };
+      const env = { RUNNER: notAi };
+      const instrumented = instrumentEnv(env);
+
+      expect(instrumented.RUNNER).toBe(notAi);
+    });
+  });
+
   describe('mTLS Fetcher bindings', () => {
     function createMtlsFetcherProxy(mockFetch: ReturnType<typeof vi.fn>) {
       return new Proxy(