From c6267101efb0fe64f8433d6cb63686dc0ceeb00a Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Wed, 8 Jul 2026 22:11:01 +0200 Subject: [PATCH] feat(cloudflare): Auto-instrument Workers AI binding via env instrumentation Detect the Workers AI binding (env.AI) in instrumentEnv via duck-typing (run + gateway + toMarkdown) and wrap it automatically, matching how D1, R2, and Queue bindings are instrumented. Manual wrapping via instrumentWorkersAiClient remains available for custom options and is now guarded against double-wrapping. Also removes the unused WORKERS_AI_INTEGRATION_NAME constant and aligns the integration test wrangler config with sibling suites (nodejs_als). Co-Authored-By: Claude Fable 5 --- .../suites/tracing/workers-ai/index.ts | 44 +++++++ .../suites/tracing/workers-ai/mocks.ts | 45 +++++++ .../suites/tracing/workers-ai/test.ts | 110 ++++++++++++++++++ .../suites/tracing/workers-ai/wrangler.jsonc | 6 + packages/cloudflare/src/index.ts | 1 + .../instrumentations/worker/instrumentEnv.ts | 17 ++- packages/cloudflare/src/utils/isBinding.ts | 16 ++- .../instrumentations/instrumentEnv.test.ts | 43 +++++++ 8 files changed, 280 insertions(+), 2 deletions(-) create mode 100644 dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/index.ts create mode 100644 dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/mocks.ts create mode 100644 dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/test.ts create mode 100644 dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/wrangler.jsonc diff --git a/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/index.ts b/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/index.ts new file mode 100644 index 000000000000..f88de43b4729 --- /dev/null +++ b/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)); + }, + }, +); diff --git a/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/mocks.ts b/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/mocks.ts new file mode 100644 index 000000000000..bce096d4ae75 --- /dev/null +++ b/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/mocks.ts @@ -0,0 +1,45 @@ +import { simulateReadableStream } from 'ai'; + +function createSseStream(events: string[]): ReadableStream { + 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): Promise { + // 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, + }, + }; + } +} diff --git a/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/test.ts b/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/test.ts new file mode 100644 index 000000000000..dffb5a88b6b8 --- /dev/null +++ b/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(); +}); diff --git a/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/wrangler.jsonc b/dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/wrangler.jsonc new file mode 100644 index 000000000000..48aaa79409a3 --- /dev/null +++ b/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"], +} diff --git a/packages/cloudflare/src/index.ts b/packages/cloudflare/src/index.ts index 20a537c5b307..7a2cce03ee7b 100644 --- a/packages/cloudflare/src/index.ts +++ b/packages/cloudflare/src/index.ts @@ -79,6 +79,7 @@ export { instrumentOpenAiClient, instrumentGoogleGenAIClient, instrumentAnthropicAiClient, + instrumentWorkersAiClient, eventFiltersIntegration, linkedErrorsIntegration, requestDataIntegration, diff --git a/packages/cloudflare/src/instrumentations/worker/instrumentEnv.ts b/packages/cloudflare/src/instrumentations/worker/instrumentEnv.ts index 393072718ab0..38176eddedeb 100644 --- a/packages/cloudflare/src/instrumentations/worker/instrumentEnv.ts +++ b/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(); * - 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: Env, opt return instrumented; } + if (isAiBinding(item)) { + const instrumented = instrumentWorkersAiClient(item); + instrumentedBindings.set(item, instrumented); + return instrumented; + } + if (!rpcPropagation) { return item; } diff --git a/packages/cloudflare/src/utils/isBinding.ts b/packages/cloudflare/src/utils/isBinding.ts index 88832f375f21..e675807b8fd7 100644 --- a/packages/cloudflare/src/utils/isBinding.ts +++ b/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. diff --git a/packages/cloudflare/test/instrumentations/instrumentEnv.test.ts b/packages/cloudflare/test/instrumentations/instrumentEnv.test.ts index 9309284abb41..4e9aa76ef204 100644 --- a/packages/cloudflare/test/instrumentations/instrumentEnv.test.ts +++ b/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) { return new Proxy(