fix(v10/cloudflare): Prevent AI provider skips

Backport of: #22719

Author: JPeer264

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<typeof makeFlushLock> | 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.
    */

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 } }

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;

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<Pick<WorkersAiOptions, 'recordInputs' | 'recordOutputs'>>,
 ): (...args: unknown[]) => Promise<unknown> {
   return function instrumentedRun(...args: unknown[]): Promise<unknown> {
+    // 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<string, unknown> | undefined];
 
     const operationName = getOperationName(inputs);

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<unknown>): Promise<void> {
+      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');
+    });
+  });
 });

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);

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<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);
+  });
+});