refactor(core): Consolidate LangChain callback merging into mergeSentryCallback

mdnanocom · mdnanocom · commit 0c2c66603309 · 2026-05-13T10:24:39.000+02:00
The langchain instrumentation's `augmentCallbackHandlers` and the
langgraph instrumentation's `mergeSentryCallback` solved the same
problem two different ways. The langgraph helper had three latent
bugs that this PR's fix already covered for langchain:

  - mutated the caller's CallbackManager (handlers accumulate across runs)
  - called `addHandler(handler)` without `inherit=true`, so the handler
    never propagated to child managers created by `getChild`
  - deduped only against `manager.handlers`, not `inheritableHandlers`

Lift the fixed implementation up to `@sentry/core` so both
instrumentations share it. The langgraph integration silently picks
up the streaming fix as a bonus. Drops the duplicate `augmentCallback`
helper and its test; behavior is covered by the expanded
`mergeSentryCallback` suite (14 cases).
diff --git a/packages/core/src/shared-exports.ts b/packages/core/src/shared-exports.ts
@@ -173,6 +173,7 @@ export { createLangChainCallbackHandler, instrumentLangChainEmbeddings } from '.
 export { LANGCHAIN_INTEGRATION_NAME } from './tracing/langchain/constants';
 export type { LangChainOptions, LangChainIntegration } from './tracing/langchain/types';
 export { instrumentStateGraphCompile, instrumentCreateReactAgent, instrumentLangGraph } from './tracing/langgraph';
+export { mergeSentryCallback } from './tracing/langgraph/utils';
 export { LANGGRAPH_INTEGRATION_NAME } from './tracing/langgraph/constants';
 export type { LangGraphOptions, LangGraphIntegration, CompiledGraph } from './tracing/langgraph/types';
 export type { OpenAiClient, OpenAiOptions, InstrumentedMethod } from './tracing/openai/types';
diff --git a/packages/core/src/tracing/langgraph/utils.ts b/packages/core/src/tracing/langgraph/utils.ts
@@ -335,7 +335,44 @@ export function setResponseAttributes(span: Span, inputMessages: LangChainMessag
   }
 }
 
-/** Merge `sentryHandler` into a langchain `callbacks` value (`BaseCallbackHandler[]` or `BaseCallbackManager`). */
+/**
+ * Duck-types a LangChain `CallbackManager` instance. We can't `instanceof`
+ * check because `@langchain/core` may be bundled, deduped, or absent from
+ * our module graph; checking the shape avoids that coupling.
+ */
+function isCallbackManager(value: unknown): value is {
+  addHandler: (handler: unknown, inherit?: boolean) => void;
+  copy: () => unknown;
+  handlers?: unknown[];
+  inheritableHandlers?: unknown[];
+} {
+  if (!value || typeof value !== 'object') {
+    return false;
+  }
+  const candidate = value as { addHandler?: unknown; copy?: unknown };
+  return typeof candidate.addHandler === 'function' && typeof candidate.copy === 'function';
+}
+
+/**
+ * Merge `sentryHandler` into a langchain `callbacks` value.
+ *
+ * Accepts the three shapes LangChain's `RunnableConfig.callbacks` can take:
+ *   - `undefined`              → wrap our handler in a fresh array
+ *   - `BaseCallbackHandler[]`  → append (idempotent)
+ *   - `BaseCallbackManager`    → `.copy()` and register as an inheritable child
+ *
+ * The `CallbackManager` case is load-bearing: when LangGraph sets up a run
+ * with `streamMode: ['messages']`, it puts a `StreamMessagesHandler` onto a
+ * `CallbackManager` and passes that manager through `options.callbacks`. If
+ * we naively wrap the manager into `[manager, sentryHandler]`, LangChain
+ * downstream treats the whole manager as a single opaque handler — its
+ * inheritable children (`StreamMessagesHandler`, the tracer, etc.) are
+ * never unpacked, and per-token streaming events silently disappear.
+ *
+ * We always operate on a copy so repeated invocations don't accumulate
+ * handlers on the caller's manager, and we add ourselves as inheritable
+ * so child callback managers created by `getChild` see us too.
+ */
 export function mergeSentryCallback(existing: unknown, sentryHandler: unknown): unknown {
   if (!existing) {
     return [sentryHandler];
@@ -348,12 +385,23 @@ export function mergeSentryCallback(existing: unknown, sentryHandler: unknown):
     return [...existing, sentryHandler];
   }
 
-  const manager = existing as { addHandler?: (h: unknown) => void; handlers?: unknown[] };
-  if (typeof manager.addHandler === 'function') {
-    const alreadyAdded = Array.isArray(manager.handlers) && manager.handlers.includes(sentryHandler);
-    if (!alreadyAdded) {
-      manager.addHandler(sentryHandler);
+  if (isCallbackManager(existing)) {
+    const copied = existing.copy() as {
+      addHandler: (handler: unknown, inherit?: boolean) => void;
+      handlers?: unknown[];
+      inheritableHandlers?: unknown[];
+    };
+    // CallbackManager keeps `inheritableHandlers ⊆ handlers` (both
+    // `addHandler` and `setHandlers` maintain the invariant), so checking
+    // `handlers` alone normally suffices — we check both as a defensive
+    // guard against externally-constructed managers that bypass `addHandler`.
+    const alreadyRegistered =
+      (copied.handlers?.includes(sentryHandler) ?? false) ||
+      (copied.inheritableHandlers?.includes(sentryHandler) ?? false);
+    if (!alreadyRegistered) {
+      copied.addHandler(sentryHandler, true);
     }
+    return copied;
   }
 
   return existing;
diff --git a/packages/core/test/lib/utils/langgraph-utils.test.ts b/packages/core/test/lib/utils/langgraph-utils.test.ts
@@ -48,6 +48,30 @@ describe('extractAgentNameFromParams', () => {
 describe('mergeSentryCallback', () => {
   const sentryHandler = { _sentry: true };
 
+  /**
+   * Minimal `CallbackManager` stand-in. Mirrors `@langchain/core`'s real
+   * semantics: `addHandler(_, inherit)` pushes to both `handlers` and
+   * `inheritableHandlers` when `inherit !== false`, and `copy()` returns
+   * a fresh manager carrying the same handlers — so we don't accidentally
+   * test against a degenerate shape that bypasses `addHandler`.
+   */
+  function makeFakeCallbackManager(existingHandlers: unknown[] = [], existingInheritableHandlers?: unknown[]) {
+    const manager = {
+      handlers: [...existingHandlers],
+      inheritableHandlers: [...(existingInheritableHandlers ?? existingHandlers)],
+      addHandler: vi.fn(function (this: any, handler: unknown, inherit?: boolean) {
+        this.handlers.push(handler);
+        if (inherit !== false) {
+          this.inheritableHandlers.push(handler);
+        }
+      }),
+      copy: vi.fn(function (this: any) {
+        return makeFakeCallbackManager(this.handlers, this.inheritableHandlers);
+      }),
+    };
+    return manager;
+  }
+
   it('returns a fresh array when no existing callbacks are present', () => {
     expect(mergeSentryCallback(undefined, sentryHandler)).toStrictEqual([sentryHandler]);
     expect(mergeSentryCallback(null, sentryHandler)).toStrictEqual([sentryHandler]);
@@ -65,19 +89,74 @@ describe('mergeSentryCallback', () => {
     expect(mergeSentryCallback(existing, sentryHandler)).toBe(existing);
   });
 
-  it('calls addHandler on a CallbackManager-like object', () => {
-    const addHandler = vi.fn();
-    const manager = { addHandler, handlers: [] as unknown[] };
-    const result = mergeSentryCallback(manager, sentryHandler);
-    expect(result).toBe(manager);
-    expect(addHandler).toHaveBeenCalledWith(sentryHandler);
-    expect(addHandler).toHaveBeenCalledTimes(1);
+  it('preserves inheritable handlers when callbacks is a CallbackManager', () => {
+    // Reproduces the LangGraph `streamMode: ['messages']` setup: a
+    // CallbackManager carrying a StreamMessagesHandler is passed via
+    // options.callbacks. Wrapping it as `[manager, sentryHandler]` would
+    // drop the manager's inheritable children — instead we register
+    // Sentry on a copy and keep the existing handler chain intact.
+    const streamMessagesHandler = {
+      name: 'StreamMessagesHandler',
+      lc_prefer_streaming: true,
+    };
+    const manager = makeFakeCallbackManager([streamMessagesHandler]);
+    const result = mergeSentryCallback(manager, sentryHandler) as {
+      handlers: unknown[];
+    };
+    expect(Array.isArray(result)).toBe(false);
+    expect(result.handlers).toEqual([streamMessagesHandler, sentryHandler]);
   });
 
-  it('does not re-add when the manager already has the sentry handler', () => {
-    const addHandler = vi.fn();
-    const manager = { addHandler, handlers: [sentryHandler] };
+  it('copies the manager rather than mutating the caller-supplied one', () => {
+    // If we mutated the original, repeated invocations would accumulate
+    // Sentry handlers (and tracers from prior runs would leak across runs).
+    const manager = makeFakeCallbackManager([]);
     mergeSentryCallback(manager, sentryHandler);
-    expect(addHandler).not.toHaveBeenCalled();
+    expect(manager.copy).toHaveBeenCalledTimes(1);
+    expect(manager.handlers).toEqual([]);
+  });
+
+  it('registers the sentry handler as inheritable so child managers see it', () => {
+    // LangChain's CallbackManager.getChild creates child managers via
+    // `setHandlers(this.inheritableHandlers)`. If we add ourselves without
+    // `inherit=true`, nested LLM calls inside an agent never receive the
+    // Sentry handler.
+    const manager = makeFakeCallbackManager([]);
+    const result = mergeSentryCallback(manager, sentryHandler) as {
+      addHandler: ReturnType<typeof vi.fn>;
+      handlers: unknown[];
+      inheritableHandlers: unknown[];
+    };
+    expect(result.addHandler).toHaveBeenCalledWith(sentryHandler, true);
+    expect(result.inheritableHandlers).toEqual([sentryHandler]);
+  });
+
+  it('does not double-register when the copied manager already contains the handler', () => {
+    const manager = makeFakeCallbackManager([sentryHandler]);
+    const result = mergeSentryCallback(manager, sentryHandler) as {
+      handlers: unknown[];
+      addHandler: ReturnType<typeof vi.fn>;
+    };
+    expect(result.handlers).toEqual([sentryHandler]);
+    expect(result.addHandler).not.toHaveBeenCalled();
+  });
+
+  it('does not double-register when the handler lives only on inheritableHandlers', () => {
+    // Defensive: a CallbackManager subclass or externally-constructed
+    // instance might keep the Sentry handler on `inheritableHandlers`
+    // without mirroring it onto `handlers`. We must still recognize it
+    // as already-registered to avoid duplicate spans on nested calls.
+    const manager = makeFakeCallbackManager([], [sentryHandler]);
+    const result = mergeSentryCallback(manager, sentryHandler) as {
+      addHandler: ReturnType<typeof vi.fn>;
+      inheritableHandlers: unknown[];
+    };
+    expect(result.addHandler).not.toHaveBeenCalled();
+    expect(result.inheritableHandlers).toEqual([sentryHandler]);
+  });
+
+  it('returns the value unchanged when it is neither an array nor a CallbackManager', () => {
+    const opaque = { name: 'NotAManager' };
+    expect(mergeSentryCallback(opaque, sentryHandler)).toBe(opaque);
   });
 });
diff --git a/packages/node/src/integrations/tracing/langchain/instrumentation.ts b/packages/node/src/integrations/tracing/langchain/instrumentation.ts
@@ -12,6 +12,7 @@ import {
   createLangChainCallbackHandler,
   GOOGLE_GENAI_INTEGRATION_NAME,
   instrumentLangChainEmbeddings,
+  mergeSentryCallback,
   OPENAI_INTEGRATION_NAME,
   SDK_VERSION,
 } from '@sentry/core';
@@ -27,92 +28,6 @@ interface PatchedLangChainExports {
   [key: string]: unknown;
 }
 
-/**
- * Duck-types a LangChain `CallbackManager` instance. We can't `instanceof`
- * check because `@langchain/core` may be bundled, deduped, or absent from
- * our module graph; checking the shape avoids that coupling.
- */
-function isCallbackManager(value: unknown): value is {
-  addHandler: (handler: unknown, inherit?: boolean) => void;
-  copy: () => unknown;
-  handlers?: unknown[];
-} {
-  if (!value || typeof value !== 'object') {
-    return false;
-  }
-  const candidate = value as { addHandler?: unknown; copy?: unknown };
-  return typeof candidate.addHandler === 'function' && typeof candidate.copy === 'function';
-}
-
-/**
- * Exported for testing.
- * @internal
- */
-export const _INTERNAL_augmentCallbackHandlers = augmentCallbackHandlers;
-
-/**
- * Augments a callback handler list with Sentry's handler if not already present.
- *
- * `options.callbacks` may be one of three shapes (per LangChain's RunnableConfig):
- *   - `undefined`           → no callbacks configured
- *   - `BaseCallbackHandler[]` → list of handler instances
- *   - `CallbackManager`     → a manager that already holds (potentially
- *                             inheritable) child handlers
- *
- * The `CallbackManager` case is the load-bearing one: when LangGraph sets up
- * a run with `streamMode: ['messages']`, it puts a `StreamMessagesHandler`
- * onto a `CallbackManager` and passes that manager through `options.callbacks`.
- * If we naively wrap the manager into `[manager, sentryHandler]`, LangChain
- * downstream treats the whole manager as a single opaque handler — its
- * inheritable children (`StreamMessagesHandler`, the tracer, etc.) are never
- * unpacked, and per-token streaming events silently disappear.
- *
- * Instead, when we receive a `CallbackManager`, we copy it (so we don't
- * mutate the caller's manager across invocations) and register Sentry's
- * handler as an inheritable child via `.addHandler()`.
- */
-function augmentCallbackHandlers(handlers: unknown, sentryHandler: unknown): unknown {
-  // Handle null/undefined - return array with just our handler
-  if (!handlers) {
-    return [sentryHandler];
-  }
-
-  // If handlers is already an array
-  if (Array.isArray(handlers)) {
-    // Check if our handler is already in the list
-    if (handlers.includes(sentryHandler)) {
-      return handlers;
-    }
-    // Add our handler to the list
-    return [...handlers, sentryHandler];
-  }
-
-  // CallbackManager: register our handler as an inheritable child on a copy
-  // so we preserve any handlers already attached (notably LangGraph's
-  // StreamMessagesHandler used by `streamMode: ['messages']`).
-  if (isCallbackManager(handlers)) {
-    const copied = handlers.copy() as {
-      addHandler: (handler: unknown, inherit?: boolean) => void;
-      handlers?: unknown[];
-      inheritableHandlers?: unknown[];
-    };
-    // Avoid double-registering on nested invocations. CallbackManager's own
-    // `addHandler` keeps `inheritableHandlers ⊆ handlers`, so checking
-    // `handlers` alone is normally enough — but we check both as a defensive
-    // guard against externally-constructed managers that bypass `addHandler`.
-    const alreadyRegistered =
-      (copied.handlers?.includes(sentryHandler) ?? false) ||
-      (copied.inheritableHandlers?.includes(sentryHandler) ?? false);
-    if (!alreadyRegistered) {
-      copied.addHandler(sentryHandler, true);
-    }
-    return copied;
-  }
-
-  // Unknown type - return original
-  return handlers;
-}
-
 /**
  * Wraps Runnable methods (invoke, stream, batch) to inject Sentry callbacks at request time
  * Uses a Proxy to intercept method calls and augment the options.callbacks
@@ -140,9 +55,7 @@ function wrapRunnableMethod(
       }
 
       // Inject our callback handler into options.callbacks (request time callbacks)
-      const existingCallbacks = options.callbacks;
-      const augmentedCallbacks = augmentCallbackHandlers(existingCallbacks, sentryHandler);
-      options.callbacks = augmentedCallbacks;
+      options.callbacks = mergeSentryCallback(options.callbacks, sentryHandler);
 
       // Call original method with augmented options
       return Reflect.apply(target, thisArg, args);
diff --git a/packages/node/test/integrations/tracing/langchain.test.ts b/packages/node/test/integrations/tracing/langchain.test.ts
