feat(core): clear up integrations on dispose

Author: JPeer264

packages/core/src/client.ts

@@ -211,6 +211,9 @@ export abstract class Client<O extends ClientOptions = ClientOptions> {
 
   protected _promiseBuffer: PromiseBuffer<unknown>;
 
+  /** Cleanup functions to call on dispose */
+  protected _disposeCallbacks: (() => void)[];
+
   /**
    * Initializes this client instance.
    *
@@ -223,6 +226,7 @@ export abstract class Client<O extends ClientOptions = ClientOptions> {
     this._outcomes = {};
     this._hooks = {};
     this._eventProcessors = [];
+    this._disposeCallbacks = [];
     this._promiseBuffer = makePromiseBuffer(options.transportOptions?.bufferSize ?? DEFAULT_TRANSPORT_BUFFER_SIZE);
 
     if (options.dsn) {
@@ -1145,6 +1149,14 @@ export abstract class Client<O extends ClientOptions = ClientOptions> {
     return {};
   }
 
+  /**
+   * Register a cleanup function to be called when the client is disposed.
+   * This is useful for integrations that need to clean up global state.
+   */
+  public registerCleanup(callback: () => void): void {
+    this._disposeCallbacks.push(callback);
+  }
+
   /**
    * Disposes of the client and releases all resources.
    *

packages/core/src/instrument/console.ts

@@ -8,14 +8,16 @@ import { addHandler, maybeInstrument, triggerHandlers } from './handlers';
 
 /**
  * Add an instrumentation handler for when a console.xxx method is called.
+ * Returns a function to remove the handler.
  *
  * Use at your own risk, this might break without changelog notice, only used internally.
  * @hidden
  */
-export function addConsoleInstrumentationHandler(handler: (data: HandlerDataConsole) => void): void {
+export function addConsoleInstrumentationHandler(handler: (data: HandlerDataConsole) => void): () => void {
   const type = 'console';
-  addHandler(type, handler);
+  const removeHandler = addHandler(type, handler);
   maybeInstrument(type, instrumentConsole);
+  return removeHandler;
 }
 
 function instrumentConsole(): void {

packages/core/src/instrument/fetch.ts

@@ -15,31 +15,35 @@ type FetchResource = string | { toString(): string } | { url: string };
  * Add an instrumentation handler for when a fetch request happens.
  * The handler function is called once when the request starts and once when it ends,
  * which can be identified by checking if it has an `endTimestamp`.
+ * Returns a function to remove the handler.
  *
  * Use at your own risk, this might break without changelog notice, only used internally.
  * @hidden
  */
 export function addFetchInstrumentationHandler(
   handler: (data: HandlerDataFetch) => void,
   skipNativeFetchCheck?: boolean,
-): void {
+): () => void {
   const type = 'fetch';
-  addHandler(type, handler);
+  const removeHandler = addHandler(type, handler);
   maybeInstrument(type, () => instrumentFetch(undefined, skipNativeFetchCheck));
+  return removeHandler;
 }
 
 /**
  * Add an instrumentation handler for long-lived fetch requests, like consuming server-sent events (SSE) via fetch.
  * The handler will resolve the request body and emit the actual `endTimestamp`, so that the
  * span can be updated accordingly.
+ * Returns a function to remove the handler.
  *
  * Only used internally
  * @hidden
  */
-export function addFetchEndInstrumentationHandler(handler: (data: HandlerDataFetch) => void): void {
+export function addFetchEndInstrumentationHandler(handler: (data: HandlerDataFetch) => void): () => void {
   const type = 'fetch-body-resolved';
-  addHandler(type, handler);
+  const removeHandler = addHandler(type, handler);
   maybeInstrument(type, () => instrumentFetch(streamHandler));
+  return removeHandler;
 }
 
 function instrumentFetch(onFetchResolved?: (response: Response) => void, skipNativeFetchCheck: boolean = false): void {

packages/core/src/instrument/handlers.ts

@@ -18,10 +18,20 @@ export type InstrumentHandlerCallback = (data: any) => void;
 const handlers: { [key in InstrumentHandlerType]?: InstrumentHandlerCallback[] } = {};
 const instrumented: { [key in InstrumentHandlerType]?: boolean } = {};
 
-/** Add a handler function. */
-export function addHandler(type: InstrumentHandlerType, handler: InstrumentHandlerCallback): void {
+/** Add a handler function. Returns a function to remove the handler. */
+export function addHandler(type: InstrumentHandlerType, handler: InstrumentHandlerCallback): () => void {
   handlers[type] = handlers[type] || [];
   handlers[type].push(handler);
+
+  return () => {
+    const typeHandlers = handlers[type];
+    if (typeHandlers) {
+      const index = typeHandlers.indexOf(handler);
+      if (index !== -1) {
+        typeHandlers.splice(index, 1);
+      }
+    }
+  };
 }
 
 /**

packages/core/src/integrations/console.ts

@@ -41,13 +41,15 @@ export const consoleIntegration = defineIntegration((options: Partial<ConsoleInt
   return {
     name: INTEGRATION_NAME,
     setup(client) {
-      addConsoleInstrumentationHandler(({ args, level }) => {
+      const unsubscribe = addConsoleInstrumentationHandler(({ args, level }) => {
         if (getClient() !== client || !levels.has(level)) {
           return;
         }
 
         addConsoleBreadcrumb(level, args);
       });
+
+      client.registerCleanup(unsubscribe);
     },
   };
 });

packages/core/src/integrations/conversationId.ts

@@ -12,7 +12,7 @@ const _conversationIdIntegration = (() => {
   return {
     name: INTEGRATION_NAME,
     setup(client: Client) {
-      client.on('spanStart', (span: Span) => {
+      const unsubscribe = client.on('spanStart', (span: Span) => {
         const scopeData = getCurrentScope().getScopeData();
         const isolationScopeData = getIsolationScope().getScopeData();
 
@@ -32,6 +32,8 @@ const _conversationIdIntegration = (() => {
           span.setAttribute(GEN_AI_CONVERSATION_ID_ATTRIBUTE, conversationId);
         }
       });
+
+      client.registerCleanup(unsubscribe);
     },
   };
 }) satisfies IntegrationFn;

packages/core/src/logs/console-integration.ts

@@ -31,7 +31,7 @@ const _consoleLoggingIntegration = ((options: Partial<CaptureConsoleOptions> = {
         return;
       }
 
-      addConsoleInstrumentationHandler(({ args, level }) => {
+      const unsubscribe = addConsoleInstrumentationHandler(({ args, level }) => {
         if (getClient() !== client || !levels.includes(level)) {
           return;
         }
@@ -66,6 +66,8 @@ const _consoleLoggingIntegration = ((options: Partial<CaptureConsoleOptions> = {
           attributes,
         });
       });
+
+      client.registerCleanup(unsubscribe);
     },
   };
 }) satisfies IntegrationFn;

packages/core/src/server-runtime-client.ts

@@ -168,6 +168,16 @@ export class ServerRuntimeClient<
   public override dispose(): void {
     DEBUG_BUILD && debug.log('Disposing client...');
 
+    // Run all registered cleanup callbacks
+    for (const callback of this._disposeCallbacks) {
+      try {
+        callback();
+      } catch {
+        // Ignore errors in cleanup callbacks
+      }
+    }
+    this._disposeCallbacks.length = 0;
+
     for (const hookName of Object.keys(this._hooks)) {
       this._hooks[hookName]?.clear();
     }

packages/core/test/lib/instrument/handlers.test.ts

@@ -1,5 +1,14 @@
-import { describe, test } from 'vitest';
-import { maybeInstrument } from '../../../src/instrument/handlers';
+import { afterEach, describe, expect, test, vi } from 'vitest';
+import {
+  addHandler,
+  maybeInstrument,
+  resetInstrumentationHandlers,
+  triggerHandlers,
+} from '../../../src/instrument/handlers';
+
+afterEach(() => {
+  resetInstrumentationHandlers();
+});
 
 describe('maybeInstrument', () => {
   test('does not throw when instrumenting fails', () => {
@@ -12,3 +21,89 @@ describe('maybeInstrument', () => {
     maybeInstrument('xhr', undefined as any);
   });
 });
+
+describe('addHandler', () => {
+  test('returns an unsubscribe function', () => {
+    const handler = vi.fn();
+    const unsubscribe = addHandler('fetch', handler);
+
+    expect(typeof unsubscribe).toBe('function');
+  });
+
+  test('handler is called when triggerHandlers is invoked', () => {
+    const handler = vi.fn();
+    addHandler('fetch', handler);
+
+    triggerHandlers('fetch', { url: 'https://example.com' });
+
+    expect(handler).toHaveBeenCalledTimes(1);
+    expect(handler).toHaveBeenCalledWith({ url: 'https://example.com' });
+  });
+
+  test('unsubscribe removes the handler', () => {
+    const handler = vi.fn();
+    const unsubscribe = addHandler('fetch', handler);
+
+    triggerHandlers('fetch', { test: 1 });
+    expect(handler).toHaveBeenCalledTimes(1);
+
+    unsubscribe();
+
+    triggerHandlers('fetch', { test: 2 });
+    expect(handler).toHaveBeenCalledTimes(1);
+  });
+
+  test('unsubscribe only removes the specific handler', () => {
+    const handler1 = vi.fn();
+    const handler2 = vi.fn();
+
+    const unsubscribe1 = addHandler('fetch', handler1);
+    addHandler('fetch', handler2);
+
+    triggerHandlers('fetch', { test: 1 });
+    expect(handler1).toHaveBeenCalledTimes(1);
+    expect(handler2).toHaveBeenCalledTimes(1);
+
+    unsubscribe1();
+
+    triggerHandlers('fetch', { test: 2 });
+    expect(handler1).toHaveBeenCalledTimes(1);
+    expect(handler2).toHaveBeenCalledTimes(2);
+  });
+
+  test('calling unsubscribe multiple times is safe', () => {
+    const handler = vi.fn();
+    const unsubscribe = addHandler('fetch', handler);
+
+    unsubscribe();
+    expect(() => unsubscribe()).not.toThrow();
+    expect(() => unsubscribe()).not.toThrow();
+  });
+
+  test('unsubscribe works with different handler types', () => {
+    const consoleHandler = vi.fn();
+    const fetchHandler = vi.fn();
+
+    const unsubscribeConsole = addHandler('console', consoleHandler);
+    const unsubscribeFetch = addHandler('fetch', fetchHandler);
+
+    triggerHandlers('console', { level: 'log' });
+    triggerHandlers('fetch', { url: 'test' });
+
+    expect(consoleHandler).toHaveBeenCalledTimes(1);
+    expect(fetchHandler).toHaveBeenCalledTimes(1);
+
+    unsubscribeConsole();
+
+    triggerHandlers('console', { level: 'warn' });
+    triggerHandlers('fetch', { url: 'test2' });
+
+    expect(consoleHandler).toHaveBeenCalledTimes(1);
+    expect(fetchHandler).toHaveBeenCalledTimes(2);
+
+    unsubscribeFetch();
+
+    triggerHandlers('fetch', { url: 'test3' });
+    expect(fetchHandler).toHaveBeenCalledTimes(2);
+  });
+});

packages/core/test/lib/server-runtime-client.test.ts

@@ -320,5 +320,90 @@ describe('ServerRuntimeClient', () => {
       // Verify it's a fresh buffer with no pending items
       expect(bufferAfterDispose.$).toEqual([]);
     });
+
+    it('calls registered cleanup callbacks on dispose', () => {
+      const options = getDefaultClientOptions({ dsn: PUBLIC_DSN });
+      client = new ServerRuntimeClient(options);
+
+      const cleanup1 = vi.fn();
+      const cleanup2 = vi.fn();
+      const cleanup3 = vi.fn();
+
+      client.registerCleanup(cleanup1);
+      client.registerCleanup(cleanup2);
+      client.registerCleanup(cleanup3);
+
+      expect(cleanup1).not.toHaveBeenCalled();
+      expect(cleanup2).not.toHaveBeenCalled();
+      expect(cleanup3).not.toHaveBeenCalled();
+
+      client.dispose();
+
+      expect(cleanup1).toHaveBeenCalledTimes(1);
+      expect(cleanup2).toHaveBeenCalledTimes(1);
+      expect(cleanup3).toHaveBeenCalledTimes(1);
+    });
+
+    it('clears cleanup callbacks after dispose', () => {
+      const options = getDefaultClientOptions({ dsn: PUBLIC_DSN });
+      client = new ServerRuntimeClient(options);
+
+      const cleanup = vi.fn();
+      client.registerCleanup(cleanup);
+
+      client.dispose();
+      expect(cleanup).toHaveBeenCalledTimes(1);
+
+      // Calling dispose again should not call cleanup again
+      client.dispose();
+      expect(cleanup).toHaveBeenCalledTimes(1);
+    });
+
+    it('continues to call other cleanup callbacks if one throws', () => {
+      const options = getDefaultClientOptions({ dsn: PUBLIC_DSN });
+      client = new ServerRuntimeClient(options);
+
+      const cleanup1 = vi.fn();
+      const throwingCleanup = vi.fn(() => {
+        throw new Error('cleanup error');
+      });
+      const cleanup2 = vi.fn();
+
+      client.registerCleanup(cleanup1);
+      client.registerCleanup(throwingCleanup);
+      client.registerCleanup(cleanup2);
+
+      expect(() => client.dispose()).not.toThrow();
+
+      expect(cleanup1).toHaveBeenCalledTimes(1);
+      expect(throwingCleanup).toHaveBeenCalledTimes(1);
+      expect(cleanup2).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('registerCleanup', () => {
+    it('accepts cleanup functions', () => {
+      const options = getDefaultClientOptions({ dsn: PUBLIC_DSN });
+      client = new ServerRuntimeClient(options);
+
+      const cleanup = vi.fn();
+
+      expect(() => client.registerCleanup(cleanup)).not.toThrow();
+    });
+
+    it('can register multiple cleanup functions', () => {
+      const options = getDefaultClientOptions({ dsn: PUBLIC_DSN });
+      client = new ServerRuntimeClient(options);
+
+      const cleanups = Array.from({ length: 10 }, () => vi.fn());
+
+      cleanups.forEach(cleanup => client.registerCleanup(cleanup));
+
+      client.dispose();
+
+      cleanups.forEach(cleanup => {
+        expect(cleanup).toHaveBeenCalledTimes(1);
+      });
+    });
   });
 });