diff --git a/packages/cloudflare/src/client.ts b/packages/cloudflare/src/client.ts index 087c1ad720d9..d3628acf51b1 100644 --- a/packages/cloudflare/src/client.ts +++ b/packages/cloudflare/src/client.ts @@ -1,5 +1,11 @@ import type { ClientOptions, Options, ServerRuntimeClientOptions } from '@sentry/core'; -import { applySdkMetadata, debug, ServerRuntimeClient, spanIsSampled } from '@sentry/core'; +import { + _INTERNAL_clearAiProviderSkips, + applySdkMetadata, + debug, + ServerRuntimeClient, + spanIsSampled, +} from '@sentry/core'; import { DEBUG_BUILD } from './debug-build'; import type { ExecutionContextCompat } from './executionContext'; import type { makeFlushLock } from './flush'; @@ -140,6 +146,16 @@ export class CloudflareClient extends ServerRuntimeClient { (this as unknown as { _flushLock: ReturnType | void })._flushLock = undefined; } + /** @inheritDoc */ + protected override _setupIntegrations(): void { + // Clear AI provider skip registrations before setting up integrations. + // The registry is module-global and Cloudflare calls `init()` per request, so without this a + // single `ai` SDK call would suppress direct `env.AI.run` spans for the rest of the isolate's + // life. Mirrors the same reset in the Node client. + _INTERNAL_clearAiProviderSkips(); + super._setupIntegrations(); + } + /** * Resets the span completion promise and resolve function. */ diff --git a/packages/core/src/tracing/vercel-ai/index.ts b/packages/core/src/tracing/vercel-ai/index.ts index c905ac980614..7d72b3de3ae7 100644 --- a/packages/core/src/tracing/vercel-ai/index.ts +++ b/packages/core/src/tracing/vercel-ai/index.ts @@ -5,7 +5,9 @@ import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from ' import { shouldEnableTruncation } from '../ai/utils'; import type { Event } from '../../types/event'; import type { Span, SpanAttributes, SpanAttributeValue, SpanJSON, StreamedSpanJSON } from '../../types/span'; +import { _INTERNAL_skipAiProviderWrapping } from '../../utils/ai/providerSkip'; import { spanToJSON } from '../../utils/spanUtils'; +import { WORKERS_AI_INTEGRATION_NAME } from '../workers-ai/constants'; import { GEN_AI_CONVERSATION_ID_ATTRIBUTE, GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, @@ -86,6 +88,12 @@ function onVercelAiSpanStart(span: Span): void { return; } + // Registered lazily here (not at `setupOnce`) so a direct `env.AI.run` call made before any `ai` + // SDK call still gets its own span. + if (SPAN_TO_OPERATION_NAME.get(name) === 'generate_content') { + _INTERNAL_skipAiProviderWrapping([WORKERS_AI_INTEGRATION_NAME]); + } + const client = getClient(); const integration = client?.getIntegrationByName('VercelAI') as | { options?: { enableTruncation?: boolean } } diff --git a/packages/core/src/tracing/workers-ai/constants.ts b/packages/core/src/tracing/workers-ai/constants.ts index 8f532a0a0372..1dfff520b185 100644 --- a/packages/core/src/tracing/workers-ai/constants.ts +++ b/packages/core/src/tracing/workers-ai/constants.ts @@ -8,3 +8,10 @@ export const WORKERS_AI_PROVIDER_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'; + +/** + * The key used to register this provider in the AI provider skip registry. + * + * @see `_INTERNAL_skipAiProviderWrapping` + */ +export const WORKERS_AI_INTEGRATION_NAME = 'WorkersAI' as const; diff --git a/packages/core/src/tracing/workers-ai/index.ts b/packages/core/src/tracing/workers-ai/index.ts index c6be4824305a..9232b52bd5f1 100644 --- a/packages/core/src/tracing/workers-ai/index.ts +++ b/packages/core/src/tracing/workers-ai/index.ts @@ -1,8 +1,10 @@ import { SPAN_STATUS_ERROR } from '../../tracing'; import { startSpan, startSpanManual } from '../../tracing/trace'; import type { Span } from '../../types/span'; +import { _INTERNAL_shouldSkipAiProviderWrapping } from '../../utils/ai/providerSkip'; import { isObjectLike } from '../../utils/is'; import { resolveAIRecordingOptions, shouldEnableTruncation } from '../ai/utils'; +import { WORKERS_AI_INTEGRATION_NAME } from './constants'; import { instrumentWorkersAiStream } from './streaming'; import type { WorkersAiOptions } from './types'; import { addRequestAttributes, addResponseAttributes, extractRequestAttributes, getOperationName } from './utils'; @@ -27,6 +29,12 @@ function instrumentRun( options: WorkersAiOptions & Required>, ): (...args: unknown[]) => Promise { return function instrumentedRun(...args: unknown[]): Promise { + // When another integration (e.g. Vercel AI via `workers-ai-provider`) is driving this binding, + // it records the spans itself and marks this provider as skipped; skip here to avoid double spans. + if (_INTERNAL_shouldSkipAiProviderWrapping(WORKERS_AI_INTEGRATION_NAME)) { + return originalRun.apply(context, args); + } + const [model, inputs, runOptions] = args as [unknown, unknown, Record | undefined]; const operationName = getOperationName(inputs); diff --git a/packages/core/test/lib/tracing/workers-ai.test.ts b/packages/core/test/lib/tracing/workers-ai.test.ts index d9c563b59962..78e5dc79fe2f 100644 --- a/packages/core/test/lib/tracing/workers-ai.test.ts +++ b/packages/core/test/lib/tracing/workers-ai.test.ts @@ -1,5 +1,11 @@ -import { describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { getCurrentScope, getGlobalScope, getIsolationScope, setCurrentClient, startSpan } from '../../../src'; +import { addVercelAiProcessors } from '../../../src/tracing/vercel-ai'; +import { AI_OPERATION_ID_ATTRIBUTE } from '../../../src/tracing/vercel-ai/vercel-ai-attributes'; import { instrumentWorkersAiClient } from '../../../src/tracing/workers-ai'; +import { _INTERNAL_clearAiProviderSkips } from '../../../src/utils/ai/providerSkip'; +import { spanToJSON } from '../../../src/utils/spanUtils'; +import { getDefaultTestClientOptions, TestClient } from '../../mocks/client'; describe('instrumentWorkersAiClient', () => { it('passes through non-run methods bound to the original client', () => { @@ -28,4 +34,85 @@ describe('instrumentWorkersAiClient', () => { expect(client.run).toHaveBeenCalledWith('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' }); expect(result).toEqual({ response: 'Paris' }); }); + + describe('when the Vercel AI SDK drives the binding', () => { + let spans: string[]; + + /** Set up a client with the Vercel AI processors registered, recording every ended span. */ + function setupClient(): void { + getCurrentScope().clear(); + getIsolationScope().clear(); + getGlobalScope().clear(); + + spans = []; + const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1 })); + client.on('spanEnd', span => { + spans.push(spanToJSON(span).description ?? ''); + }); + setCurrentClient(client); + addVercelAiProcessors(client); + } + + beforeEach(() => { + _INTERNAL_clearAiProviderSkips(); + setupClient(); + }); + + afterEach(() => { + _INTERNAL_clearAiProviderSkips(); + }); + + /** + * Emit the span the `ai` SDK creates for a model call. Its `spanStart` handler is what marks + * Workers AI as skipped, exactly as it would at runtime. + */ + async function withVercelAiModelCall(callback: () => Promise): Promise { + await startSpan( + { name: 'ai.streamText.doStream', attributes: { [AI_OPERATION_ID_ATTRIBUTE]: 'ai.streamText.doStream' } }, + async () => { + await callback(); + }, + ); + } + + it('does not create a duplicate span for the nested `run` call', async () => { + const client = { run: vi.fn().mockResolvedValue({ response: 'Paris' }) }; + const instrumented = instrumentWorkersAiClient(client); + + await withVercelAiModelCall(() => instrumented.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' })); + + // The call is forwarded, but no duplicate `gen_ai.chat` span is emitted — only the + // Vercel AI model-call span remains. + expect(client.run).toHaveBeenCalledWith('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' }); + expect(spans).not.toContain('chat @cf/meta/llama-3.1-8b-instruct'); + expect(spans).toEqual(['streamText.doStream']); + }); + + it('still creates a span for a direct `run` call made before any Vercel AI call', async () => { + const client = { run: vi.fn().mockResolvedValue({ response: 'Paris' }) }; + const instrumented = instrumentWorkersAiClient(client); + + await startSpan({ name: 'root' }, () => instrumented.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' })); + + expect(spans).toContain('chat @cf/meta/llama-3.1-8b-instruct'); + }); + + it('clears the skip between clients so a later isolate reuse is unaffected', async () => { + const client = { run: vi.fn().mockResolvedValue({ response: 'Paris' }) }; + const instrumented = instrumentWorkersAiClient(client); + + // First request: the `ai` SDK runs and marks Workers AI as skipped. + await withVercelAiModelCall(() => instrumented.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' })); + expect(spans).not.toContain('chat @cf/meta/llama-3.1-8b-instruct'); + + // Second request on the same isolate: `_setupIntegrations` resets the registry, so a direct + // `env.AI.run` call must get its span back. Without the reset this would stay suppressed. + _INTERNAL_clearAiProviderSkips(); + setupClient(); + + await startSpan({ name: 'root' }, () => instrumented.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' })); + + expect(spans).toContain('chat @cf/meta/llama-3.1-8b-instruct'); + }); + }); }); diff --git a/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts b/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts index 0b9e7f39e101..fead1c746383 100644 --- a/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts +++ b/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts @@ -28,6 +28,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'; // 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); diff --git a/packages/server-utils/test/vercel-ai/skip-workers-ai.test.ts b/packages/server-utils/test/vercel-ai/skip-workers-ai.test.ts new file mode 100644 index 000000000000..b42565d595ee --- /dev/null +++ b/packages/server-utils/test/vercel-ai/skip-workers-ai.test.ts @@ -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 { + public eventFromException(): PromiseLike { + return resolvedSyncPromise({}); + } + + public eventFromMessage(): PromiseLike { + 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[0], + {} as Parameters[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[0], + {} as Parameters[1], + ); + span?.end(); + + expect(_INTERNAL_shouldSkipAiProviderWrapping(WORKERS_AI_INTEGRATION_NAME)).toBe(false); + }); +});