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..3c887bb33fa5 100644 --- a/packages/node/src/integrations/tracing/dataloader.ts +++ b/packages/node/src/integrations/tracing/dataloader.ts @@ -6,7 +6,7 @@ import { spanToJSON, } from '@sentry/core'; import type { IntegrationFn } from '@sentry/core'; -import { generateInstrumentOnce } from '../../otel/instrument'; +import { instrumentWhenWrapped, generateInstrumentOnce } from '../../otel/instrument'; const INTEGRATION_NAME = 'Dataloader'; @@ -19,31 +19,37 @@ export const instrumentDataloader = generateInstrumentOnce( ); const _dataloaderIntegration = (() => { + let instrumentationWrappedCallback: undefined | ((callback: () => void) => void); + return { name: INTEGRATION_NAME, setupOnce() { - instrumentDataloader(); + const instrumentation = instrumentDataloader(); + instrumentationWrappedCallback = instrumentWhenWrapped(instrumentation); }, 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'); - } + // 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'); + } - // 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. - } + // 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/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..06b75f0c707e 100644 --- a/packages/node/src/integrations/tracing/prisma.ts +++ b/packages/node/src/integrations/tracing/prisma.ts @@ -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; @@ -119,6 +125,12 @@ export const prismaIntegration = defineIntegration( instrumentPrisma({ prismaInstrumentation }); }, setup(client) { + // 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:')) { 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..26bf57ed38b2 100644 --- a/packages/node/src/integrations/tracing/vercelai/index.ts +++ b/packages/node/src/integrations/tracing/vercelai/index.ts @@ -3,152 +3,153 @@ import { SEMANTIC_ATTRIBUTE_SENTRY_OP, defineIntegration, spanToJSON } from '@se import type { IntegrationFn } from '@sentry/core'; 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 instrumentation: undefined | SentryVercelAiInstrumentation; + 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; + instrumentation = instrumentVercelAi(); }, 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')); + instrumentation?.callWhenPatched(() => { + client.on('spanStart', span => { + 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/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) => { diff --git a/packages/node/src/otel/instrument.ts b/packages/node/src/otel/instrument.ts index c5bd7500a739..6f8b10db2ba7 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,55 @@ 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 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 instrumentWhenWrapped(instrumentation: T): (callback: () => void) => void { + let isWrapped = false; + let callbacks: (() => void)[] = []; + + if (!hasWrap(instrumentation)) { + isWrapped = true; + } else { + const originalWrap = instrumentation['_wrap']; + + instrumentation['_wrap'] = (...args: Parameters) => { + isWrapped = true; + callbacks.forEach(callback => callback()); + callbacks = []; + return originalWrap(...args); + }; + } + + const registerCallback = (callback: () => void): void => { + if (isWrapped) { + callback(); + } else { + callbacks.push(callback); + } + }; + + return registerCallback; +} + +function hasWrap( + instrumentation: T, +): instrumentation is T & { _wrap: (...args: unknown[]) => unknown } { + return typeof (instrumentation as T & { _wrap?: (...args: unknown[]) => unknown })['_wrap'] === 'function'; +} diff --git a/packages/node/test/utils/instrument.test.ts b/packages/node/test/utils/instrument.test.ts new file mode 100644 index 000000000000..3cc7fedb4d22 --- /dev/null +++ b/packages/node/test/utils/instrument.test.ts @@ -0,0 +1,96 @@ +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; + + const registerCallback = instrumentWhenWrapped(instrumentation); + + // Call _wrap first + instrumentation._wrap(); + + 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); + }); +});