ref: Use top-level export patching for Mastra

.changeset/large-gifts-go.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+ref: Fork import-in-the-middle and require-in-the-middle

js/src/auto-instrumentations/bundler/plugin.ts

@@ -1,11 +1,20 @@
 import { createUnplugin } from "unplugin";
 import { create, type InstrumentationConfig } from "../orchestrion-js";
-import { extname, join, sep } from "path";
+import { dirname, extname, join, resolve, sep } from "path";
 import { readFileSync } from "fs";
 import { fileURLToPath } from "url";
 import moduleDetailsFromPath from "module-details-from-path";
 import { getDefaultInstrumentationConfigs } from "../configs/all";
 import { applySpecialCasePatch } from "../loader/special-case-patches";
+import {
+  buildTopLevelImportHookSourceWrapper,
+  getDefaultTopLevelImportHooks,
+  type TopLevelImportHookTarget,
+} from "../loader/top-level-export-patches";
+import { readDisabledInstrumentationEnvConfig } from "../../instrumentation/config";
+
+const TOP_LEVEL_ORIGINAL_IMPORT_PREFIX = "braintrust-top-level-original:";
+const TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX = `\0${TOP_LEVEL_ORIGINAL_IMPORT_PREFIX}`;
 
 export interface LegacyBundlerPluginOptions {
   /**
@@ -75,6 +84,18 @@ export const unplugin = createUnplugin<LegacyBundlerPluginOptions>(
     const allInstrumentations = getDefaultInstrumentationConfigs({
       additionalInstrumentations: options.instrumentations,
     });
+    const disabledIntegrationConfig = readDisabledInstrumentationEnvConfig(
+      process.env.BRAINTRUST_DISABLE_INSTRUMENTATION,
+    ).integrations;
+    const topLevelImportHookTarget: TopLevelImportHookTarget =
+      options.browser === true ? "browser" : "node";
+    const topLevelImportHooks = getDefaultTopLevelImportHooks({
+      disabledIntegrationConfig,
+      target: topLevelImportHookTarget,
+    });
+    const originalSources = new Map<string, string>();
+    const originalDirectories = new Map<string, string>();
+    let nextOriginalId = 0;
 
     // Default to browser build, use polyfill unless explicitly disabled
     const dcModule = options.browser === false ? undefined : "dc-browser";
@@ -85,11 +106,37 @@ export const unplugin = createUnplugin<LegacyBundlerPluginOptions>(
     return {
       name: "code-transformer",
       enforce: "pre",
+      resolveId(id: string, importer?: string) {
+        if (id.startsWith(TOP_LEVEL_ORIGINAL_IMPORT_PREFIX)) {
+          return `${TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX}${id.slice(
+            TOP_LEVEL_ORIGINAL_IMPORT_PREFIX.length,
+          )}`;
+        }
+        if (
+          id.startsWith(".") &&
+          importer?.startsWith(TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX)
+        ) {
+          const originalDirectory = originalDirectories.get(importer);
+          if (originalDirectory) {
+            return resolve(originalDirectory, id);
+          }
+        }
+        return null;
+      },
+      load(id: string) {
+        if (id.startsWith(TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX)) {
+          return originalSources.get(id) ?? null;
+        }
+        return null;
+      },
       transform(code: string, id: string) {
         if (!id) {
           // Some modules apparently don't have an id?
           return null;
         }
+        if (id.startsWith(TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX)) {
+          return null;
+        }
 
         // Convert file:// URLs to regular paths at entry point
         // Node.js ESM loader hooks provide file:// URLs, but downstream code expects paths
@@ -123,25 +170,52 @@ export const unplugin = createUnplugin<LegacyBundlerPluginOptions>(
         const normalizedModulePath = moduleDetails.path.replace(/\\/g, "/");
         const moduleVersion = getModuleVersion(moduleDetails.basedir);
 
+        let nextCode = code;
+        let didPatch = false;
+
         // Per-package source patches (see loader/special-case-patches.ts).
-        // Same anti-pattern fallback the runtime loader uses — mirrored here
-        // so bundled apps get the patches without relying on hook.mjs.
-        // Skipped for browser bundles since the wrapper templates use
-        // `node:module`/`require` to resolve `@mastra/observability`.
+        // Skipped for browser bundles to preserve the historical behavior of
+        // the OpenAI APIPromise patch path.
         if (options.browser !== true) {
           const patched = applySpecialCasePatch({
             packageName: moduleName,
             modulePath: normalizedModulePath,
-            source: code,
+            source: nextCode,
             format: isModule ? "esm" : "cjs",
           });
           if (patched !== null) {
-            return { code: patched, map: null };
+            nextCode = patched;
+            didPatch = true;
           }
         }
 
+        const originalModuleSpecifier = `${TOP_LEVEL_ORIGINAL_IMPORT_PREFIX}${nextOriginalId++}`;
+        const originalModuleId = `${TOP_LEVEL_ORIGINAL_RESOLVED_PREFIX}${originalModuleSpecifier.slice(
+          TOP_LEVEL_ORIGINAL_IMPORT_PREFIX.length,
+        )}`;
+        const topLevelWrapper = buildTopLevelImportHookSourceWrapper(
+          topLevelImportHooks,
+          {
+            format: isModule ? "esm" : "cjs",
+            modulePath: normalizedModulePath,
+            originalModuleSpecifier,
+            packageName: moduleName,
+            source: nextCode,
+            target: topLevelImportHookTarget,
+          },
+        );
+        if (topLevelWrapper !== null) {
+          originalSources.set(originalModuleId, nextCode);
+          originalDirectories.set(originalModuleId, dirname(filePath));
+          nextCode = topLevelWrapper;
+          didPatch = true;
+        }
+
         // If no version found
         if (!moduleVersion) {
+          if (didPatch) {
+            return { code: nextCode, map: null };
+          }
           console.warn(
             `No 'package.json' version found for module ${moduleName} at ${moduleDetails.basedir}. Skipping transformation.`,
           );
@@ -157,13 +231,13 @@ export const unplugin = createUnplugin<LegacyBundlerPluginOptions>(
 
         if (!transformer) {
           // No instrumentations match this file
-          return null;
+          return didPatch ? { code: nextCode, map: null } : null;
         }
 
         try {
           // Transform the code
           const moduleType = isModule ? "esm" : "cjs";
-          const result = transformer.transform(code, moduleType);
+          const result = transformer.transform(nextCode, moduleType);
           const transformedCode = result.code.replace(
             /const \{tracingChannel: ([A-Za-z_$][\w$]*)\} = ([A-Za-z_$][\w$]*);/g,
             "const $1 = $2.tracingChannel;",

js/src/auto-instrumentations/bundler/webpack-loader.ts

@@ -27,6 +27,15 @@ import { readFileSync } from "fs";
 import moduleDetailsFromPath from "module-details-from-path";
 import { getDefaultInstrumentationConfigs } from "../configs/all";
 import { type LegacyBundlerPluginOptions } from "./plugin";
+import { applySpecialCasePatch } from "../loader/special-case-patches";
+import {
+  buildTopLevelImportHookSourceWrapper,
+  getDefaultTopLevelImportHooks,
+  type TopLevelImportHookTarget,
+} from "../loader/top-level-export-patches";
+import { readDisabledInstrumentationEnvConfig } from "../../instrumentation/config";
+
+const TOP_LEVEL_ORIGINAL_QUERY = "braintrust-top-level-original";
 
 /**
  * Helper function to get module version from package.json
@@ -94,11 +103,15 @@ function codeTransformerLoader(
   const callback = this.async();
   const options: LegacyBundlerPluginOptions = this.getOptions() ?? {};
   const resourcePath: string = this.resourcePath;
+  const resourceQuery: string = this.resourceQuery ?? "";
 
   // Skip virtual modules (e.g. Next.js loaders pass query-string URLs with no real path)
   if (!resourcePath) {
     return callback(null, code, inputSourceMap);
   }
+  if (resourceQuery.includes(TOP_LEVEL_ORIGINAL_QUERY)) {
+    return callback(null, code, inputSourceMap);
+  }
 
   // Determine if this is an ES module using multiple methods for accurate detection
   const ext = extname(resourcePath);
@@ -122,12 +135,54 @@ function codeTransformerLoader(
   const moduleName = moduleDetails.name;
   const moduleVersion = getModuleVersion(moduleDetails.basedir);
 
-  if (!moduleVersion) {
-    return callback(null, code, inputSourceMap);
-  }
-
   // Normalize the module path for Windows compatibility (WASM transformer expects forward slashes)
   const normalizedModulePath = moduleDetails.path.replace(/\\/g, "/");
+  const moduleType: ModuleType = isModule ? "esm" : "cjs";
+  const target: TopLevelImportHookTarget =
+    options.browser === true ? "browser" : "node";
+  const disabledIntegrationConfig = readDisabledInstrumentationEnvConfig(
+    process.env.BRAINTRUST_DISABLE_INSTRUMENTATION,
+  ).integrations;
+  const topLevelImportHooks = getDefaultTopLevelImportHooks({
+    disabledIntegrationConfig,
+    target,
+  });
+
+  let nextCode = code;
+  let didPatch = false;
+
+  if (options.browser !== true) {
+    const patched = applySpecialCasePatch({
+      format: moduleType,
+      modulePath: normalizedModulePath,
+      packageName: moduleName,
+      source: nextCode,
+    });
+    if (patched !== null) {
+      nextCode = patched;
+      didPatch = true;
+    }
+  }
+
+  const topLevelWrapper = buildTopLevelImportHookSourceWrapper(
+    topLevelImportHooks,
+    {
+      format: moduleType,
+      modulePath: normalizedModulePath,
+      originalModuleSpecifier: `${resourcePath}?${TOP_LEVEL_ORIGINAL_QUERY}`,
+      packageName: moduleName,
+      source: nextCode,
+      target,
+    },
+  );
+  if (topLevelWrapper !== null) {
+    nextCode = topLevelWrapper;
+    didPatch = true;
+  }
+
+  if (!moduleVersion) {
+    return callback(null, nextCode, inputSourceMap);
+  }
 
   const matcher = getMatcher(options);
   const transformer = matcher.getTransformer(
@@ -137,19 +192,18 @@ function codeTransformerLoader(
   );
 
   if (!transformer) {
-    return callback(null, code, inputSourceMap);
+    return callback(null, nextCode, inputSourceMap);
   }
 
   try {
-    const moduleType: ModuleType = isModule ? "esm" : "cjs";
-    const result = transformer.transform(code, moduleType);
+    const result = transformer.transform(nextCode, moduleType);
     callback(null, result.code, result.map ?? undefined);
   } catch (error) {
     console.warn(
       `[code-transformer-loader] Error transforming ${resourcePath}:`,
       error,
     );
-    callback(null, code, inputSourceMap);
+    callback(null, didPatch ? nextCode : code, inputSourceMap);
   }
 }
 

js/src/auto-instrumentations/configs/all.ts

@@ -96,12 +96,9 @@ const defaultInstrumentationConfigGroups: readonly InstrumentationConfigGroup[]
       configs: flueConfigs,
     },
     // Note: `@mastra/core` is not listed here because its instrumentation
-    // doesn't go through the AST `code-transformer` matcher — Mastra's
-    // content-hashed chunks make `filePath`-based matching too brittle.
-    // Instead it's handled by the source-replacement entry in
-    // `loader/special-case-patches.ts`, which both the runtime loader
-    // (`hook.mjs` → `cjs-patch.ts`/`esm-hook.mts`) and the bundler plugin
-    // (`bundler/plugin.ts`) call. The `mastra` env-var disable still works.
+    // doesn't go through the AST `code-transformer` matcher. It is handled by
+    // the internal top-level import hook registry in
+    // `loader/top-level-export-patches.ts`.
   ];
 
 export function getDefaultInstrumentationConfigs({

js/src/auto-instrumentations/hook.mts

@@ -13,7 +13,7 @@
  * - CJS modules: Transformed via ModulePatch monkey-patching Module._compile
  */
 
-import { register } from "node:module";
+import { register as registerModule } from "node:module";
 import {
   isInstrumentationIntegrationDisabled,
   readDisabledInstrumentationEnvConfig,
@@ -23,6 +23,28 @@ import { installMastraExporterFactory } from "./loader/mastra-observability-patc
 import { getDefaultAutoInstrumentationConfigs } from "./configs/all.js";
 import { ModulePatch } from "./loader/cjs-patch.js";
 import { patchTracingChannel } from "./patch-tracing-channel.js";
+import registerState from "./import-in-the-middle/lib/register.js";
+import { createHook as createImportInTheMiddleHook } from "./import-in-the-middle/create-hook.mjs";
+import {
+  getDefaultTopLevelImportHooks,
+  installTopLevelImportHookRunner,
+} from "./loader/top-level-export-patches.js";
+import { installNodeTopLevelExportPatches } from "./loader/top-level-export-patches-node.js";
+
+const BRAINTRUST_IITM_LOADER_PARAM = "braintrust-iitm-loader";
+const registryImportUrl = getCanonicalHookUrl(import.meta.url);
+const asyncImportHookUrl = getImportInTheMiddleLoaderUrl(registryImportUrl);
+const isImportInTheMiddleLoader = hasImportInTheMiddleLoaderParam(
+  import.meta.url,
+);
+const importInTheMiddleHook = createImportInTheMiddleHook(import.meta, {
+  registerUrl: registryImportUrl,
+});
+
+export const initialize = importInTheMiddleHook.initialize;
+export const resolve = importInTheMiddleHook.resolve;
+export const load = importInTheMiddleHook.load;
+export const register = registerState.register;
 
 const state = ((globalThis as any)[
   Symbol.for("braintrust.applyAutoInstrumentation")
@@ -32,13 +54,13 @@ const alreadyApplied = state.applied;
 // Patch diagnostics_channel.tracePromise to handle APIPromise correctly.
 // MUST be done here (before any SDK code runs) to fix Anthropic APIPromise incompatibility.
 // Construct the module path dynamically to prevent build from stripping "node:" prefix.
-if (!alreadyApplied) {
+if (!isImportInTheMiddleLoader && !alreadyApplied) {
   const dcPath = ["node", "diagnostics_channel"].join(":");
   const dc: any = await import(/* @vite-ignore */ dcPath as any);
   patchTracingChannel(dc.tracingChannel);
 }
 
-if (!alreadyApplied) {
+if (!isImportInTheMiddleLoader && !alreadyApplied) {
   const allConfigs = getDefaultAutoInstrumentationConfigs();
 
   // Expose the Mastra exporter factory on globalThis so the loader patches
@@ -52,8 +74,19 @@ if (!alreadyApplied) {
     installMastraExporterFactory(() => new BraintrustObservabilityExporter());
   }
 
+  const topLevelImportHooks = getDefaultTopLevelImportHooks({
+    disabledIntegrationConfig: disabled,
+    target: "node",
+  });
+  installTopLevelImportHookRunner(topLevelImportHooks);
+  installNodeTopLevelExportPatches({
+    asyncImportHookUrl,
+    hooks: topLevelImportHooks,
+    registryImportUrl,
+  });
+
   // 1. Register ESM loader for ESM modules
-  register("./loader/esm-hook.mjs", {
+  registerModule("./loader/esm-hook.mjs", {
     parentURL: import.meta.url,
     data: { instrumentations: allConfigs },
   } as any);
@@ -82,3 +115,24 @@ if (!alreadyApplied) {
     }
   }
 }
+
+function hasImportInTheMiddleLoaderParam(url: string): boolean {
+  try {
+    return new URL(url).searchParams.has(BRAINTRUST_IITM_LOADER_PARAM);
+  } catch {
+    return false;
+  }
+}
+
+function getCanonicalHookUrl(url: string): string {
+  const parsed = new URL(url);
+  parsed.searchParams.delete(BRAINTRUST_IITM_LOADER_PARAM);
+  parsed.hash = "";
+  return parsed.href;
+}
+
+function getImportInTheMiddleLoaderUrl(canonicalUrl: string): string {
+  const parsed = new URL(canonicalUrl);
+  parsed.searchParams.set(BRAINTRUST_IITM_LOADER_PARAM, "true");
+  return parsed.href;
+}