Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion packages/cloudflare/src/client.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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.
*/
Expand Down
8 changes: 8 additions & 0 deletions packages/core/src/tracing/vercel-ai/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 } }
Expand Down
7 changes: 7 additions & 0 deletions packages/core/src/tracing/workers-ai/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
8 changes: 8 additions & 0 deletions packages/core/src/tracing/workers-ai/index.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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);
Expand Down
89 changes: 88 additions & 1 deletion packages/core/test/lib/tracing/workers-ai.test.ts
Original file line number Diff line number Diff line change
@@ -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', () => {
Expand Down Expand Up @@ -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');
});
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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';
Expand Down Expand Up @@ -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);
Expand Down
81 changes: 81 additions & 0 deletions packages/server-utils/test/vercel-ai/skip-workers-ai.test.ts
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);
});
});
Loading