-
-
Notifications
You must be signed in to change notification settings - Fork 1.8k
feat(cloudflare): Auto-instrument Workers AI binding via env instrumentation #22126
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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)); | ||
| }, | ||
| }, | ||
| ); |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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, | ||
| }, | ||
| }; | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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( | ||
|
Check failure on line 28 in dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/test.ts
|
||
| 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( | ||
|
Check failure on line 73 in dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/test.ts
|
||
| 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(); | ||
| }); | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| { | ||
| "name": "worker-name", | ||
| "compatibility_date": "2025-06-17", | ||
| "main": "index.ts", | ||
| "compatibility_flags": ["nodejs_als"], | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Bug: The
MockAiclass in integration tests is missinggatewayandtoMarkdownmethods, which means the new auto-instrumentation logic for AI bindings is not being tested.Severity: LOW
Suggested Fix
Update the
MockAiclass indev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/mocks.tsto include mock implementations for thegatewayandtoMarkdownmethods. Modify the integration test to rely oninstrumentEnvto automatically detect and instrument theenv.AIbinding, which will properly test the new auto-instrumentation code path.Prompt for AI Agent
Also affects:
dev-packages/cloudflare-integration-tests/suites/tracing/workers-ai/index.ts:10~10packages/cloudflare/src/utils/isBinding.ts:94~96packages/cloudflare/src/instrumentations/worker/instrumentEnv.ts:91~95