Split patch for main and prebuilt langgraph modules

Author: nicohrubec

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

@@ -35,19 +35,19 @@ export class SentryLangGraphInstrumentation extends InstrumentationBase<LangGrap
     const mainModule = new InstrumentationNodeModuleDefinition(
       '@langchain/langgraph',
       supportedVersions,
-      this._patch.bind(this),
+      this._patchMainModule.bind(this),
       exports => exports,
       [
         new InstrumentationNodeModuleFile(
           /**
            * In CJS, LangGraph packages re-export from dist/index.cjs files.
            * Patching only the root module sometimes misses the real implementation or
            * gets overwritten when that file is loaded. We add a file-level patch so that
-           * _patch runs again on the concrete implementation
-           */
+           * _patchMainModule runs again on the concrete implementation
+          */
           '@langchain/langgraph/dist/index.cjs',
           supportedVersions,
-          this._patch.bind(this),
+          this._patchMainModule.bind(this),
           exports => exports,
         ),
       ],
@@ -56,17 +56,19 @@ export class SentryLangGraphInstrumentation extends InstrumentationBase<LangGrap
     const prebuiltModule = new InstrumentationNodeModuleDefinition(
       '@langchain/langgraph/prebuilt',
       supportedVersions,
-      this._patch.bind(this),
+      this._patchPrebuiltModule.bind(this),
       exports => exports,
       [
         new InstrumentationNodeModuleFile(
           /**
-           * Patch the prebuilt subpath exports for CJS.
-           * The @langchain/langgraph/prebuilt entry point re-exports from dist/prebuilt/index.cjs
-           */
+           * In CJS, LangGraph packages re-export from dist/prebuilt/index.cjs files.
+           * Patching only the root module sometimes misses the real implementation or
+           * gets overwritten when that file is loaded. We add a file-level patch so that
+           * _patchPrebuiltModule runs again on the concrete implementation
+          */
           '@langchain/langgraph/dist/prebuilt/index.cjs',
           supportedVersions,
-          this._patch.bind(this),
+          this._patchPrebuiltModule.bind(this),
           exports => exports,
         ),
       ],
@@ -76,20 +78,10 @@ export class SentryLangGraphInstrumentation extends InstrumentationBase<LangGrap
   }
 
   /**
-   * Core patch logic applying instrumentation to the LangGraph module.
+   * Patch logic applying instrumentation to the LangGraph main module.
    */
-  private _patch(exports: PatchedModuleExports): PatchedModuleExports | void {
-    const client = getClient();
-    const defaultPii = Boolean(client?.getOptions().sendDefaultPii);
-
-    const config = this.getConfig();
-    const recordInputs = config.recordInputs ?? defaultPii;
-    const recordOutputs = config.recordOutputs ?? defaultPii;
-
-    const options: LangGraphOptions = {
-      recordInputs,
-      recordOutputs,
-    };
+  private _patchMainModule(exports: PatchedModuleExports): PatchedModuleExports | void {
+    const options = this._getOptions();
 
     // Patch StateGraph.compile to instrument both compile() and invoke()
     if (exports.StateGraph && typeof exports.StateGraph === 'function') {
@@ -103,6 +95,15 @@ export class SentryLangGraphInstrumentation extends InstrumentationBase<LangGrap
       );
     }
 
+    return exports;
+  }
+
+  /**
+   * Patch logic applying instrumentation to the LangGraph prebuilt module.
+   */
+  private _patchPrebuiltModule(exports: PatchedModuleExports): PatchedModuleExports | void {
+    const options = this._getOptions();
+
     // Patch createReactAgent to instrument the agent creation and invocation
     if (exports.createReactAgent && typeof exports.createReactAgent === 'function') {
       exports.createReactAgent = instrumentCreateReactAgent(
@@ -113,4 +114,18 @@ export class SentryLangGraphInstrumentation extends InstrumentationBase<LangGrap
 
     return exports;
   }
+
+  /**
+   * Helper to get instrumentation options
+   */
+  private _getOptions(): LangGraphOptions {
+    const client = getClient();
+    const defaultPii = Boolean(client?.getOptions().sendDefaultPii);
+    const config = this.getConfig();
+
+    return {
+      recordInputs: config.recordInputs ?? defaultPii,
+      recordOutputs: config.recordOutputs ?? defaultPii,
+    };
+  }
 }