ref: Patch Mastra using top level export patching

e2e/scenarios/mastra-instrumentation/scenario.mjs

@@ -60,12 +60,9 @@ const travelWorkflow = createWorkflow({
   .commit();
 
 // The scenario constructs `new Mastra({})` with no observability config.
-// Under `node --import braintrust/hook.mjs`, the loader patches:
-//   - `@mastra/core` Mastra constructor → injects a default Observability when
-//     `observability` is missing
-//   - `@mastra/observability` Observability constructor → injects
-//     BraintrustObservabilityExporter when not already present
-// So this snippet exercises the truly zero-line integration path.
+// `braintrust/hook.mjs` wraps the Mastra constructors in tracing channels, and
+// the Braintrust import in the shared provider runtime subscribes the Mastra
+// plugin that injects observability when it is missing.
 const mastra = new Mastra({
   agents: { "weather-agent": weatherAgent },
   workflows: { "travel-flow": travelWorkflow },

js/src/auto-instrumentations/README.md

@@ -137,7 +137,8 @@ The plugin system follows the OpenTelemetry pattern:
 
 - **Integration Configs**: `src/configs/` in this package
   - SDK-specific instrumentation configurations
-  - Consumed by orchestrion-js for transformation
+  - Orchestrion configs describe function transformations
+  - Module export patch configs describe constructor tracing channels
 
 ### Auto-Enable Mechanism
 
@@ -167,7 +168,7 @@ Configuration is merged in the following order (later overrides earlier):
 ```typescript
 import type { InstrumentationConfig } from "braintrust/vite";
 
-export const anthropicConfigs: InstrumentationConfig[] = [
+export const anthropicOrchestrionConfigs: InstrumentationConfig[] = [
   {
     channelName: "messages.create",
     module: {
@@ -187,7 +188,7 @@ export const anthropicConfigs: InstrumentationConfig[] = [
 2. **Export the configs** from `src/index.ts`:
 
 ```typescript
-export { anthropicConfigs } from "./configs/anthropic";
+export { anthropicOrchestrionConfigs } from "./configs/anthropic";
 ```
 
 3. **Add channel handlers** in `BraintrustPlugin` (in the main `braintrust` package):

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

@@ -1,11 +1,22 @@
 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 {
+  getDefaultModuleExportPatchConfigs,
+  getDefaultOrchestrionConfigs,
+} from "../configs/all";
+import { applySourcePatch } from "../loader/source-patches";
+import {
+  buildModuleExportSourceWrapper,
+  type ModuleExportPatchTarget,
+} from "../loader/module-hooks/registry";
+import { readDisabledInstrumentationEnvConfig } from "../../instrumentation/config";
+
+const MODULE_EXPORT_ORIGINAL_IMPORT_PREFIX = "braintrust-top-level-original:";
+const MODULE_EXPORT_ORIGINAL_RESOLVED_PREFIX = `\0${MODULE_EXPORT_ORIGINAL_IMPORT_PREFIX}`;
 
 export interface LegacyBundlerPluginOptions {
   /**
@@ -14,7 +25,7 @@ export interface LegacyBundlerPluginOptions {
   debug?: boolean;
 
   /**
-   * Additional instrumentation configs to apply
+   * Additional Orchestrion configs to apply
    */
   instrumentations?: InstrumentationConfig[];
 
@@ -35,7 +46,7 @@ export interface BundlerPluginOptions {
   debug?: boolean;
 
   /**
-   * Additional instrumentation configs to apply
+   * Additional Orchestrion configs to apply
    */
   instrumentations?: InstrumentationConfig[];
 
@@ -72,24 +83,66 @@ function getModuleVersion(basedir: string): string | undefined {
 
 export const unplugin = createUnplugin<LegacyBundlerPluginOptions>(
   (options = {}) => {
-    const allInstrumentations = getDefaultInstrumentationConfigs({
-      additionalInstrumentations: options.instrumentations,
+    const disabledIntegrationConfig = readDisabledInstrumentationEnvConfig(
+      process.env.BRAINTRUST_DISABLE_INSTRUMENTATION,
+    ).integrations;
+    const orchestrionConfigs = getDefaultOrchestrionConfigs({
+      additionalOrchestrionConfigs: options.instrumentations,
+      disabledIntegrationConfig,
+    });
+    const moduleExportPatchTarget: ModuleExportPatchTarget =
+      options.browser === false ? "node" : "browser";
+    const moduleExportPatchConfigs = getDefaultModuleExportPatchConfigs({
+      disabledIntegrationConfig,
+      target: moduleExportPatchTarget,
     });
+    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";
 
     // Create the code transformer instrumentor
-    const instrumentationMatcher = create(allInstrumentations, dcModule);
+    const instrumentationMatcher = create(orchestrionConfigs, dcModule);
 
     return {
       name: "code-transformer",
       enforce: "pre",
+      resolveId(id: string, importer?: string) {
+        if (id.startsWith(MODULE_EXPORT_ORIGINAL_IMPORT_PREFIX)) {
+          return `${MODULE_EXPORT_ORIGINAL_RESOLVED_PREFIX}${id.slice(
+            MODULE_EXPORT_ORIGINAL_IMPORT_PREFIX.length,
+          )}`;
+        }
+        if (
+          id.startsWith(".") &&
+          importer?.startsWith(MODULE_EXPORT_ORIGINAL_RESOLVED_PREFIX)
+        ) {
+          const originalDirectory = originalDirectories.get(importer);
+          if (originalDirectory) {
+            return resolve(originalDirectory, id);
+          }
+        }
+        return null;
+      },
+      loadInclude(id: string) {
+        return id.startsWith(MODULE_EXPORT_ORIGINAL_RESOLVED_PREFIX);
+      },
+      load(id: string) {
+        if (id.startsWith(MODULE_EXPORT_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(MODULE_EXPORT_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
@@ -122,26 +175,54 @@ export const unplugin = createUnplugin<LegacyBundlerPluginOptions>(
         // Normalize the module path for Windows compatibility (WASM transformer expects forward slashes)
         const normalizedModulePath = moduleDetails.path.replace(/\\/g, "/");
         const moduleVersion = getModuleVersion(moduleDetails.basedir);
+        const moduleType = isModule ? "esm" : "cjs";
+        let nextCode = code;
+        let didPatch = false;
 
-        // Per-package source patches (see loader/special-case-patches.ts).
+        // Per-package source patches (see loader/source-patches/).
         // 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`.
         if (options.browser !== true) {
-          const patched = applySpecialCasePatch({
+          const patched = applySourcePatch({
             packageName: moduleName,
             modulePath: normalizedModulePath,
-            source: code,
-            format: isModule ? "esm" : "cjs",
+            source: nextCode,
+            format: moduleType,
           });
           if (patched !== null) {
-            return { code: patched, map: null };
+            nextCode = patched;
+            didPatch = true;
           }
         }
 
+        const originalModuleSpecifier = `${MODULE_EXPORT_ORIGINAL_IMPORT_PREFIX}${nextOriginalId++}`;
+        const originalModuleId = `${MODULE_EXPORT_ORIGINAL_RESOLVED_PREFIX}${originalModuleSpecifier.slice(
+          MODULE_EXPORT_ORIGINAL_IMPORT_PREFIX.length,
+        )}`;
+        const moduleExportWrapper = buildModuleExportSourceWrapper(
+          moduleExportPatchConfigs,
+          {
+            format: moduleType,
+            modulePath: normalizedModulePath,
+            moduleVersion,
+            originalModuleSpecifier,
+            packageName: moduleName,
+            source: nextCode,
+            target: moduleExportPatchTarget,
+          },
+        );
+        if (moduleExportWrapper !== null) {
+          originalSources.set(originalModuleId, nextCode);
+          originalDirectories.set(originalModuleId, dirname(filePath));
+          nextCode = moduleExportWrapper;
+          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 +238,12 @@ 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;",
@@ -176,7 +256,7 @@ export const unplugin = createUnplugin<LegacyBundlerPluginOptions>(
         } catch (error) {
           // If transformation fails, warn and return original code
           console.warn(`Code transformation failed for ${id}: ${error}`);
-          return null;
+          return didPatch ? { code: nextCode, map: null } : null;
         }
       },
     };

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

@@ -25,9 +25,23 @@ import { create } from "../orchestrion-js";
 import { extname, join, sep } from "path";
 import { readFileSync } from "fs";
 import moduleDetailsFromPath from "module-details-from-path";
-import { getDefaultInstrumentationConfigs } from "../configs/all";
+import {
+  getDefaultModuleExportPatchConfigs,
+  getDefaultOrchestrionConfigs,
+} from "../configs/all";
+import { applySourcePatch } from "../loader/source-patches";
+import {
+  buildModuleExportSourceWrapper,
+  type ModuleExportPatchTarget,
+} from "../loader/module-hooks/registry";
+import { readDisabledInstrumentationEnvConfig } from "../../instrumentation/config";
 import { type LegacyBundlerPluginOptions } from "./plugin";
 
+const MODULE_EXPORT_ORIGINAL_QUERY = "braintrust-top-level-original";
+const disabledIntegrationConfig = readDisabledInstrumentationEnvConfig(
+  process.env.BRAINTRUST_DISABLE_INSTRUMENTATION,
+).integrations;
+
 /**
  * Helper function to get module version from package.json
  */
@@ -55,11 +69,12 @@ const matcherCache = new Map<string, Matcher>();
  * Get or create a matcher instance, caching by config hash
  */
 function getMatcher(options: LegacyBundlerPluginOptions): Matcher {
-  const allInstrumentations = getDefaultInstrumentationConfigs({
-    additionalInstrumentations: options.instrumentations,
+  const orchestrionConfigs = getDefaultOrchestrionConfigs({
+    additionalOrchestrionConfigs: options.instrumentations,
+    disabledIntegrationConfig,
   });
   const dcModule = options.browser ? "dc-browser" : undefined;
-  const configHash = JSON.stringify({ allInstrumentations, dcModule });
+  const configHash = JSON.stringify({ orchestrionConfigs, dcModule });
 
   if (matcherCache.has(configHash)) {
     return matcherCache.get(configHash)!;
@@ -71,7 +86,7 @@ function getMatcher(options: LegacyBundlerPluginOptions): Matcher {
     }
   }
 
-  const matcher = create(allInstrumentations, dcModule ?? null);
+  const matcher = create(orchestrionConfigs, dcModule ?? null);
   matcherCache.set(configHash, matcher);
   return matcher;
 }
@@ -94,11 +109,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(MODULE_EXPORT_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 +141,52 @@ 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: ModuleExportPatchTarget =
+    options.browser === true ? "browser" : "node";
+  const moduleExportPatchConfigs = getDefaultModuleExportPatchConfigs({
+    disabledIntegrationConfig,
+    target,
+  });
+
+  let nextCode = code;
+  let didPatch = false;
+
+  if (options.browser !== true) {
+    const patched = applySourcePatch({
+      format: moduleType,
+      modulePath: normalizedModulePath,
+      packageName: moduleName,
+      source: nextCode,
+    });
+    if (patched !== null) {
+      nextCode = patched;
+      didPatch = true;
+    }
+  }
+
+  const moduleExportWrapper = buildModuleExportSourceWrapper(
+    moduleExportPatchConfigs,
+    {
+      format: moduleType,
+      modulePath: normalizedModulePath,
+      moduleVersion,
+      originalModuleSpecifier: `${resourcePath}?${MODULE_EXPORT_ORIGINAL_QUERY}`,
+      packageName: moduleName,
+      source: nextCode,
+      target,
+    },
+  );
+  if (moduleExportWrapper !== null) {
+    nextCode = moduleExportWrapper;
+    didPatch = true;
+  }
+
+  if (!moduleVersion) {
+    return callback(null, nextCode, inputSourceMap);
+  }
 
   const matcher = getMatcher(options);
   const transformer = matcher.getTransformer(
@@ -137,19 +196,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);
   }
 }