feat: auto-register ESM loader hooks in TuskDrift.initialize() (#146)

Author: sohankshirsagar

CONTRIBUTING.md

@@ -65,22 +65,40 @@ The SDK supports both CommonJS and ESM module systems, using different intercept
 
 #### CommonJS Module Interception
 
-- **Package**: `require-in-the-middle`
-- **How it works**: Hooks into `Module.prototype.require` globally
-- **When it activates**: When `require()` is called
-- **Setup**: Automatic - no special flags needed
-- **Use case**: Works for all CommonJS modules
+- **Package**: `require-in-the-middle` (RITM)
+- **How it works**: CJS modules are loaded through a single JavaScript function (`Module._load`). RITM monkey-patches this function so that every `require()` call passes through the patch, giving the SDK a chance to intercept and wrap module exports.
+- **Setup**: Automatic -- no special flags or loader registration needed. Just calling `TuskDrift.initialize()` before other `require()` calls is sufficient.
 
 #### ESM Module Interception
 
-- **Package**: `import-in-the-middle`, created by [Datadog](https://opensource.datadoghq.com/projects/node/#the-import-in-the-middle-library)
-- **How it works**: Uses Node.js loader hooks to intercept imports before they're cached
-- **When it activates**: During module resolution/loading phase
-- **Setup**: Requires `--import` flag or `module.register()` call
-- **Use case**: Required for ESM modules
-- **Loader file**: `hook.mjs` - re-exports loader hooks from `import-in-the-middle`
+- **Package**: `import-in-the-middle` (IITM), created by [Datadog](https://opensource.datadoghq.com/projects/node/#the-import-in-the-middle-library)
+- **How it works**: Unlike CJS, ESM module loading is handled by Node.js internals (C++), not a patchable JavaScript function. The only way to intercept ESM imports is through Node's official [customization hooks API](https://nodejs.org/api/module.html#customization-hooks) (`module.register`), which runs hook code in a separate loader thread.
+- **Setup**: The SDK automatically registers ESM loader hooks inside `TuskDrift.initialize()` via `module.register()` (see `src/core/esmLoader.ts`). ESM applications must still use `--import` to ensure the init file runs before the application's import graph is resolved. The `hook.mjs` file at the package root is kept for backward compatibility but is no longer required for manual registration.
 
-**Key difference**: CommonJS's `require()` is synchronous and sequential, so you can control order. ESM's `import` is hoisted and parallel, requiring loader hooks to intercept before evaluation.
+#### How ESM instrumentation works end-to-end
+
+1. **Loader registration**: `initializeEsmLoader()` (called from `TuskDrift.initialize()`) uses `createAddHookMessageChannel()` from IITM to set up a `MessagePort` between the main thread and the loader thread, then calls `module.register('import-in-the-middle/hook.mjs', ...)` to install the loader hooks.
+2. **Module wrapping**: When any ESM module is imported, IITM's `load` hook transforms its source code on the fly, replacing all named exports with getter/setter proxies. The module works normally, but exports now pass through a proxy layer.
+3. **Hook registration**: `TdInstrumentationBase.enable()` creates `new HookImport(['pg'], {}, hookFn)` for each instrumented module. This registers a callback and sends the module name to the loader thread via the `MessagePort` so the loader knows to watch for it.
+4. **Interception at runtime**: When application code accesses a wrapped module's exports (e.g., `import { Client } from 'pg'`), the getter proxy fires, the `hookFn` callback runs, and the SDK patches the export with its instrumented version.
+
+For CJS, steps 1-2 are unnecessary -- RITM patches `Module._load` directly in the main thread, and the rest works the same way.
+
+#### Why `--import` is still needed for ESM
+
+In CJS, `require()` is synchronous and imperative -- putting `require('./tuskDriftInit')` first guarantees it runs before other modules. In ESM, all `import` declarations are hoisted and the entire module graph is resolved before any module-level code executes. The `--import` flag runs the init file in a pre-evaluation phase, ensuring `TuskDrift.initialize()` (and the loader registration) happens before the application's imports are resolved.
+
+#### Node.js built-in modules are always CJS
+
+Node.js built-in modules (`http`, `https`, `net`, `fs`, etc.) are loaded through the CJS `require()` path internally, even when imported via ESM `import` syntax. This means RITM can intercept them regardless of the application's module system, and the ESM loader hooks are not required for built-in module instrumentation.
+
+#### The `registerEsmLoaderHooks` opt-out
+
+Because we pass `include: []` during `module.register()`, IITM starts with an empty allowlist and only wraps modules that are explicitly registered via `new Hook([...])` on the main thread (sent to the loader thread over the `MessagePort`). This means only modules the SDK actually instruments get their exports wrapped with getter/setter proxies -- unrelated modules are left untouched. In rare cases, the wrapping can still conflict with non-standard export patterns, native/WASM bindings, or bundler-generated ESM in the instrumented modules themselves. Users can disable this with `registerEsmLoaderHooks: false` in `TuskDrift.initialize()`, which means only CJS-loaded modules will be instrumentable. See `docs/initialization.md` for the user-facing documentation.
+
+#### Compatibility with other IITM consumers (Sentry, OpenTelemetry)
+
+Multiple SDKs can each call `module.register()` with their own IITM loader instance and `MessagePort`. IITM detects the duplicate initialization (`global.__import_in_the_middle_initialized__`) and logs a warning, but both SDKs' hooks will fire correctly. Patches layer on top of each other -- if Sentry wraps `pg.Client.query` and Drift also wraps it, the final export passes through both wrappers.
 
 ### When Does an Instrumentation Need Special ESM Handling?
 

docs/initialization.md

@@ -20,26 +20,7 @@ Create a separate file (e.g. `tuskDriftInit.ts` or `tuskDriftInit.js`) to initia
 
 **IMPORTANT**: Ensure that `TuskDrift` is initialized before any other telemetry providers (e.g. OpenTelemetry, Sentry, etc.). If not, your existing telemetry may not work properly.
 
-### Determining Your Module System
-
-Before proceeding, you need to determine whether your application uses **CommonJS** or **ESM** (ECMAScript Modules).
-
-The easiest way to determine this is by looking at your import syntax.
-
-**If your application uses `require()`:**
-
-- Your application is CommonJS (use the CommonJS setup below)
-
-**If your application uses `import` statements:**
-
-- This could be either CommonJS or ESM, depending on your build configuration
-- Check your compiled output (if you compile to a directory like `dist/`):
-  - If the compiled code contains `require()` statements → CommonJS application
-  - If the compiled code contains `import` statements → ESM application
-- If you don't compile your code (running source files directly):
-  - It is an ESM application
-
-### For CommonJS Applications
+The initialization file is the same for both CommonJS and ESM applications. The SDK automatically registers ESM loader hooks when running in an ESM environment (Node.js >= 18.19.0 or >= 20.6.0).
 
 ```typescript
 // tuskDriftInit.ts or tuskDriftInit.js
@@ -54,31 +35,7 @@ TuskDrift.initialize({
 export { TuskDrift };
 ```
 
-### For ESM Applications
-
-ESM applications require additional setup to properly intercept module imports:
-
-```typescript
-// tuskDriftInit.ts
-import { register } from "node:module";
-import { pathToFileURL } from "node:url";
-
-// Register the ESM loader
-// This enables interception of ESM module imports
-register("@use-tusk/drift-node-sdk/hook.mjs", pathToFileURL("./"));
-
-import { TuskDrift } from "@use-tusk/drift-node-sdk";
-
-// Initialize SDK immediately
-TuskDrift.initialize({
-  apiKey: process.env.TUSK_API_KEY,
-  env: process.env.NODE_ENV,
-});
-
-export { TuskDrift };
-```
-
-**Why the ESM loader is needed**: ESM imports are statically analyzed and hoisted, meaning all imports are resolved before any code runs. The `register()` call sets up Node.js loader hooks that intercept module imports, allowing the SDK to instrument packages like `postgres`, `http`, etc. Without this, the SDK cannot patch ESM modules.
+> **Note:** ESM applications still require the `--import` flag when starting Node.js. See [Step 2](#2-import-sdk-at-application-entry-point) for details.
 
 ### Initialization Parameters
 
@@ -116,13 +73,36 @@ export { TuskDrift };
       <td><code>1.0</code></td>
       <td>Override sampling rate (0.0 - 1.0) for recording. Takes precedence over <code>TUSK_SAMPLING_RATE</code> env var and config file.</td>
     </tr>
+    <tr>
+      <td><code>registerEsmLoaderHooks</code></td>
+      <td><code>boolean</code></td>
+      <td><code>true</code></td>
+      <td>Automatically register ESM loader hooks for module interception. Set to <code>false</code> to disable if <code>import-in-the-middle</code> causes issues with certain packages. See <a href="#troubleshooting-esm">Troubleshooting ESM</a>.</td>
+    </tr>
   </tbody>
 </table>
 
 > **See also:** [Environment Variables guide](./environment-variables.md) for detailed information about environment variables.
 
 ## 2. Import SDK at Application Entry Point
 
+### Determining Your Module System
+
+You need to know whether your application uses **CommonJS** or **ESM** (ECMAScript Modules) because the entry point setup differs.
+
+**If your application uses `require()`:**
+
+- Your application is CommonJS
+
+**If your application uses `import` statements:**
+
+- This could be either CommonJS or ESM, depending on your build configuration
+- Check your compiled output (if you compile to a directory like `dist/`):
+  - If the compiled code contains `require()` statements → CommonJS application
+  - If the compiled code contains `import` statements → ESM application
+- If you don't compile your code (running source files directly):
+  - It is an ESM application
+
 ### For CommonJS Applications
 
 In your main server file (e.g., `server.ts`, `index.ts`, `app.ts`), require the initialized SDK **at the very top**, before any other requires:
@@ -153,7 +133,7 @@ For ESM applications, you **cannot** control import order within your applicatio
 }
 ```
 
-**Why `--import` is required for ESM**: In ESM, all `import` statements are hoisted and evaluated before any code runs, making it impossible to control import order within a file. The `--import` flag ensures the SDK initialization (including loader registration) happens in a separate phase before your application code loads, guaranteeing proper module interception.
+**Why `--import` is required for ESM**: In ESM, all `import` statements are hoisted and evaluated before any code runs, making it impossible to control import order within a file. The `--import` flag ensures the SDK initialization happens in a separate phase before your application code loads, guaranteeing proper module interception.
 
 ### 3. Configure Sampling Rate
 
@@ -255,3 +235,33 @@ app.listen(8000, () => {
   console.log("Server started and ready for Tusk Drift");
 });
 ```
+
+## Troubleshooting ESM
+
+The SDK automatically registers ESM loader hooks via [`import-in-the-middle`](https://www.npmjs.com/package/import-in-the-middle) to intercept ES module imports. Only modules that the SDK instruments have their exports wrapped with getter/setter proxies -- unrelated modules are left untouched.
+
+In rare cases, the wrapping can cause issues with instrumented packages:
+
+- **Non-standard export patterns**: Some packages use dynamic `export *` re-exports or conditional exports that the wrapper's static analysis cannot parse, resulting in runtime syntax errors.
+- **Native or WASM bindings**: Packages with native addons loaded via ESM can conflict with the proxy wrapping mechanism.
+- **Bundler-generated ESM**: Code that was bundled (e.g., by esbuild or webpack) into ESM sometimes produces patterns the wrapper does not handle correctly.
+- **Circular ESM imports**: The proxy layer can interact badly with circular ESM import graphs in some edge cases.
+
+If you encounter errors like:
+
+```
+SyntaxError: The requested module '...' does not provide an export named '...'
+(node:1234) Error: 'import-in-the-middle' failed to wrap 'file://../../path/to/file.js'
+```
+
+You can disable the automatic ESM hook registration:
+
+```typescript
+TuskDrift.initialize({
+  apiKey: process.env.TUSK_API_KEY,
+  env: process.env.NODE_ENV,
+  registerEsmLoaderHooks: false,
+});
+```
+
+> **Note:** Disabling ESM loader hooks means the SDK will only be able to instrument packages loaded via CommonJS (`require()`). Packages loaded purely through ESM `import` statements will not be intercepted. Node.js built-in modules (like `http`, `https`, `net`) are always loaded through the CJS path internally, so they will continue to be instrumented regardless of this setting.

src/core/TuskDrift.ts

@@ -40,18 +40,22 @@ import {
   loadTuskConfig,
   TuskConfig,
   OriginalGlobalUtils,
+  isCommonJS,
 } from "./utils";
 import { TransformConfigs } from "../instrumentation/libraries/types";
 import { ATTR_SERVICE_NAME } from "@opentelemetry/semantic-conventions";
 import { Resource } from "@opentelemetry/resources";
 import { getRustCoreStartupStatus } from "./rustCoreBinding";
+import { initializeEsmLoader } from "./esmLoader";
 
 export interface InitParams {
   apiKey?: string;
   env?: string;
   logLevel?: LogLevel;
   transforms?: TransformConfigs;
   samplingRate?: number;
+  /** Set to `false` to disable automatic ESM loader hook registration. Defaults to `true`. */
+  registerEsmLoaderHooks?: boolean;
 }
 
 export enum TuskDriftMode {
@@ -83,15 +87,6 @@ export class TuskDriftCore {
     this.config = loadTuskConfig() || {};
   }
 
-  private isCommonJS(): boolean {
-    return (
-      typeof module !== "undefined" &&
-      "exports" in module &&
-      typeof require !== "undefined" &&
-      typeof require.cache !== "undefined"
-    );
-  }
-
   private getPackageName(modulePath: string): string | null {
     let dir = path.dirname(modulePath);
     while (dir) {
@@ -117,7 +112,7 @@ export class TuskDriftCore {
   private alreadyRequiredModules(): Set<string> {
     const alreadyRequiredModuleNames = new Set<string>();
 
-    if (this.isCommonJS()) {
+    if (isCommonJS()) {
       const requireCache = Object.keys(require.cache);
       for (const modulePath of requireCache) {
         if (modulePath.includes("node_modules")) {
@@ -478,6 +473,10 @@ export class TuskDriftCore {
       return;
     }
 
+    if (initParams.registerEsmLoaderHooks !== false) {
+      initializeEsmLoader();
+    }
+
     this.logRustCoreStartupStatus();
     logger.debug(`Initializing in ${this.mode} mode`);
 

src/core/esmLoader.ts

@@ -0,0 +1,70 @@
+import { createAddHookMessageChannel } from "import-in-the-middle";
+import * as moduleModule from "module";
+import { logger } from "./utils";
+import { isCommonJS } from "./utils/runtimeDetectionUtils";
+
+const NODE_MAJOR = parseInt(process.versions.node.split(".")[0]!, 10);
+const NODE_MINOR = parseInt(process.versions.node.split(".")[1]!, 10);
+
+function supportsModuleRegister(): boolean {
+  return (
+    NODE_MAJOR >= 21 ||
+    (NODE_MAJOR === 20 && NODE_MINOR >= 6) ||
+    (NODE_MAJOR === 18 && NODE_MINOR >= 19)
+  );
+}
+
+/**
+ * Automatically register ESM loader hooks via `import-in-the-middle` so that
+ * ESM imports can be intercepted for instrumentation.
+ *
+ * In CJS mode this is a no-op because `require-in-the-middle` handles
+ * interception. On Node versions that lack `module.register` support
+ * (< 18.19, < 20.6) we log a warning and skip.
+ *
+ * https://nodejs.org/api/module.html#moduleregisterspecifier-parenturl-options
+ */
+export function initializeEsmLoader(): void {
+  if (isCommonJS()) {
+    return;
+  }
+
+  if (!supportsModuleRegister()) {
+    logger.warn(
+      `Node.js ${process.versions.node} does not support module.register(). ` +
+        `ESM loader hooks will not be registered automatically. ` +
+        `Upgrade to Node.js >= 18.19.0 or >= 20.6.0, or register the hooks manually.`,
+    );
+    return;
+  }
+
+  if ((globalThis as any).__tuskDriftEsmLoaderRegistered) {
+    return;
+  }
+  (globalThis as any).__tuskDriftEsmLoaderRegistered = true;
+
+  try {
+    // createAddHookMessageChannel sets up a MessagePort so the main thread can
+    // send new hook registrations (from `new Hook(...)` calls in userland) to
+    // the loader thread, which runs in a separate context.
+    const { addHookMessagePort } = createAddHookMessageChannel();
+
+    // The IITM loader hook module that intercepts ESM imports.
+    // Resolved relative to this SDK package (import.meta.url) so the hook
+    // module is found from node_modules regardless of the user's cwd.
+    // @ts-expect-error register exists on module in supported Node versions
+    moduleModule.register("import-in-the-middle/hook.mjs", import.meta.url, {
+      // Payload sent to the loader hook's initialize() function:
+      // - addHookMessagePort: the MessagePort for main↔loader communication
+      // - include: [] starts with an empty allowlist; only modules registered
+      //   via new Hook([...]) on the main thread get added dynamically through
+      //   the MessagePort, so only instrumented modules are wrapped.
+      data: { addHookMessagePort, include: [] },
+      // Transfer (not clone) the port — a MessagePort can only be owned by one thread
+      transferList: [addHookMessagePort],
+    });
+    logger.debug("ESM loader hooks registered successfully");
+  } catch (error) {
+    logger.warn("Failed to register ESM loader hooks:", error);
+  }
+}

src/core/utils/runtimeDetectionUtils.ts

@@ -1,3 +1,16 @@
+export function isCommonJS(): boolean {
+  try {
+    return (
+      typeof module !== "undefined" &&
+      "exports" in module &&
+      typeof require !== "undefined" &&
+      typeof require.cache !== "undefined"
+    );
+  } catch {
+    return false;
+  }
+}
+
 export function isNextJsRuntime(): boolean {
   return (
     process.env.NEXT_RUNTIME !== undefined || typeof (global as any).__NEXT_DATA__ !== "undefined"

src/instrumentation/libraries/fetch/e2e-tests/esm-fetch/src/tdInit.ts

@@ -1,11 +1,4 @@
-import { register } from 'node:module';
-import { pathToFileURL } from 'node:url';
-
-// Register the ESM loader
-// This enables interception of ESM module imports
-register('@use-tusk/drift-node-sdk/hook.mjs', pathToFileURL('./'));
-
-import { TuskDrift } from '@use-tusk/drift-node-sdk';
+import { TuskDrift } from "@use-tusk/drift-node-sdk";
 
 TuskDrift.initialize({
   apiKey: "api-key",

src/instrumentation/libraries/firestore/e2e-tests/esm-firestore/src/tdInit.ts

@@ -1,11 +1,4 @@
-import { register } from 'node:module';
-import { pathToFileURL } from 'node:url';
-
-// Register the ESM loader
-// This enables interception of ESM module imports
-register('@use-tusk/drift-node-sdk/hook.mjs', pathToFileURL('./'));
-
-import { TuskDrift } from '@use-tusk/drift-node-sdk';
+import { TuskDrift } from "@use-tusk/drift-node-sdk";
 
 TuskDrift.initialize({
   apiKey: "api-key",

src/instrumentation/libraries/grpc/e2e-tests/esm-grpc/src/tdInit.ts

@@ -1,11 +1,4 @@
-import { register } from 'node:module';
-import { pathToFileURL } from 'node:url';
-
-// Register the ESM loader
-// This enables interception of ESM module imports
-register('@use-tusk/drift-node-sdk/hook.mjs', pathToFileURL('./'));
-
-import { TuskDrift } from '@use-tusk/drift-node-sdk';
+import { TuskDrift } from "@use-tusk/drift-node-sdk";
 
 TuskDrift.initialize({
   apiKey: "api-key",

src/instrumentation/libraries/http/e2e-tests/esm-http/src/tdInit.ts

@@ -1,11 +1,4 @@
-import { register } from 'node:module';
-import { pathToFileURL } from 'node:url';
-
-// Register the ESM loader
-// This enables interception of ESM module imports
-register('@use-tusk/drift-node-sdk/hook.mjs', pathToFileURL('./'));
-
-import { TuskDrift } from '@use-tusk/drift-node-sdk';
+import { TuskDrift } from "@use-tusk/drift-node-sdk";
 
 TuskDrift.initialize({
   apiKey: "api-key",