fix it for ai

Author: mydea

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

packages/node/src/integrations/tracing/vercelai/instrumentation.ts

@@ -23,15 +23,16 @@ type MethodArgs = [MethodFirstArg, ...unknown[]];
 type PatchedModuleExports = Record<(typeof INSTRUMENTED_METHODS)[number], (...args: MethodArgs) => unknown> &
   Record<string, unknown>;
 
-export let sentryVercelAiPatched = false;
-
 /**
  * This detects is added by the Sentry Vercel AI Integration to detect if the integration should
  * be enabled.
  *
  * 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) => {