From 3431a27aec7cb168e09fb51b018352f25e24a73f Mon Sep 17 00:00:00 2001 From: Francesco Novy Date: Mon, 24 Mar 2025 14:52:42 +0100 Subject: [PATCH 1/4] feat(node): WIP Only add span listeners conditionally --- .../node/src/integrations/tracing/connect.ts | 2 +- .../src/integrations/tracing/dataloader.ts | 48 +++++++++++-------- packages/node/src/otel/instrument.ts | 41 ++++++++++++++-- 3 files changed, 65 insertions(+), 26 deletions(-) diff --git a/packages/node/src/integrations/tracing/connect.ts b/packages/node/src/integrations/tracing/connect.ts index 8cfe2bb74050..ae7958c7da83 100644 --- a/packages/node/src/integrations/tracing/connect.ts +++ b/packages/node/src/integrations/tracing/connect.ts @@ -104,7 +104,7 @@ function addConnectSpanAttributes(span: Span): void { [SEMANTIC_ATTRIBUTE_SENTRY_OP]: `${type}.connect`, }); - // Also update the name, we don't need to "middleware - " prefix + // Also update the name, we don't need the "middleware - " prefix const name = attributes['connect.name']; if (typeof name === 'string') { span.updateName(name); diff --git a/packages/node/src/integrations/tracing/dataloader.ts b/packages/node/src/integrations/tracing/dataloader.ts index 6339515c6f28..26ca83851209 100644 --- a/packages/node/src/integrations/tracing/dataloader.ts +++ b/packages/node/src/integrations/tracing/dataloader.ts @@ -3,10 +3,11 @@ import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, defineIntegration, + getClient, spanToJSON, } from '@sentry/core'; import type { IntegrationFn } from '@sentry/core'; -import { generateInstrumentOnce } from '../../otel/instrument'; +import { callWhenWrapped, generateInstrumentOnce } from '../../otel/instrument'; const INTEGRATION_NAME = 'Dataloader'; @@ -19,31 +20,38 @@ export const instrumentDataloader = generateInstrumentOnce( ); const _dataloaderIntegration = (() => { + let hookCallback: undefined | (() => void); + return { name: INTEGRATION_NAME, setupOnce() { - instrumentDataloader(); - }, + const instrumentation = instrumentDataloader(); - setup(client) { - client.on('spanStart', span => { - const spanJSON = spanToJSON(span); - if (spanJSON.description?.startsWith('dataloader')) { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.dataloader'); + callWhenWrapped(instrumentation, () => { + const client = getClient(); + if (hookCallback || !client) { + return; } - // These are all possible dataloader span descriptions - // Still checking for the future versions - // in case they add support for `clear` and `prime` - if ( - spanJSON.description === 'dataloader.load' || - spanJSON.description === 'dataloader.loadMany' || - spanJSON.description === 'dataloader.batch' - ) { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'cache.get'); - // TODO: We can try adding `key` to the `data` attribute upstream. - // Or alternatively, we can add `requestHook` to the dataloader instrumentation. - } + hookCallback = client.on('spanStart', span => { + const spanJSON = spanToJSON(span); + if (spanJSON.description?.startsWith('dataloader')) { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.dataloader'); + } + + // These are all possible dataloader span descriptions + // Still checking for the future versions + // in case they add support for `clear` and `prime` + if ( + spanJSON.description === 'dataloader.load' || + spanJSON.description === 'dataloader.loadMany' || + spanJSON.description === 'dataloader.batch' + ) { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'cache.get'); + // TODO: We can try adding `key` to the `data` attribute upstream. + // Or alternatively, we can add `requestHook` to the dataloader instrumentation. + } + }); }); }, }; diff --git a/packages/node/src/otel/instrument.ts b/packages/node/src/otel/instrument.ts index c5bd7500a739..cd7b7f0c8583 100644 --- a/packages/node/src/otel/instrument.ts +++ b/packages/node/src/otel/instrument.ts @@ -7,19 +7,22 @@ export const INSTRUMENTED: Record = {}; * Instrument an OpenTelemetry instrumentation once. * This will skip running instrumentation again if it was already instrumented. */ -export function generateInstrumentOnce( +export function generateInstrumentOnce< + Options = unknown, + InstrumentationInstance extends Instrumentation = Instrumentation, +>( name: string, - creator: (options?: Options) => Instrumentation, -): ((options?: Options) => void) & { id: string } { + creator: (options?: Options) => InstrumentationInstance, +): ((options?: Options) => InstrumentationInstance) & { id: string } { return Object.assign( (options?: Options) => { - const instrumented = INSTRUMENTED[name]; + const instrumented = INSTRUMENTED[name] as InstrumentationInstance | undefined; if (instrumented) { // If options are provided, ensure we update them if (options) { instrumented.setConfig(options); } - return; + return instrumented; } const instrumentation = creator(options); @@ -28,7 +31,35 @@ export function generateInstrumentOnce( registerInstrumentations({ instrumentations: [instrumentation], }); + + return instrumentation; }, { id: name }, ); } + +/** + * Ensure a given callback is called when the instrumentation is actually wrapping something. + * This can be used to ensure some logic is only called when the instrumentation is actually active. + * This depends on wrapping `_wrap` (inception!). If this is not possible (e.g. the property name is mangled, ...) + * the callback will be called immediately. + */ +export function callWhenWrapped(instrumentation: T, callback: () => void): void { + if (!hasWrap(instrumentation)) { + callback(); + return; + } + + const originalWrap = instrumentation['_wrap']; + + instrumentation['_wrap'] = (...args: Parameters) => { + callback(); + return originalWrap(...args); + }; +} + +function hasWrap( + instrumentation: T, +): instrumentation is T & { _wrap: (...args: unknown[]) => unknown } { + return typeof (instrumentation as T & { _wrap?: (...args: unknown[]) => unknown })['_wrap'] === 'function'; +} From ee303399a0b441e257ab45358d4c37a3872ca7a7 Mon Sep 17 00:00:00 2001 From: Francesco Novy Date: Tue, 25 Mar 2025 10:55:50 +0100 Subject: [PATCH 2/4] ref add stuff --- .../src/integrations/tracing/dataloader.ts | 18 +- .../src/integrations/tracing/genericPool.ts | 29 +- .../node/src/integrations/tracing/knex.ts | 25 +- .../node/src/integrations/tracing/prisma.ts | 39 +-- .../node/src/integrations/tracing/tedious.ts | 33 ++- .../integrations/tracing/vercelai/index.ts | 276 +++++++++--------- packages/node/src/otel/instrument.ts | 40 ++- packages/node/test/utils/instrument.test.ts | 95 ++++++ 8 files changed, 347 insertions(+), 208 deletions(-) create mode 100644 packages/node/test/utils/instrument.test.ts diff --git a/packages/node/src/integrations/tracing/dataloader.ts b/packages/node/src/integrations/tracing/dataloader.ts index 26ca83851209..3c887bb33fa5 100644 --- a/packages/node/src/integrations/tracing/dataloader.ts +++ b/packages/node/src/integrations/tracing/dataloader.ts @@ -3,11 +3,10 @@ import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, defineIntegration, - getClient, spanToJSON, } from '@sentry/core'; import type { IntegrationFn } from '@sentry/core'; -import { callWhenWrapped, generateInstrumentOnce } from '../../otel/instrument'; +import { instrumentWhenWrapped, generateInstrumentOnce } from '../../otel/instrument'; const INTEGRATION_NAME = 'Dataloader'; @@ -20,20 +19,19 @@ export const instrumentDataloader = generateInstrumentOnce( ); const _dataloaderIntegration = (() => { - let hookCallback: undefined | (() => void); + let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); return { name: INTEGRATION_NAME, setupOnce() { const instrumentation = instrumentDataloader(); + instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); + }, - callWhenWrapped(instrumentation, () => { - const client = getClient(); - if (hookCallback || !client) { - return; - } - - hookCallback = client.on('spanStart', span => { + setup(client) { + // This is called either immediately or when the instrumentation is wrapped + instrumentationWrappedCallback?.(() => { + client.on('spanStart', span => { const spanJSON = spanToJSON(span); if (spanJSON.description?.startsWith('dataloader')) { span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.dataloader'); diff --git a/packages/node/src/integrations/tracing/genericPool.ts b/packages/node/src/integrations/tracing/genericPool.ts index 6c4169d66b99..d6f736c6f9a5 100644 --- a/packages/node/src/integrations/tracing/genericPool.ts +++ b/packages/node/src/integrations/tracing/genericPool.ts @@ -1,33 +1,38 @@ import { GenericPoolInstrumentation } from '@opentelemetry/instrumentation-generic-pool'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, defineIntegration, spanToJSON } from '@sentry/core'; import type { IntegrationFn } from '@sentry/core'; -import { generateInstrumentOnce } from '../../otel/instrument'; +import { generateInstrumentOnce, instrumentWhenWrapped } from '../../otel/instrument'; const INTEGRATION_NAME = 'GenericPool'; export const instrumentGenericPool = generateInstrumentOnce(INTEGRATION_NAME, () => new GenericPoolInstrumentation({})); const _genericPoolIntegration = (() => { + let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); + return { name: INTEGRATION_NAME, setupOnce() { - instrumentGenericPool(); + const instrumentation = instrumentGenericPool(); + instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); }, setup(client) { - client.on('spanStart', span => { - const spanJSON = spanToJSON(span); + instrumentationWrappedCallback?.(() => + client.on('spanStart', span => { + const spanJSON = spanToJSON(span); - const spanDescription = spanJSON.description; + const spanDescription = spanJSON.description; - // typo in emitted span for version <= 0.38.0 of @opentelemetry/instrumentation-generic-pool - const isGenericPoolSpan = - spanDescription === 'generic-pool.aquire' || spanDescription === 'generic-pool.acquire'; + // typo in emitted span for version <= 0.38.0 of @opentelemetry/instrumentation-generic-pool + const isGenericPoolSpan = + spanDescription === 'generic-pool.aquire' || spanDescription === 'generic-pool.acquire'; - if (isGenericPoolSpan) { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.generic_pool'); - } - }); + if (isGenericPoolSpan) { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.generic_pool'); + } + }), + ); }, }; }) satisfies IntegrationFn; diff --git a/packages/node/src/integrations/tracing/knex.ts b/packages/node/src/integrations/tracing/knex.ts index 86c728d115bd..da70d8b3c8b7 100644 --- a/packages/node/src/integrations/tracing/knex.ts +++ b/packages/node/src/integrations/tracing/knex.ts @@ -1,7 +1,7 @@ import { KnexInstrumentation } from '@opentelemetry/instrumentation-knex'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, defineIntegration, spanToJSON } from '@sentry/core'; import type { IntegrationFn } from '@sentry/core'; -import { generateInstrumentOnce } from '../../otel/instrument'; +import { generateInstrumentOnce, instrumentWhenWrapped } from '../../otel/instrument'; const INTEGRATION_NAME = 'Knex'; @@ -11,21 +11,26 @@ export const instrumentKnex = generateInstrumentOnce( ); const _knexIntegration = (() => { + let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); + return { name: INTEGRATION_NAME, setupOnce() { - instrumentKnex(); + const instrumentation = instrumentKnex(); + instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); }, setup(client) { - client.on('spanStart', span => { - const { data } = spanToJSON(span); - // knex.version is always set in the span data - // https://github.com/open-telemetry/opentelemetry-js-contrib/blob/0309caeafc44ac9cb13a3345b790b01b76d0497d/plugins/node/opentelemetry-instrumentation-knex/src/instrumentation.ts#L138 - if ('knex.version' in data) { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.knex'); - } - }); + instrumentationWrappedCallback?.(() => + client.on('spanStart', span => { + const { data } = spanToJSON(span); + // knex.version is always set in the span data + // https://github.com/open-telemetry/opentelemetry-js-contrib/blob/0309caeafc44ac9cb13a3345b790b01b76d0497d/plugins/node/opentelemetry-instrumentation-knex/src/instrumentation.ts#L138 + if ('knex.version' in data) { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.knex'); + } + }), + ); }, }; }) satisfies IntegrationFn; diff --git a/packages/node/src/integrations/tracing/prisma.ts b/packages/node/src/integrations/tracing/prisma.ts index 58516671c9a3..2e1a08bb9ca6 100644 --- a/packages/node/src/integrations/tracing/prisma.ts +++ b/packages/node/src/integrations/tracing/prisma.ts @@ -2,7 +2,7 @@ import type { Instrumentation } from '@opentelemetry/instrumentation'; // When importing CJS modules into an ESM module, we cannot import the named exports directly. import * as prismaInstrumentation from '@prisma/instrumentation'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, consoleSandbox, defineIntegration, spanToJSON } from '@sentry/core'; -import { generateInstrumentOnce } from '../../otel/instrument'; +import { generateInstrumentOnce, instrumentWhenWrapped } from '../../otel/instrument'; import type { PrismaV5TracingHelper } from './prisma/vendor/v5-tracing-helper'; import type { PrismaV6TracingHelper } from './prisma/vendor/v6-tracing-helper'; @@ -113,29 +113,34 @@ export const prismaIntegration = defineIntegration( */ prismaInstrumentation?: Instrumentation; } = {}) => { + let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); + return { name: INTEGRATION_NAME, setupOnce() { - instrumentPrisma({ prismaInstrumentation }); + const instrumentation = instrumentPrisma({ prismaInstrumentation }); + instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); }, setup(client) { - client.on('spanStart', span => { - const spanJSON = spanToJSON(span); - if (spanJSON.description?.startsWith('prisma:')) { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.prisma'); - } + instrumentationWrappedCallback?.(() => + client.on('spanStart', span => { + const spanJSON = spanToJSON(span); + if (spanJSON.description?.startsWith('prisma:')) { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.prisma'); + } - // Make sure we use the query text as the span name, for ex. SELECT * FROM "User" WHERE "id" = $1 - if (spanJSON.description === 'prisma:engine:db_query' && spanJSON.data['db.query.text']) { - span.updateName(spanJSON.data['db.query.text'] as string); - } + // Make sure we use the query text as the span name, for ex. SELECT * FROM "User" WHERE "id" = $1 + if (spanJSON.description === 'prisma:engine:db_query' && spanJSON.data['db.query.text']) { + span.updateName(spanJSON.data['db.query.text'] as string); + } - // In Prisma v5.22+, the `db.system` attribute is automatically set - // On older versions, this is missing, so we add it here - if (spanJSON.description === 'prisma:engine:db_query' && !spanJSON.data['db.system']) { - span.setAttribute('db.system', 'prisma'); - } - }); + // In Prisma v5.22+, the `db.system` attribute is automatically set + // On older versions, this is missing, so we add it here + if (spanJSON.description === 'prisma:engine:db_query' && !spanJSON.data['db.system']) { + span.setAttribute('db.system', 'prisma'); + } + }), + ); }, }; }, diff --git a/packages/node/src/integrations/tracing/tedious.ts b/packages/node/src/integrations/tracing/tedious.ts index 644601986a7e..dd74c095a566 100644 --- a/packages/node/src/integrations/tracing/tedious.ts +++ b/packages/node/src/integrations/tracing/tedious.ts @@ -1,7 +1,7 @@ import { TediousInstrumentation } from '@opentelemetry/instrumentation-tedious'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, defineIntegration, spanToJSON } from '@sentry/core'; import type { IntegrationFn } from '@sentry/core'; -import { generateInstrumentOnce } from '../../otel/instrument'; +import { generateInstrumentOnce, instrumentWhenWrapped } from '../../otel/instrument'; const TEDIUS_INSTRUMENTED_METHODS = new Set([ 'callProcedure', @@ -17,25 +17,30 @@ const INTEGRATION_NAME = 'Tedious'; export const instrumentTedious = generateInstrumentOnce(INTEGRATION_NAME, () => new TediousInstrumentation({})); const _tediousIntegration = (() => { + let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); + return { name: INTEGRATION_NAME, setupOnce() { - instrumentTedious(); + const instrumentation = instrumentTedious(); + instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); }, setup(client) { - client.on('spanStart', span => { - const { description, data } = spanToJSON(span); - // Tedius integration always set a span name and `db.system` attribute to `mssql`. - if (!description || data['db.system'] !== 'mssql') { - return; - } - - const operation = description.split(' ')[0] || ''; - if (TEDIUS_INSTRUMENTED_METHODS.has(operation)) { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.tedious'); - } - }); + instrumentationWrappedCallback?.(() => + client.on('spanStart', span => { + const { description, data } = spanToJSON(span); + // Tedius integration always set a span name and `db.system` attribute to `mssql`. + if (!description || data['db.system'] !== 'mssql') { + return; + } + + const operation = description.split(' ')[0] || ''; + if (TEDIUS_INSTRUMENTED_METHODS.has(operation)) { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.tedious'); + } + }), + ); }, }; }) satisfies IntegrationFn; diff --git a/packages/node/src/integrations/tracing/vercelai/index.ts b/packages/node/src/integrations/tracing/vercelai/index.ts index a8198a8458ec..78d28fcba97d 100644 --- a/packages/node/src/integrations/tracing/vercelai/index.ts +++ b/packages/node/src/integrations/tracing/vercelai/index.ts @@ -1,7 +1,7 @@ /* eslint-disable complexity */ import { SEMANTIC_ATTRIBUTE_SENTRY_OP, defineIntegration, spanToJSON } from '@sentry/core'; import type { IntegrationFn } from '@sentry/core'; -import { generateInstrumentOnce } from '../../../otel/instrument'; +import { generateInstrumentOnce, instrumentWhenWrapped } from '../../../otel/instrument'; import { addOriginToSpan } from '../../../utils/addOriginToSpan'; import { SentryVercelAiInstrumentation, sentryVercelAiPatched } from './instrumentation'; @@ -10,145 +10,151 @@ const INTEGRATION_NAME = 'VercelAI'; export const instrumentVercelAi = generateInstrumentOnce(INTEGRATION_NAME, () => new SentryVercelAiInstrumentation({})); const _vercelAIIntegration = (() => { + let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); + return { name: INTEGRATION_NAME, setupOnce() { - instrumentVercelAi(); - }, - processEvent(event) { - if (event.type === 'transaction' && event.spans?.length) { - for (const span of event.spans) { - const { data: attributes, description: name } = span; - - if (!name || span.origin !== 'auto.vercelai.otel') { - continue; - } - - if (attributes['ai.usage.completionTokens'] != undefined) { - attributes['ai.completion_tokens.used'] = attributes['ai.usage.completionTokens']; - } - if (attributes['ai.usage.promptTokens'] != undefined) { - attributes['ai.prompt_tokens.used'] = attributes['ai.usage.promptTokens']; - } - if ( - typeof attributes['ai.usage.completionTokens'] == 'number' && - typeof attributes['ai.usage.promptTokens'] == 'number' - ) { - attributes['ai.total_tokens.used'] = - attributes['ai.usage.completionTokens'] + attributes['ai.usage.promptTokens']; - } - } - } - - return event; + const instrumentation = instrumentVercelAi(); + instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); }, setup(client) { - client.on('spanStart', span => { - if (!sentryVercelAiPatched) { - return; - } - - const { data: attributes, description: name } = spanToJSON(span); - - if (!name) { - return; - } - - // The id of the model - const aiModelId = attributes['ai.model.id']; - - // the provider of the model - const aiModelProvider = attributes['ai.model.provider']; - - // both of these must be defined for the integration to work - if (typeof aiModelId !== 'string' || typeof aiModelProvider !== 'string' || !aiModelId || !aiModelProvider) { - return; - } - - let isPipelineSpan = false; - - switch (name) { - case 'ai.generateText': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.generateText'); - isPipelineSpan = true; - break; - } - case 'ai.generateText.doGenerate': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run.doGenerate'); - break; - } - case 'ai.streamText': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.streamText'); - isPipelineSpan = true; - break; - } - case 'ai.streamText.doStream': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run.doStream'); - break; - } - case 'ai.generateObject': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.generateObject'); - isPipelineSpan = true; - break; - } - case 'ai.generateObject.doGenerate': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run.doGenerate'); - break; - } - case 'ai.streamObject': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.streamObject'); - isPipelineSpan = true; - break; - } - case 'ai.streamObject.doStream': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run.doStream'); - break; - } - case 'ai.embed': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.embed'); - isPipelineSpan = true; - break; - } - case 'ai.embed.doEmbed': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.embeddings'); - break; - } - case 'ai.embedMany': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.embedMany'); - isPipelineSpan = true; - break; - } - case 'ai.embedMany.doEmbed': { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.embeddings'); - break; - } - case 'ai.toolCall': - case 'ai.stream.firstChunk': - case 'ai.stream.finish': - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run'); - break; - } - - addOriginToSpan(span, 'auto.vercelai.otel'); - - const nameWthoutAi = name.replace('ai.', ''); - span.setAttribute('ai.pipeline.name', nameWthoutAi); - span.updateName(nameWthoutAi); - - // If a Telemetry name is set and it is a pipeline span, use that as the operation name - const functionId = attributes['ai.telemetry.functionId']; - if (functionId && typeof functionId === 'string' && isPipelineSpan) { - span.updateName(functionId); - span.setAttribute('ai.pipeline.name', functionId); - } - - if (attributes['ai.prompt']) { - span.setAttribute('ai.input_messages', attributes['ai.prompt']); - } - if (attributes['ai.model.id']) { - span.setAttribute('ai.model_id', attributes['ai.model.id']); - } - span.setAttribute('ai.streaming', name.includes('stream')); + instrumentationWrappedCallback?.(() => { + client.on('spanStart', span => { + if (!sentryVercelAiPatched) { + return; + } + + const { data: attributes, description: name } = spanToJSON(span); + + if (!name) { + return; + } + + // The id of the model + const aiModelId = attributes['ai.model.id']; + + // the provider of the model + const aiModelProvider = attributes['ai.model.provider']; + + // both of these must be defined for the integration to work + if (typeof aiModelId !== 'string' || typeof aiModelProvider !== 'string' || !aiModelId || !aiModelProvider) { + return; + } + + let isPipelineSpan = false; + + switch (name) { + case 'ai.generateText': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.generateText'); + isPipelineSpan = true; + break; + } + case 'ai.generateText.doGenerate': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run.doGenerate'); + break; + } + case 'ai.streamText': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.streamText'); + isPipelineSpan = true; + break; + } + case 'ai.streamText.doStream': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run.doStream'); + break; + } + case 'ai.generateObject': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.generateObject'); + isPipelineSpan = true; + break; + } + case 'ai.generateObject.doGenerate': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run.doGenerate'); + break; + } + case 'ai.streamObject': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.streamObject'); + isPipelineSpan = true; + break; + } + case 'ai.streamObject.doStream': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run.doStream'); + break; + } + case 'ai.embed': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.embed'); + isPipelineSpan = true; + break; + } + case 'ai.embed.doEmbed': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.embeddings'); + break; + } + case 'ai.embedMany': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.pipeline.embedMany'); + isPipelineSpan = true; + break; + } + case 'ai.embedMany.doEmbed': { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.embeddings'); + break; + } + case 'ai.toolCall': + case 'ai.stream.firstChunk': + case 'ai.stream.finish': + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run'); + break; + } + + addOriginToSpan(span, 'auto.vercelai.otel'); + + const nameWthoutAi = name.replace('ai.', ''); + span.setAttribute('ai.pipeline.name', nameWthoutAi); + span.updateName(nameWthoutAi); + + // If a Telemetry name is set and it is a pipeline span, use that as the operation name + const functionId = attributes['ai.telemetry.functionId']; + if (functionId && typeof functionId === 'string' && isPipelineSpan) { + span.updateName(functionId); + span.setAttribute('ai.pipeline.name', functionId); + } + + if (attributes['ai.prompt']) { + span.setAttribute('ai.input_messages', attributes['ai.prompt']); + } + if (attributes['ai.model.id']) { + span.setAttribute('ai.model_id', attributes['ai.model.id']); + } + span.setAttribute('ai.streaming', name.includes('stream')); + }); + + client.addEventProcessor(event => { + if (event.type === 'transaction' && event.spans?.length) { + for (const span of event.spans) { + const { data: attributes, description: name } = span; + + if (!name || span.origin !== 'auto.vercelai.otel') { + continue; + } + + if (attributes['ai.usage.completionTokens'] != undefined) { + attributes['ai.completion_tokens.used'] = attributes['ai.usage.completionTokens']; + } + if (attributes['ai.usage.promptTokens'] != undefined) { + attributes['ai.prompt_tokens.used'] = attributes['ai.usage.promptTokens']; + } + if ( + typeof attributes['ai.usage.completionTokens'] == 'number' && + typeof attributes['ai.usage.promptTokens'] == 'number' + ) { + attributes['ai.total_tokens.used'] = + attributes['ai.usage.completionTokens'] + attributes['ai.usage.promptTokens']; + } + } + } + + return event; + }); }); }, }; diff --git a/packages/node/src/otel/instrument.ts b/packages/node/src/otel/instrument.ts index cd7b7f0c8583..6f8b10db2ba7 100644 --- a/packages/node/src/otel/instrument.ts +++ b/packages/node/src/otel/instrument.ts @@ -41,21 +41,41 @@ export function generateInstrumentOnce< /** * Ensure a given callback is called when the instrumentation is actually wrapping something. * This can be used to ensure some logic is only called when the instrumentation is actually active. - * This depends on wrapping `_wrap` (inception!). If this is not possible (e.g. the property name is mangled, ...) - * the callback will be called immediately. + * + * This function returns a function that can be invoked with a callback. + * This callback will either be invoked immediately + * (e.g. if the instrumentation was already wrapped, or if _wrap could not be patched), + * or once the instrumentation is actually wrapping something. + * + * Make sure to call this function right after adding the instrumentation, otherwise it may be too late! + * The returned callback can be used any time, and also multiple times. */ -export function callWhenWrapped(instrumentation: T, callback: () => void): void { +export function instrumentWhenWrapped(instrumentation: T): (callback: () => void) => void { + let isWrapped = false; + let callbacks: (() => void)[] = []; + if (!hasWrap(instrumentation)) { - callback(); - return; - } + isWrapped = true; + } else { + const originalWrap = instrumentation['_wrap']; - const originalWrap = instrumentation['_wrap']; + instrumentation['_wrap'] = (...args: Parameters) => { + isWrapped = true; + callbacks.forEach(callback => callback()); + callbacks = []; + return originalWrap(...args); + }; + } - instrumentation['_wrap'] = (...args: Parameters) => { - callback(); - return originalWrap(...args); + const registerCallback = (callback: () => void): void => { + if (isWrapped) { + callback(); + } else { + callbacks.push(callback); + } }; + + return registerCallback; } function hasWrap( diff --git a/packages/node/test/utils/instrument.test.ts b/packages/node/test/utils/instrument.test.ts new file mode 100644 index 000000000000..6d905b715bc0 --- /dev/null +++ b/packages/node/test/utils/instrument.test.ts @@ -0,0 +1,95 @@ +import { describe, test, vi, expect } from 'vitest'; +import { instrumentWhenWrapped } from '../../src/otel/instrument'; + +describe('instrumentWhenWrapped', () => { + test('calls callback immediately when instrumentation has no _wrap method', () => { + const callback = vi.fn(); + const instrumentation = {} as any; + + const registerCallback = instrumentWhenWrapped(instrumentation); + registerCallback(callback); + + expect(callback).toHaveBeenCalledTimes(1); + }); + + test('calls callback when _wrap is called', () => { + const callback = vi.fn(); + const originalWrap = vi.fn(); + const instrumentation = { + _wrap: originalWrap, + } as any; + + const registerCallback = instrumentWhenWrapped(instrumentation); + registerCallback(callback); + + // Callback should not be called yet + expect(callback).not.toHaveBeenCalled(); + + // Call _wrap + instrumentation._wrap(); + + // Callback should be called once + expect(callback).toHaveBeenCalledTimes(1); + expect(originalWrap).toHaveBeenCalled(); + }); + + test('calls multiple callbacks when _wrap is called', () => { + const callback1 = vi.fn(); + const callback2 = vi.fn(); + const originalWrap = vi.fn(); + const instrumentation = { + _wrap: originalWrap, + } as any; + + const registerCallback = instrumentWhenWrapped(instrumentation); + registerCallback(callback1); + registerCallback(callback2); + + // Callbacks should not be called yet + expect(callback1).not.toHaveBeenCalled(); + expect(callback2).not.toHaveBeenCalled(); + + // Call _wrap + instrumentation._wrap(); + + // Both callbacks should be called once + expect(callback1).toHaveBeenCalledTimes(1); + expect(callback2).toHaveBeenCalledTimes(1); + expect(originalWrap).toHaveBeenCalled(); + }); + + test('calls callback immediately if already wrapped', () => { + const callback = vi.fn(); + const originalWrap = vi.fn(); + const instrumentation = { + _wrap: originalWrap, + } as any; + + // Call _wrap first + instrumentation._wrap(); + + const registerCallback = instrumentWhenWrapped(instrumentation); + registerCallback(callback); + + // Callback should be called immediately + expect(callback).toHaveBeenCalledTimes(1); + expect(originalWrap).toHaveBeenCalled(); + }); + + test('passes through arguments to original _wrap', () => { + const callback = vi.fn(); + const originalWrap = vi.fn(); + const instrumentation = { + _wrap: originalWrap, + } as any; + + const registerCallback = instrumentWhenWrapped(instrumentation); + registerCallback(callback); + + // Call _wrap with arguments + const args = ['arg1', 'arg2']; + instrumentation._wrap(...args); + + expect(originalWrap).toHaveBeenCalledWith(...args); + }); +}); From d0af51d2b93b6be1e6a429f51aed83250d774153 Mon Sep 17 00:00:00 2001 From: Francesco Novy Date: Tue, 25 Mar 2025 11:50:48 +0100 Subject: [PATCH 3/4] fixes --- .../node/src/integrations/tracing/prisma.ts | 69 ++++++++++--------- packages/node/test/utils/instrument.test.ts | 3 +- 2 files changed, 40 insertions(+), 32 deletions(-) diff --git a/packages/node/src/integrations/tracing/prisma.ts b/packages/node/src/integrations/tracing/prisma.ts index 2e1a08bb9ca6..06b75f0c707e 100644 --- a/packages/node/src/integrations/tracing/prisma.ts +++ b/packages/node/src/integrations/tracing/prisma.ts @@ -2,7 +2,7 @@ import type { Instrumentation } from '@opentelemetry/instrumentation'; // When importing CJS modules into an ESM module, we cannot import the named exports directly. import * as prismaInstrumentation from '@prisma/instrumentation'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, consoleSandbox, defineIntegration, spanToJSON } from '@sentry/core'; -import { generateInstrumentOnce, instrumentWhenWrapped } from '../../otel/instrument'; +import { generateInstrumentOnce } from '../../otel/instrument'; import type { PrismaV5TracingHelper } from './prisma/vendor/v5-tracing-helper'; import type { PrismaV6TracingHelper } from './prisma/vendor/v6-tracing-helper'; @@ -19,6 +19,18 @@ function isPrismaV6TracingHelper(helper: unknown): helper is PrismaV6TracingHelp return !!helper && typeof helper === 'object' && 'dispatchEngineSpans' in helper; } +function getPrismaTracingHelper(): unknown | undefined { + const prismaInstrumentationObject = (globalThis as Record).PRISMA_INSTRUMENTATION; + const prismaTracingHelper = + prismaInstrumentationObject && + typeof prismaInstrumentationObject === 'object' && + 'helper' in prismaInstrumentationObject + ? prismaInstrumentationObject.helper + : undefined; + + return prismaTracingHelper; +} + class SentryPrismaInteropInstrumentation extends EsmInteropPrismaInstrumentation { public constructor() { super(); @@ -30,13 +42,7 @@ class SentryPrismaInteropInstrumentation extends EsmInteropPrismaInstrumentation // The PrismaIntegration (super class) defines a global variable `global["PRISMA_INSTRUMENTATION"]` when `enable()` is called. This global variable holds a "TracingHelper" which Prisma uses internally to create tracing data. It's their way of not depending on OTEL with their main package. The sucky thing is, prisma broke the interface of the tracing helper with the v6 major update. This means that if you use Prisma 5 with the v6 instrumentation (or vice versa) Prisma just blows up, because tries to call methods on the helper that no longer exist. // Because we actually want to use the v6 instrumentation and not blow up in Prisma 5 user's faces, what we're doing here is backfilling the v5 method (`createEngineSpan`) with a noop so that no longer crashes when it attempts to call that function. // We still won't fully emit all the spans, but this could potentially be implemented in the future. - const prismaInstrumentationObject = (globalThis as Record).PRISMA_INSTRUMENTATION; - const prismaTracingHelper = - prismaInstrumentationObject && - typeof prismaInstrumentationObject === 'object' && - 'helper' in prismaInstrumentationObject - ? prismaInstrumentationObject.helper - : undefined; + const prismaTracingHelper = getPrismaTracingHelper(); let emittedWarning = false; @@ -113,34 +119,35 @@ export const prismaIntegration = defineIntegration( */ prismaInstrumentation?: Instrumentation; } = {}) => { - let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); - return { name: INTEGRATION_NAME, setupOnce() { - const instrumentation = instrumentPrisma({ prismaInstrumentation }); - instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); + instrumentPrisma({ prismaInstrumentation }); }, setup(client) { - instrumentationWrappedCallback?.(() => - client.on('spanStart', span => { - const spanJSON = spanToJSON(span); - if (spanJSON.description?.startsWith('prisma:')) { - span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.prisma'); - } - - // Make sure we use the query text as the span name, for ex. SELECT * FROM "User" WHERE "id" = $1 - if (spanJSON.description === 'prisma:engine:db_query' && spanJSON.data['db.query.text']) { - span.updateName(spanJSON.data['db.query.text'] as string); - } - - // In Prisma v5.22+, the `db.system` attribute is automatically set - // On older versions, this is missing, so we add it here - if (spanJSON.description === 'prisma:engine:db_query' && !spanJSON.data['db.system']) { - span.setAttribute('db.system', 'prisma'); - } - }), - ); + // If no tracing helper exists, we skip any work here + // this means that prisma is not being used + if (!getPrismaTracingHelper()) { + return; + } + + client.on('spanStart', span => { + const spanJSON = spanToJSON(span); + if (spanJSON.description?.startsWith('prisma:')) { + span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.prisma'); + } + + // Make sure we use the query text as the span name, for ex. SELECT * FROM "User" WHERE "id" = $1 + if (spanJSON.description === 'prisma:engine:db_query' && spanJSON.data['db.query.text']) { + span.updateName(spanJSON.data['db.query.text'] as string); + } + + // In Prisma v5.22+, the `db.system` attribute is automatically set + // On older versions, this is missing, so we add it here + if (spanJSON.description === 'prisma:engine:db_query' && !spanJSON.data['db.system']) { + span.setAttribute('db.system', 'prisma'); + } + }); }, }; }, diff --git a/packages/node/test/utils/instrument.test.ts b/packages/node/test/utils/instrument.test.ts index 6d905b715bc0..3cc7fedb4d22 100644 --- a/packages/node/test/utils/instrument.test.ts +++ b/packages/node/test/utils/instrument.test.ts @@ -65,10 +65,11 @@ describe('instrumentWhenWrapped', () => { _wrap: originalWrap, } as any; + const registerCallback = instrumentWhenWrapped(instrumentation); + // Call _wrap first instrumentation._wrap(); - const registerCallback = instrumentWhenWrapped(instrumentation); registerCallback(callback); // Callback should be called immediately From 4f540e507f08547ba0a123b8211bbafc4cde27fd Mon Sep 17 00:00:00 2001 From: Francesco Novy Date: Tue, 25 Mar 2025 12:18:43 +0100 Subject: [PATCH 4/4] fix it for ai --- .../integrations/tracing/vercelai/index.ts | 15 +++++-------- .../tracing/vercelai/instrumentation.ts | 22 ++++++++++++++++--- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/packages/node/src/integrations/tracing/vercelai/index.ts b/packages/node/src/integrations/tracing/vercelai/index.ts index 78d28fcba97d..26bf57ed38b2 100644 --- a/packages/node/src/integrations/tracing/vercelai/index.ts +++ b/packages/node/src/integrations/tracing/vercelai/index.ts @@ -1,30 +1,25 @@ /* eslint-disable complexity */ import { SEMANTIC_ATTRIBUTE_SENTRY_OP, defineIntegration, spanToJSON } from '@sentry/core'; import type { IntegrationFn } from '@sentry/core'; -import { generateInstrumentOnce, instrumentWhenWrapped } from '../../../otel/instrument'; +import { generateInstrumentOnce } from '../../../otel/instrument'; import { addOriginToSpan } from '../../../utils/addOriginToSpan'; -import { SentryVercelAiInstrumentation, sentryVercelAiPatched } from './instrumentation'; +import { SentryVercelAiInstrumentation } from './instrumentation'; const INTEGRATION_NAME = 'VercelAI'; export const instrumentVercelAi = generateInstrumentOnce(INTEGRATION_NAME, () => new SentryVercelAiInstrumentation({})); const _vercelAIIntegration = (() => { - let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); + let instrumentation: undefined | SentryVercelAiInstrumentation; return { name: INTEGRATION_NAME, setupOnce() { - const instrumentation = instrumentVercelAi(); - instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); + instrumentation = instrumentVercelAi(); }, setup(client) { - instrumentationWrappedCallback?.(() => { + instrumentation?.callWhenPatched(() => { client.on('spanStart', span => { - if (!sentryVercelAiPatched) { - return; - } - const { data: attributes, description: name } = spanToJSON(span); if (!name) { diff --git a/packages/node/src/integrations/tracing/vercelai/instrumentation.ts b/packages/node/src/integrations/tracing/vercelai/instrumentation.ts index 97721eaee15d..2cddd175c0d8 100644 --- a/packages/node/src/integrations/tracing/vercelai/instrumentation.ts +++ b/packages/node/src/integrations/tracing/vercelai/instrumentation.ts @@ -23,8 +23,6 @@ type MethodArgs = [MethodFirstArg, ...unknown[]]; type PatchedModuleExports = Record<(typeof INSTRUMENTED_METHODS)[number], (...args: MethodArgs) => unknown> & Record; -export let sentryVercelAiPatched = false; - /** * This detects is added by the Sentry Vercel AI Integration to detect if the integration should * be enabled. @@ -32,6 +30,9 @@ export let sentryVercelAiPatched = false; * It also patches the `ai` module to enable Vercel AI telemetry automatically for all methods. */ export class SentryVercelAiInstrumentation extends InstrumentationBase { + private _isPatched = false; + private _callbacks: (() => void)[] = []; + public constructor(config: InstrumentationConfig = {}) { super('@sentry/instrumentation-vercel-ai', SDK_VERSION, config); } @@ -44,11 +45,26 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase { return module; } + /** + * Call the provided callback when the module is patched. + * If it has already been patched, the callback will be called immediately. + */ + public callWhenPatched(callback: () => void): void { + if (this._isPatched) { + callback(); + } else { + this._callbacks.push(callback); + } + } + /** * Patches module exports to enable Vercel AI telemetry. */ private _patch(moduleExports: PatchedModuleExports): unknown { - sentryVercelAiPatched = true; + this._isPatched = true; + + this._callbacks.forEach(callback => callback()); + this._callbacks = []; function generatePatch(name: string) { return (...args: MethodArgs) => {