From 2f041bad21f62a45ad502fe6b28a6b779d7553bf Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 13 Jul 2026 16:03:07 +0200 Subject: [PATCH 1/2] ref: Patch Mastra using top level export patching --- js/src/auto-instrumentations/README.md | 7 +- .../auto-instrumentations/bundler/plugin.ts | 115 +++- .../bundler/webpack-loader.ts | 84 ++- .../auto-instrumentations/configs/ai-sdk.ts | 4 +- .../auto-instrumentations/configs/all.test.ts | 58 ++ js/src/auto-instrumentations/configs/all.ts | 225 ++++---- .../configs/anthropic.ts | 4 +- .../configs/bedrock-runtime.test.ts | 104 ---- .../configs/bedrock-runtime.ts | 2 +- .../configs/claude-agent-sdk.ts | 4 +- .../auto-instrumentations/configs/cohere.ts | 2 +- .../configs/cursor-sdk.ts | 2 +- .../configs/flue.test.ts | 63 --- js/src/auto-instrumentations/configs/flue.ts | 2 +- .../auto-instrumentations/configs/genkit.ts | 4 +- .../configs/github-copilot.ts | 2 +- .../configs/google-adk.ts | 4 +- .../configs/google-genai.ts | 4 +- js/src/auto-instrumentations/configs/groq.ts | 2 +- .../configs/huggingface.ts | 2 +- .../configs/langchain.ts | 2 +- .../auto-instrumentations/configs/mastra.ts | 53 ++ .../auto-instrumentations/configs/mistral.ts | 2 +- .../configs/openai-agents.ts | 2 +- .../configs/openai-codex.ts | 2 +- .../auto-instrumentations/configs/openai.ts | 4 +- .../configs/openrouter-agent.ts | 2 +- .../configs/openrouter.ts | 2 +- .../configs/pi-coding-agent.test.ts | 52 -- .../configs/pi-coding-agent.ts | 2 +- .../configs/strands-agent-sdk.test.ts | 78 --- .../configs/strands-agent-sdk.ts | 2 +- js/src/auto-instrumentations/hook.mts | 112 ++-- .../import-in-the-middle/create-hook.mts | 22 +- .../lib/get-esm-exports.mts | 18 +- js/src/auto-instrumentations/index.ts | 44 +- .../loader/{cjs-patch.ts => cjs.ts} | 8 +- .../loader/{esm-hook.mts => esm.mts} | 30 +- .../loader/mastra-observability-patch.test.ts | 117 ---- .../loader/mastra-observability-patch.ts | 288 ---------- .../loader/module-hooks/iitm.ts | 150 +++++ .../loader/module-hooks/node-runtime.ts | 27 + .../loader/module-hooks/node.test.ts | 310 ++++++++++ .../loader/module-hooks/node.ts | 123 ++++ .../loader/module-hooks/registry.test.ts | 308 ++++++++++ .../loader/module-hooks/registry.ts | 531 ++++++++++++++++++ .../loader/module-hooks/ritm.ts | 274 +++++++++ ...-package-version.ts => package-version.ts} | 8 +- .../index.ts} | 39 +- .../openai.ts} | 0 .../require-in-the-middle/index.ts | 239 -------- .../require-in-the-middle/tsconfig.json | 4 - js/src/imports.test.ts | 12 +- .../instrumentation/braintrust-plugin.test.ts | 33 ++ js/src/instrumentation/braintrust-plugin.ts | 17 +- .../plugins/mastra-channels.ts | 5 + .../plugins/mastra-plugin.test.ts | 231 ++++++++ .../instrumentation/plugins/mastra-plugin.ts | 114 ++++ .../node/apply-auto-instrumentation-entry.ts | 45 +- js/src/node/config.ts | 28 +- js/src/wrappers/mastra.ts | 14 +- .../auto-instrumentations/fixtures/.gitignore | 5 + .../fixtures/import-hook-query-mode.mjs | 15 + .../node_modules/@mastra/core/dist/index.cjs | 3 + .../node_modules/@mastra/core/dist/index.js | 1 + .../node_modules/@mastra/core/dist/mastra.cjs | 8 + .../node_modules/@mastra/core/dist/mastra.js | 6 + .../@mastra/core/dist/mastra/index.cjs | 3 + .../@mastra/core/dist/mastra/index.js | 1 + .../node_modules/@mastra/core/package.json | 15 + .../@mastra/observability/dist/index.cjs | 7 + .../@mastra/observability/dist/index.js | 5 + .../@mastra/observability/package.json | 11 + ...untime-apply-auto-mastra-top-level-esm.mjs | 3 + .../fixtures/test-mastra-bundler.js | 3 + .../fixtures/test-mastra-top-level-cjs.cjs | 41 ++ .../fixtures/test-mastra-top-level-esm.mjs | 45 ++ .../fixtures/vendor-hooks/ritm-app.cjs | 4 +- .../vendor-hooks/ritm-coexist-app.cjs | 4 +- .../import-require-in-the-middle.test.ts | 2 +- .../auto-instrumentations/loader-hook.test.ts | 144 +++++ .../transformation.test.ts | 134 +++++ js/tsconfig.vendor-hooks.json | 2 +- js/tsup.config.ts | 6 +- 84 files changed, 3227 insertions(+), 1289 deletions(-) create mode 100644 js/src/auto-instrumentations/configs/all.test.ts delete mode 100644 js/src/auto-instrumentations/configs/bedrock-runtime.test.ts delete mode 100644 js/src/auto-instrumentations/configs/flue.test.ts create mode 100644 js/src/auto-instrumentations/configs/mastra.ts delete mode 100644 js/src/auto-instrumentations/configs/pi-coding-agent.test.ts delete mode 100644 js/src/auto-instrumentations/configs/strands-agent-sdk.test.ts rename js/src/auto-instrumentations/loader/{cjs-patch.ts => cjs.ts} (92%) rename js/src/auto-instrumentations/loader/{esm-hook.mts => esm.mts} (85%) delete mode 100644 js/src/auto-instrumentations/loader/mastra-observability-patch.test.ts delete mode 100644 js/src/auto-instrumentations/loader/mastra-observability-patch.ts create mode 100644 js/src/auto-instrumentations/loader/module-hooks/iitm.ts create mode 100644 js/src/auto-instrumentations/loader/module-hooks/node-runtime.ts create mode 100644 js/src/auto-instrumentations/loader/module-hooks/node.test.ts create mode 100644 js/src/auto-instrumentations/loader/module-hooks/node.ts create mode 100644 js/src/auto-instrumentations/loader/module-hooks/registry.test.ts create mode 100644 js/src/auto-instrumentations/loader/module-hooks/registry.ts create mode 100644 js/src/auto-instrumentations/loader/module-hooks/ritm.ts rename js/src/auto-instrumentations/loader/{get-package-version.ts => package-version.ts} (90%) rename js/src/auto-instrumentations/loader/{special-case-patches.ts => source-patches/index.ts} (60%) rename js/src/auto-instrumentations/loader/{openai-api-promise-patch.ts => source-patches/openai.ts} (100%) delete mode 100644 js/src/auto-instrumentations/require-in-the-middle/index.ts delete mode 100644 js/src/auto-instrumentations/require-in-the-middle/tsconfig.json create mode 100644 js/src/instrumentation/plugins/mastra-channels.ts create mode 100644 js/src/instrumentation/plugins/mastra-plugin.test.ts create mode 100644 js/src/instrumentation/plugins/mastra-plugin.ts create mode 100644 js/tests/auto-instrumentations/fixtures/import-hook-query-mode.mjs create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/index.cjs create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/index.js create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra.cjs create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra.js create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra/index.cjs create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra/index.js create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/package.json create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/dist/index.cjs create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/dist/index.js create mode 100644 js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json create mode 100644 js/tests/auto-instrumentations/fixtures/runtime-apply-auto-mastra-top-level-esm.mjs create mode 100644 js/tests/auto-instrumentations/fixtures/test-mastra-bundler.js create mode 100644 js/tests/auto-instrumentations/fixtures/test-mastra-top-level-cjs.cjs create mode 100644 js/tests/auto-instrumentations/fixtures/test-mastra-top-level-esm.mjs diff --git a/js/src/auto-instrumentations/README.md b/js/src/auto-instrumentations/README.md index f2a319210..65c570e62 100644 --- a/js/src/auto-instrumentations/README.md +++ b/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): diff --git a/js/src/auto-instrumentations/bundler/plugin.ts b/js/src/auto-instrumentations/bundler/plugin.ts index dda1b7490..cb890b067 100644 --- a/js/src/auto-instrumentations/bundler/plugin.ts +++ b/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,63 @@ function getModuleVersion(basedir: string): string | undefined { export const unplugin = createUnplugin( (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(); + const originalDirectories = new Map(); + 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; + }, + 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 +172,54 @@ export const unplugin = createUnplugin( // 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 +235,12 @@ export const unplugin = createUnplugin( 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 +253,7 @@ export const unplugin = createUnplugin( } 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; } }, }; diff --git a/js/src/auto-instrumentations/bundler/webpack-loader.ts b/js/src/auto-instrumentations/bundler/webpack-loader.ts index c218c21a6..22e1c759c 100644 --- a/js/src/auto-instrumentations/bundler/webpack-loader.ts +++ b/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(); * 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); } } diff --git a/js/src/auto-instrumentations/configs/ai-sdk.ts b/js/src/auto-instrumentations/configs/ai-sdk.ts index b7e3bfa7e..b4b58e0db 100644 --- a/js/src/auto-instrumentations/configs/ai-sdk.ts +++ b/js/src/auto-instrumentations/configs/ai-sdk.ts @@ -2,7 +2,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { aiSDKChannels } from "../../instrumentation/plugins/ai-sdk-channels"; /** - * Instrumentation configurations for the Vercel AI SDK. + * Orchestrion configurations for the Vercel AI SDK. * * These configs define which functions to instrument and what channel * to emit events on. They are used by orchestrion-js to perform AST @@ -12,7 +12,7 @@ import { aiSDKChannels } from "../../instrumentation/plugins/ai-sdk-channels"; * will prepend "orchestrion:ai-sdk:" to these names, resulting in final channel names like: * "orchestrion:ai-sdk:generateText" */ -export const aiSDKConfigs: InstrumentationConfig[] = [ +export const aiSDKOrchestrionConfigs: InstrumentationConfig[] = [ // generateText - async function { channelName: aiSDKChannels.generateText.channelName, diff --git a/js/src/auto-instrumentations/configs/all.test.ts b/js/src/auto-instrumentations/configs/all.test.ts new file mode 100644 index 000000000..99ec1d631 --- /dev/null +++ b/js/src/auto-instrumentations/configs/all.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, it } from "vitest"; +import { mastraChannels } from "../../instrumentation/plugins/mastra-channels"; +import { getModuleExportPatchSpecifiers } from "../loader/module-hooks/registry"; +import { getDefaultModuleExportPatchConfigs } from "./all"; +import { mastraModuleExportPatchConfigs } from "./mastra"; + +describe("module export patch configs", () => { + it("registers versioned Mastra constructor channels for Node", () => { + const configs = getDefaultModuleExportPatchConfigs({ target: "node" }); + + expect(configs).toEqual(mastraModuleExportPatchConfigs); + expect(getModuleExportPatchSpecifiers(configs)).toEqual([ + "@mastra/core", + "@mastra/core/mastra", + "@mastra/observability", + ]); + expect(configs[0].modules.map((module) => module.versionRange)).toEqual([ + ">=1.20.0", + ">=1.20.0", + ">=1.20.0", + ]); + expect(configs[0].modules.map((module) => module.patches)).toEqual([ + [ + { + channelName: mastraChannels.mastraConstructor, + exportName: "Mastra", + kind: "constructor", + }, + ], + [ + { + channelName: mastraChannels.mastraConstructor, + exportName: "Mastra", + kind: "constructor", + }, + ], + [ + { + channelName: mastraChannels.observabilityConstructor, + exportName: "Observability", + kind: "constructor", + }, + ], + ]); + }); + + it("filters Mastra configs for browser and disabled integrations", () => { + expect(getDefaultModuleExportPatchConfigs({ target: "browser" })).toEqual( + [], + ); + expect( + getDefaultModuleExportPatchConfigs({ + disabledIntegrationConfig: { mastra: false }, + target: "node", + }), + ).toEqual([]); + }); +}); diff --git a/js/src/auto-instrumentations/configs/all.ts b/js/src/auto-instrumentations/configs/all.ts index b4642974a..db6a2aa61 100644 --- a/js/src/auto-instrumentations/configs/all.ts +++ b/js/src/auto-instrumentations/configs/all.ts @@ -1,115 +1,120 @@ import type { InstrumentationConfig } from "../orchestrion-js"; +import { + filterModuleExportPatchConfigs, + type ModuleExportPatchConfig, + type ModuleExportPatchTarget, +} from "../loader/module-hooks/registry"; import { isInstrumentationIntegrationDisabled, readDisabledInstrumentationEnvConfig, type InstrumentationIntegrationsConfig, } from "../../instrumentation/config"; -import { aiSDKConfigs } from "./ai-sdk"; -import { anthropicConfigs } from "./anthropic"; -import { bedrockRuntimeConfigs } from "./bedrock-runtime"; -import { claudeAgentSDKConfigs } from "./claude-agent-sdk"; -import { cohereConfigs } from "./cohere"; -import { cursorSDKConfigs } from "./cursor-sdk"; -import { flueConfigs } from "./flue"; -import { genkitConfigs } from "./genkit"; -import { gitHubCopilotConfigs } from "./github-copilot"; -import { googleADKConfigs } from "./google-adk"; -import { googleGenAIConfigs } from "./google-genai"; -import { groqConfigs } from "./groq"; -import { huggingFaceConfigs } from "./huggingface"; -import { langchainConfigs } from "./langchain"; -import { mistralConfigs } from "./mistral"; -import { openAIAgentsCoreConfigs } from "./openai-agents"; -import { openaiConfigs } from "./openai"; -import { openAICodexConfigs } from "./openai-codex"; -import { openRouterConfigs } from "./openrouter"; -import { openRouterAgentConfigs } from "./openrouter-agent"; -import { piCodingAgentConfigs } from "./pi-coding-agent"; -import { strandsAgentSDKConfigs } from "./strands-agent-sdk"; +import { aiSDKOrchestrionConfigs } from "./ai-sdk"; +import { anthropicOrchestrionConfigs } from "./anthropic"; +import { bedrockRuntimeOrchestrionConfigs } from "./bedrock-runtime"; +import { claudeAgentSDKOrchestrionConfigs } from "./claude-agent-sdk"; +import { cohereOrchestrionConfigs } from "./cohere"; +import { cursorSDKOrchestrionConfigs } from "./cursor-sdk"; +import { flueOrchestrionConfigs } from "./flue"; +import { genkitOrchestrionConfigs } from "./genkit"; +import { gitHubCopilotOrchestrionConfigs } from "./github-copilot"; +import { googleADKOrchestrionConfigs } from "./google-adk"; +import { googleGenAIOrchestrionConfigs } from "./google-genai"; +import { groqOrchestrionConfigs } from "./groq"; +import { huggingFaceOrchestrionConfigs } from "./huggingface"; +import { langchainOrchestrionConfigs } from "./langchain"; +import { mistralOrchestrionConfigs } from "./mistral"; +import { mastraModuleExportPatchConfigs } from "./mastra"; +import { openAIAgentsCoreOrchestrionConfigs } from "./openai-agents"; +import { openaiOrchestrionConfigs } from "./openai"; +import { openAICodexOrchestrionConfigs } from "./openai-codex"; +import { openRouterOrchestrionConfigs } from "./openrouter"; +import { openRouterAgentOrchestrionConfigs } from "./openrouter-agent"; +import { piCodingAgentOrchestrionConfigs } from "./pi-coding-agent"; +import { strandsAgentSDKOrchestrionConfigs } from "./strands-agent-sdk"; -interface InstrumentationConfigGroup { +interface OrchestrionConfigGroup { integrations: readonly (keyof InstrumentationIntegrationsConfig)[]; configs: readonly InstrumentationConfig[]; } -const defaultInstrumentationConfigGroups: readonly InstrumentationConfigGroup[] = - [ - { integrations: ["openai"], configs: openaiConfigs }, - { - integrations: ["openaiCodexSDK"], - configs: openAICodexConfigs, - }, - { integrations: ["anthropic"], configs: anthropicConfigs }, - { - integrations: ["bedrock", "awsBedrock", "awsBedrockRuntime"], - configs: bedrockRuntimeConfigs, - }, - { - integrations: ["aisdk", "vercel"], - configs: aiSDKConfigs, - }, - { - integrations: ["claudeAgentSDK"], - configs: claudeAgentSDKConfigs, - }, - { integrations: ["cursor", "cursorSDK"], configs: cursorSDKConfigs }, - { - integrations: ["openAIAgents"], - configs: openAIAgentsCoreConfigs, - }, - { - integrations: ["google", "googleGenAI"], - configs: googleGenAIConfigs, - }, - { integrations: ["huggingface"], configs: huggingFaceConfigs }, - { - integrations: ["langchain", "langgraph"], - configs: langchainConfigs, - }, - { integrations: ["openrouter"], configs: openRouterConfigs }, - { - integrations: ["openrouterAgent"], - configs: openRouterAgentConfigs, - }, - { integrations: ["mistral"], configs: mistralConfigs }, - { integrations: ["googleADK"], configs: googleADKConfigs }, - { integrations: ["cohere"], configs: cohereConfigs }, - { integrations: ["groq"], configs: groqConfigs }, - { - integrations: ["genkit"], - configs: genkitConfigs, - }, - { - integrations: ["gitHubCopilot"], - configs: gitHubCopilotConfigs, - }, - { - integrations: ["piCodingAgent"], - configs: piCodingAgentConfigs, - }, - { - integrations: ["strandsAgentSDK"], - configs: strandsAgentSDKConfigs, - }, - { - integrations: ["flue"], - 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. - ]; +const defaultOrchestrionConfigGroups: readonly OrchestrionConfigGroup[] = [ + { integrations: ["openai"], configs: openaiOrchestrionConfigs }, + { + integrations: ["openaiCodexSDK"], + configs: openAICodexOrchestrionConfigs, + }, + { integrations: ["anthropic"], configs: anthropicOrchestrionConfigs }, + { + integrations: ["bedrock", "awsBedrock", "awsBedrockRuntime"], + configs: bedrockRuntimeOrchestrionConfigs, + }, + { + integrations: ["aisdk", "vercel"], + configs: aiSDKOrchestrionConfigs, + }, + { + integrations: ["claudeAgentSDK"], + configs: claudeAgentSDKOrchestrionConfigs, + }, + { + integrations: ["cursor", "cursorSDK"], + configs: cursorSDKOrchestrionConfigs, + }, + { + integrations: ["openAIAgents"], + configs: openAIAgentsCoreOrchestrionConfigs, + }, + { + integrations: ["google", "googleGenAI"], + configs: googleGenAIOrchestrionConfigs, + }, + { integrations: ["huggingface"], configs: huggingFaceOrchestrionConfigs }, + { + integrations: ["langchain", "langgraph"], + configs: langchainOrchestrionConfigs, + }, + { integrations: ["openrouter"], configs: openRouterOrchestrionConfigs }, + { + integrations: ["openrouterAgent"], + configs: openRouterAgentOrchestrionConfigs, + }, + { integrations: ["mistral"], configs: mistralOrchestrionConfigs }, + { integrations: ["googleADK"], configs: googleADKOrchestrionConfigs }, + { integrations: ["cohere"], configs: cohereOrchestrionConfigs }, + { integrations: ["groq"], configs: groqOrchestrionConfigs }, + { + integrations: ["genkit"], + configs: genkitOrchestrionConfigs, + }, + { + integrations: ["gitHubCopilot"], + configs: gitHubCopilotOrchestrionConfigs, + }, + { + integrations: ["piCodingAgent"], + configs: piCodingAgentOrchestrionConfigs, + }, + { + integrations: ["strandsAgentSDK"], + configs: strandsAgentSDKOrchestrionConfigs, + }, + { + integrations: ["flue"], + configs: flueOrchestrionConfigs, + }, +]; + +const defaultModuleExportPatchConfigs: readonly ModuleExportPatchConfig[] = [ + ...mastraModuleExportPatchConfigs, +]; -export function getDefaultInstrumentationConfigs({ - additionalInstrumentations, +export function getDefaultOrchestrionConfigs({ + additionalOrchestrionConfigs, disabledIntegrationConfig, disabledIntegrations, }: { - additionalInstrumentations?: readonly InstrumentationConfig[]; + additionalOrchestrionConfigs?: readonly InstrumentationConfig[]; disabledIntegrationConfig?: InstrumentationIntegrationsConfig; disabledIntegrations?: ReadonlySet; } = {}): InstrumentationConfig[] { @@ -122,20 +127,24 @@ export function getDefaultInstrumentationConfigs({ : undefined); return [ - ...defaultInstrumentationConfigGroups.flatMap( - ({ configs, integrations }) => - isInstrumentationIntegrationDisabled(disabledConfig, ...integrations) - ? [] - : configs, + ...defaultOrchestrionConfigGroups.flatMap(({ configs, integrations }) => + isInstrumentationIntegrationDisabled(disabledConfig, ...integrations) + ? [] + : configs, ), - ...(additionalInstrumentations ?? []), + ...(additionalOrchestrionConfigs ?? []), ]; } -export function getDefaultAutoInstrumentationConfigs(): InstrumentationConfig[] { - return getDefaultInstrumentationConfigs({ - disabledIntegrationConfig: readDisabledInstrumentationEnvConfig( - process.env.BRAINTRUST_DISABLE_INSTRUMENTATION, - ).integrations, +export function getDefaultModuleExportPatchConfigs({ + disabledIntegrationConfig, + target, +}: { + disabledIntegrationConfig?: InstrumentationIntegrationsConfig; + target: ModuleExportPatchTarget; +}): ModuleExportPatchConfig[] { + return filterModuleExportPatchConfigs(defaultModuleExportPatchConfigs, { + disabledIntegrationConfig, + target, }); } diff --git a/js/src/auto-instrumentations/configs/anthropic.ts b/js/src/auto-instrumentations/configs/anthropic.ts index 61d4c48db..8969ab78c 100644 --- a/js/src/auto-instrumentations/configs/anthropic.ts +++ b/js/src/auto-instrumentations/configs/anthropic.ts @@ -2,7 +2,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { anthropicChannels } from "../../instrumentation/plugins/anthropic-channels"; /** - * Instrumentation configurations for the Anthropic SDK. + * Orchestrion configurations for the Anthropic SDK. * * These configs define which functions to instrument and what channel * to emit events on. They are used by orchestrion-js to perform AST @@ -12,7 +12,7 @@ import { anthropicChannels } from "../../instrumentation/plugins/anthropic-chann * will prepend "orchestrion:" + module.name + ":" to these names, resulting in final channel names like: * "orchestrion:@anthropic-ai/sdk:messages.create" */ -export const anthropicConfigs: InstrumentationConfig[] = [ +export const anthropicOrchestrionConfigs: InstrumentationConfig[] = [ // Each logical target is listed for both published module formats: // `.mjs` covers ESM imports, while `.js` covers CJS requires. The Bedrock // SDK delegates CJS `messages.create` calls through these Anthropic SDK diff --git a/js/src/auto-instrumentations/configs/bedrock-runtime.test.ts b/js/src/auto-instrumentations/configs/bedrock-runtime.test.ts deleted file mode 100644 index 77b3997b1..000000000 --- a/js/src/auto-instrumentations/configs/bedrock-runtime.test.ts +++ /dev/null @@ -1,104 +0,0 @@ -import { create, type ModuleType } from "@apm-js-collab/code-transformer"; -import { describe, expect, it } from "vitest"; -import { bedrockRuntimeConfigs } from "./bedrock-runtime"; - -const CJS_CLIENT_SOURCE = ` -class Client { - send(command, optionsOrCb, cb) { - if (typeof optionsOrCb === "function" || typeof cb === "function") { - return undefined; - } - return Promise.resolve(command); - } -} -module.exports = { Client }; -`; - -const ESM_CLIENT_SOURCE = ` -export class Client { - send(command, optionsOrCb, cb) { - if (typeof optionsOrCb === "function" || typeof cb === "function") { - return undefined; - } - return Promise.resolve(command); - } -} -`; - -describe("bedrockRuntimeConfigs", () => { - it("matches current and legacy Smithy Client.send entry files", () => { - const matcher = create(bedrockRuntimeConfigs); - - try { - for (const entry of [ - { - channelName: "orchestrion:@smithy/core:client.send", - moduleName: "@smithy/core", - moduleType: "cjs" as ModuleType, - path: "dist-cjs/submodules/client/index.js", - source: CJS_CLIENT_SOURCE, - version: "3.25.1", - }, - { - channelName: "orchestrion:@smithy/core:client.send", - moduleName: "@smithy/core", - moduleType: "esm" as ModuleType, - path: "dist-es/submodules/client/smithy-client/client.js", - source: ESM_CLIENT_SOURCE, - version: "3.25.1", - }, - { - channelName: "orchestrion:@smithy/smithy-client:client.send", - moduleName: "@smithy/smithy-client", - moduleType: "cjs" as ModuleType, - path: "dist-cjs/index.js", - source: CJS_CLIENT_SOURCE, - version: "3.0.0", - }, - { - channelName: "orchestrion:@smithy/smithy-client:client.send", - moduleName: "@smithy/smithy-client", - moduleType: "esm" as ModuleType, - path: "dist-es/client.js", - source: ESM_CLIENT_SOURCE, - version: "3.0.0", - }, - { - channelName: "orchestrion:@smithy/smithy-client:client.send", - moduleName: "@smithy/smithy-client", - moduleType: "cjs" as ModuleType, - path: "dist-cjs/index.js", - source: CJS_CLIENT_SOURCE, - version: "4.8.0", - }, - { - channelName: "orchestrion:@smithy/smithy-client:client.send", - moduleName: "@smithy/smithy-client", - moduleType: "esm" as ModuleType, - path: "dist-es/client.js", - source: ESM_CLIENT_SOURCE, - version: "4.8.0", - }, - ]) { - const transformer = matcher.getTransformer( - entry.moduleName, - entry.version, - entry.path, - ); - - try { - expect(transformer).toBeDefined(); - const transformed = transformer!.transform( - entry.source, - entry.moduleType, - ); - expect(transformed.code).toContain(entry.channelName); - } finally { - transformer?.free(); - } - } - } finally { - matcher.free(); - } - }); -}); diff --git a/js/src/auto-instrumentations/configs/bedrock-runtime.ts b/js/src/auto-instrumentations/configs/bedrock-runtime.ts index 29b4d0f0d..822fefaa4 100644 --- a/js/src/auto-instrumentations/configs/bedrock-runtime.ts +++ b/js/src/auto-instrumentations/configs/bedrock-runtime.ts @@ -4,7 +4,7 @@ import { smithyCoreChannels, } from "../../instrumentation/plugins/bedrock-runtime-channels"; -export const bedrockRuntimeConfigs: InstrumentationConfig[] = [ +export const bedrockRuntimeOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: smithyCoreChannels.clientSend.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/claude-agent-sdk.ts b/js/src/auto-instrumentations/configs/claude-agent-sdk.ts index 1cf2c6537..074886c1c 100644 --- a/js/src/auto-instrumentations/configs/claude-agent-sdk.ts +++ b/js/src/auto-instrumentations/configs/claude-agent-sdk.ts @@ -2,7 +2,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { claudeAgentSDKChannels } from "../../instrumentation/plugins/claude-agent-sdk-channels"; /** - * Instrumentation configuration for the Claude Agent SDK. + * Orchestrion configuration for the Claude Agent SDK. * * This config defines which functions to instrument and what channel * to emit events on. It is used by orchestrion-js to perform AST @@ -12,7 +12,7 @@ import { claudeAgentSDKChannels } from "../../instrumentation/plugins/claude-age * will prepend "orchestrion:claude-agent-sdk:" to these names, resulting in final channel * names like: "orchestrion:claude-agent-sdk:query" */ -export const claudeAgentSDKConfigs: InstrumentationConfig[] = [ +export const claudeAgentSDKOrchestrionConfigs: InstrumentationConfig[] = [ // query - Main entry point for agent interactions. The SDK returns an async // iterable, but the exported query function itself is synchronous. { diff --git a/js/src/auto-instrumentations/configs/cohere.ts b/js/src/auto-instrumentations/configs/cohere.ts index 106980884..caf19b7be 100644 --- a/js/src/auto-instrumentations/configs/cohere.ts +++ b/js/src/auto-instrumentations/configs/cohere.ts @@ -1,7 +1,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { cohereChannels } from "../../instrumentation/plugins/cohere-channels"; -export const cohereConfigs: InstrumentationConfig[] = [ +export const cohereOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: cohereChannels.chat.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/cursor-sdk.ts b/js/src/auto-instrumentations/configs/cursor-sdk.ts index ecb845b4b..4f4389cb5 100644 --- a/js/src/auto-instrumentations/configs/cursor-sdk.ts +++ b/js/src/auto-instrumentations/configs/cursor-sdk.ts @@ -5,7 +5,7 @@ const cursorSDKVersionRange = ">=1.0.7 <2.0.0"; const cursorSDKEntrypoints = ["dist/esm/index.js", "dist/cjs/index.js"]; -export const cursorSDKConfigs: InstrumentationConfig[] = +export const cursorSDKOrchestrionConfigs: InstrumentationConfig[] = cursorSDKEntrypoints.flatMap((filePath) => [ { channelName: cursorSDKChannels.create.channelName, diff --git a/js/src/auto-instrumentations/configs/flue.test.ts b/js/src/auto-instrumentations/configs/flue.test.ts deleted file mode 100644 index 456b5110f..000000000 --- a/js/src/auto-instrumentations/configs/flue.test.ts +++ /dev/null @@ -1,63 +0,0 @@ -import { create } from "@apm-js-collab/code-transformer"; -import { describe, expect, it } from "vitest"; -import { flueConfigs, flueVersionRange } from "./flue"; - -describe("flue auto-instrumentation configs", () => { - it("targets the stable Flue 0.8 context factory", () => { - const matcher = create(flueConfigs); - const flue08Transformer = matcher.getTransformer( - "@flue/runtime", - "0.8.0", - "dist/internal.mjs", - ); - const flue10BetaTransformer = matcher.getTransformer( - "@flue/runtime", - "1.0.0-beta.3", - "dist/internal.mjs", - ); - const flue10Transformer = matcher.getTransformer( - "@flue/runtime", - "1.0.0", - "dist/internal.mjs", - ); - - expect(flueVersionRange).toBe(">=0.8.0 <1.0.0"); - expect(flue08Transformer).toBeDefined(); - expect(flue10BetaTransformer).toBeUndefined(); - expect(flue10Transformer).toBeUndefined(); - const transformed = flue08Transformer!.transform( - ` -function createFlueContext(config) { - return { config }; -} -`, - "esm", - ).code; - - expect(transformed).toContain( - "orchestrion:@flue/runtime:createFlueContext", - ); - }); - - it("does not target Flue content-hashed workflow or tool chunks", () => { - const matcher = create(flueConfigs); - - expect( - matcher.getTransformer( - "@flue/runtime", - "0.8.0", - "dist/handle-agent-DcUclCE2.mjs", - ), - ).toBeUndefined(); - expect( - matcher.getTransformer( - "@flue/runtime", - "0.8.0", - "dist/sandbox-DNEJXjr_.mjs", - ), - ).toBeUndefined(); - expect( - matcher.getTransformer("@flue/runtime", "1.0.0-beta.3", "dist/index.mjs"), - ).toBeUndefined(); - }); -}); diff --git a/js/src/auto-instrumentations/configs/flue.ts b/js/src/auto-instrumentations/configs/flue.ts index a04c02bf2..5f911bee0 100644 --- a/js/src/auto-instrumentations/configs/flue.ts +++ b/js/src/auto-instrumentations/configs/flue.ts @@ -3,7 +3,7 @@ import { flueChannels } from "../../instrumentation/plugins/flue-channels"; export const flueVersionRange = ">=0.8.0 <1.0.0"; -export const flueConfigs: InstrumentationConfig[] = [ +export const flueOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: flueChannels.createContext.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/genkit.ts b/js/src/auto-instrumentations/configs/genkit.ts index 349635b01..c411d951b 100644 --- a/js/src/auto-instrumentations/configs/genkit.ts +++ b/js/src/auto-instrumentations/configs/genkit.ts @@ -7,13 +7,13 @@ import { const genkitVersionRange = ">=1.0.0 <2.0.0"; /** - * Instrumentation configurations for Genkit's JavaScript SDK. + * Orchestrion configurations for Genkit's JavaScript SDK. * * Genkit's public instance methods live on the GenkitAI base class in * @genkit-ai/ai. The top-level `genkit` package subclasses that class, so * targeting these methods instruments regular `genkit({ ... })` instances. */ -export const genkitConfigs: InstrumentationConfig[] = [ +export const genkitOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: genkitChannels.generate.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/github-copilot.ts b/js/src/auto-instrumentations/configs/github-copilot.ts index 901f9e682..2ad17c30d 100644 --- a/js/src/auto-instrumentations/configs/github-copilot.ts +++ b/js/src/auto-instrumentations/configs/github-copilot.ts @@ -1,7 +1,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { gitHubCopilotChannels } from "../../instrumentation/plugins/github-copilot-channels"; -export const gitHubCopilotConfigs: InstrumentationConfig[] = [ +export const gitHubCopilotOrchestrionConfigs: InstrumentationConfig[] = [ // ESM: CopilotClient.createSession { channelName: gitHubCopilotChannels.createSession.channelName, diff --git a/js/src/auto-instrumentations/configs/google-adk.ts b/js/src/auto-instrumentations/configs/google-adk.ts index aeb72c44b..d8622147e 100644 --- a/js/src/auto-instrumentations/configs/google-adk.ts +++ b/js/src/auto-instrumentations/configs/google-adk.ts @@ -6,7 +6,7 @@ const googleADKBundledIndexV06VersionRange = ">=0.6.1 <0.7.0"; const googleADKBundledIndexV1VersionRange = ">=1.0.0 <2.0.0"; /** - * Instrumentation configurations for the Google ADK (@google/adk). + * Orchestrion configurations for the Google ADK (@google/adk). * * Runner.runAsync and BaseAgent.runAsync are async generators (`async *`). * They synchronously return an AsyncGenerator object, so we use kind "Sync" @@ -16,7 +16,7 @@ const googleADKBundledIndexV1VersionRange = ">=1.0.0 <2.0.0"; * FunctionTool.runAsync is a regular async method (returns Promise) * and uses kind "Async". */ -export const googleADKConfigs: InstrumentationConfig[] = [ +export const googleADKOrchestrionConfigs: InstrumentationConfig[] = [ // --- Runner.runAsync --- async generator, kind "Sync" + sync-stream channel // Runner.runAsync — ESM individual module file diff --git a/js/src/auto-instrumentations/configs/google-genai.ts b/js/src/auto-instrumentations/configs/google-genai.ts index 16edd1d5c..d944845c3 100644 --- a/js/src/auto-instrumentations/configs/google-genai.ts +++ b/js/src/auto-instrumentations/configs/google-genai.ts @@ -2,7 +2,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { googleGenAIChannels } from "../../instrumentation/plugins/google-genai-channels"; /** - * Instrumentation configurations for the Google GenAI SDK. + * Orchestrion configurations for the Google GenAI SDK. * * These configs define which functions to instrument and what channel * to emit events on. They are used by orchestrion-js to perform AST @@ -12,7 +12,7 @@ import { googleGenAIChannels } from "../../instrumentation/plugins/google-genai- * will prepend "orchestrion:google-genai:" to these names, resulting in final channel names like: * "orchestrion:google-genai:models.generateContent" */ -export const googleGenAIConfigs: InstrumentationConfig[] = [ +export const googleGenAIOrchestrionConfigs: InstrumentationConfig[] = [ // Models.generateContentInternal - The actual class method (Node.js entry point) // Note: generateContent is an arrow function property that calls this internal method { diff --git a/js/src/auto-instrumentations/configs/groq.ts b/js/src/auto-instrumentations/configs/groq.ts index 864b2a166..050b550a2 100644 --- a/js/src/auto-instrumentations/configs/groq.ts +++ b/js/src/auto-instrumentations/configs/groq.ts @@ -1,7 +1,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { groqChannels } from "../../instrumentation/plugins/groq-channels"; -export const groqConfigs: InstrumentationConfig[] = [ +export const groqOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: groqChannels.chatCompletionsCreate.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/huggingface.ts b/js/src/auto-instrumentations/configs/huggingface.ts index f4c9bec06..12596e6a9 100644 --- a/js/src/auto-instrumentations/configs/huggingface.ts +++ b/js/src/auto-instrumentations/configs/huggingface.ts @@ -1,7 +1,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { huggingFaceChannels } from "../../instrumentation/plugins/huggingface-channels"; -export const huggingFaceConfigs: InstrumentationConfig[] = [ +export const huggingFaceOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: huggingFaceChannels.chatCompletion.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/langchain.ts b/js/src/auto-instrumentations/configs/langchain.ts index e8108ed02..a5942f9d0 100644 --- a/js/src/auto-instrumentations/configs/langchain.ts +++ b/js/src/auto-instrumentations/configs/langchain.ts @@ -4,7 +4,7 @@ import { langChainChannels } from "../../instrumentation/plugins/langchain-chann const langChainCoreVersionRange = ">=0.3.42"; const langChainCallbackManagerFilePath = "dist/callbacks/manager.js"; -export const langchainConfigs: InstrumentationConfig[] = [ +export const langchainOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: langChainChannels.configure.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/mastra.ts b/js/src/auto-instrumentations/configs/mastra.ts new file mode 100644 index 000000000..587165993 --- /dev/null +++ b/js/src/auto-instrumentations/configs/mastra.ts @@ -0,0 +1,53 @@ +import { mastraChannels } from "../../instrumentation/plugins/mastra-channels.js"; +import type { ModuleExportPatchConfig } from "../loader/module-hooks/registry.js"; + +const mastraConstructorPatches = [ + { + channelName: mastraChannels.mastraConstructor, + exportName: "Mastra", + kind: "constructor", + }, +] as const; + +export const mastraModuleExportPatchConfigs: readonly ModuleExportPatchConfig[] = + [ + { + integrations: ["mastra"], + modules: [ + { + packageName: "@mastra/core", + patches: mastraConstructorPatches, + source: { + modulePaths: ["dist/index.js", "dist/index.cjs"], + }, + specifier: "@mastra/core", + versionRange: ">=1.20.0", + }, + { + packageName: "@mastra/core", + patches: mastraConstructorPatches, + source: { + modulePaths: ["dist/mastra/index.js", "dist/mastra/index.cjs"], + }, + specifier: "@mastra/core/mastra", + versionRange: ">=1.20.0", + }, + { + packageName: "@mastra/observability", + patches: [ + { + channelName: mastraChannels.observabilityConstructor, + exportName: "Observability", + kind: "constructor", + }, + ], + source: { + modulePaths: ["dist/index.js", "dist/index.cjs"], + }, + specifier: "@mastra/observability", + versionRange: ">=1.20.0", + }, + ], + targets: ["node"], + }, + ]; diff --git a/js/src/auto-instrumentations/configs/mistral.ts b/js/src/auto-instrumentations/configs/mistral.ts index e61fc24be..04471721a 100644 --- a/js/src/auto-instrumentations/configs/mistral.ts +++ b/js/src/auto-instrumentations/configs/mistral.ts @@ -1,7 +1,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { mistralChannels } from "../../instrumentation/plugins/mistral-channels"; -export const mistralConfigs: InstrumentationConfig[] = [ +export const mistralOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: mistralChannels.chatComplete.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/openai-agents.ts b/js/src/auto-instrumentations/configs/openai-agents.ts index 3647edbc5..fda4674bc 100644 --- a/js/src/auto-instrumentations/configs/openai-agents.ts +++ b/js/src/auto-instrumentations/configs/openai-agents.ts @@ -8,7 +8,7 @@ const lifecycleMethods = [ ["onSpanEnd", openAIAgentsCoreChannels.onSpanEnd.channelName], ] as const; -export const openAIAgentsCoreConfigs: InstrumentationConfig[] = +export const openAIAgentsCoreOrchestrionConfigs: InstrumentationConfig[] = lifecycleMethods.flatMap(([methodName, channelName]) => ["dist/tracing/processor.mjs", "dist/tracing/processor.js"].map( (filePath) => ({ diff --git a/js/src/auto-instrumentations/configs/openai-codex.ts b/js/src/auto-instrumentations/configs/openai-codex.ts index 40ee60784..2570480b2 100644 --- a/js/src/auto-instrumentations/configs/openai-codex.ts +++ b/js/src/auto-instrumentations/configs/openai-codex.ts @@ -3,7 +3,7 @@ import { openAICodexChannels } from "../../instrumentation/plugins/openai-codex- const openAICodexVersionRange = ">=0.128.0 <1.0.0"; -export const openAICodexConfigs: InstrumentationConfig[] = [ +export const openAICodexOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: openAICodexChannels.run.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/openai.ts b/js/src/auto-instrumentations/configs/openai.ts index 28d85b7f6..fda5e8a32 100644 --- a/js/src/auto-instrumentations/configs/openai.ts +++ b/js/src/auto-instrumentations/configs/openai.ts @@ -2,7 +2,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { openAIChannels } from "../../instrumentation/plugins/openai-channels"; /** - * Instrumentation configurations for the OpenAI SDK. + * Orchestrion configurations for the OpenAI SDK. * * These configs define which functions to instrument and what channel * to emit events on. They are used by orchestrion-js to perform AST @@ -12,7 +12,7 @@ import { openAIChannels } from "../../instrumentation/plugins/openai-channels"; * will prepend "orchestrion:" + module.name + ":" to these names, resulting in final channel names like: * "orchestrion:openai:chat.completions.create" */ -export const openaiConfigs: InstrumentationConfig[] = [ +export const openaiOrchestrionConfigs: InstrumentationConfig[] = [ // Chat Completions { channelName: openAIChannels.chatCompletionsCreate.channelName, diff --git a/js/src/auto-instrumentations/configs/openrouter-agent.ts b/js/src/auto-instrumentations/configs/openrouter-agent.ts index f9f6191e1..ea32a46bf 100644 --- a/js/src/auto-instrumentations/configs/openrouter-agent.ts +++ b/js/src/auto-instrumentations/configs/openrouter-agent.ts @@ -1,7 +1,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { openRouterAgentChannels } from "../../instrumentation/plugins/openrouter-agent-channels"; -export const openRouterAgentConfigs: InstrumentationConfig[] = [ +export const openRouterAgentOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: openRouterAgentChannels.callModel.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/openrouter.ts b/js/src/auto-instrumentations/configs/openrouter.ts index 15bf47919..ff472d64e 100644 --- a/js/src/auto-instrumentations/configs/openrouter.ts +++ b/js/src/auto-instrumentations/configs/openrouter.ts @@ -1,7 +1,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { openRouterChannels } from "../../instrumentation/plugins/openrouter-channels"; -export const openRouterConfigs: InstrumentationConfig[] = [ +export const openRouterOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: openRouterChannels.chatSend.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/pi-coding-agent.test.ts b/js/src/auto-instrumentations/configs/pi-coding-agent.test.ts deleted file mode 100644 index b35b470cc..000000000 --- a/js/src/auto-instrumentations/configs/pi-coding-agent.test.ts +++ /dev/null @@ -1,52 +0,0 @@ -import { describe, expect, it } from "vitest"; -import { readDisabledInstrumentationEnvConfig } from "../../instrumentation/config"; -import { getDefaultInstrumentationConfigs } from "./all"; -import { piCodingAgentConfigs } from "./pi-coding-agent"; - -const piCodingAgentChannelName = "AgentSession.prompt"; - -describe("piCodingAgentConfigs", () => { - it("targets the Pi Coding Agent library AgentSession.prompt entrypoint", () => { - expect(piCodingAgentConfigs).toHaveLength(1); - expect(piCodingAgentConfigs[0]).toMatchObject({ - channelName: piCodingAgentChannelName, - module: { - name: "@earendil-works/pi-coding-agent", - versionRange: ">=0.79.0 <0.80.0", - filePath: "dist/core/agent-session.js", - }, - functionQuery: { - className: "AgentSession", - methodName: "prompt", - kind: "Async", - }, - }); - expect(piCodingAgentConfigs[0].module.filePath).not.toContain("cli"); - }); - - it("is included by default and disabled by Pi Coding Agent env aliases", () => { - expect( - getDefaultInstrumentationConfigs().some( - (config) => config.channelName === piCodingAgentChannelName, - ), - ).toBe(true); - - for (const alias of [ - "pi-coding-agent", - "pi-coding-agent-sdk", - "picodingagent", - "picodingagentsdk", - "@earendil-works/pi-coding-agent", - ]) { - const disabledConfig = - readDisabledInstrumentationEnvConfig(alias).integrations; - - expect(disabledConfig).toMatchObject({ piCodingAgent: false }); - expect( - getDefaultInstrumentationConfigs({ - disabledIntegrationConfig: disabledConfig, - }).some((config) => config.channelName === piCodingAgentChannelName), - ).toBe(false); - } - }); -}); diff --git a/js/src/auto-instrumentations/configs/pi-coding-agent.ts b/js/src/auto-instrumentations/configs/pi-coding-agent.ts index 660143369..ca4f2f289 100644 --- a/js/src/auto-instrumentations/configs/pi-coding-agent.ts +++ b/js/src/auto-instrumentations/configs/pi-coding-agent.ts @@ -3,7 +3,7 @@ import { piCodingAgentChannels } from "../../instrumentation/plugins/pi-coding-a const piCodingAgentVersionRange = ">=0.79.0 <0.80.0"; -export const piCodingAgentConfigs: InstrumentationConfig[] = [ +export const piCodingAgentOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: piCodingAgentChannels.prompt.channelName, module: { diff --git a/js/src/auto-instrumentations/configs/strands-agent-sdk.test.ts b/js/src/auto-instrumentations/configs/strands-agent-sdk.test.ts deleted file mode 100644 index 58dcff867..000000000 --- a/js/src/auto-instrumentations/configs/strands-agent-sdk.test.ts +++ /dev/null @@ -1,78 +0,0 @@ -import { describe, expect, it } from "vitest"; -import { readDisabledInstrumentationEnvConfig } from "../../instrumentation/config"; -import { getDefaultInstrumentationConfigs } from "./all"; -import { strandsAgentSDKConfigs } from "./strands-agent-sdk"; - -const strandsChannelNames = [ - "Agent.stream", - "Graph.stream", - "Swarm.stream", -] as const; - -describe("strandsAgentSDKConfigs", () => { - it("targets the Strands Agent SDK stream entrypoints", () => { - expect(strandsAgentSDKConfigs).toHaveLength(3); - expect(strandsAgentSDKConfigs.map((config) => config.channelName)).toEqual( - strandsChannelNames, - ); - expect(strandsAgentSDKConfigs).toEqual([ - expect.objectContaining({ - module: { - name: "@strands-agents/sdk", - versionRange: ">=1.0.0 <2.0.0", - filePath: "dist/src/agent/agent.js", - }, - functionQuery: { - className: "Agent", - methodName: "stream", - kind: "Sync", - }, - }), - expect.objectContaining({ - module: expect.objectContaining({ - filePath: "dist/src/multiagent/graph.js", - }), - functionQuery: expect.objectContaining({ - className: "Graph", - methodName: "stream", - }), - }), - expect.objectContaining({ - module: expect.objectContaining({ - filePath: "dist/src/multiagent/swarm.js", - }), - functionQuery: expect.objectContaining({ - className: "Swarm", - methodName: "stream", - }), - }), - ]); - }); - - it("is included by default and disabled by Strands Agent SDK integration keys", () => { - expect( - getDefaultInstrumentationConfigs().some((config) => - strandsChannelNames.includes(config.channelName as never), - ), - ).toBe(true); - - for (const alias of [ - "strandsAgentSDK", - "strandsagentsdk", - "strands-agent-sdk", - "@strands-agents/sdk", - ]) { - const disabledConfig = - readDisabledInstrumentationEnvConfig(alias).integrations; - - expect(disabledConfig).toMatchObject({ strandsAgentSDK: false }); - expect( - getDefaultInstrumentationConfigs({ - disabledIntegrationConfig: disabledConfig, - }).some((config) => - strandsChannelNames.includes(config.channelName as never), - ), - ).toBe(false); - } - }); -}); diff --git a/js/src/auto-instrumentations/configs/strands-agent-sdk.ts b/js/src/auto-instrumentations/configs/strands-agent-sdk.ts index 7a7f7a5b9..e21fe7764 100644 --- a/js/src/auto-instrumentations/configs/strands-agent-sdk.ts +++ b/js/src/auto-instrumentations/configs/strands-agent-sdk.ts @@ -3,7 +3,7 @@ import { strandsAgentSDKChannels } from "../../instrumentation/plugins/strands-a const strandsAgentSDKVersionRange = ">=1.0.0 <2.0.0"; -export const strandsAgentSDKConfigs: InstrumentationConfig[] = [ +export const strandsAgentSDKOrchestrionConfigs: InstrumentationConfig[] = [ { channelName: strandsAgentSDKChannels.agentStream.channelName, module: { diff --git a/js/src/auto-instrumentations/hook.mts b/js/src/auto-instrumentations/hook.mts index 10626cebd..ee8fc9d21 100644 --- a/js/src/auto-instrumentations/hook.mts +++ b/js/src/auto-instrumentations/hook.mts @@ -3,82 +3,118 @@ * * Usage: * node --import @braintrust/auto-instrumentations/hook.mjs app.js - * - * This hook performs AST transformation at load-time for BOTH ESM and CJS modules, - * injecting TracingChannel calls into AI SDK functions. - * - * Many modern apps use a mix of ESM and CJS modules, so this single hook - * handles both: - * - ESM modules: Transformed via register() loader hook - * - CJS modules: Transformed via ModulePatch monkey-patching Module._compile */ -import { register } from "node:module"; +import { register as registerModule } from "node:module"; +import { readDisabledInstrumentationEnvConfig } from "../instrumentation/config.js"; +import { createHook as createImportInTheMiddleHook } from "./import-in-the-middle/create-hook.mjs"; +import registerState from "./import-in-the-middle/lib/register.mjs"; import { - isInstrumentationIntegrationDisabled, - readDisabledInstrumentationEnvConfig, -} from "../instrumentation/config.js"; -import { BraintrustObservabilityExporter } from "../wrappers/mastra.js"; -import { installMastraExporterFactory } from "./loader/mastra-observability-patch.js"; -import { getDefaultAutoInstrumentationConfigs } from "./configs/all.js"; -import { ModulePatch } from "./loader/cjs-patch.js"; + getDefaultModuleExportPatchConfigs, + getDefaultOrchestrionConfigs, +} from "./configs/all.js"; +import { ModulePatch } from "./loader/cjs.js"; +import { nodeModuleExportPatchRuntime } from "./loader/module-hooks/node-runtime.js"; +import { installModuleExportPatchRunner } from "./loader/module-hooks/registry.js"; +import { installNodeModuleExportHooks } from "./loader/module-hooks/node.js"; import { patchTracingChannel } from "./patch-tracing-channel.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; +export default registerState; + const state = ((globalThis as any)[ Symbol.for("braintrust.applyAutoInstrumentation") ] ??= {}) as { applied?: boolean }; 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) { +// Query-mode imports expose IITM's loader API without bootstrapping Braintrust +// again when module.register() loads this file with its private query marker. +if (!isImportInTheMiddleLoader && !alreadyApplied) { const dcPath = ["node", "diagnostics_channel"].join(":"); const dc: any = await import(/* @vite-ignore */ dcPath as any); patchTracingChannel(dc.tracingChannel); } -if (!alreadyApplied) { - const allConfigs = getDefaultAutoInstrumentationConfigs(); - - // Expose the Mastra exporter factory on globalThis so the loader patches - // for `@mastra/core` / `@mastra/observability` can find it without having - // to resolve `braintrust` from the user's module graph. Skipped when the - // user opts out via `BRAINTRUST_DISABLE_INSTRUMENTATION=mastra`. +if (!isImportInTheMiddleLoader && !alreadyApplied) { const disabled = readDisabledInstrumentationEnvConfig( process.env.BRAINTRUST_DISABLE_INSTRUMENTATION, ).integrations; - if (!isInstrumentationIntegrationDisabled(disabled, "mastra")) { - installMastraExporterFactory(() => new BraintrustObservabilityExporter()); - } + const orchestrionConfigs = getDefaultOrchestrionConfigs({ + disabledIntegrationConfig: disabled, + }); + const moduleExportPatchConfigs = getDefaultModuleExportPatchConfigs({ + disabledIntegrationConfig: disabled, + target: "node", + }); - // 1. Register ESM loader for ESM modules - register("./loader/esm-hook.mjs", { + installModuleExportPatchRunner( + moduleExportPatchConfigs, + nodeModuleExportPatchRuntime, + ); + installNodeModuleExportHooks({ + asyncImportHookUrl, + configs: moduleExportPatchConfigs, + registryImportUrl, + }); + + registerModule("./loader/esm.mjs", { parentURL: import.meta.url, - data: { instrumentations: allConfigs }, + data: { instrumentations: orchestrionConfigs }, } as any); state.applied = true; - // 2. Also load CJS register for CJS modules (many apps use mixed ESM/CJS) try { - const patch = new ModulePatch({ instrumentations: allConfigs }); + const patch = new ModulePatch({ instrumentations: orchestrionConfigs }); patch.patch(); if (process.env.DEBUG === "@braintrust*" || process.env.DEBUG === "*") { console.log( "[Braintrust] Auto-instrumentation active (ESM + CJS) for:", - allConfigs.map((c) => c.channelName).join(", "), + orchestrionConfigs.map((config) => config.channelName).join(", "), ); } } catch (err) { - // CJS patch failed, but ESM hook is still active if (process.env.DEBUG === "@braintrust*" || process.env.DEBUG === "*") { console.log( "[Braintrust] Auto-instrumentation active (ESM only) for:", - allConfigs.map((c) => c.channelName).join(", "), + orchestrionConfigs.map((config) => config.channelName).join(", "), ); console.error("[Braintrust] CJS patch failed:", err); } } } + +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; +} diff --git a/js/src/auto-instrumentations/import-in-the-middle/create-hook.mts b/js/src/auto-instrumentations/import-in-the-middle/create-hook.mts index debbfa6f3..e465da088 100644 --- a/js/src/auto-instrumentations/import-in-the-middle/create-hook.mts +++ b/js/src/auto-instrumentations/import-in-the-middle/create-hook.mts @@ -35,8 +35,10 @@ const HANDLED_FORMATS = new Set([ const TRACE_WARNINGS = process.execArgv.includes("--trace-warnings"); type LoaderMeta = { url: string }; +type HookOptions = { registerUrl?: string }; type HookData = { addHookMessagePort?: MessagePort; + include?: readonly string[]; }; type AsyncLoadFunction = ( url: string, @@ -59,7 +61,7 @@ type ProcessModuleResult = { setters: SetterMap; }; -interface ImportInTheMiddleHook { +export interface ImportInTheMiddleHook { applyOptions(data: HookData): void; initialize(data?: HookData): Promise; load( @@ -360,13 +362,18 @@ function addIitm(url: string): string { return urlObj.href; } -export function createHook(meta: LoaderMeta): ImportInTheMiddleHook { +export function createHook( + meta: LoaderMeta, + options: HookOptions = {}, +): ImportInTheMiddleHook { let cachedAsyncResolve: AsyncResolveFunction | undefined; let cachedSyncResolve: SyncResolveFunction | undefined; - const iitmURL = new URL( - meta.url.endsWith(".mts") ? "lib/register.mts" : "lib/register.mjs", - meta.url, - ).toString(); + const iitmURL = + options.registerUrl ?? + new URL( + meta.url.endsWith(".mts") ? "lib/register.mts" : "lib/register.mjs", + meta.url, + ).toString(); const loaderThreadHookedModules = new Set(); // Track CJS module URLs that IITM has wrapped. On Node 24+, CJS modules loaded @@ -433,7 +440,8 @@ export function createHook(meta: LoaderMeta): ImportInTheMiddleHook { } function applyOptions(data: HookData): void { - const { addHookMessagePort } = data; + const { addHookMessagePort, include } = data; + addExplicitHookModules(include); if (addHookMessagePort) { addHookMessagePort .on("message", (modules: unknown) => { diff --git a/js/src/auto-instrumentations/import-in-the-middle/lib/get-esm-exports.mts b/js/src/auto-instrumentations/import-in-the-middle/lib/get-esm-exports.mts index 3f6137282..1f69c58fe 100644 --- a/js/src/auto-instrumentations/import-in-the-middle/lib/get-esm-exports.mts +++ b/js/src/auto-instrumentations/import-in-the-middle/lib/get-esm-exports.mts @@ -2,9 +2,11 @@ import { readFileSync } from "node:fs"; import { createRequire, Module } from "node:module"; +import { join } from "node:path"; +import { pathToFileURL } from "node:url"; import { initSync, parse as parseWasm } from "es-module-lexer"; -const require = createRequire(import.meta.url); +const require = createRequire(getCurrentModuleUrl()); type LexerImport = { n: string; @@ -36,6 +38,20 @@ if (disallowCodegen) { parse = loadAsmParse(); } +function getCurrentModuleUrl(): string { + if (typeof __filename !== "undefined") { + return pathToFileURL(__filename).href; + } + + const stack = new Error().stack ?? ""; + const match = + stack.match(/\((file:\/\/[^)]+)\)/) ?? stack.match(/\s(file:\/\/\S+)/); + return ( + match?.[1].replace(/:\d+:\d+$/, "") ?? + pathToFileURL(join(process.cwd(), "package.json")).href + ); +} + function loadAsmParse(): typeof parseWasm { const asmPath = require.resolve("es-module-lexer/js"); const source = diff --git a/js/src/auto-instrumentations/index.ts b/js/src/auto-instrumentations/index.ts index 26a97fa0f..c7e2dcbe2 100644 --- a/js/src/auto-instrumentations/index.ts +++ b/js/src/auto-instrumentations/index.ts @@ -4,7 +4,8 @@ * Auto-instrumentation for AI SDKs using orchestrion-js and diagnostics_channel. * * This package provides: - * - Instrumentation configs for orchestrion-js + * - Orchestrion configs + * - Declarative module export patch configs * - ESM loader hooks for load-time instrumentation * - CJS register for CommonJS instrumentation * - Bundler plugins for build-time instrumentation @@ -28,26 +29,27 @@ * ``` */ -export { openaiConfigs } from "./configs/openai"; -export { openAICodexConfigs } from "./configs/openai-codex"; -export { anthropicConfigs } from "./configs/anthropic"; -export { bedrockRuntimeConfigs } from "./configs/bedrock-runtime"; -export { aiSDKConfigs } from "./configs/ai-sdk"; -export { claudeAgentSDKConfigs } from "./configs/claude-agent-sdk"; -export { cursorSDKConfigs } from "./configs/cursor-sdk"; -export { openAIAgentsCoreConfigs } from "./configs/openai-agents"; -export { googleGenAIConfigs } from "./configs/google-genai"; -export { huggingFaceConfigs } from "./configs/huggingface"; -export { openRouterAgentConfigs } from "./configs/openrouter-agent"; -export { openRouterConfigs } from "./configs/openrouter"; -export { mistralConfigs } from "./configs/mistral"; -export { googleADKConfigs } from "./configs/google-adk"; -export { cohereConfigs } from "./configs/cohere"; -export { groqConfigs } from "./configs/groq"; -export { genkitConfigs } from "./configs/genkit"; -export { gitHubCopilotConfigs } from "./configs/github-copilot"; -export { langchainConfigs } from "./configs/langchain"; -export { piCodingAgentConfigs } from "./configs/pi-coding-agent"; +export { openaiOrchestrionConfigs } from "./configs/openai"; +export { openAICodexOrchestrionConfigs } from "./configs/openai-codex"; +export { anthropicOrchestrionConfigs } from "./configs/anthropic"; +export { bedrockRuntimeOrchestrionConfigs } from "./configs/bedrock-runtime"; +export { aiSDKOrchestrionConfigs } from "./configs/ai-sdk"; +export { claudeAgentSDKOrchestrionConfigs } from "./configs/claude-agent-sdk"; +export { cursorSDKOrchestrionConfigs } from "./configs/cursor-sdk"; +export { openAIAgentsCoreOrchestrionConfigs } from "./configs/openai-agents"; +export { googleGenAIOrchestrionConfigs } from "./configs/google-genai"; +export { huggingFaceOrchestrionConfigs } from "./configs/huggingface"; +export { openRouterAgentOrchestrionConfigs } from "./configs/openrouter-agent"; +export { openRouterOrchestrionConfigs } from "./configs/openrouter"; +export { mistralOrchestrionConfigs } from "./configs/mistral"; +export { googleADKOrchestrionConfigs } from "./configs/google-adk"; +export { cohereOrchestrionConfigs } from "./configs/cohere"; +export { groqOrchestrionConfigs } from "./configs/groq"; +export { genkitOrchestrionConfigs } from "./configs/genkit"; +export { gitHubCopilotOrchestrionConfigs } from "./configs/github-copilot"; +export { langchainOrchestrionConfigs } from "./configs/langchain"; +export { piCodingAgentOrchestrionConfigs } from "./configs/pi-coding-agent"; +export { mastraModuleExportPatchConfigs } from "./configs/mastra"; // Re-export orchestrion configuration types from the internal fork. export type { InstrumentationConfig } from "./orchestrion-js"; diff --git a/js/src/auto-instrumentations/loader/cjs-patch.ts b/js/src/auto-instrumentations/loader/cjs.ts similarity index 92% rename from js/src/auto-instrumentations/loader/cjs-patch.ts rename to js/src/auto-instrumentations/loader/cjs.ts index 670a6a5b1..8583b326d 100644 --- a/js/src/auto-instrumentations/loader/cjs-patch.ts +++ b/js/src/auto-instrumentations/loader/cjs.ts @@ -7,8 +7,8 @@ import { create, type InstrumentationConfig } from "../orchestrion-js"; import * as NodeModule from "node:module"; import { sep } from "node:path"; import moduleDetailsFromPath from "module-details-from-path"; -import { getPackageName, getPackageVersion } from "./get-package-version.js"; -import { applySpecialCasePatch } from "./special-case-patches.js"; +import { getPackageName, getPackageVersion } from "./package-version.js"; +import { applySourcePatch } from "./source-patches/index.js"; export class ModulePatch { private packages: Set; @@ -48,11 +48,11 @@ export class ModulePatch { const normalizedModulePath = resolvedModule.path.replace(/\\/g, "/"); const version = getPackageVersion(resolvedModule.basedir); - // Per-package source patches (see loader/special-case-patches.ts). + // Per-package source patches (see loader/source-patches/). // Anti-pattern intentionally isolated in its own module — do not // expand inline here; new integrations belong in the standard plugin // pipeline. - const patched = applySpecialCasePatch({ + const patched = applySourcePatch({ packageName, modulePath: normalizedModulePath, source: String(content), diff --git a/js/src/auto-instrumentations/loader/esm-hook.mts b/js/src/auto-instrumentations/loader/esm.mts similarity index 85% rename from js/src/auto-instrumentations/loader/esm-hook.mts rename to js/src/auto-instrumentations/loader/esm.mts index 09bf03139..70fb4097f 100644 --- a/js/src/auto-instrumentations/loader/esm-hook.mts +++ b/js/src/auto-instrumentations/loader/esm.mts @@ -8,19 +8,19 @@ import { fileURLToPath } from "node:url"; import { extname, sep } from "node:path"; import { create, type InstrumentationConfig } from "../orchestrion-js"; import moduleDetailsFromPath from "module-details-from-path"; -import { getPackageName, getPackageVersion } from "./get-package-version.js"; +import { getPackageName, getPackageVersion } from "./package-version.js"; import { - applySpecialCasePatch, - isSpecialCaseTarget, -} from "./special-case-patches.js"; + applySourcePatch, + isSourcePatchTarget, +} from "./source-patches/index.js"; let instrumentator: any; let packages: Set; let transformers: Map = new Map(); -// URLs that need a per-package source patch (see special-case-patches.ts). +// URLs that need a per-package source patch (see source-patches/). // Resolve records the (packageName, modulePath) it already computed so load // doesn't have to redo the moduleDetailsFromPath roundtrip. -const specialCaseUrls: Map< +const sourcePatchUrls: Map< string, { packageName: string; modulePath: string } > = new Map(); @@ -79,11 +79,11 @@ export async function resolve( const version = getPackageVersion(resolvedModule.basedir); // Track files that need per-package source patches (see - // loader/special-case-patches.ts). Anti-pattern fallback for SDKs we + // loader/source-patches/). Anti-pattern fallback for SDKs we // can't instrument via the standard pipeline; the load step retrieves // the recorded (packageName, modulePath) and applies the patch. - if (isSpecialCaseTarget(packageName, normalizedModulePath)) { - specialCaseUrls.set(url.url, { + if (isSourcePatchTarget(packageName, normalizedModulePath)) { + sourcePatchUrls.set(url.url, { packageName, modulePath: normalizedModulePath, }); @@ -110,18 +110,18 @@ export async function resolve( export async function load(url: string, context: any, nextLoad: Function) { const result = await nextLoad(url, context); - // Per-package source patches (see loader/special-case-patches.ts). + // Per-package source patches (see loader/source-patches/). // Anti-pattern fallback — keep this branch and the helper module narrow. - const specialCase = specialCaseUrls.get(url); - if (specialCase) { + const sourcePatch = sourcePatchUrls.get(url); + if (sourcePatch) { if (result.format === "commonjs") { const parsedUrl = new URL(result.responseURL ?? url); result.source ??= await readFile(parsedUrl); } if (result.source) { - const patched = applySpecialCasePatch({ - packageName: specialCase.packageName, - modulePath: specialCase.modulePath, + const patched = applySourcePatch({ + packageName: sourcePatch.packageName, + modulePath: sourcePatch.modulePath, source: result.source.toString("utf8"), format: result.format === "commonjs" ? "cjs" : "esm", }); diff --git a/js/src/auto-instrumentations/loader/mastra-observability-patch.test.ts b/js/src/auto-instrumentations/loader/mastra-observability-patch.test.ts deleted file mode 100644 index cda12afdf..000000000 --- a/js/src/auto-instrumentations/loader/mastra-observability-patch.test.ts +++ /dev/null @@ -1,117 +0,0 @@ -import { describe, expect, it } from "vitest"; -import { - classifyMastraTarget, - patchMastraSource, -} from "./mastra-observability-patch"; - -describe("classifyMastraTarget", () => { - it("identifies @mastra/core main and submodule entries", () => { - expect(classifyMastraTarget("@mastra/core", "dist/index.js")).toBe("core"); - expect(classifyMastraTarget("@mastra/core", "dist/index.cjs")).toBe("core"); - expect(classifyMastraTarget("@mastra/core", "dist/mastra/index.js")).toBe( - "core", - ); - expect(classifyMastraTarget("@mastra/core", "dist/mastra/index.cjs")).toBe( - "core", - ); - }); - - it("identifies @mastra/observability entry", () => { - expect(classifyMastraTarget("@mastra/observability", "dist/index.js")).toBe( - "observability", - ); - expect( - classifyMastraTarget("@mastra/observability", "dist/index.cjs"), - ).toBe("observability"); - }); - - it("returns null for unrelated paths", () => { - expect( - classifyMastraTarget("@mastra/core", "dist/agent/index.js"), - ).toBeNull(); - expect( - classifyMastraTarget("@mastra/core", "dist/chunk-XYZ.js"), - ).toBeNull(); - expect(classifyMastraTarget("openai", "dist/index.js")).toBeNull(); - }); -}); - -describe("patchMastraSource — @mastra/core ESM entry", () => { - it("rewrites a thin re-export into a Proxy-wrapped class", () => { - const original = `export { Mastra } from './chunk-PLCLLPJL.js';\n`; - const patched = patchMastraSource(original, "core", "esm"); - - // The rewritten source must import the original class from the same chunk - expect(patched).toContain( - `import { Mastra as __braintrustOrigMastra } from "./chunk-PLCLLPJL.js"`, - ); - // It must wrap in a Proxy with a `construct` trap - expect(patched).toContain("new Proxy(__braintrustOrigMastra,"); - expect(patched).toContain("construct(target, args, newTarget)"); - // It must re-export `Mastra` so consumers' bindings are unchanged - expect(patched).toContain("export { Mastra }"); - // It must pull Observability via createRequire so it resolves from the - // user's node_modules tree (not our SDK's) - expect(patched).toContain("createRequire"); - expect(patched).toContain(`__braintrustRequire("@mastra/observability")`); - }); - - it("returns the original source when the re-export shape doesn't match", () => { - const original = `// arbitrary code that doesn't re-export Mastra\n`; - const patched = patchMastraSource(original, "core", "esm"); - expect(patched).toBe(original); - }); - - it("preserves the exact chunk path Mastra references", () => { - const original = `export { Mastra } from '../chunk-DIFFERENT.js';\n`; - const patched = patchMastraSource(original, "core", "esm"); - expect(patched).toContain("../chunk-DIFFERENT.js"); - }); -}); - -describe("patchMastraSource — @mastra/core CJS entry", () => { - it("rewrites the require + defineProperty shape into a Proxy", () => { - const original = `'use strict'; -var chunkVOP4TUHG_cjs = require('./chunk-VOP4TUHG.cjs'); -Object.defineProperty(exports, "Mastra", { - enumerable: true, - get: function () { return chunkVOP4TUHG_cjs.Mastra; } -}); -`; - const patched = patchMastraSource(original, "core", "cjs"); - - expect(patched).toContain(`require("./chunk-VOP4TUHG.cjs")`); - expect(patched).toContain(`require("@mastra/observability")`); - expect(patched).toContain("__braintrustChunk.Mastra"); - expect(patched).toContain("new Proxy("); - expect(patched).toContain(`Object.defineProperty(exports, "Mastra"`); - }); -}); - -describe("patchMastraSource — @mastra/observability entry", () => { - it("appends a Proxy wrap, leaving the original source intact above", () => { - const original = `var Observability = class extends MastraBase {};\nexport { Observability };\n`; - const patched = patchMastraSource(original, "observability", "esm"); - - // Original source still starts the file - expect(patched.startsWith(original)).toBe(true); - // Append wraps the Observability binding via Proxy - expect(patched).toContain("function __braintrustWrapObservability"); - expect(patched).toContain( - `if (typeof Observability === "undefined") return`, - ); - expect(patched).toContain("new Proxy(__OriginalObservability"); - expect(patched).toContain("factory()"); - expect(patched).toContain("__braintrustWrapped"); - }); - - it("doesn't depend on chunk path extraction for observability", () => { - // The Observability entry is one big inline file, not a re-export. - // patchMastraSource should still produce a valid append even when no - // chunk-shaped pattern is present in the source. - const original = `// big inline bundle\nvar Observability = class {};\nexport { Observability };\n`; - const patched = patchMastraSource(original, "observability", "esm"); - expect(patched).not.toBe(original); - expect(patched.length).toBeGreaterThan(original.length); - }); -}); diff --git a/js/src/auto-instrumentations/loader/mastra-observability-patch.ts b/js/src/auto-instrumentations/loader/mastra-observability-patch.ts deleted file mode 100644 index 95b8f1aa1..000000000 --- a/js/src/auto-instrumentations/loader/mastra-observability-patch.ts +++ /dev/null @@ -1,288 +0,0 @@ -/** - * Runtime patches that auto-install the Braintrust observability exporter for - * Mastra users. - * - * Two entry shapes need handling, and they're each different enough that we - * pattern-match them separately rather than trying to share one strategy: - * - * 1. **`@mastra/core` `Mastra` entries** (`dist/index.{js,cjs}` and - * `dist/mastra/index.{js,cjs}`) are thin re-exports from a content-hashed - * chunk that changes filename every release: - * - * // dist/index.js - * export { Mastra } from "./chunk-PLCLLPJL.js"; - * - * The entry filename itself is stable (pinned by the package's `exports` - * map), so we patch it — but a re-export has no local binding to mutate, - * so we have to rewrite the line into an explicit import + Proxy + export. - * That means extracting the chunk path with a regex. The regex pattern is - * stable; only the chunk hash inside the matched string changes. - * - * 2. **`@mastra/observability` `Observability` entry** (`dist/index.{js,cjs}`) - * inlines its class declaration in the entry itself: - * - * var Observability = class extends MastraBase { ... }; - * export { ..., Observability, ... }; - * - * Because the local `Observability` is a `var` binding, we can *append* a - * wrap (reassign `Observability` to a Proxy) and ESM's live-binding - * semantics propagate it to importers. No source rewrite needed, no chunk - * path involved. - * - * In both cases the generated wrapper is a `Proxy` with a `construct` trap. - * On the Mastra side the trap injects an `Observability` instance into the - * config when the user didn't pass one (via a factory registered on - * `globalThis` by `hook.mjs`). On the Observability side the trap walks the - * `configs` map and appends our exporter to any instance config that doesn't - * already have one with `name === "braintrust"`. - * - * The patch is a no-op if the regex doesn't match (e.g., Mastra restructures - * its build) or if our globalThis factory isn't present (e.g., user disabled - * via `BRAINTRUST_DISABLE_INSTRUMENTATION=mastra`). - */ - -/** - * Name of the `globalThis` property the generated patches read to look up the - * Braintrust exporter factory at runtime. Set by `hook.mjs` (loader path) and - * by `configureNode()` (main braintrust import — covers the bundler-plugin - * path where the patched module is built into the user's bundle rather than - * loaded via the ESM hook). Both call sites use `installMastraExporterFactory` - * below so the `??=` lives in one place. Kept module-private — external - * callers should go through `installMastraExporterFactory`. - */ -const MASTRA_EXPORTER_FACTORY_GLOBAL = "__braintrustMastraExporterFactory"; - -/** - * Idempotently install the Braintrust Mastra exporter factory on - * `globalThis`. Callers are responsible for honoring - * `BRAINTRUST_DISABLE_INSTRUMENTATION=mastra`; this helper is just the - * `??=` placement. - */ -export function installMastraExporterFactory(factory: () => unknown): void { - const globals = globalThis as Record; - globals[MASTRA_EXPORTER_FACTORY_GLOBAL] ??= factory; -} - -const MASTRA_CORE_PACKAGE = "@mastra/core"; -const MASTRA_OBSERVABILITY_PACKAGE = "@mastra/observability"; - -// Entrypoints pinned by each package's `exports` map. Both Mastra entries -// (`dist/index.js` and `dist/mastra/index.js`) re-export `Mastra` from a -// chunk, so they both need the rewrite treatment. -const MASTRA_CORE_ENTRY_PATHS = new Set([ - "dist/index.js", - "dist/index.cjs", - "dist/mastra/index.js", - "dist/mastra/index.cjs", -]); -const MASTRA_OBSERVABILITY_ENTRY_PATHS = new Set([ - "dist/index.js", - "dist/index.cjs", -]); - -type MastraTargetKind = "core" | "observability"; -export type MastraModuleFormat = "esm" | "cjs"; - -export function classifyMastraTarget( - packageName: string, - modulePath: string, -): MastraTargetKind | null { - if ( - packageName === MASTRA_CORE_PACKAGE && - MASTRA_CORE_ENTRY_PATHS.has(modulePath) - ) { - return "core"; - } - if ( - packageName === MASTRA_OBSERVABILITY_PACKAGE && - MASTRA_OBSERVABILITY_ENTRY_PATHS.has(modulePath) - ) { - return "observability"; - } - return null; -} - -// Extract the chunk path from a Mastra entry source. -// ESM shape: `export { Mastra } from "../chunk-XYZ.js";` -// CJS shape: `var chunkXYZ_cjs = require("../chunk-XYZ.cjs"); ...` -// Returns null when the shape isn't what we expect — caller falls through and -// emits the original source unchanged so user code keeps working. -function extractChunkPath(source: string): string | null { - const esmMatch = source.match( - /export\s*\{\s*Mastra(?:\s+as\s+\w+)?\s*\}\s*from\s*['"]([^'"]+)['"]/, - ); - if (esmMatch) return esmMatch[1]; - - const cjsMatch = source.match( - /require\s*\(\s*['"]([^'"]+chunk-[^'"]+)['"]\s*\)/, - ); - if (cjsMatch) return cjsMatch[1]; - - return null; -} - -// All wrapper templates below are emitted as runtime JS in the target module's -// scope. They never reference Braintrust SDK code directly — instead they look -// up a factory function on `globalThis` (registered by `hook.mjs`). That keeps -// the patch independent of how the user's module graph resolves `braintrust`, -// and lets us cleanly no-op when the user disables our integration via env. - -const EXPORTER_FACTORY_KEY = JSON.stringify(MASTRA_EXPORTER_FACTORY_GLOBAL); - -// Construct-trap body shared between ESM and CJS wrappers. The wrapper loads -// `@mastra/observability` at module-evaluation time from the Mastra entry's -// own location (so the resolver finds the user's `node_modules`, not our -// SDK's). When the user constructs `new Mastra(...)` without an observability -// config, the trap injects a default `Observability` — which, because it -// loaded through our patched ESM hook / CJS patch, already auto-installs the -// Braintrust exporter via its own constructor wrap. -const MASTRA_PROXY_HANDLER_BODY = ` -{ - construct(target, args, newTarget) { - var firstArg = args[0]; - if ( - (!firstArg || typeof firstArg !== "object" || !firstArg.observability) && - __braintrustObservabilityClass - ) { - try { - // serviceName is required by Mastra's Observability validator; pass - // something sensible by default. Users who want a different name - // should construct Observability themselves. - var observability = new __braintrustObservabilityClass({ - configs: { default: { serviceName: "mastra" } }, - }); - args = args.slice(); - args[0] = Object.assign({}, firstArg, { observability: observability }); - } catch (e) { - // Fall through. Mastra will use its own NoOp; user code still works, - // just without auto-instrumentation. - } - } - return Reflect.construct(target, args, newTarget); - }, -}`; - -function buildMastraEsmWrapper(chunkPath: string): string { - const chunk = JSON.stringify(chunkPath); - return `import { Mastra as __braintrustOrigMastra } from ${chunk}; -import { createRequire as __braintrustCreateRequire } from "node:module"; - -let __braintrustObservabilityClass = null; -try { - // Resolve @mastra/observability relative to this module (the Mastra entry), - // so it's looked up from the user's node_modules tree. - const __braintrustRequire = __braintrustCreateRequire(import.meta.url); - __braintrustObservabilityClass = - __braintrustRequire("@mastra/observability").Observability; -} catch (e) { - // @mastra/observability isn't installed; the construct trap will skip the - // auto-construct branch and Mastra falls back to its own NoOp. -} -const Mastra = new Proxy(__braintrustOrigMastra, ${MASTRA_PROXY_HANDLER_BODY}); -export { Mastra }; -`; -} - -function buildMastraCjsWrapper(chunkPath: string): string { - const chunk = JSON.stringify(chunkPath); - return `'use strict'; -const __braintrustChunk = require(${chunk}); - -let __braintrustObservabilityClass = null; -try { - __braintrustObservabilityClass = require("@mastra/observability").Observability; -} catch (e) { - // @mastra/observability isn't installed; same fallback as the ESM wrapper. -} - -const __braintrustWrappedMastra = new Proxy( - __braintrustChunk.Mastra, - ${MASTRA_PROXY_HANDLER_BODY}, -); -Object.defineProperty(exports, "Mastra", { - enumerable: true, - configurable: true, - get: function () { return __braintrustWrappedMastra; } -}); -`; -} - -// Appended to the end of `@mastra/observability/dist/index.{js,cjs}`. The -// entry declares `var Observability = class extends MastraBase { ... }`, -// which means we can reassign the binding from inside the same module after -// the class is created and the export line has run. ESM live-binding -// semantics make external importers see the new value; CJS lookups go -// through `exports.Observability`, which we redefine to mirror the wrap. -const OBSERVABILITY_APPEND_BODY = ` -;(function __braintrustWrapObservability() { - // Top-level so we can both read and reassign the var binding the original - // entry declared. - if (typeof Observability === "undefined") return; - if (Observability.__braintrustWrapped) return; - function __braintrustEnsureExporter(rawConfig) { - try { - var factory = globalThis[${EXPORTER_FACTORY_KEY}]; - if (typeof factory !== "function") return rawConfig; - var config = rawConfig && typeof rawConfig === "object" ? rawConfig : {}; - var configsIn = config.configs && typeof config.configs === "object" ? config.configs : null; - var configsOut = {}; - var hadEntries = false; - if (configsIn) { - for (var name in configsIn) { - if (!Object.prototype.hasOwnProperty.call(configsIn, name)) continue; - hadEntries = true; - var inst = configsIn[name] || {}; - var existing = Array.isArray(inst.exporters) ? inst.exporters : []; - var hasOurs = existing.some(function (e) { return e && e.name === "braintrust"; }); - configsOut[name] = Object.assign({}, inst, { - exporters: hasOurs ? existing : existing.concat([factory()]), - }); - } - } - if (!hadEntries) { - configsOut.default = { - serviceName: "mastra", - exporters: [factory()], - }; - } - return Object.assign({}, config, { configs: configsOut }); - } catch (e) { - return rawConfig; - } - } - var __OriginalObservability = Observability; - Observability = new Proxy(__OriginalObservability, { - construct: function (target, args, newTarget) { - var nextArgs = args.slice(); - nextArgs[0] = __braintrustEnsureExporter(nextArgs[0]); - return Reflect.construct(target, nextArgs, newTarget); - }, - }); - Observability.__braintrustWrapped = true; - if (typeof exports !== "undefined" && exports && typeof exports === "object") { - try { - Object.defineProperty(exports, "Observability", { - enumerable: true, - configurable: true, - get: function () { return Observability; }, - }); - } catch (e) {} - } -})(); -`; - -export function patchMastraSource( - source: string, - target: MastraTargetKind, - format: MastraModuleFormat, -): string { - if (target === "core") { - const chunkPath = extractChunkPath(source); - if (!chunkPath) return source; - return format === "esm" - ? buildMastraEsmWrapper(chunkPath) - : buildMastraCjsWrapper(chunkPath); - } - // Observability: append-wrap, leaving original source intact above. - return source + OBSERVABILITY_APPEND_BODY; -} diff --git a/js/src/auto-instrumentations/loader/module-hooks/iitm.ts b/js/src/auto-instrumentations/loader/module-hooks/iitm.ts new file mode 100644 index 000000000..bf8d30a0b --- /dev/null +++ b/js/src/auto-instrumentations/loader/module-hooks/iitm.ts @@ -0,0 +1,150 @@ +import moduleDetailsFromPath from "module-details-from-path"; +import { isBuiltin } from "node:module"; +import { fileURLToPath } from "node:url"; +import registerState, { + type ImportHook, + type Namespace, +} from "../../import-in-the-middle/lib/register.mjs"; + +const { + addHookedModules, + deleteHookedModules, + importHooks, + specifiers, + toHook, +} = registerState; + +type HookFn = (exported: Namespace, name: string, baseDir?: string) => unknown; + +function addHook(hook: ImportHook): void { + importHooks.push(hook); + toHook.forEach(([name, namespace, specifier]) => + hook(name, namespace, specifier), + ); +} + +function removeHook(hook: ImportHook): void { + const index = importHooks.indexOf(hook); + if (index > -1) { + importHooks.splice(index, 1); + } +} + +function callHookFn( + hookFn: HookFn, + namespace: Namespace, + name: string, + baseDir?: string, +): void { + const newDefault = hookFn(namespace, name, baseDir); + if (newDefault && newDefault !== namespace && "default" in namespace) { + namespace.default = newDefault; + } +} + +function normalizeModules(modules: string[]): string[] { + if (!Array.isArray(modules) || modules.length === 0) { + throw new TypeError( + "Braintrust import-in-the-middle requires a non-empty modules array", + ); + } + + for (const each of modules) { + if (typeof each !== "string") { + throw new TypeError( + "Braintrust import-in-the-middle only supports string module names or file URLs", + ); + } + } + + return modules; +} + +function moduleMatches( + matchArg: string, + name: string, + filePath: string | undefined, + specifier: string | undefined, + baseDir: string | undefined, + loadUrl: string, +): boolean { + if (filePath && matchArg === filePath) { + return true; + } + + if (matchArg === specifier || matchArg === loadUrl || matchArg === name) { + return true; + } + + if (!baseDir) { + return false; + } + + return matchArg === name && baseDir.endsWith(specifiers.get(loadUrl) ?? ""); +} + +export default class Hook { + private readonly modules: string[]; + private readonly iitmHook: ImportHook; + + constructor(modules: string[], hookFn: HookFn) { + modules = normalizeModules(modules); + if (typeof hookFn !== "function") { + throw new TypeError( + "Braintrust import-in-the-middle requires a hook function", + ); + } + + addHookedModules(modules); + + this.modules = modules; + this.iitmHook = (name, namespace, specifier) => { + const loadUrl = name; + let filePath: string | undefined; + let baseDir: string | undefined; + + if (loadUrl.startsWith("node:")) { + const unprefixed = name.slice(5); + if (isBuiltin(unprefixed)) { + name = unprefixed; + } + } else if (loadUrl.startsWith("file://")) { + const stackTraceLimit = Error.stackTraceLimit; + Error.stackTraceLimit = 0; + try { + filePath = fileURLToPath(name); + name = filePath; + } catch {} + Error.stackTraceLimit = stackTraceLimit; + + if (filePath) { + const details = moduleDetailsFromPath(filePath); + if (details) { + name = details.name; + baseDir = details.basedir; + } + } + } + + for (const matchArg of modules) { + if ( + moduleMatches(matchArg, name, filePath, specifier, baseDir, loadUrl) + ) { + callHookFn( + hookFn, + namespace, + filePath && matchArg === filePath ? filePath : matchArg, + baseDir, + ); + } + } + }; + + addHook(this.iitmHook); + } + + unhook(): void { + removeHook(this.iitmHook); + deleteHookedModules(this.modules); + } +} diff --git a/js/src/auto-instrumentations/loader/module-hooks/node-runtime.ts b/js/src/auto-instrumentations/loader/module-hooks/node-runtime.ts new file mode 100644 index 000000000..ebe464af8 --- /dev/null +++ b/js/src/auto-instrumentations/loader/module-hooks/node-runtime.ts @@ -0,0 +1,27 @@ +import * as diagnosticsChannel from "node:diagnostics_channel"; +import { createRequire } from "node:module"; +import { join } from "node:path"; +import type { + ModuleExportConstructorEvent, + ModuleExportPatchRuntime, +} from "./registry.js"; + +export const nodeModuleExportPatchRuntime: ModuleExportPatchRuntime = { + resolveModule(specifier, context) { + try { + return createRequire( + context.resolutionBase ?? + (context.baseDir + ? join(context.baseDir, "package.json") + : join(process.cwd(), "package.json")), + )(specifier); + } catch { + return undefined; + } + }, + traceConstructor(channelName, event, construct) { + return diagnosticsChannel + .tracingChannel(channelName) + .traceSync(construct, event); + }, +}; diff --git a/js/src/auto-instrumentations/loader/module-hooks/node.test.ts b/js/src/auto-instrumentations/loader/module-hooks/node.test.ts new file mode 100644 index 000000000..05d22f3fe --- /dev/null +++ b/js/src/auto-instrumentations/loader/module-hooks/node.test.ts @@ -0,0 +1,310 @@ +import { describe, expect, it, vi } from "vitest"; +import { debugLogger } from "../../../debug-logger"; +import { createHook as createImportInTheMiddleHook } from "../../import-in-the-middle/create-hook.mjs"; +import type { ModuleExportPatchConfig } from "./registry"; +import { + installNodeModuleExportHooks, + type NodeModuleExportHookRuntime, +} from "./node"; + +describe("installNodeModuleExportHooks", () => { + it("uses in-process registerHooks when synchronous hooks are supported", () => { + const events: string[] = []; + const importHook = makeNoopImportHook({ + initialize(data: unknown) { + events.push(`initialize:${JSON.stringify(data)}`); + return Promise.resolve(); + }, + }); + + installNodeModuleExportHooks({ + asyncImportHookUrl: + "file:///braintrust/hook.mjs?braintrust-iitm-loader=true", + configs: [fakeConfig()], + registryImportUrl: "file:///braintrust/hook.mjs", + runtime: fakeRuntime(events, { + createImportHook(meta, options) { + events.push(`create:${meta.url}:${options.registerUrl}`); + return importHook as ReturnType< + NodeModuleExportHookRuntime["createImportHook"] + >; + }, + register() { + events.push("register"); + }, + registerHooks(hooks) { + events.push( + `registerHooks:${typeof hooks.resolve}:${typeof hooks.load}`, + ); + }, + supportsSyncHooks: () => true, + }), + }); + + expect(events).toEqual([ + "create:file:///braintrust/hook.mjs:file:///braintrust/hook.mjs", + 'initialize:{"include":["pkg"]}', + "registerHooks:function:function", + "importHook:pkg", + "importHookPatched:patched", + "requireHook:pkg", + "requireHookPatched:patched", + ]); + }); + + it("falls back to self-registering hook.mjs as the async loader", () => { + const events: string[] = []; + const asyncImportHookUrl = + "file:///braintrust/hook.mjs?braintrust-iitm-loader=true"; + + installNodeModuleExportHooks({ + asyncImportHookUrl, + configs: [fakeConfig()], + registryImportUrl: "file:///braintrust/hook.mjs", + runtime: fakeRuntime(events, { + register(specifier, options) { + events.push(`register:${specifier}:${JSON.stringify(options)}`); + }, + registerHooks() { + events.push("registerHooks"); + }, + supportsSyncHooks: () => false, + }), + }); + + expect(events).toEqual([ + `register:${asyncImportHookUrl}:{"data":{"include":["pkg"]}}`, + "importHook:pkg", + "importHookPatched:patched", + "requireHook:pkg", + "requireHookPatched:patched", + ]); + }); + + it("passes the loaded package version to constrained hooks", () => { + const events: string[] = []; + + installNodeModuleExportHooks({ + asyncImportHookUrl: + "file:///braintrust/hook.mjs?braintrust-iitm-loader=true", + configs: [fakeConfig(">=2.0.0")], + registryImportUrl: "file:///braintrust/hook.mjs", + runtime: fakeRuntime(events, { + baseDir: "/pkg", + getPackageVersion(baseDir) { + events.push(`version:${baseDir}`); + return "1.0.0"; + }, + supportsSyncHooks: () => false, + }), + }); + + expect(events).toEqual([ + "register", + "importHook:pkg", + "version:/pkg", + "importHookPatched:original", + "requireHook:pkg", + "version:/pkg", + "requireHookPatched:original", + ]); + }); + + it("does not throw when ESM or CJS hook installation fails", () => { + const warnSpy = vi.spyOn(debugLogger, "warn").mockImplementation(() => {}); + + expect(() => + installNodeModuleExportHooks({ + asyncImportHookUrl: + "file:///braintrust/hook.mjs?braintrust-iitm-loader=true", + configs: [fakeConfig()], + registryImportUrl: "file:///braintrust/hook.mjs", + runtime: { + createImportHook: (() => { + throw new Error("esm"); + }) as NodeModuleExportHookRuntime["createImportHook"], + importHookConstructor: class {}, + moduleApi: { + register() { + throw new Error("register"); + }, + }, + requireHookConstructor: class { + constructor() { + throw new Error("cjs"); + } + }, + supportsSyncHooks: () => true, + }, + }), + ).not.toThrow(); + + expect(warnSpy).toHaveBeenCalledWith( + "Failed to install ESM module export hooks", + expect.any(Error), + ); + expect(warnSpy).toHaveBeenCalledWith( + "Failed to install CJS module export hooks", + expect.any(Error), + ); + + warnSpy.mockRestore(); + }); +}); + +describe("merged import-in-the-middle wrapper generation", () => { + it("imports register state from the canonical Braintrust hook URL", async () => { + const hook = createImportInTheMiddleHook( + { url: "file:///braintrust/hook.mjs?braintrust-iitm-loader=true" }, + { registerUrl: "file:///braintrust/hook.mjs" }, + ); + await hook.initialize({ include: ["target"] }); + + const resolved = (await hook.resolve( + "target", + { conditions: ["import"], parentURL: "file:///app.mjs" }, + async () => ({ format: "module", url: "file:///target.mjs" }), + )) as { url: string }; + const loaded = (await hook.load( + resolved.url, + { format: "module" }, + async () => ({ + format: "module", + source: 'export const value = "original";', + }), + )) as { source: string }; + + expect(loaded.source).toContain( + "import registerState from 'file:///braintrust/hook.mjs'", + ); + expect(loaded.source).not.toContain("lib/register."); + }); +}); + +function fakeConfig(versionRange?: string): ModuleExportPatchConfig { + return { + integrations: ["mastra"], + modules: [ + { + packageName: "pkg", + patches: [ + { + channelName: "orchestrion:pkg:Target.constructor", + exportName: "Target", + kind: "constructor", + }, + ], + specifier: "pkg", + versionRange, + }, + ], + targets: ["node"], + }; +} + +function fakeRuntime( + events: string[], + { + createImportHook, + baseDir, + getPackageVersion, + register, + registerHooks, + supportsSyncHooks, + }: { + baseDir?: string; + createImportHook?: NodeModuleExportHookRuntime["createImportHook"]; + getPackageVersion?: NodeModuleExportHookRuntime["getPackageVersion"]; + register?: (specifier: string, options?: unknown) => void; + registerHooks?: (hooks: { load: unknown; resolve: unknown }) => void; + supportsSyncHooks: () => boolean; + }, +): Partial { + class FakeImportHook { + constructor( + modules: string[], + hookFn: ( + exportsValue: unknown, + name: string, + baseDir?: string, + ) => unknown, + ) { + events.push(`importHook:${modules.join(",")}`); + const namespace = { + Target: class { + constructor(public readonly value: string) {} + }, + }; + hookFn(namespace, "pkg", baseDir); + events.push( + `importHookPatched:${new namespace.Target("original").value}`, + ); + } + } + class FakeRequireHook { + constructor( + modules: string[], + hookFn: ( + exportsValue: unknown, + name: string, + baseDir?: string, + ) => unknown, + ) { + events.push(`requireHook:${modules.join(",")}`); + const namespace = { + Target: class { + constructor(public readonly value: string) {} + }, + }; + hookFn(namespace, "pkg", baseDir); + events.push( + `requireHookPatched:${new namespace.Target("original").value}`, + ); + } + } + return { + createImportHook: createImportHook ?? (() => makeNoopImportHook()), + ...(getPackageVersion ? { getPackageVersion } : {}), + importHookConstructor: FakeImportHook, + moduleApi: { + register: + register ?? + (() => { + events.push("register"); + }), + registerHooks, + }, + patchRuntime: { + resolveModule: () => undefined, + traceConstructor(_channelName, event, construct) { + event.arguments[0] = "patched"; + return construct(); + }, + }, + requireHookConstructor: FakeRequireHook, + supportsSyncHooks, + }; +} + +function makeNoopImportHook({ + initialize, +}: { + initialize?: (data: unknown) => Promise; +} = {}): ReturnType { + return { + applyOptions() {}, + initialize: initialize ?? (() => Promise.resolve()), + async load() { + return { format: "module", source: "" }; + }, + loadSync() { + return { format: "module", source: "" }; + }, + async resolve() { + return { url: "file:///noop.mjs" }; + }, + resolveSync() { + return { url: "file:///noop.mjs" }; + }, + }; +} diff --git a/js/src/auto-instrumentations/loader/module-hooks/node.ts b/js/src/auto-instrumentations/loader/module-hooks/node.ts new file mode 100644 index 000000000..d0395343e --- /dev/null +++ b/js/src/auto-instrumentations/loader/module-hooks/node.ts @@ -0,0 +1,123 @@ +import * as module from "node:module"; +import { debugLogger } from "../../../debug-logger"; +import { + createHook as createImportInTheMiddleHook, + type ImportInTheMiddleHook as ImportInTheMiddleHookApi, +} from "../../import-in-the-middle/create-hook.mjs"; +import { supportsSyncHooks } from "../../import-in-the-middle/supports-sync-hooks.mjs"; +import ImportInTheMiddleRuntimeHook from "./iitm.js"; +import RequireInTheMiddleHook from "./ritm.js"; +import { getPackageVersionIfAvailable } from "../package-version.js"; +import { nodeModuleExportPatchRuntime } from "./node-runtime.js"; +import { + getModuleExportPatchSpecifiers, + runModuleExportPatches, + type ModuleExportPatchConfig, + type ModuleExportPatchRuntime, +} from "./registry.js"; + +export interface InstallNodeModuleExportHooksOptions { + configs: readonly ModuleExportPatchConfig[]; + asyncImportHookUrl: string; + registryImportUrl: string; + runtime?: Partial; +} + +type HookCallback = ( + exportsValue: unknown, + name: string, + baseDir?: string, +) => unknown; + +type HookConstructor = new (modules: string[], hookFn: HookCallback) => unknown; + +interface NodeModuleApi { + register: (specifier: string, options?: unknown) => void; + registerHooks?: (hooks: { load: unknown; resolve: unknown }) => void; +} + +export interface NodeModuleExportHookRuntime { + createImportHook: ( + meta: { url: string }, + options: { registerUrl: string }, + ) => ImportInTheMiddleHookApi; + importHookConstructor: HookConstructor; + moduleApi: NodeModuleApi; + patchRuntime: ModuleExportPatchRuntime; + requireHookConstructor: HookConstructor; + supportsSyncHooks: () => boolean; + getPackageVersion: (baseDir: string) => string | undefined; +} + +export function installNodeModuleExportHooks({ + asyncImportHookUrl, + configs, + registryImportUrl, + runtime, +}: InstallNodeModuleExportHooksOptions): void { + const specifiers = getModuleExportPatchSpecifiers(configs); + if (specifiers.length === 0) return; + + const effectiveRuntime = getRuntime(runtime); + const hookCallback: HookCallback = (exportsValue, name, baseDir) => + runModuleExportPatches( + configs, + exportsValue, + { + baseDir, + moduleName: name, + moduleVersion: baseDir + ? effectiveRuntime.getPackageVersion(baseDir) + : undefined, + }, + effectiveRuntime.patchRuntime, + ); + + try { + if ( + effectiveRuntime.moduleApi.registerHooks && + effectiveRuntime.supportsSyncHooks() + ) { + const importHook = effectiveRuntime.createImportHook( + { url: registryImportUrl }, + { registerUrl: registryImportUrl }, + ); + void importHook.initialize({ include: specifiers }); + effectiveRuntime.moduleApi.registerHooks({ + load: importHook.loadSync, + resolve: importHook.resolveSync, + }); + } else { + effectiveRuntime.moduleApi.register(asyncImportHookUrl, { + data: { include: specifiers }, + }); + } + + new effectiveRuntime.importHookConstructor(specifiers, hookCallback); + } catch (err) { + debugLogger.warn("Failed to install ESM module export hooks", err); + } + + try { + new effectiveRuntime.requireHookConstructor(specifiers, hookCallback); + } catch (err) { + debugLogger.warn("Failed to install CJS module export hooks", err); + } +} + +function getRuntime( + overrides: Partial | undefined, +): NodeModuleExportHookRuntime { + return { + createImportHook: createImportInTheMiddleHook, + getPackageVersion: getPackageVersionIfAvailable, + importHookConstructor: + ImportInTheMiddleRuntimeHook as unknown as HookConstructor, + moduleApi: module as unknown as NodeModuleApi, + patchRuntime: nodeModuleExportPatchRuntime, + requireHookConstructor: + RequireInTheMiddleHook as unknown as HookConstructor, + supportsSyncHooks, + ...overrides, + }; +} diff --git a/js/src/auto-instrumentations/loader/module-hooks/registry.test.ts b/js/src/auto-instrumentations/loader/module-hooks/registry.test.ts new file mode 100644 index 000000000..0dcf515cf --- /dev/null +++ b/js/src/auto-instrumentations/loader/module-hooks/registry.test.ts @@ -0,0 +1,308 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + buildModuleExportSourceWrapper, + filterModuleExportPatchConfigs, + getModuleExportPatchSpecifiers, + installModuleExportPatchRunner, + runModuleExportPatches, + type ModuleExportConstructorEvent, + type ModuleExportPatchConfig, + type ModuleExportPatchRuntime, + type ModuleExportPatchTarget, +} from "./registry"; + +const runnerKey = "__braintrustTopLevelImportHookRunner"; + +describe("module export patch registry", () => { + afterEach(() => { + delete (globalThis as Record)[runnerKey]; + }); + + it("filters configs by target and disabled integrations", () => { + const nodeConfig = fakeConfig("node", { integrations: ["mastra"] }); + const browserConfig = fakeConfig("browser", { + integrations: ["openai"], + }); + + expect( + filterModuleExportPatchConfigs([nodeConfig, browserConfig], { + target: "node", + }), + ).toEqual([nodeConfig]); + expect( + filterModuleExportPatchConfigs([nodeConfig, browserConfig], { + disabledIntegrationConfig: { mastra: false }, + target: "node", + }), + ).toEqual([]); + }); + + it("collects unique runtime specifiers in first-seen order", () => { + expect( + getModuleExportPatchSpecifiers([ + fakeConfig("node", { moduleSpecifiers: ["pkg-a", "pkg-b"] }), + fakeConfig("node", { moduleSpecifiers: ["pkg-b", "pkg-c"] }), + ]), + ).toEqual(["pkg-a", "pkg-b", "pkg-c"]); + }); + + it("wraps constructors and lets channel subscribers mutate arguments", () => { + class Target { + constructor(public readonly value: string) {} + } + const namespace = { Target }; + const runtime = fakeRuntime({ + traceConstructor(_channelName, event, construct) { + event.arguments[0] = "patched"; + return construct(); + }, + }); + + const result = runModuleExportPatches( + [fakeConfig("node")], + namespace, + { moduleName: "pkg" }, + runtime, + ) as typeof namespace; + + expect(result).toBe(namespace); + expect(new result.Target("original").value).toBe("patched"); + }); + + it("passes module context and a safe resolver through constructor events", () => { + class Target {} + let seen: ModuleExportConstructorEvent | undefined; + const resolveModule = vi.fn(() => ({ value: true })); + const runtime = fakeRuntime({ + resolveModule, + traceConstructor(_channelName, event, construct) { + seen = event; + expect(event.resolveModule("dependency")).toEqual({ value: true }); + return construct(); + }, + }); + const result = runModuleExportPatches( + [fakeConfig("node")], + { Target }, + { + baseDir: "/pkg", + moduleName: "pkg", + moduleVersion: "2.1.0", + resolutionBase: "file:///app.mjs", + }, + runtime, + ) as { Target: typeof Target }; + + new result.Target(); + + expect(seen).toMatchObject({ + baseDir: "/pkg", + moduleName: "pkg", + moduleVersion: "2.1.0", + resolutionBase: "file:///app.mjs", + }); + expect(resolveModule).toHaveBeenCalledWith("dependency", { + baseDir: "/pkg", + moduleName: "pkg", + moduleVersion: "2.1.0", + resolutionBase: "file:///app.mjs", + }); + }); + + it("applies only matching module versions", () => { + class Target {} + const runtime = fakeRuntime(); + const config = fakeConfig("node", { versionRange: ">=2.0.0 <3.0.0" }); + const matching = { Target }; + const nonMatching = { Target }; + + runModuleExportPatches( + [config], + matching, + { moduleName: "pkg", moduleVersion: "2.1.0" }, + runtime, + ); + runModuleExportPatches( + [config], + nonMatching, + { moduleName: "pkg", moduleVersion: "3.0.0" }, + runtime, + ); + + expect(matching.Target).not.toBe(Target); + expect(nonMatching.Target).toBe(Target); + }); + + it("wraps each namespace once without changing constructor semantics", () => { + let constructionCount = 0; + class Target { + constructor() { + constructionCount++; + throw new Error("user error"); + } + } + const namespace = { Target }; + const runtime = fakeRuntime(); + const config = fakeConfig("node"); + + runModuleExportPatches([config], namespace, { moduleName: "pkg" }, runtime); + runModuleExportPatches([config], namespace, { moduleName: "pkg" }, runtime); + + expect(() => new namespace.Target()).toThrow("user error"); + expect(constructionCount).toBe(1); + }); + + it("skips missing and non-constructor exports and contains lookup errors", () => { + const runtime = fakeRuntime(); + const missing = { other: true }; + const throwing = Object.defineProperty({}, "Target", { + get() { + throw new Error("lookup failed"); + }, + }); + + expect( + runModuleExportPatches( + [fakeConfig("node")], + missing, + { moduleName: "pkg" }, + runtime, + ), + ).toBe(missing); + expect(() => + runModuleExportPatches( + [fakeConfig("node")], + throwing, + { moduleName: "pkg" }, + runtime, + ), + ).not.toThrow(); + }); + + it("clones immutable namespaces when installing a patched export", () => { + class Target {} + const namespace = Object.freeze({ Target }); + const result = runModuleExportPatches( + [fakeConfig("node")], + namespace, + { moduleName: "pkg" }, + fakeRuntime(), + ) as typeof namespace; + + expect(result).not.toBe(namespace); + expect(result.Target).not.toBe(Target); + }); + + it("installs a non-enumerable global runner", () => { + class Target {} + installModuleExportPatchRunner([fakeConfig("node")], fakeRuntime()); + + const descriptor = Object.getOwnPropertyDescriptor(globalThis, runnerKey); + const namespace = { Target }; + const result = ( + descriptor?.value as ( + exportsValue: unknown, + name: string, + baseDir?: string, + resolutionBase?: string, + moduleVersion?: string, + ) => unknown + )(namespace, "pkg", undefined, undefined, "1.0.0") as typeof namespace; + + expect(descriptor).toMatchObject({ + configurable: true, + enumerable: false, + writable: true, + }); + expect(result.Target).not.toBe(Target); + }); + + it("generates ESM and CJS wrappers for matching source modules", () => { + const config = fakeConfig("node"); + const esm = buildModuleExportSourceWrapper([config], { + format: "esm", + modulePath: "index.js", + originalModuleSpecifier: "braintrust-top-level-original:0", + packageName: "pkg", + source: "export { Target } from './target.js';", + target: "node", + }); + const cjs = buildModuleExportSourceWrapper([config], { + format: "cjs", + modulePath: "index.js", + originalModuleSpecifier: "/pkg/index.js?braintrust-top-level-original", + packageName: "pkg", + source: "exports.Target = Target;", + target: "node", + }); + + expect(esm).toContain("__braintrustTopLevelImportHookRunner"); + expect(esm).toContain("__braintrustExports"); + expect(esm).toContain(" as Target"); + expect(cjs).toContain("__braintrustTopLevelImportHookRunner"); + expect(cjs).toContain('Object.defineProperty(exports, "Target"'); + }); + + it("returns null for target and version mismatches", () => { + const config = fakeConfig("node", { versionRange: ">=2.0.0" }); + const input = { + format: "esm" as const, + modulePath: "index.js", + moduleVersion: "1.0.0", + originalModuleSpecifier: "original", + packageName: "pkg", + source: "export class Target {}", + }; + + expect( + buildModuleExportSourceWrapper([config], { + ...input, + target: "node", + }), + ).toBeNull(); + expect( + buildModuleExportSourceWrapper([config], { + ...input, + moduleVersion: "2.0.0", + target: "browser", + }), + ).toBeNull(); + }); +}); + +function fakeConfig( + target: ModuleExportPatchTarget, + options: { + integrations?: ModuleExportPatchConfig["integrations"]; + moduleSpecifiers?: string[]; + versionRange?: string; + } = {}, +): ModuleExportPatchConfig { + return { + integrations: options.integrations ?? ["mastra"], + modules: (options.moduleSpecifiers ?? ["pkg"]).map((specifier) => ({ + packageName: specifier, + patches: [ + { + channelName: `orchestrion:${specifier}:Target.constructor`, + exportName: "Target", + kind: "constructor", + }, + ], + source: { modulePaths: ["index.js"] }, + specifier, + versionRange: options.versionRange, + })), + targets: [target], + }; +} + +function fakeRuntime( + overrides: Partial = {}, +): ModuleExportPatchRuntime { + return { + resolveModule: () => undefined, + traceConstructor: (_channelName, _event, construct) => construct(), + ...overrides, + }; +} diff --git a/js/src/auto-instrumentations/loader/module-hooks/registry.ts b/js/src/auto-instrumentations/loader/module-hooks/registry.ts new file mode 100644 index 000000000..881962193 --- /dev/null +++ b/js/src/auto-instrumentations/loader/module-hooks/registry.ts @@ -0,0 +1,531 @@ +import type { InstrumentationIntegrationsConfig } from "../../../instrumentation/config"; +import { isInstrumentationIntegrationDisabled } from "../../../instrumentation/config"; +import semifies from "semifies"; + +export type ModuleExportPatchTarget = "node" | "browser"; +type ModuleExportPatchFormat = "esm" | "cjs"; +type MutableExportNamespace = Record; + +type RuntimeConstructor = new (...args: unknown[]) => object; + +interface ModuleExportSource { + modulePaths: readonly string[]; +} + +interface ModuleExportModule { + packageName: string; + patches: readonly ModuleExportConstructorPatch[]; + specifier: string; + versionRange?: string; + source?: ModuleExportSource; +} + +interface ModuleExportConstructorPatch { + channelName: string; + exportName: string; + kind: "constructor"; +} + +export interface ModuleExportPatchConfig { + integrations: readonly (keyof InstrumentationIntegrationsConfig)[]; + modules: readonly ModuleExportModule[]; + targets: readonly ModuleExportPatchTarget[]; +} + +export interface ModuleExportPatchContext { + moduleName: string; + moduleVersion?: string; + baseDir?: string; + resolutionBase?: string; +} + +export interface ModuleExportConstructorEvent extends ModuleExportPatchContext { + arguments: unknown[]; + resolveModule(specifier: string): unknown; +} + +export interface ModuleExportPatchRuntime { + resolveModule(specifier: string, context: ModuleExportPatchContext): unknown; + traceConstructor( + channelName: string, + event: ModuleExportConstructorEvent, + construct: () => object, + ): object; +} + +interface ModuleExportSourceWrapperInput { + baseDir?: string; + format: ModuleExportPatchFormat; + modulePath: string; + moduleVersion?: string; + originalModuleSpecifier: string; + packageName: string; + source: string; + target: ModuleExportPatchTarget; +} + +const MODULE_EXPORT_HOOK_RUNNER_GLOBAL = "__braintrustTopLevelImportHookRunner"; + +export function filterModuleExportPatchConfigs( + configs: readonly ModuleExportPatchConfig[], + { + disabledIntegrationConfig, + target, + }: { + disabledIntegrationConfig?: InstrumentationIntegrationsConfig; + target: ModuleExportPatchTarget; + }, +): ModuleExportPatchConfig[] { + return configs.filter( + (config) => + config.targets.includes(target) && + !isInstrumentationIntegrationDisabled( + disabledIntegrationConfig, + ...config.integrations, + ), + ); +} + +export function getModuleExportPatchSpecifiers( + configs: readonly ModuleExportPatchConfig[], +): string[] { + return [ + ...new Set( + configs.flatMap((config) => + config.modules.map((module) => module.specifier), + ), + ), + ]; +} + +export function runModuleExportPatches( + configs: readonly ModuleExportPatchConfig[], + exportsValue: unknown, + context: ModuleExportPatchContext, + runtime: ModuleExportPatchRuntime, +): unknown { + let namespace = asMutableNamespace(exportsValue); + for (const config of configs) { + for (const module of config.modules) { + if ( + module.specifier !== context.moduleName || + !matchesVersion(context.moduleVersion, module.versionRange) + ) { + continue; + } + + for (const patch of module.patches) { + try { + const original = namespace[patch.exportName]; + if ( + typeof original !== "function" || + isConstructorWrapped(original, patch.channelName) + ) { + continue; + } + + const wrapped = wrapConstructor( + original as RuntimeConstructor, + patch.channelName, + context, + runtime, + ); + namespace = setNamespaceExport(namespace, patch.exportName, wrapped); + } catch { + // Export patches are best-effort and must not break imports. + } + } + } + } + return namespace; +} + +export function installModuleExportPatchRunner( + configs: readonly ModuleExportPatchConfig[], + runtime: ModuleExportPatchRuntime, +): void { + Object.defineProperty(globalThis, MODULE_EXPORT_HOOK_RUNNER_GLOBAL, { + configurable: true, + enumerable: false, + value( + exportsValue: unknown, + name: string, + baseDir?: string, + resolutionBase?: string, + moduleVersion?: string, + ) { + return runModuleExportPatches( + configs, + exportsValue, + { + baseDir, + moduleName: name, + moduleVersion, + resolutionBase, + }, + runtime, + ); + }, + writable: true, + }); +} + +export function buildModuleExportSourceWrapper( + configs: readonly ModuleExportPatchConfig[], + input: ModuleExportSourceWrapperInput, +): string | null { + const sourceModules = configs + .filter((config) => config.targets.includes(input.target)) + .flatMap((config) => config.modules) + .filter( + (module) => + module.packageName === input.packageName && + module.source?.modulePaths.includes(input.modulePath) && + matchesVersion(input.moduleVersion, module.versionRange), + ); + + if (sourceModules.length === 0) { + return null; + } + + const specifier = sourceModules[0].specifier; + const exportNames = [ + ...new Set([ + ...sourceModules.flatMap((module) => + module.patches.map((patch) => patch.exportName), + ), + ...collectStaticExportNames(input.source, input.format), + ]), + ].sort((a, b) => Number(a === "default") - Number(b === "default")); + + if (exportNames.length === 0) { + return null; + } + + return input.format === "esm" + ? buildEsmSourceWrapper({ + baseDir: input.baseDir, + exportNames, + moduleName: specifier, + moduleVersion: input.moduleVersion, + originalModuleSpecifier: input.originalModuleSpecifier, + }) + : buildCjsSourceWrapper({ + baseDir: input.baseDir, + exportNames, + moduleName: specifier, + moduleVersion: input.moduleVersion, + originalModuleSpecifier: input.originalModuleSpecifier, + }); +} + +function matchesVersion( + moduleVersion: string | undefined, + versionRange: string | undefined, +): boolean { + if (!versionRange) return true; + if (!moduleVersion) return false; + + try { + return semifies(moduleVersion, versionRange); + } catch { + return false; + } +} + +function asMutableNamespace(exportsValue: unknown): MutableExportNamespace { + if (exportsValue && typeof exportsValue === "object") { + return exportsValue as MutableExportNamespace; + } + + return { default: exportsValue }; +} + +function wrapConstructor( + constructor: RuntimeConstructor, + channelName: string, + context: ModuleExportPatchContext, + runtime: ModuleExportPatchRuntime, +): RuntimeConstructor { + const marker = Symbol.for( + `braintrust.module-export-constructor.${channelName}`, + ); + return new Proxy(constructor, { + construct(target, args, newTarget) { + const event: ModuleExportConstructorEvent = { + ...context, + arguments: args.slice(), + resolveModule(specifier) { + return runtime.resolveModule(specifier, context); + }, + }; + return runtime.traceConstructor(channelName, event, () => + Reflect.construct(target, event.arguments, newTarget), + ); + }, + get(target, property, receiver) { + if (property === marker) return true; + return Reflect.get(target, property, receiver); + }, + }); +} + +function isConstructorWrapped(value: Function, channelName: string): boolean { + try { + return ( + (value as unknown as Record)[ + Symbol.for(`braintrust.module-export-constructor.${channelName}`) + ] === true + ); + } catch { + return false; + } +} + +function setNamespaceExport( + namespace: MutableExportNamespace, + key: string, + value: RuntimeConstructor, +): MutableExportNamespace { + try { + namespace[key] = value; + if (namespace[key] === value) return namespace; + } catch { + // ESM namespaces are immutable; clone below. + } + + return Object.defineProperties( + Object.create(Object.getPrototypeOf(namespace)), + { + ...Object.getOwnPropertyDescriptors(namespace), + [key]: { + configurable: true, + enumerable: true, + value, + writable: true, + }, + }, + ) as MutableExportNamespace; +} + +function collectStaticExportNames( + source: string, + format: ModuleExportPatchFormat, +): string[] { + if (format === "cjs") { + const names = new Set(); + for (const match of source.matchAll( + /\bexports\.([A-Za-z_$][\w$]*)\s*=|\bmodule\.exports\.([A-Za-z_$][\w$]*)\s*=|Object\.defineProperty\s*\(\s*exports\s*,\s*["']([^"']+)["']/g, + )) { + names.add(match[1] ?? match[2] ?? match[3]); + } + return [...names]; + } + + const names = new Set(); + + for (const match of source.matchAll( + /\bexport\s+(?:async\s+)?(?:class|function|const|let|var)\s+([A-Za-z_$][\w$]*)/g, + )) { + names.add(match[1]); + } + + if (/\bexport\s+default\b/.test(source)) { + names.add("default"); + } + + for (const match of source.matchAll( + /\bexport\s*\{([^}]+)\}(?:\s*from\s*["'][^"']+["'])?/g, + )) { + for (const part of match[1].split(",")) { + const trimmed = part.trim(); + if (!trimmed || trimmed.startsWith("type ")) continue; + const aliasMatch = trimmed.match(/\bas\s+([A-Za-z_$][\w$]*|default)$/); + const directMatch = trimmed.match(/^([A-Za-z_$][\w$]*|default)$/); + const name = aliasMatch?.[1] ?? directMatch?.[1]; + if (name) names.add(name); + } + } + + return [...names]; +} + +function buildEsmSourceWrapper({ + baseDir, + exportNames, + moduleName, + moduleVersion, + originalModuleSpecifier, +}: { + baseDir?: string; + exportNames: readonly string[]; + moduleName: string; + moduleVersion?: string; + originalModuleSpecifier: string; +}): string { + const locals = exportNames.map((name, index) => ({ + exportName: name, + localName: + name === "default" + ? "__braintrustDefaultExport" + : toSafeLocalName(name, index), + })); + + return `import * as __braintrustOriginal from ${JSON.stringify(originalModuleSpecifier)}; + +${locals + .map( + ({ exportName, localName }) => + `let ${localName} = __braintrustOriginal[${JSON.stringify(exportName)}];`, + ) + .join("\n")} + +const __braintrustExports = Object.create(null); +${locals.map(({ exportName, localName }) => buildMutableExportDescriptor(exportName, localName)).join("\n")} + +try { + const __braintrustHookRunner = globalThis[${JSON.stringify(MODULE_EXPORT_HOOK_RUNNER_GLOBAL)}]; + if (typeof __braintrustHookRunner === "function") { + const __braintrustPatched = __braintrustHookRunner( + __braintrustExports, + ${JSON.stringify(moduleName)}, + ${JSON.stringify(baseDir)}, + import.meta.url, + ${JSON.stringify(moduleVersion)}, + ); + if (__braintrustPatched && __braintrustPatched !== __braintrustExports) { +${locals + .map( + ({ exportName, localName }) => + ` if (Object.prototype.hasOwnProperty.call(__braintrustPatched, ${JSON.stringify(exportName)})) ${localName} = __braintrustPatched[${JSON.stringify(exportName)}];`, + ) + .join("\n")} + } + } +} catch (e) { + // Module export hooks are best-effort in bundled output. +} + +export * from ${JSON.stringify(originalModuleSpecifier)}; +${buildEsmExports(locals)} +`; +} + +function buildCjsSourceWrapper({ + baseDir, + exportNames, + moduleName, + moduleVersion, + originalModuleSpecifier, +}: { + baseDir?: string; + exportNames: readonly string[]; + moduleName: string; + moduleVersion?: string; + originalModuleSpecifier: string; +}): string { + const locals = exportNames.map((name, index) => ({ + exportName: name, + localName: + name === "default" + ? "__braintrustDefaultExport" + : toSafeLocalName(name, index), + })); + const cjsRequire = "require"; + + return `"use strict"; +const __braintrustOriginal = ${cjsRequire}(${JSON.stringify(originalModuleSpecifier)}); +const __braintrustPatchedExportNames = new Set(${JSON.stringify(exportNames)}); +try { + const __braintrustOriginalDescriptors = Object.getOwnPropertyDescriptors(__braintrustOriginal); +${locals.map(({ exportName }) => ` delete __braintrustOriginalDescriptors[${JSON.stringify(exportName)}];`).join("\n")} + Object.defineProperties(exports, __braintrustOriginalDescriptors); +} catch (e) { + for (const __braintrustKey in __braintrustOriginal) { + if (!__braintrustPatchedExportNames.has(__braintrustKey)) { + exports[__braintrustKey] = __braintrustOriginal[__braintrustKey]; + } + } +} + +${locals + .map( + ({ exportName, localName }) => + `let ${localName} = __braintrustOriginal[${JSON.stringify(exportName)}];`, + ) + .join("\n")} + +const __braintrustExports = Object.create(null); +${locals.map(({ exportName, localName }) => buildMutableExportDescriptor(exportName, localName)).join("\n")} + +try { + const __braintrustHookRunner = globalThis[${JSON.stringify(MODULE_EXPORT_HOOK_RUNNER_GLOBAL)}]; + if (typeof __braintrustHookRunner === "function") { + const __braintrustPatched = __braintrustHookRunner( + __braintrustExports, + ${JSON.stringify(moduleName)}, + ${JSON.stringify(baseDir)}, + typeof __filename === "string" ? __filename : undefined, + ${JSON.stringify(moduleVersion)}, + ); + if (__braintrustPatched && __braintrustPatched !== __braintrustExports) { +${locals + .map( + ({ exportName, localName }) => + ` if (Object.prototype.hasOwnProperty.call(__braintrustPatched, ${JSON.stringify(exportName)})) ${localName} = __braintrustPatched[${JSON.stringify(exportName)}];`, + ) + .join("\n")} + } + } +} catch (e) { + // Module export hooks are best-effort in bundled output. +} + +${locals.map(({ exportName, localName }) => buildCjsExportDescriptor(exportName, localName)).join("\n")} +`; +} + +function buildMutableExportDescriptor( + exportName: string, + localName: string, +): string { + return `Object.defineProperty(__braintrustExports, ${JSON.stringify(exportName)}, { + configurable: true, + enumerable: true, + get() { return ${localName}; }, + set(value) { ${localName} = value; return true; }, +});`; +} + +function buildEsmExports( + locals: readonly { exportName: string; localName: string }[], +): string { + const named = locals.filter(({ exportName }) => exportName !== "default"); + const defaultExport = locals.find( + ({ exportName }) => exportName === "default", + ); + const lines = named.length + ? [ + `export { ${named.map(({ localName, exportName }) => `${localName} as ${exportName}`).join(", ")} };`, + ] + : []; + if (defaultExport) { + lines.push(`export { ${defaultExport.localName} as default };`); + } + return lines.join("\n"); +} + +function buildCjsExportDescriptor( + exportName: string, + localName: string, +): string { + return `Object.defineProperty(exports, ${JSON.stringify(exportName)}, { + configurable: true, + enumerable: true, + get() { return ${localName}; }, +});`; +} + +function toSafeLocalName(exportName: string, index: number): string { + return `__braintrust_export_${index}_${exportName.replace(/\W/g, "_")}`; +} diff --git a/js/src/auto-instrumentations/loader/module-hooks/ritm.ts b/js/src/auto-instrumentations/loader/module-hooks/ritm.ts new file mode 100644 index 000000000..a7de07bad --- /dev/null +++ b/js/src/auto-instrumentations/loader/module-hooks/ritm.ts @@ -0,0 +1,274 @@ +import moduleDetailsFromPath from "module-details-from-path"; +import Module, { builtinModules, createRequire } from "node:module"; +import { parse as parsePath } from "node:path"; +import { pathToFileURL } from "node:url"; + +type OnRequireFn = (exports: T, name: string, basedir?: string) => T; +type GetBuiltinModuleFn = (this: unknown, id: string) => unknown; +type HookedRequire = (this: NodeJS.Module, id: string) => unknown; +type ModuleWithInternals = typeof Module & { + _cache: Record) | undefined>; + _resolveFilename(id: string, parent: NodeJS.Module): string; +}; + +let builtinModulesSet: Set | undefined; +const ModuleInternals = Module as ModuleWithInternals; +const requireForResolve = createRequire( + pathToFileURL(process.argv[1] ?? process.cwd()).href, +); + +const isCore = + typeof Module.isBuiltin === "function" + ? Module.isBuiltin + : (moduleName: string) => { + if (moduleName.startsWith("node:")) { + return true; + } + + builtinModulesSet ??= new Set(builtinModules); + return builtinModulesSet.has(moduleName); + }; + +const normalize = /([/\\]index)?(\.js|\.cjs)?$/; + +class ExportsCache { + private readonly localCache = new Map(); + private readonly kRitmExports = Symbol("RitmExports"); + + has(filename: string, isBuiltin: boolean): boolean { + if (this.localCache.has(filename)) { + return true; + } else if (!isBuiltin) { + const mod = ModuleInternals._cache[filename] as + | (NodeJS.Module & Record) + | undefined; + return !!(mod && this.kRitmExports in mod); + } else { + return false; + } + } + + get(filename: string, isBuiltin: boolean): unknown { + const cachedExports = this.localCache.get(filename); + if (cachedExports !== undefined) { + return cachedExports; + } else if (!isBuiltin) { + const mod = ModuleInternals._cache[filename] as + | (NodeJS.Module & Record) + | undefined; + return mod && mod[this.kRitmExports]; + } + } + + set(filename: string, exports: unknown, isBuiltin: boolean): void { + if (isBuiltin) { + this.localCache.set(filename, exports); + } else if (filename in ModuleInternals._cache) { + ( + ModuleInternals._cache[filename] as NodeJS.Module & + Record + )[this.kRitmExports] = exports; + } else { + this.localCache.set(filename, exports); + } + } +} + +function normalizeModules(modules: string[]): string[] { + if (!Array.isArray(modules) || modules.length === 0) { + throw new TypeError( + "Braintrust require-in-the-middle requires a non-empty modules array", + ); + } + + for (const each of modules) { + if (typeof each !== "string") { + throw new TypeError( + "Braintrust require-in-the-middle only supports string module names or absolute paths", + ); + } + } + + return modules; +} + +export default class Hook { + private readonly cache = new ExportsCache(); + private readonly hookedRequire: HookedRequire; + private hookedGetBuiltinModule?: GetBuiltinModuleFn; + private unhooked = false; + private readonly origRequire = Module.prototype.require; + private readonly origGetBuiltinModule: GetBuiltinModuleFn | undefined = ( + process as unknown as { getBuiltinModule?: GetBuiltinModuleFn } + ).getBuiltinModule; + + constructor(modules: string[], onrequire: OnRequireFn) { + modules = normalizeModules(modules); + if (typeof onrequire !== "function") { + throw new TypeError( + "Braintrust require-in-the-middle requires an onrequire function", + ); + } + + if (typeof ModuleInternals._resolveFilename !== "function") { + throw new Error( + `Expected Module._resolveFilename to be a function, got ${typeof ModuleInternals._resolveFilename}`, + ); + } + + const self = this; + const patching = new Set(); + + this.hookedRequire = Module.prototype.require = function (id: string) { + if (self.unhooked === true) { + return self.origRequire.apply(this, arguments as any); + } + + return patchedRequire.call(this, arguments, false); + }; + + if (typeof this.origGetBuiltinModule === "function") { + this.hookedGetBuiltinModule = ( + process as unknown as { getBuiltinModule: GetBuiltinModuleFn } + ).getBuiltinModule = function (id: string) { + if (self.unhooked === true) { + return self.origGetBuiltinModule!.apply(process, arguments as any); + } + + return patchedRequire.call( + process as unknown as NodeJS.Module, + arguments, + true, + ); + }; + } + + function patchedRequire( + this: NodeJS.Module, + args: IArguments, + coreOnly: boolean, + ): unknown { + const id = args[0] as string; + const core = isCore(id); + let filename: string; + if (core) { + filename = id; + if (id.startsWith("node:")) { + const idWithoutPrefix = id.slice(5); + if (isCore(idWithoutPrefix)) { + filename = idWithoutPrefix; + } + } + } else if (coreOnly) { + return self.origGetBuiltinModule!.call(process, id); + } else { + try { + filename = ModuleInternals._resolveFilename(id, this); + } catch { + return self.origRequire.apply(this, args as any); + } + } + + if (self.cache.has(filename, core) === true) { + return self.cache.get(filename, core); + } + + const isPatching = patching.has(filename); + if (isPatching === false) { + patching.add(filename); + } + + const exports = coreOnly + ? self.origGetBuiltinModule!.call(process, id) + : self.origRequire.apply(this, args as any); + + if (isPatching === true) { + return exports; + } + + patching.delete(filename); + + let moduleName: string; + let basedir: string | undefined; + + if (core === true) { + if (modules.includes(filename) === false) { + return exports; + } + moduleName = filename; + } else if (modules.includes(filename)) { + const parsedPath = parsePath(filename); + moduleName = parsedPath.name; + basedir = parsedPath.dir; + } else { + const stat = moduleDetailsFromPath(filename); + if (!stat) { + return exports; + } + moduleName = stat.name; + basedir = stat.basedir; + + const fullModuleName = resolveModuleName(stat); + let matchFound = false; + if (!id.startsWith(".") && modules.includes(id)) { + moduleName = id; + matchFound = true; + } + + if ( + !modules.includes(moduleName) && + !modules.includes(fullModuleName) + ) { + return exports; + } + + if (modules.includes(fullModuleName) && fullModuleName !== moduleName) { + moduleName = fullModuleName; + matchFound = true; + } + + if (!matchFound) { + let res: string; + try { + res = requireForResolve.resolve(moduleName, { + paths: basedir ? [basedir] : undefined, + }); + } catch { + self.cache.set(filename, exports, core); + return exports; + } + + if (res !== filename) { + self.cache.set(filename, exports, core); + return exports; + } + } + } + + const patchedExports = onrequire(exports, moduleName, basedir); + self.cache.set(filename, patchedExports, core); + return patchedExports; + } + } + + unhook(): void { + this.unhooked = true; + if (Module.prototype.require === this.hookedRequire) { + Module.prototype.require = this.origRequire; + } + const processWithGetBuiltinModule = process as unknown as { + getBuiltinModule?: GetBuiltinModuleFn; + }; + if ( + this.hookedGetBuiltinModule && + processWithGetBuiltinModule.getBuiltinModule === + this.hookedGetBuiltinModule + ) { + processWithGetBuiltinModule.getBuiltinModule = this.origGetBuiltinModule; + } + } +} + +function resolveModuleName(stat: { name: string; path: string }): string { + return `${stat.name}/${stat.path.replace(normalize, "")}`; +} diff --git a/js/src/auto-instrumentations/loader/get-package-version.ts b/js/src/auto-instrumentations/loader/package-version.ts similarity index 90% rename from js/src/auto-instrumentations/loader/get-package-version.ts rename to js/src/auto-instrumentations/loader/package-version.ts index d2470060b..7f4f917fd 100644 --- a/js/src/auto-instrumentations/loader/get-package-version.ts +++ b/js/src/auto-instrumentations/loader/package-version.ts @@ -44,6 +44,12 @@ function readPackageJsonWithFallback( } export function getPackageVersion(baseDir: string): string { + return getPackageVersionIfAvailable(baseDir) ?? process.version.slice(1); +} + +export function getPackageVersionIfAvailable( + baseDir: string, +): string | undefined { if (packageVersions.has(baseDir)) { return packageVersions.get(baseDir)!; } @@ -54,7 +60,7 @@ export function getPackageVersion(baseDir: string): string { return packageJson.version; } - return process.version.slice(1); + return undefined; } export function getPackageName(baseDir: string): string | undefined { diff --git a/js/src/auto-instrumentations/loader/special-case-patches.ts b/js/src/auto-instrumentations/loader/source-patches/index.ts similarity index 60% rename from js/src/auto-instrumentations/loader/special-case-patches.ts rename to js/src/auto-instrumentations/loader/source-patches/index.ts index 4b984c158..97721d05f 100644 --- a/js/src/auto-instrumentations/loader/special-case-patches.ts +++ b/js/src/auto-instrumentations/loader/source-patches/index.ts @@ -21,31 +21,19 @@ * .parse()` doesn't double-read the response body. Removable once OpenAI * stops sharing the same `APIPromise` between `create()` and * `_thenUnwrap()`. - * - `@mastra/core` and `@mastra/observability` entries: Mastra ships - * code-split bundles with content-hashed chunk filenames, so we patch - * the stable submodule entries to install the - * `BraintrustObservabilityExporter` automatically. Removable when Mastra - * adopts a NPM-installable Braintrust exporter package directly, or when - * `import-in-the-middle` is reliable enough across Node versions to use - * for the same job. */ -import { - classifyMastraTarget, - patchMastraSource, - type MastraModuleFormat, -} from "./mastra-observability-patch.js"; -import { OPENAI_API_PROMISE_PATCH } from "./openai-api-promise-patch.js"; +import { OPENAI_API_PROMISE_PATCH } from "./openai.js"; -type SpecialCaseFormat = MastraModuleFormat; +type SourcePatchFormat = "esm" | "cjs"; -interface SpecialCaseInput { +interface SourcePatchInput { packageName: string; /** Forward-slash-normalized path inside the package, e.g. `dist/index.js`. */ modulePath: string; /** Original module source as a string. */ source: string; - format: SpecialCaseFormat; + format: SourcePatchFormat; } /** @@ -53,7 +41,7 @@ interface SpecialCaseInput { * Returns the patched source, or `null` if no entry matched (caller falls * through to the standard transformation pipeline). */ -export function applySpecialCasePatch(input: SpecialCaseInput): string | null { +export function applySourcePatch(input: SourcePatchInput): string | null { // OpenAI: append idempotent .then() wrap to dist/api-promise.{m,c}js. if ( input.packageName === "openai" && @@ -62,33 +50,20 @@ export function applySpecialCasePatch(input: SpecialCaseInput): string | null { return input.source + OPENAI_API_PROMISE_PATCH; } - // Mastra: rewrite the stable submodule entries (@mastra/core) or append a - // Proxy wrap to the inline class binding (@mastra/observability). - const mastraTarget = classifyMastraTarget( - input.packageName, - input.modulePath, - ); - if (mastraTarget) { - return patchMastraSource(input.source, mastraTarget, input.format); - } - return null; } /** * Synchronous predicate variant used by the ESM resolve hook to decide * up-front whether a URL needs to be remembered for later patching. The - * load step still calls `applySpecialCasePatch` against the actual source. + * load step still calls `applySourcePatch` against the actual source. */ -export function isSpecialCaseTarget( +export function isSourcePatchTarget( packageName: string, modulePath: string, ): boolean { if (packageName === "openai" && modulePath.includes("api-promise")) { return true; } - if (classifyMastraTarget(packageName, modulePath) !== null) { - return true; - } return false; } diff --git a/js/src/auto-instrumentations/loader/openai-api-promise-patch.ts b/js/src/auto-instrumentations/loader/source-patches/openai.ts similarity index 100% rename from js/src/auto-instrumentations/loader/openai-api-promise-patch.ts rename to js/src/auto-instrumentations/loader/source-patches/openai.ts diff --git a/js/src/auto-instrumentations/require-in-the-middle/index.ts b/js/src/auto-instrumentations/require-in-the-middle/index.ts deleted file mode 100644 index 416e22644..000000000 --- a/js/src/auto-instrumentations/require-in-the-middle/index.ts +++ /dev/null @@ -1,239 +0,0 @@ -import Module from "node:module"; -import moduleDetailsFromPath from "module-details-from-path"; -import path from "node:path"; - -type OnRequireFn = ( - exports: unknown, - name: string, - basedir?: string, -) => unknown; -type GetBuiltinModuleFn = (this: unknown, id: string) => unknown; -type ProcessWithGetBuiltinModule = Omit & { - getBuiltinModule?: GetBuiltinModuleFn; -}; -type ModuleWithInternals = typeof Module & { - _resolveFilename(id: string, parent: NodeJS.Module): string; -}; -type CachedModule = NodeJS.Module & Record; -type HookedRequire = (this: NodeJS.Module, id: string) => unknown; - -const ModuleInternals = Module as ModuleWithInternals; -const isCore = Module.isBuiltin; - -const normalize = /([/\\]index)?(\.js|\.cjs)?$/; - -class ExportsCache { - private readonly localCache = new Map(); - private readonly kRitmExports = Symbol("RitmExports"); - - has(filename: string, isBuiltin: boolean): boolean { - if (this.localCache.has(filename)) { - return true; - } else if (!isBuiltin) { - const mod = require.cache[filename] as CachedModule | undefined; - return !!(mod && this.kRitmExports in mod); - } else { - return false; - } - } - - get(filename: string, isBuiltin: boolean): unknown { - const cachedExports = this.localCache.get(filename); - if (cachedExports !== undefined) { - return cachedExports; - } else if (!isBuiltin) { - const mod = require.cache[filename] as CachedModule | undefined; - return mod && mod[this.kRitmExports]; - } - } - - set(filename: string, exports: unknown, isBuiltin: boolean): void { - if (isBuiltin) { - this.localCache.set(filename, exports); - } else if (filename in require.cache) { - (require.cache[filename] as CachedModule)[this.kRitmExports] = exports; - } else { - this.localCache.set(filename, exports); - } - } -} - -class Hook { - private readonly _cache = new ExportsCache(); - private readonly _origRequire: HookedRequire; - private readonly _patching = new Set(); - private readonly _normalizedModules: string[]; - private _getBuiltinModule?: GetBuiltinModuleFn; - private _origGetBuiltinModule?: GetBuiltinModuleFn; - private _require: HookedRequire; - private _unhooked = false; - - constructor( - modules: readonly string[], - private readonly _onRequireFn: OnRequireFn, - ) { - this._normalizedModules = Array.from(modules); - this._origRequire = Module.prototype.require; - const self = this; - - this._require = Module.prototype.require = function ( - this: NodeJS.Module, - id: string, - ) { - if (self._unhooked === true) { - return self._origRequire.call(this, id); - } - - return self._patchedRequire(this, id, false); - }; - - const processWithGetBuiltinModule: ProcessWithGetBuiltinModule = process; - if (typeof processWithGetBuiltinModule.getBuiltinModule === "function") { - this._origGetBuiltinModule = processWithGetBuiltinModule.getBuiltinModule; - this._getBuiltinModule = processWithGetBuiltinModule.getBuiltinModule = - function (this: unknown, id: string) { - if (self._unhooked === true) { - return self._origGetBuiltinModule!.call(this, id); - } - - return self._patchedRequire(this as NodeJS.Module, id, true); - }; - } - } - - unhook(): void { - this._unhooked = true; - - if (this._require === Module.prototype.require) { - Module.prototype.require = this._origRequire; - } - - const processWithGetBuiltinModule: ProcessWithGetBuiltinModule = process; - if ( - this._getBuiltinModule && - this._getBuiltinModule === processWithGetBuiltinModule.getBuiltinModule - ) { - processWithGetBuiltinModule.getBuiltinModule = this._origGetBuiltinModule; - } - } - - private _patchedRequire( - requireThis: NodeJS.Module, - id: string, - coreOnly: boolean, - ): unknown { - const core = isCore(id); - let filename: string; - if (core) { - filename = id; - if (id.startsWith("node:")) { - const idWithoutPrefix = id.slice(5); - if (isCore(idWithoutPrefix)) { - filename = idWithoutPrefix; - } - } - } else if (coreOnly) { - return this._origGetBuiltinModule!.call(requireThis, id); - } else { - try { - filename = ModuleInternals._resolveFilename(id, requireThis); - } catch { - return this._origRequire.call(requireThis, id); - } - } - - if (this._cache.has(filename, core) === true) { - return this._cache.get(filename, core); - } - - const isPatching = this._patching.has(filename); - if (isPatching === false) { - this._patching.add(filename); - } - - const exports = coreOnly - ? this._origGetBuiltinModule!.call(requireThis, id) - : this._origRequire.call(requireThis, id); - - if (isPatching === true) { - return exports; - } - - this._patching.delete(filename); - - let moduleName: string; - let basedir: string | undefined; - - if (core === true) { - if (this._normalizedModules.includes(filename) === false) { - return exports; - } - moduleName = filename; - } else if (this._normalizedModules.includes(filename)) { - const parsedPath = path.parse(filename); - moduleName = parsedPath.name; - basedir = parsedPath.dir; - } else { - const stat = moduleDetailsFromPath(filename); - if (stat === undefined || stat === null) { - return exports; - } - moduleName = stat.name; - basedir = stat.basedir; - - const fullModuleName = resolveModuleName(stat); - let matchFound = false; - if (!id.startsWith(".") && this._normalizedModules.includes(id)) { - moduleName = id; - matchFound = true; - } - - if ( - !this._normalizedModules.includes(moduleName) && - !this._normalizedModules.includes(fullModuleName) - ) { - return exports; - } - - if ( - this._normalizedModules.includes(fullModuleName) && - fullModuleName !== moduleName - ) { - moduleName = fullModuleName; - matchFound = true; - } - - if (!matchFound) { - let res: string; - try { - res = require.resolve(moduleName, { paths: [basedir] }); - } catch { - this._cache.set(filename, exports, core); - return exports; - } - - if (res !== filename) { - this._cache.set(filename, exports, core); - return exports; - } - } - } - - this._cache.set(filename, exports, core); - const patchedExports = this._onRequireFn(exports, moduleName, basedir); - this._cache.set(filename, patchedExports, core); - - return patchedExports; - } -} - -module.exports = Hook; -module.exports.Hook = Hook; - -function resolveModuleName(stat: { name: string; path: string }): string { - const normalizedPath = - path.sep !== "/" ? stat.path.split(path.sep).join("/") : stat.path; - return path.posix.join(stat.name, normalizedPath).replace(normalize, ""); -} - -export {}; diff --git a/js/src/auto-instrumentations/require-in-the-middle/tsconfig.json b/js/src/auto-instrumentations/require-in-the-middle/tsconfig.json deleted file mode 100644 index 1046b1ae6..000000000 --- a/js/src/auto-instrumentations/require-in-the-middle/tsconfig.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "extends": "../../../tsconfig.vendor-hooks.json", - "include": ["../types/module-details-from-path.d.ts", "**/*.ts"] -} diff --git a/js/src/imports.test.ts b/js/src/imports.test.ts index 96a4851d5..e093463cb 100644 --- a/js/src/imports.test.ts +++ b/js/src/imports.test.ts @@ -204,19 +204,9 @@ describe("CLI import restrictions", () => { const content = fs.readFileSync(filePath, "utf-8"); const lines = content.split("\n"); - // Exception: the Mastra observability patch *generates* CJS source code - // (a string of JavaScript that is appended to @mastra/core's - // dist/mastra/index.cjs). That source legitimately contains - // `require(...)` because it runs inside a CommonJS module's evaluation - // — it isn't a require() call from our SDK code itself. Normalize the - // separator so the exemption matches on Windows (`\\`) too. - const isMastraPatchGenerator = relativePath - .replaceAll(path.sep, "/") - .includes("auto-instrumentations/loader/mastra-observability-patch.ts"); - lines.forEach((line, index) => { // Check for require() calls - if (/\brequire\s*\(/.test(line) && !isMastraPatchGenerator) { + if (/\brequire\s*\(/.test(line)) { violations.push( `${relativePath}:${index + 1} - Found require() statement: "${line.trim()}"`, ); diff --git a/js/src/instrumentation/braintrust-plugin.test.ts b/js/src/instrumentation/braintrust-plugin.test.ts index 067c1f85f..bcbf67b30 100644 --- a/js/src/instrumentation/braintrust-plugin.test.ts +++ b/js/src/instrumentation/braintrust-plugin.test.ts @@ -17,6 +17,7 @@ import { GitHubCopilotPlugin } from "./plugins/github-copilot-plugin"; import { LangChainPlugin } from "./plugins/langchain-plugin"; import { PiCodingAgentPlugin } from "./plugins/pi-coding-agent-plugin"; import { StrandsAgentSDKPlugin } from "./plugins/strands-agent-sdk-plugin"; +import { MastraPlugin } from "./plugins/mastra-plugin"; function createPluginClassMock() { return vi.fn(function MockPlugin(this: { @@ -106,6 +107,10 @@ vi.mock("./plugins/strands-agent-sdk-plugin", () => ({ StrandsAgentSDKPlugin: createPluginClassMock(), })); +vi.mock("./plugins/mastra-plugin", () => ({ + MastraPlugin: createPluginClassMock(), +})); + describe("BraintrustPlugin", () => { beforeEach(() => { vi.clearAllMocks(); @@ -260,6 +265,15 @@ describe("BraintrustPlugin", () => { expect(mockInstance.enable).toHaveBeenCalledTimes(1); }); + it("should create and enable Mastra plugin by default", () => { + const plugin = new BraintrustPlugin(); + plugin.enable(); + + expect(MastraPlugin).toHaveBeenCalledTimes(1); + const mockInstance = vi.mocked(MastraPlugin).mock.results[0].value; + expect(mockInstance.enable).toHaveBeenCalledTimes(1); + }); + it("should create all plugins when enabled with no config", () => { const plugin = new BraintrustPlugin(); plugin.enable(); @@ -280,6 +294,7 @@ describe("BraintrustPlugin", () => { expect(GitHubCopilotPlugin).toHaveBeenCalledTimes(1); expect(StrandsAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(LangChainPlugin).toHaveBeenCalledTimes(1); + expect(MastraPlugin).toHaveBeenCalledTimes(1); }); it("should create all plugins when enabled with empty config", () => { @@ -302,6 +317,7 @@ describe("BraintrustPlugin", () => { expect(GitHubCopilotPlugin).toHaveBeenCalledTimes(1); expect(StrandsAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(LangChainPlugin).toHaveBeenCalledTimes(1); + expect(MastraPlugin).toHaveBeenCalledTimes(1); }); it("should create all plugins when enabled with empty integrations config", () => { @@ -323,6 +339,7 @@ describe("BraintrustPlugin", () => { expect(GroqPlugin).toHaveBeenCalledTimes(1); expect(StrandsAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(LangChainPlugin).toHaveBeenCalledTimes(1); + expect(MastraPlugin).toHaveBeenCalledTimes(1); }); }); @@ -582,6 +599,7 @@ describe("BraintrustPlugin", () => { langchain: false, piCodingAgent: false, strandsAgentSDK: false, + mastra: false, }, }); plugin.enable(); @@ -603,6 +621,7 @@ describe("BraintrustPlugin", () => { expect(LangChainPlugin).not.toHaveBeenCalled(); expect(PiCodingAgentPlugin).not.toHaveBeenCalled(); expect(StrandsAgentSDKPlugin).not.toHaveBeenCalled(); + expect(MastraPlugin).not.toHaveBeenCalled(); }); it("should not create Pi Coding Agent plugin when piCodingAgent: false", () => { @@ -626,6 +645,16 @@ describe("BraintrustPlugin", () => { expect(OpenAIPlugin).toHaveBeenCalledTimes(1); }); + it("should not create Mastra plugin when mastra: false", () => { + const plugin = new BraintrustPlugin({ + integrations: { mastra: false }, + }); + plugin.enable(); + + expect(MastraPlugin).not.toHaveBeenCalled(); + expect(OpenAIPlugin).toHaveBeenCalledTimes(1); + }); + it("should allow selective enabling of plugins", () => { const plugin = new BraintrustPlugin({ integrations: { @@ -784,6 +813,7 @@ describe("BraintrustPlugin", () => { const strandsAgentSDKMock = vi.mocked(StrandsAgentSDKPlugin).mock .results[0].value; const langChainMock = vi.mocked(LangChainPlugin).mock.results[0].value; + const mastraMock = vi.mocked(MastraPlugin).mock.results[0].value; expect(openaiMock.enable).toHaveBeenCalledTimes(1); expect(openAICodexMock.enable).toHaveBeenCalledTimes(1); @@ -801,6 +831,7 @@ describe("BraintrustPlugin", () => { expect(piCodingAgentMock.enable).toHaveBeenCalledTimes(1); expect(strandsAgentSDKMock.enable).toHaveBeenCalledTimes(1); expect(langChainMock.enable).toHaveBeenCalledTimes(1); + expect(mastraMock.enable).toHaveBeenCalledTimes(1); }); it("should disable and nullify all sub-plugins when disabled", () => { @@ -831,6 +862,7 @@ describe("BraintrustPlugin", () => { const strandsAgentSDKMock = vi.mocked(StrandsAgentSDKPlugin).mock .results[0].value; const langChainMock = vi.mocked(LangChainPlugin).mock.results[0].value; + const mastraMock = vi.mocked(MastraPlugin).mock.results[0].value; plugin.disable(); @@ -850,6 +882,7 @@ describe("BraintrustPlugin", () => { expect(piCodingAgentMock.disable).toHaveBeenCalledTimes(1); expect(strandsAgentSDKMock.disable).toHaveBeenCalledTimes(1); expect(langChainMock.disable).toHaveBeenCalledTimes(1); + expect(mastraMock.disable).toHaveBeenCalledTimes(1); }); it("should be idempotent on multiple enable calls", () => { diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index a1fe0e352..2266f9c17 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -21,6 +21,7 @@ import { FluePlugin } from "./plugins/flue-plugin"; import { LangChainPlugin } from "./plugins/langchain-plugin"; import { PiCodingAgentPlugin } from "./plugins/pi-coding-agent-plugin"; import { StrandsAgentSDKPlugin } from "./plugins/strands-agent-sdk-plugin"; +import { MastraPlugin } from "./plugins/mastra-plugin"; import type { InstrumentationIntegrationsConfig } from "./config"; export interface BraintrustPluginConfig { @@ -40,6 +41,7 @@ export interface BraintrustPluginConfig { * - LangChain.js and LangGraph * - Mistral SDK * - Cohere SDK + * - Mastra observability * * The plugin is automatically enabled when the Braintrust library is loaded. * Individual integrations can be disabled via configuration. @@ -68,6 +70,7 @@ export class BraintrustPlugin extends BasePlugin { private langChainPlugin: LangChainPlugin | null = null; private piCodingAgentPlugin: PiCodingAgentPlugin | null = null; private strandsAgentSDKPlugin: StrandsAgentSDKPlugin | null = null; + private mastraPlugin: MastraPlugin | null = null; constructor(config: BraintrustPluginConfig = {}) { super(); @@ -200,11 +203,10 @@ export class BraintrustPlugin extends BasePlugin { this.langChainPlugin.enable(); } - // Mastra is intentionally not wired here: `@mastra/core` ships its own - // ObservabilityExporter contract, and `BraintrustObservabilityExporter` - // (wrappers/mastra.ts) is auto-installed by the loader patch in - // `auto-instrumentations/loader/mastra-observability-patch.ts` rather than - // by a BasePlugin / tracingChannel subscription. + if (integrations.mastra !== false) { + this.mastraPlugin = new MastraPlugin(); + this.mastraPlugin.enable(); + } } protected onDisable(): void { @@ -317,6 +319,11 @@ export class BraintrustPlugin extends BasePlugin { this.langChainPlugin.disable(); this.langChainPlugin = null; } + + if (this.mastraPlugin) { + this.mastraPlugin.disable(); + this.mastraPlugin = null; + } } } diff --git a/js/src/instrumentation/plugins/mastra-channels.ts b/js/src/instrumentation/plugins/mastra-channels.ts new file mode 100644 index 000000000..6fcb01235 --- /dev/null +++ b/js/src/instrumentation/plugins/mastra-channels.ts @@ -0,0 +1,5 @@ +export const mastraChannels = { + mastraConstructor: "orchestrion:@mastra/core:Mastra.constructor", + observabilityConstructor: + "orchestrion:@mastra/observability:Observability.constructor", +} as const; diff --git a/js/src/instrumentation/plugins/mastra-plugin.test.ts b/js/src/instrumentation/plugins/mastra-plugin.test.ts new file mode 100644 index 000000000..54b15023b --- /dev/null +++ b/js/src/instrumentation/plugins/mastra-plugin.test.ts @@ -0,0 +1,231 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { debugLogger } from "../../debug-logger"; +import iso, { type IsoChannelHandlers } from "../../isomorph"; +import type { ModuleExportConstructorEvent } from "../../auto-instrumentations/loader/module-hooks/registry"; +import { MastraPlugin } from "./mastra-plugin"; +import { mastraChannels } from "./mastra-channels"; + +vi.mock("../../wrappers/mastra", () => ({ + BraintrustObservabilityExporter: class { + readonly name = "braintrust"; + }, +})); + +const originalNewTracingChannel = iso.newTracingChannel; +const channels = new Map(); + +beforeEach(() => { + channels.clear(); + iso.newTracingChannel = vi.fn((name: string | object) => { + const channel = new FakeTracingChannel(); + channels.set(String(name), channel); + return channel; + }) as typeof iso.newTracingChannel; +}); + +afterEach(() => { + iso.newTracingChannel = originalNewTracingChannel; + vi.restoreAllMocks(); +}); + +describe("MastraPlugin", () => { + it("preserves Observability configs and appends one Braintrust exporter", () => { + const plugin = new MastraPlugin(); + plugin.enable(); + const event = constructorEvent([ + { + custom: "kept", + configs: { + default: { + exporters: [{ name: "other" }], + serviceName: "service", + }, + secondary: { customValue: true, serviceName: "secondary" }, + }, + }, + ]); + + start(mastraChannels.observabilityConstructor, event); + + expect(event.arguments[0]).toEqual({ + custom: "kept", + configs: { + default: { + exporters: [{ name: "other" }, { name: "braintrust" }], + serviceName: "service", + }, + secondary: { + customValue: true, + exporters: [{ name: "braintrust" }], + serviceName: "secondary", + }, + }, + }); + }); + + it("does not duplicate an existing Braintrust exporter", () => { + const plugin = new MastraPlugin(); + plugin.enable(); + const existing = { name: "braintrust" }; + const event = constructorEvent([ + { + configs: { + default: { exporters: [existing], serviceName: "service" }, + }, + }, + ]); + + start(mastraChannels.observabilityConstructor, event); + + expect( + ( + event.arguments[0] as { + configs: { default: { exporters: unknown[] } }; + } + ).configs.default.exporters, + ).toEqual([existing]); + }); + + it("creates a default Observability config", () => { + const plugin = new MastraPlugin(); + plugin.enable(); + const event = constructorEvent([undefined]); + + start(mastraChannels.observabilityConstructor, event); + + expect(event.arguments[0]).toEqual({ + configs: { + default: { + exporters: [{ name: "braintrust" }], + serviceName: "mastra", + }, + }, + }); + }); + + it("injects Observability into Mastra constructor arguments", () => { + class Observability { + constructor(public readonly config: unknown) {} + } + const plugin = new MastraPlugin(); + plugin.enable(); + const event = constructorEvent([{ custom: "kept" }], () => ({ + Observability, + })); + + start(mastraChannels.mastraConstructor, event); + + const config = event.arguments[0] as { + custom: string; + observability: Observability; + }; + expect(config.custom).toBe("kept"); + expect(config.observability).toBeInstanceOf(Observability); + expect(config.observability.config).toEqual({ + configs: { default: { serviceName: "mastra" } }, + }); + }); + + it.each([false, undefined, { user: true }])( + "preserves an explicitly provided observability value: %j", + (observability) => { + const plugin = new MastraPlugin(); + plugin.enable(); + const event = constructorEvent([{ observability }], () => { + throw new Error("should not resolve"); + }); + + start(mastraChannels.mastraConstructor, event); + + expect(event.arguments[0]).toEqual({ observability }); + }, + ); + + it("leaves Mastra arguments unchanged when Observability cannot be resolved", () => { + const plugin = new MastraPlugin(); + plugin.enable(); + const original = { custom: true }; + const event = constructorEvent([original], () => undefined); + + start(mastraChannels.mastraConstructor, event); + + expect(event.arguments[0]).toBe(original); + }); + + it("contains subscriber errors and logs a warning", () => { + const warn = vi.spyOn(debugLogger, "warn").mockImplementation(() => {}); + const plugin = new MastraPlugin(); + plugin.enable(); + const event = constructorEvent([{}], () => { + throw new Error("resolution failed"); + }); + + expect(() => start(mastraChannels.mastraConstructor, event)).not.toThrow(); + expect(warn).toHaveBeenCalledWith( + "Failed to configure Mastra observability", + expect.any(Error), + ); + }); + + it("unsubscribes both channels when disabled", () => { + const plugin = new MastraPlugin(); + plugin.enable(); + plugin.disable(); + + expect(channels.get(mastraChannels.mastraConstructor)?.handlers.size).toBe( + 0, + ); + expect( + channels.get(mastraChannels.observabilityConstructor)?.handlers.size, + ).toBe(0); + }); +}); + +class FakeTracingChannel { + readonly handlers = new Set< + IsoChannelHandlers + >(); + + subscribe(handlers: IsoChannelHandlers): void { + this.handlers.add(handlers); + } + + unsubscribe( + handlers: IsoChannelHandlers, + ): boolean { + return this.handlers.delete(handlers); + } + + get hasSubscribers(): boolean { + return this.handlers.size > 0; + } + + traceSync(fn: () => T): T { + return fn(); + } + + tracePromise(fn: () => PromiseLike): PromiseLike { + return fn(); + } + + traceCallback(fn: () => T): T { + return fn(); + } +} + +function constructorEvent( + args: unknown[], + resolveModule: (specifier: string) => unknown = () => undefined, +): ModuleExportConstructorEvent { + return { + arguments: args, + moduleName: "pkg", + resolveModule, + }; +} + +function start(channelName: string, event: ModuleExportConstructorEvent): void { + for (const handlers of channels.get(channelName)?.handlers ?? []) { + handlers.start?.(event, channelName); + } +} diff --git a/js/src/instrumentation/plugins/mastra-plugin.ts b/js/src/instrumentation/plugins/mastra-plugin.ts new file mode 100644 index 000000000..6954cf3e9 --- /dev/null +++ b/js/src/instrumentation/plugins/mastra-plugin.ts @@ -0,0 +1,114 @@ +import { debugLogger } from "../../debug-logger"; +import iso, { type IsoChannelHandlers } from "../../isomorph"; +import { BraintrustObservabilityExporter } from "../../wrappers/mastra"; +import { isObject } from "../../../util/index"; +import type { ModuleExportConstructorEvent } from "../../auto-instrumentations/loader/module-hooks/registry"; +import { BasePlugin } from "../core"; +import { unsubscribeAll } from "../core/channel-tracing"; +import { mastraChannels } from "./mastra-channels"; + +export class MastraPlugin extends BasePlugin { + protected onEnable(): void { + const mastraChannel = iso.newTracingChannel( + mastraChannels.mastraConstructor, + ); + const mastraHandlers: IsoChannelHandlers = { + start(event) { + try { + const config = event.arguments[0]; + if (isObject(config) && "observability" in config) { + return; + } + + const observabilityModule = event.resolveModule( + "@mastra/observability", + ); + if (!isObject(observabilityModule)) { + return; + } + const Observability = observabilityModule.Observability; + if (typeof Observability !== "function") { + return; + } + + const observability = Reflect.construct(Observability, [ + { configs: { default: { serviceName: "mastra" } } }, + ]); + event.arguments[0] = isObject(config) + ? { ...config, observability } + : { observability }; + } catch (error) { + debugLogger.warn("Failed to configure Mastra observability", error); + } + }, + }; + mastraChannel.subscribe(mastraHandlers); + this.unsubscribers.push(() => mastraChannel.unsubscribe(mastraHandlers)); + + const observabilityChannel = + iso.newTracingChannel( + mastraChannels.observabilityConstructor, + ); + const observabilityHandlers: IsoChannelHandlers = + { + start(event) { + try { + const rawConfig = event.arguments[0]; + const config = isObject(rawConfig) ? rawConfig : {}; + const configsIn = isObject(config.configs) + ? config.configs + : undefined; + const configsOut: Record = {}; + + if (configsIn && Object.keys(configsIn).length > 0) { + for (const [name, rawInstanceConfig] of Object.entries( + configsIn, + )) { + const instanceConfig = isObject(rawInstanceConfig) + ? rawInstanceConfig + : {}; + const existingExporters = Array.isArray( + instanceConfig.exporters, + ) + ? instanceConfig.exporters + : []; + const hasBraintrustExporter = existingExporters.some( + (exporter) => + isObject(exporter) && exporter.name === "braintrust", + ); + configsOut[name] = { + ...instanceConfig, + exporters: hasBraintrustExporter + ? existingExporters + : [ + ...existingExporters, + new BraintrustObservabilityExporter(), + ], + }; + } + } else { + configsOut.default = { + exporters: [new BraintrustObservabilityExporter()], + serviceName: "mastra", + }; + } + + event.arguments[0] = { ...config, configs: configsOut }; + } catch (error) { + debugLogger.warn( + "Failed to configure the Braintrust Mastra exporter", + error, + ); + } + }, + }; + observabilityChannel.subscribe(observabilityHandlers); + this.unsubscribers.push(() => + observabilityChannel.unsubscribe(observabilityHandlers), + ); + } + + protected onDisable(): void { + this.unsubscribers = unsubscribeAll(this.unsubscribers); + } +} diff --git a/js/src/node/apply-auto-instrumentation-entry.ts b/js/src/node/apply-auto-instrumentation-entry.ts index 0328688dd..7fc024f5e 100644 --- a/js/src/node/apply-auto-instrumentation-entry.ts +++ b/js/src/node/apply-auto-instrumentation-entry.ts @@ -1,9 +1,16 @@ import * as diagnostics_channel from "node:diagnostics_channel"; import { register } from "node:module"; import { pathToFileURL } from "node:url"; -import { getDefaultAutoInstrumentationConfigs } from "../auto-instrumentations/configs/all"; -import { ModulePatch } from "../auto-instrumentations/loader/cjs-patch"; +import { + getDefaultModuleExportPatchConfigs, + getDefaultOrchestrionConfigs, +} from "../auto-instrumentations/configs/all"; +import { ModulePatch } from "../auto-instrumentations/loader/cjs"; +import { nodeModuleExportPatchRuntime } from "../auto-instrumentations/loader/module-hooks/node-runtime"; +import { installModuleExportPatchRunner } from "../auto-instrumentations/loader/module-hooks/registry"; +import { installNodeModuleExportHooks } from "../auto-instrumentations/loader/module-hooks/node"; import { patchTracingChannel } from "../auto-instrumentations/patch-tracing-channel"; +import { readDisabledInstrumentationEnvConfig } from "../instrumentation/config"; interface ApplyAutoInstrumentationState { applied?: boolean; @@ -32,18 +39,44 @@ if (state !== existingState) { if (!state.applied) { patchTracingChannel(diagnostics_channel.tracingChannel); - const allConfigs = getDefaultAutoInstrumentationConfigs(); + const disabled = readDisabledInstrumentationEnvConfig( + process.env.BRAINTRUST_DISABLE_INSTRUMENTATION, + ).integrations; + const orchestrionConfigs = getDefaultOrchestrionConfigs({ + disabledIntegrationConfig: disabled, + }); + const moduleExportPatchConfigs = getDefaultModuleExportPatchConfigs({ + disabledIntegrationConfig: disabled, + target: "node", + }); + + installModuleExportPatchRunner( + moduleExportPatchConfigs, + nodeModuleExportPatchRuntime, + ); const currentModuleUrl = getCurrentModuleUrl(); - register("./auto-instrumentations/loader/esm-hook.mjs", { + const autoInstrumentationHookUrl = new URL( + "./auto-instrumentations/hook.mjs", + currentModuleUrl, + ).href; + const asyncImportHookUrl = new URL(autoInstrumentationHookUrl); + asyncImportHookUrl.searchParams.set("braintrust-iitm-loader", "true"); + installNodeModuleExportHooks({ + asyncImportHookUrl: asyncImportHookUrl.href, + configs: moduleExportPatchConfigs, + registryImportUrl: autoInstrumentationHookUrl, + }); + + register("./auto-instrumentations/loader/esm.mjs", { parentURL: currentModuleUrl, - data: { instrumentations: allConfigs }, + data: { instrumentations: orchestrionConfigs }, }); state.applied = true; try { - const patch = new ModulePatch({ instrumentations: allConfigs }); + const patch = new ModulePatch({ instrumentations: orchestrionConfigs }); patch.patch(); } catch { // ESM instrumentation is already active; keep user code running if CJS patching fails. diff --git a/js/src/node/config.ts b/js/src/node/config.ts index 09c27b807..48c183984 100644 --- a/js/src/node/config.ts +++ b/js/src/node/config.ts @@ -15,12 +15,10 @@ import { getRepoInfo, getPastNAncestors } from "../gitutil"; import { getCallerLocation } from "../stackutil"; import { _internalSetInitialState } from "../logger"; import { registry } from "../instrumentation/registry"; -import { - isInstrumentationIntegrationDisabled, - readDisabledInstrumentationEnvConfig, -} from "../instrumentation/config"; -import { installMastraExporterFactory } from "../auto-instrumentations/loader/mastra-observability-patch"; -import { BraintrustObservabilityExporter } from "../wrappers/mastra"; +import { readDisabledInstrumentationEnvConfig } from "../instrumentation/config"; +import { getDefaultModuleExportPatchConfigs } from "../auto-instrumentations/configs/all"; +import { nodeModuleExportPatchRuntime } from "../auto-instrumentations/loader/module-hooks/node-runtime"; +import { installModuleExportPatchRunner } from "../auto-instrumentations/loader/module-hooks/registry"; const BRAINTRUST_ENV_SEARCH_PARENT_LIMIT = 64; @@ -136,18 +134,18 @@ export function configureNode() { _internalSetInitialState(); - // The Mastra patches rewritten by the bundler plugin and the loader hook - // both look up an exporter factory on `globalThis` at runtime. Register it - // here too so a bundled app (where neither hook.mjs nor the loader runs - // against `@mastra/core`) still gets the exporter installed when its code - // imports `braintrust`. The loader path (hook.mjs) also calls this; the - // `??=` makes double-registration a no-op. + // Bundled wrappers call the module export runner through globalThis. Installing + // it here covers applications that import Braintrust without hook.mjs. const disabled = readDisabledInstrumentationEnvConfig( iso.getEnv("BRAINTRUST_DISABLE_INSTRUMENTATION"), ).integrations; - if (!isInstrumentationIntegrationDisabled(disabled, "mastra")) { - installMastraExporterFactory(() => new BraintrustObservabilityExporter()); - } + installModuleExportPatchRunner( + getDefaultModuleExportPatchConfigs({ + disabledIntegrationConfig: disabled, + target: "node", + }), + nodeModuleExportPatchRuntime, + ); // Enable auto-instrumentation registry.enable(); diff --git a/js/src/wrappers/mastra.ts b/js/src/wrappers/mastra.ts index 9946a9643..44a3740ee 100644 --- a/js/src/wrappers/mastra.ts +++ b/js/src/wrappers/mastra.ts @@ -10,9 +10,9 @@ * Two integration paths: * - **Manual**: `new Mastra({ observability: new Observability({ configs: { * default: { exporters: [new BraintrustObservabilityExporter()] } } }) })` - * - **Auto** (under `node --import braintrust/hook.mjs`): the loader patches - * `@mastra/core`'s `dist/mastra/index.{js,cjs}` to wrap `Mastra` so it - * calls `defaultInstance.registerExporter(exporter)` after construction. + * - **Auto**: declarative module export patches wrap Mastra's public + * constructors in tracing channels. `MastraPlugin` subscribes when the + * Braintrust SDK is imported and adds this exporter to Observability. * * Minimum supported Mastra version: 1.20.0 (when `Mastra.prototype.register` * `Exporter` and `ObservabilityInstance.registerExporter` were added). The @@ -368,10 +368,10 @@ function logExporterError(err: unknown): void { * * To capture Mastra spans in Braintrust, do one of: * - * - **Auto-instrumentation**: run your app with - * `node --import braintrust/hook.mjs`. The loader installs - * `BraintrustObservabilityExporter` into every `new Mastra(...)` - * automatically. + * - **Auto-instrumentation**: import Braintrust and run your app with + * `node --import braintrust/hook.mjs`. Declarative constructor patches emit + * tracing-channel events that `MastraPlugin` uses to install + * `BraintrustObservabilityExporter` into Mastra Observability configs. * - **Manual wiring**: pass the exporter yourself: * * ```ts diff --git a/js/tests/auto-instrumentations/fixtures/.gitignore b/js/tests/auto-instrumentations/fixtures/.gitignore index b09bbb816..51d008aa1 100644 --- a/js/tests/auto-instrumentations/fixtures/.gitignore +++ b/js/tests/auto-instrumentations/fixtures/.gitignore @@ -2,5 +2,10 @@ # This negates the global node_modules ignore !/node_modules +# Mastra's mock package uses dist/ entrypoints, which are otherwise ignored by +# js/.gitignore. +!/node_modules/@mastra/**/dist/ +!/node_modules/@mastra/**/dist/** + # Exclude dynamically generated test files /test-files diff --git a/js/tests/auto-instrumentations/fixtures/import-hook-query-mode.mjs b/js/tests/auto-instrumentations/fixtures/import-hook-query-mode.mjs new file mode 100644 index 000000000..3ad020cf0 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/import-hook-query-mode.mjs @@ -0,0 +1,15 @@ +import { parentPort } from "node:worker_threads"; + +const imported = await import(process.env.BRAINTRUST_QUERY_HOOK_URL); +const state = globalThis[Symbol.for("braintrust.applyAutoInstrumentation")]; + +parentPort?.postMessage({ + result: { + applied: state?.applied === true, + hasInitialize: typeof imported.initialize === "function", + hasLoad: typeof imported.load === "function", + hasRegister: typeof imported.register === "function", + hasResolve: typeof imported.resolve === "function", + }, + type: "hook-query-result", +}); diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/index.cjs b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/index.cjs new file mode 100644 index 000000000..065e586cd --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/index.cjs @@ -0,0 +1,3 @@ +const { Mastra } = require("./mastra.cjs"); + +module.exports = { Mastra }; diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/index.js b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/index.js new file mode 100644 index 000000000..0e52385ad --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/index.js @@ -0,0 +1 @@ +export { Mastra } from "./mastra.js"; diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra.cjs b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra.cjs new file mode 100644 index 000000000..2cb2090d1 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra.cjs @@ -0,0 +1,8 @@ +class Mastra { + constructor(config = {}) { + this.config = config; + this.observability = config.observability; + } +} + +module.exports = { Mastra }; diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra.js b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra.js new file mode 100644 index 000000000..294de01a6 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra.js @@ -0,0 +1,6 @@ +export class Mastra { + constructor(config = {}) { + this.config = config; + this.observability = config.observability; + } +} diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra/index.cjs b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra/index.cjs new file mode 100644 index 000000000..023ffc0a1 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra/index.cjs @@ -0,0 +1,3 @@ +const { Mastra } = require("../mastra.cjs"); + +module.exports = { Mastra }; diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra/index.js b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra/index.js new file mode 100644 index 000000000..3323f7941 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/dist/mastra/index.js @@ -0,0 +1 @@ +export { Mastra } from "../mastra.js"; diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/package.json b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/package.json new file mode 100644 index 000000000..84b0d6964 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/core/package.json @@ -0,0 +1,15 @@ +{ + "name": "@mastra/core", + "version": "1.26.0", + "type": "module", + "exports": { + ".": { + "import": "./dist/index.js", + "require": "./dist/index.cjs" + }, + "./mastra": { + "import": "./dist/mastra/index.js", + "require": "./dist/mastra/index.cjs" + } + } +} diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/dist/index.cjs b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/dist/index.cjs new file mode 100644 index 000000000..ba90f04e1 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/dist/index.cjs @@ -0,0 +1,7 @@ +class Observability { + constructor(config) { + this.config = config; + } +} + +module.exports = { Observability }; diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/dist/index.js b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/dist/index.js new file mode 100644 index 000000000..cd799c455 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/dist/index.js @@ -0,0 +1,5 @@ +export class Observability { + constructor(config) { + this.config = config; + } +} diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json new file mode 100644 index 000000000..ed647752e --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json @@ -0,0 +1,11 @@ +{ + "name": "@mastra/observability", + "version": "1.26.0", + "type": "module", + "exports": { + ".": { + "import": "./dist/index.js", + "require": "./dist/index.cjs" + } + } +} diff --git a/js/tests/auto-instrumentations/fixtures/runtime-apply-auto-mastra-top-level-esm.mjs b/js/tests/auto-instrumentations/fixtures/runtime-apply-auto-mastra-top-level-esm.mjs new file mode 100644 index 000000000..94c9ab6c0 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/runtime-apply-auto-mastra-top-level-esm.mjs @@ -0,0 +1,3 @@ +import "braintrust/apply-auto-instrumentation"; + +await import("./test-mastra-top-level-esm.mjs"); diff --git a/js/tests/auto-instrumentations/fixtures/test-mastra-bundler.js b/js/tests/auto-instrumentations/fixtures/test-mastra-bundler.js new file mode 100644 index 000000000..26a50077c --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/test-mastra-bundler.js @@ -0,0 +1,3 @@ +import { Mastra } from "@mastra/core"; + +new Mastra({}); diff --git a/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-cjs.cjs b/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-cjs.cjs new file mode 100644 index 000000000..a80e303b8 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-cjs.cjs @@ -0,0 +1,41 @@ +const { parentPort } = require("node:worker_threads"); +require("braintrust"); +const { Mastra } = require("@mastra/core"); +const { Mastra: SubpathMastra } = require("@mastra/core/mastra"); +const { Observability } = require("@mastra/observability"); + +const root = new Mastra({}); +const subpath = new SubpathMastra({}); +const userObservability = new Observability({ + custom: "kept", + configs: { + default: { + exporters: [{ name: "other" }], + serviceName: "user-service", + }, + }, +}); +const withUserObservability = new Mastra({ + observability: userObservability, +}); + +parentPort?.postMessage({ + result: { + root: summarize(root), + subpath: summarize(subpath), + userConfig: userObservability.config, + userObservabilityPreserved: + withUserObservability.observability === userObservability, + }, + type: "mastra-result", +}); + +function summarize(mastra) { + return { + exporters: + mastra.observability?.config?.configs?.default?.exporters?.map( + (exporter) => exporter.name, + ) ?? [], + hasObservability: Boolean(mastra.observability), + }; +} diff --git a/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-esm.mjs b/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-esm.mjs new file mode 100644 index 000000000..4fd7c5f78 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/test-mastra-top-level-esm.mjs @@ -0,0 +1,45 @@ +import { parentPort } from "node:worker_threads"; + +if (process.env.BRAINTRUST_TEST_ENABLE_MASTRA_PLUGIN !== "false") { + await import("braintrust"); +} + +const { Mastra } = await import("@mastra/core"); +const { Mastra: SubpathMastra } = await import("@mastra/core/mastra"); +const { Observability } = await import("@mastra/observability"); + +const root = new Mastra({}); +const subpath = new SubpathMastra({}); +const userObservability = new Observability({ + custom: "kept", + configs: { + default: { + exporters: [{ name: "other" }], + serviceName: "user-service", + }, + }, +}); +const withUserObservability = new Mastra({ + observability: userObservability, +}); + +parentPort?.postMessage({ + result: { + root: summarize(root), + subpath: summarize(subpath), + userConfig: userObservability.config, + userObservabilityPreserved: + withUserObservability.observability === userObservability, + }, + type: "mastra-result", +}); + +function summarize(mastra) { + return { + exporters: + mastra.observability?.config?.configs?.default?.exporters?.map( + (exporter) => exporter.name, + ) ?? [], + hasObservability: Boolean(mastra.observability), + }; +} diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/ritm-app.cjs b/js/tests/auto-instrumentations/fixtures/vendor-hooks/ritm-app.cjs index 7ee8d6d87..4b9b6652a 100644 --- a/js/tests/auto-instrumentations/fixtures/vendor-hooks/ritm-app.cjs +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/ritm-app.cjs @@ -1,7 +1,7 @@ const assert = require("node:assert"); const { - Hook, -} = require("../../../../src/auto-instrumentations/require-in-the-middle/index.ts"); + default: Hook, +} = require("../../../../src/auto-instrumentations/loader/module-hooks/ritm.ts"); let calls = 0; const hook = new Hook(["ritm-target"], (exports, name) => { diff --git a/js/tests/auto-instrumentations/fixtures/vendor-hooks/ritm-coexist-app.cjs b/js/tests/auto-instrumentations/fixtures/vendor-hooks/ritm-coexist-app.cjs index 2c30a0d36..6afd7cfd8 100644 --- a/js/tests/auto-instrumentations/fixtures/vendor-hooks/ritm-coexist-app.cjs +++ b/js/tests/auto-instrumentations/fixtures/vendor-hooks/ritm-coexist-app.cjs @@ -1,8 +1,8 @@ const assert = require("node:assert"); const Module = require("node:module"); const { - Hook, -} = require("../../../../src/auto-instrumentations/require-in-the-middle/index.ts"); + default: Hook, +} = require("../../../../src/auto-instrumentations/loader/module-hooks/ritm.ts"); const originalRequire = Module.prototype.require; diff --git a/js/tests/auto-instrumentations/import-require-in-the-middle.test.ts b/js/tests/auto-instrumentations/import-require-in-the-middle.test.ts index 35df2cdc1..955549bc8 100644 --- a/js/tests/auto-instrumentations/import-require-in-the-middle.test.ts +++ b/js/tests/auto-instrumentations/import-require-in-the-middle.test.ts @@ -7,7 +7,7 @@ const __dirname = path.dirname(fileURLToPath(import.meta.url)); const fixturesDir = path.join(__dirname, "fixtures", "vendor-hooks"); const iitmSrc = "../../src/auto-instrumentations/import-in-the-middle/"; -describe("vendored import-in-the-middle and require-in-the-middle", () => { +describe("import-in-the-middle and require-in-the-middle runtimes", () => { it("only wraps explicitly hooked ESM imports through the async loader", async () => { await runNode({ args: [ diff --git a/js/tests/auto-instrumentations/loader-hook.test.ts b/js/tests/auto-instrumentations/loader-hook.test.ts index 51a20be41..27dbe74dd 100644 --- a/js/tests/auto-instrumentations/loader-hook.test.ts +++ b/js/tests/auto-instrumentations/loader-hook.test.ts @@ -20,6 +20,10 @@ const helperPromisePath = path.join( fixturesDir, "test-api-promise-preservation.mjs", ); +const importHookQueryModePath = path.join( + fixturesDir, + "import-hook-query-mode.mjs", +); const runtimeApplyAutoSideEffectEsmPath = path.join( fixturesDir, "runtime-apply-auto-side-effect-esm.mjs", @@ -28,11 +32,35 @@ const runtimeApplyAutoSideEffectCjsPath = path.join( fixturesDir, "runtime-apply-auto-side-effect-cjs.cjs", ); +const mastraTopLevelEsmPath = path.join( + fixturesDir, + "test-mastra-top-level-esm.mjs", +); +const mastraTopLevelCjsPath = path.join( + fixturesDir, + "test-mastra-top-level-cjs.cjs", +); +const runtimeApplyAutoMastraTopLevelEsmPath = path.join( + fixturesDir, + "runtime-apply-auto-mastra-top-level-esm.mjs", +); interface TestResult { events: { start: any[]; end: any[]; error: any[] }; } +interface MastraResult { + root: { exporters: string[]; hasObservability: boolean }; + subpath: { exporters: string[]; hasObservability: boolean }; + userConfig: { + custom: string; + configs: { + default: { exporters: { name: string }[]; serviceName: string }; + }; + }; + userObservabilityPreserved: boolean; +} + describe("Unified Loader Hook Integration Tests", () => { beforeAll(() => { // No setup needed - test/fixtures/node_modules/openai is committed to the repo @@ -83,6 +111,93 @@ describe("Unified Loader Hook Integration Tests", () => { expect(result.withResponseOk).toBe(true); expect(result.constructorName).toBe("HelperPromise"); }); + + it("should expose import hook exports in query mode without bootstrapping", async () => { + const result = await runWithWorkerMessage<{ + applied: boolean; + hasInitialize: boolean; + hasLoad: boolean; + hasRegister: boolean; + hasResolve: boolean; + }>({ + env: { + BRAINTRUST_QUERY_HOOK_URL: `${pathToFileURL(hookPath).href}?braintrust-iitm-loader=true`, + }, + execArgv: [], + messageType: "hook-query-result", + script: importHookQueryModePath, + }); + + expect(result).toEqual({ + applied: false, + hasInitialize: true, + hasLoad: true, + hasRegister: true, + hasResolve: true, + }); + }); + + it("should patch Mastra ESM exports", async () => { + const result = await runWithWorkerMessage({ + execArgv: ["--import", hookPath], + messageType: "mastra-result", + script: mastraTopLevelEsmPath, + }); + + expectMastraEnabled(result); + }); + + it("should leave Mastra constructor channels passive until Braintrust enables its plugins", async () => { + const result = await runWithWorkerMessage({ + env: { BRAINTRUST_TEST_ENABLE_MASTRA_PLUGIN: "false" }, + execArgv: ["--import", hookPath], + messageType: "mastra-result", + script: mastraTopLevelEsmPath, + }); + + expect(result.root).toEqual({ exporters: [], hasObservability: false }); + expect(result.subpath).toEqual({ + exporters: [], + hasObservability: false, + }); + expect(result.userObservabilityPreserved).toBe(true); + expect( + result.userConfig.configs.default.exporters.map( + (exporter) => exporter.name, + ), + ).toEqual(["other"]); + }); + + it("should patch Mastra CJS exports", async () => { + const result = await runWithWorkerMessage({ + execArgv: ["--import", hookPath], + messageType: "mastra-result", + script: mastraTopLevelCjsPath, + }); + + expectMastraEnabled(result); + }); + + it("should respect Mastra disable config for module export hooks", async () => { + const result = await runWithWorkerMessage({ + env: { BRAINTRUST_DISABLE_INSTRUMENTATION: "mastra" }, + execArgv: ["--import", hookPath], + messageType: "mastra-result", + script: mastraTopLevelEsmPath, + }); + + expect(result.root).toEqual({ exporters: [], hasObservability: false }); + expect(result.subpath).toEqual({ + exporters: [], + hasObservability: false, + }); + expect(result.userObservabilityPreserved).toBe(true); + expect( + result.userConfig.configs.default.exporters.map( + (exporter) => exporter.name, + ), + ).toEqual(["other"]); + }); }); describe("apply-auto-instrumentation side-effect runtime setup", () => { @@ -126,9 +241,38 @@ describe("Unified Loader Hook Integration Tests", () => { expect(result.events.start.length).toBe(1); expect(result.events.end.length).toBe(1); }); + + it("should apply Mastra patches through the side-effect export", async () => { + const result = await runWithWorkerMessage({ + execArgv: [], + messageType: "mastra-result", + script: runtimeApplyAutoMastraTopLevelEsmPath, + }); + + expectMastraEnabled(result); + }); }); }); +function expectMastraEnabled(result: MastraResult): void { + expect(result.root).toEqual({ + exporters: ["braintrust"], + hasObservability: true, + }); + expect(result.subpath).toEqual({ + exporters: ["braintrust"], + hasObservability: true, + }); + expect(result.userObservabilityPreserved).toBe(true); + expect(result.userConfig.custom).toBe("kept"); + expect(result.userConfig.configs.default.serviceName).toBe("user-service"); + expect( + result.userConfig.configs.default.exporters.map( + (exporter) => exporter.name, + ), + ).toEqual(["other", "braintrust"]); +} + async function runWithWorker(options: { env?: NodeJS.ProcessEnv; execArgv: string[]; diff --git a/js/tests/auto-instrumentations/transformation.test.ts b/js/tests/auto-instrumentations/transformation.test.ts index fc402096b..1c9640fa8 100644 --- a/js/tests/auto-instrumentations/transformation.test.ts +++ b/js/tests/auto-instrumentations/transformation.test.ts @@ -23,6 +23,10 @@ const __dirname = path.dirname(fileURLToPath(import.meta.url)); const fixturesDir = path.join(__dirname, "fixtures"); const outputDir = path.join(__dirname, "output-transformation"); const nodeModulesDir = path.join(fixturesDir, "node_modules"); +const mastraBundlerEntryPoint = path.join( + fixturesDir, + "test-mastra-bundler.js", +); function testConfig( functionQuery: InstrumentationConfig["functionQuery"], @@ -342,6 +346,105 @@ describe("Orchestrion Transformation Tests", () => { // Should NOT import from external diagnostics_channel expect(output).not.toMatch(/from\s+["']diagnostics_channel["']/); }); + + it("should apply node module export hooks", async () => { + const { braintrustEsbuildPlugin } = + await import("../../src/auto-instrumentations/bundler/esbuild.js"); + + const outfile = path.join(outputDir, "esbuild-mastra-bundle.js"); + + const result = await esbuild.build({ + absWorkingDir: fixturesDir, + bundle: true, + entryPoints: [mastraBundlerEntryPoint], + format: "esm", + logLevel: "error", + outfile, + platform: "node", + plugins: [braintrustEsbuildPlugin()], + preserveSymlinks: true, + write: true, + }); + + expect(result.errors).toHaveLength(0); + expect(fs.existsSync(outfile)).toBe(true); + + const output = fs.readFileSync(outfile, "utf-8"); + + expect(output).toContain("__braintrustTopLevelImportHookRunner"); + expect(output).toContain("__braintrustExports"); + expect(output.replace(/\\/g, "/")).not.toContain( + JSON.stringify( + path + .join(fixturesDir, "node_modules", "@mastra", "core") + .replace(/\\/g, "/"), + ), + ); + }); + + it("should skip node-only module export wrappers for browser bundles", async () => { + const { braintrustEsbuildPlugin } = + await import("../../src/auto-instrumentations/bundler/esbuild.js"); + + const outfile = path.join( + outputDir, + "esbuild-module-export-wrapper-browser.js", + ); + + const result = await esbuild.build({ + absWorkingDir: fixturesDir, + bundle: true, + entryPoints: [mastraBundlerEntryPoint], + format: "esm", + logLevel: "error", + outfile, + platform: "browser", + plugins: [ + braintrustEsbuildPlugin({ useDiagnosticChannelCompatShim: true }), + ], + preserveSymlinks: true, + write: true, + }); + + expect(result.errors).toHaveLength(0); + expect(fs.existsSync(outfile)).toBe(true); + + const output = fs.readFileSync(outfile, "utf-8"); + + expect(output).not.toContain("__braintrustTopLevelImportHookRunner"); + expect(output).not.toContain("__braintrustOriginal"); + }); + + it("should keep legacy default bundles browser-safe for module export wrappers", async () => { + const { esbuildPlugin } = + await import("../../src/auto-instrumentations/bundler/esbuild.js"); + + const outfile = path.join( + outputDir, + "esbuild-module-export-wrapper-legacy-browser.js", + ); + + const result = await esbuild.build({ + absWorkingDir: fixturesDir, + bundle: true, + entryPoints: [mastraBundlerEntryPoint], + format: "esm", + logLevel: "error", + outfile, + platform: "browser", + plugins: [esbuildPlugin({})], + preserveSymlinks: true, + write: true, + }); + + expect(result.errors).toHaveLength(0); + expect(fs.existsSync(outfile)).toBe(true); + + const output = fs.readFileSync(outfile, "utf-8"); + + expect(output).not.toContain("__braintrustTopLevelImportHookRunner"); + expect(output).not.toContain("__braintrustOriginal"); + }); }); describe("vite", () => { @@ -602,6 +705,37 @@ describe("Orchestrion Transformation Tests", () => { expect(output).toContain("TracingChannel"); expect(output).not.toMatch(/require\(["']diagnostics_channel["']\)/); }); + + it("should apply module export wrappers (turbopack loader-only mode)", async () => { + const { errors, output } = await runWebpackWithLoader({ + entry: mastraBundlerEntryPoint, + output: { + path: outputDir, + filename: "turbopack-module-export-wrapper-bundle.js", + library: { type: "module" }, + }, + experiments: { outputModule: true }, + mode: "development", + resolve: { modules: [nodeModulesDir, "node_modules"] }, + externals: { "node:module": "module node:module" }, + module: { + rules: [ + { + use: [ + { + loader: webpackLoaderPath, + options: { browser: false }, + }, + ], + }, + ], + }, + }); + + expect(errors).toHaveLength(0); + expect(output).toContain("__braintrustTopLevelImportHookRunner"); + expect(output).toContain("__braintrustExports"); + }); }); describe("rollup", () => { diff --git a/js/tsconfig.vendor-hooks.json b/js/tsconfig.vendor-hooks.json index 79c106eea..d496a3704 100644 --- a/js/tsconfig.vendor-hooks.json +++ b/js/tsconfig.vendor-hooks.json @@ -12,7 +12,7 @@ "src/auto-instrumentations/types/module-details-from-path.d.ts", "src/auto-instrumentations/orchestrion-js/vendor-deps.d.ts", "src/auto-instrumentations/import-in-the-middle/**/*.mts", - "src/auto-instrumentations/require-in-the-middle/**/*.ts" + "src/auto-instrumentations/loader/module-hooks/ritm.ts" ], "exclude": ["node_modules/**", "**/dist/**"] } diff --git a/js/tsup.config.ts b/js/tsup.config.ts index a56751745..8e6fa0aa9 100644 --- a/js/tsup.config.ts +++ b/js/tsup.config.ts @@ -96,8 +96,8 @@ export default defineConfig([ { entry: [ "src/auto-instrumentations/index.ts", - "src/auto-instrumentations/loader/cjs-patch.ts", - "src/auto-instrumentations/loader/get-package-version.ts", + "src/auto-instrumentations/loader/cjs.ts", + "src/auto-instrumentations/loader/package-version.ts", "src/auto-instrumentations/bundler/vite.ts", "src/auto-instrumentations/bundler/webpack.ts", "src/auto-instrumentations/bundler/next.ts", @@ -133,7 +133,7 @@ export default defineConfig([ { entry: [ "src/auto-instrumentations/hook.mts", - "src/auto-instrumentations/loader/esm-hook.mts", + "src/auto-instrumentations/loader/esm.mts", ], format: ["esm"], outDir: "dist/auto-instrumentations", From f0e4ca56bdf2fa573402b99f371fc52716c5917c Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 13 Jul 2026 16:38:03 +0200 Subject: [PATCH 2/2] fix failures --- e2e/scenarios/mastra-instrumentation/scenario.mjs | 9 +++------ js/src/auto-instrumentations/bundler/plugin.ts | 3 +++ js/src/auto-instrumentations/configs/all.test.ts | 2 +- js/src/auto-instrumentations/configs/mastra.ts | 2 +- .../node_modules/@mastra/observability/package.json | 2 +- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/e2e/scenarios/mastra-instrumentation/scenario.mjs b/e2e/scenarios/mastra-instrumentation/scenario.mjs index 6b10057df..93854b1fe 100644 --- a/e2e/scenarios/mastra-instrumentation/scenario.mjs +++ b/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 }, diff --git a/js/src/auto-instrumentations/bundler/plugin.ts b/js/src/auto-instrumentations/bundler/plugin.ts index cb890b067..00bbf876a 100644 --- a/js/src/auto-instrumentations/bundler/plugin.ts +++ b/js/src/auto-instrumentations/bundler/plugin.ts @@ -126,6 +126,9 @@ export const unplugin = createUnplugin( } 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; diff --git a/js/src/auto-instrumentations/configs/all.test.ts b/js/src/auto-instrumentations/configs/all.test.ts index 99ec1d631..332b6f04c 100644 --- a/js/src/auto-instrumentations/configs/all.test.ts +++ b/js/src/auto-instrumentations/configs/all.test.ts @@ -17,7 +17,7 @@ describe("module export patch configs", () => { expect(configs[0].modules.map((module) => module.versionRange)).toEqual([ ">=1.20.0", ">=1.20.0", - ">=1.20.0", + ">=1.0.0", ]); expect(configs[0].modules.map((module) => module.patches)).toEqual([ [ diff --git a/js/src/auto-instrumentations/configs/mastra.ts b/js/src/auto-instrumentations/configs/mastra.ts index 587165993..908116879 100644 --- a/js/src/auto-instrumentations/configs/mastra.ts +++ b/js/src/auto-instrumentations/configs/mastra.ts @@ -45,7 +45,7 @@ export const mastraModuleExportPatchConfigs: readonly ModuleExportPatchConfig[] modulePaths: ["dist/index.js", "dist/index.cjs"], }, specifier: "@mastra/observability", - versionRange: ">=1.20.0", + versionRange: ">=1.0.0", }, ], targets: ["node"], diff --git a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json index ed647752e..ea851632b 100644 --- a/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json +++ b/js/tests/auto-instrumentations/fixtures/node_modules/@mastra/observability/package.json @@ -1,6 +1,6 @@ { "name": "@mastra/observability", - "version": "1.26.0", + "version": "1.13.0", "type": "module", "exports": { ".": {