-
-
Notifications
You must be signed in to change notification settings - Fork 1.8k
fix(cloudflare): Prevent AI provider skips #22719
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 |
|---|---|---|
|
|
@@ -29,6 +29,7 @@ import { | |
| import { GEN_AI_EXECUTE_TOOL_SPAN_OP, GEN_AI_INVOKE_AGENT_SPAN_OP } from '@sentry/conventions/op'; | ||
| import type { Span, SpanAttributes } from '@sentry/core'; | ||
| import { | ||
| _INTERNAL_skipAiProviderWrapping, | ||
| captureException, | ||
| GEN_AI_CONVERSATION_ID_ATTRIBUTE, | ||
| GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, | ||
|
|
@@ -66,6 +67,8 @@ const GEN_AI_RERANK_OPERATION = 'rerank'; | |
| // The model-call op matches the Vercel AI OTel integration (`gen_ai.generate_content`) rather than | ||
| // the generic `gen_ai.chat`, so v6 (OTel) and v7 (channel) produce the same spans. | ||
| const GEN_AI_GENERATE_CONTENT_OPERATION = 'generate_content'; | ||
| // TODO(v11): export the constant from server-utils and import it here instead. | ||
| const WORKERS_AI_INTEGRATION_NAME = 'WorkersAI'; | ||
|
Comment on lines
+70
to
+71
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. q: Why is this a todo? Can't we do that now? We are in v11.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes we can. I just wanted to keep things easier for backports for now.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In a follow up I can move everything over to server-utils |
||
|
|
||
| // Subset of the `vercel.ai.*` passthrough attributes the OTel integration emits that we reproduce. | ||
| const VERCEL_AI_OPERATION_ID_ATTRIBUTE = 'vercel.ai.operationId'; | ||
|
|
@@ -395,6 +398,8 @@ export function createSpanFromMessage( | |
| // the OTel path derives from the SDK's Zod schema is not reconstructed on the channel path. | ||
| return buildInvokeAgentSpan(event, baseAttributes, recordInputs, enableTruncation, callId, type === 'streamText'); | ||
| case 'languageModelCall': | ||
| _INTERNAL_skipAiProviderWrapping([WORKERS_AI_INTEGRATION_NAME]); | ||
|
|
||
| return buildModelCallSpan(event, baseAttributes, recordInputs, enableTruncation, callId, modelId); | ||
| case 'executeTool': | ||
| return buildToolSpan(event, recordInputs); | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| import { | ||
| _INTERNAL_clearAiProviderSkips, | ||
| _INTERNAL_shouldSkipAiProviderWrapping, | ||
| Client, | ||
| createTransport, | ||
| getCurrentScope, | ||
| getGlobalScope, | ||
| getIsolationScope, | ||
| initAndBind, | ||
| resolvedSyncPromise, | ||
| } from '@sentry/core'; | ||
| import { afterEach, beforeEach, describe, expect, it } from 'vitest'; | ||
| import { createSpanFromMessage } from '../../src/vercel-ai/vercel-ai-dc-subscriber'; | ||
|
|
||
| // Must match `WORKERS_AI_INTEGRATION_NAME` in core's `tracing/workers-ai/constants`. | ||
| const WORKERS_AI_INTEGRATION_NAME = 'WorkersAI'; | ||
|
|
||
| class TestClient extends Client<any> { | ||
| public eventFromException(): PromiseLike<any> { | ||
| return resolvedSyncPromise({}); | ||
| } | ||
|
|
||
| public eventFromMessage(): PromiseLike<any> { | ||
| return resolvedSyncPromise({}); | ||
| } | ||
| } | ||
|
|
||
| function initTestClient(): void { | ||
| initAndBind(TestClient, { | ||
| dsn: 'https://username@domain/123', | ||
| integrations: [], | ||
| sendClientReports: false, | ||
| stackParser: () => [], | ||
| tracesSampleRate: 1, | ||
| transport: () => createTransport({ recordDroppedEvent: () => undefined }, () => resolvedSyncPromise({})), | ||
| }); | ||
| } | ||
|
|
||
| describe('vercel ai tracing channel: workers-ai dedup', () => { | ||
| beforeEach(() => { | ||
| _INTERNAL_clearAiProviderSkips(); | ||
| getCurrentScope().clear(); | ||
| getIsolationScope().clear(); | ||
| getGlobalScope().clear(); | ||
| initTestClient(); | ||
| }); | ||
|
|
||
| afterEach(() => { | ||
| _INTERNAL_clearAiProviderSkips(); | ||
| }); | ||
|
|
||
| it('marks Workers AI as skipped on a model call, so the binding does not double-instrument', () => { | ||
| expect(_INTERNAL_shouldSkipAiProviderWrapping(WORKERS_AI_INTEGRATION_NAME)).toBe(false); | ||
|
|
||
| const span = createSpanFromMessage( | ||
| { | ||
| type: 'languageModelCall', | ||
| event: { provider: 'workers-ai', modelId: '@cf/meta/llama-3.1-8b-instruct' }, | ||
| } as Parameters<typeof createSpanFromMessage>[0], | ||
| {} as Parameters<typeof createSpanFromMessage>[1], | ||
| ); | ||
| span?.end(); | ||
|
|
||
| expect(_INTERNAL_shouldSkipAiProviderWrapping(WORKERS_AI_INTEGRATION_NAME)).toBe(true); | ||
| }); | ||
|
|
||
| it('does not mark Workers AI as skipped for tool calls', () => { | ||
| // A tool calling `env.AI.run` itself is a genuine separate inference the `ai` SDK does not | ||
| // instrument, so it must keep its own span. | ||
| const span = createSpanFromMessage( | ||
| { | ||
| type: 'executeTool', | ||
| event: { toolName: 'getWeather', toolCallId: 'call_1' }, | ||
| } as Parameters<typeof createSpanFromMessage>[0], | ||
| {} as Parameters<typeof createSpanFromMessage>[1], | ||
| ); | ||
| span?.end(); | ||
|
|
||
| expect(_INTERNAL_shouldSkipAiProviderWrapping(WORKERS_AI_INTEGRATION_NAME)).toBe(false); | ||
| }); | ||
| }); |
Uh oh!
There was an error while loading. Please reload this page.