braintrust 3.20.0 → 3.22.0

package/dist/apply-auto-instrumentation.mjs

@@ -28,7 +28,7 @@ import {
   smithyClientChannels,
   smithyCoreChannels,
   strandsAgentSDKChannels
-} from "./chunk-Y7W7WP6H.mjs";
+} from "./chunk-KMGUTPB7.mjs";
 
 // src/node/apply-auto-instrumentation-entry.ts
 import * as diagnostics_channel from "node:diagnostics_channel";
@@ -227,6 +227,18 @@ var aiSDKConfigs = [
       kind: "Sync"
     }
   },
+  {
+    channelName: aiSDKChannels.v7CreateTelemetryDispatcher.channelName,
+    module: {
+      name: "ai",
+      versionRange: ">=7.0.0-0 <8.0.0",
+      filePath: "dist/internal/index.js"
+    },
+    functionQuery: {
+      functionName: "createTelemetryDispatcher",
+      kind: "Sync"
+    }
+  },
   // streamObject - async function (v3 only, before the sync refactor in v4)
   {
     channelName: aiSDKChannels.streamObject.channelName,
@@ -2548,10 +2560,582 @@ function getDefaultAutoInstrumentationConfigs() {
   });
 }
 
+// src/auto-instrumentations/orchestrion-js/matcher.ts
+import semifies from "semifies";
+
+// src/auto-instrumentations/orchestrion-js/transformer.ts
+import esquery2 from "esquery";
+import { generate } from "astring";
+import { parse as parse2 } from "meriyah";
+import { SourceMapGenerator } from "source-map";
+
+// src/auto-instrumentations/orchestrion-js/transforms.ts
+import esquery from "esquery";
+import { parse } from "meriyah";
+var tracingChannelPredicate = (node) => node.declarations?.[0]?.id?.properties?.[0]?.value?.name === "tr_ch_apm_tracingChannel";
+var CHANNEL_REGEX = /[^\w]/g;
+function formatChannelVariable(channelName) {
+  return `tr_ch_apm$${channelName.replace(CHANNEL_REGEX, "_")}`;
+}
+var transforms = {
+  tracingChannelImport({ dcModule, moduleType }, node) {
+    if (node.body.some(tracingChannelPredicate)) {
+      return;
+    }
+    const options = { module: moduleType === "esm" };
+    const index = node.body.findIndex(
+      (child) => child.directive === "use strict"
+    );
+    const dc = moduleType === "esm" ? `import tr_ch_apm_dc from "${dcModule}"` : `const tr_ch_apm_dc = ${"require"}("${dcModule}")`;
+    const tracingChannel2 = "const { tracingChannel: tr_ch_apm_tracingChannel } = tr_ch_apm_dc";
+    const hasSubscribers = `const tr_ch_apm_hasSubscribers = ch => ch.start.hasSubscribers
+      || ch.end.hasSubscribers
+      || ch.asyncStart.hasSubscribers
+      || ch.asyncEnd.hasSubscribers
+      || ch.error.hasSubscribers`;
+    node.body.splice(
+      index + 1,
+      0,
+      parse(dc, options).body[0],
+      parse(tracingChannel2, options).body[0],
+      parse(hasSubscribers, options).body[0]
+    );
+  },
+  tracingChannelDeclaration(state2, node) {
+    const {
+      channelName,
+      module: { name }
+    } = state2;
+    const channelVariable = formatChannelVariable(channelName);
+    if (node.body.some(
+      (child) => child.declarations?.[0]?.id?.name === channelVariable
+    )) {
+      return;
+    }
+    transforms.tracingChannelImport(state2, node, null, []);
+    const index = node.body.findIndex(tracingChannelPredicate);
+    const code = `
+      const ${channelVariable} = tr_ch_apm_tracingChannel("orchestrion:${name}:${channelName}")
+    `;
+    node.body.splice(index + 1, 0, parse(code).body[0]);
+  },
+  traceCallback: traceAny,
+  tracePromise: traceAny,
+  traceSync: traceAny
+};
+function traceAny(state2, node, _parent, ancestry) {
+  const program = ancestry[ancestry.length - 1];
+  if (node.type === "ClassDeclaration" || node.type === "ClassExpression") {
+    traceInstanceMethod(state2, node, program);
+  } else {
+    traceFunction(state2, node, program);
+  }
+}
+function traceFunction(state2, node, program) {
+  transforms.tracingChannelDeclaration(state2, program, null, []);
+  const { functionQuery } = state2;
+  const methodName = "methodName" in functionQuery ? functionQuery.methodName : void 0;
+  const privateMethodName = "privateMethodName" in functionQuery ? functionQuery.privateMethodName : void 0;
+  const functionName = "functionName" in functionQuery ? functionQuery.functionName : void 0;
+  const isConstructor = methodName === "constructor" || !methodName && !privateMethodName && !functionName;
+  const type = isConstructor ? "ArrowFunctionExpression" : "FunctionExpression";
+  node.body = wrap(
+    state2,
+    {
+      type,
+      params: node.params,
+      body: node.body,
+      async: node.async,
+      expression: false,
+      generator: node.generator
+    },
+    program
+  );
+  node.generator = false;
+  node.async = false;
+  wrapSuper(node);
+}
+function traceInstanceMethod(state2, node, program) {
+  const { functionQuery, operator } = state2;
+  const { methodName } = functionQuery;
+  if (!methodName) {
+    return;
+  }
+  const classBody = node.body;
+  if (classBody.body.some(({ key }) => key.name === methodName)) {
+    return;
+  }
+  let ctor = classBody.body.find(({ kind }) => kind === "constructor");
+  transforms.tracingChannelDeclaration(state2, program, null, []);
+  if (!ctor) {
+    ctor = parse(
+      node.superClass ? "class A extends Object { constructor (...args) { super(...args) } }" : "class A { constructor () {} }"
+    ).body[0].body.body[0];
+    classBody.body.unshift(ctor);
+  }
+  const ctorBody = parse(`
+    const __apm$${methodName} = this["${methodName}"]
+    this["${methodName}"] = function () {}
+  `).body;
+  const fn = ctorBody[1].expression.right;
+  fn.async = operator === "tracePromise";
+  fn.body = wrap(
+    state2,
+    { type: "Identifier", name: `__apm$${methodName}` },
+    program
+  );
+  wrapSuper(fn);
+  ctor.value.body.body.push(...ctorBody);
+}
+function wrap(state2, node, program) {
+  const { operator, moduleVersion } = state2;
+  const wrapper = operator === "traceCallback" ? wrapCallback(state2) : operator === "tracePromise" ? wrapPromise(state2) : wrapSync(state2);
+  const block = wrapper.body[0].body;
+  const common = parse(
+    node.type === "ArrowFunctionExpression" ? `
+    const __apm$ctx = {
+      arguments,
+      moduleVersion: ${JSON.stringify(moduleVersion)}
+    };
+    const __apm$traced = () => {
+      const __apm$wrapped = () => {};
+      return __apm$wrapped(...arguments);
+    };
+  ` : `
+    const __apm$ctx = {
+      arguments,
+      self: this,
+      moduleVersion: ${JSON.stringify(moduleVersion)}
+    };
+    const __apm$traced = () => {
+      const __apm$wrapped = () => {};
+      return __apm$wrapped.apply(this, arguments);
+    };
+  `
+  ).body;
+  block.body.unshift(...common);
+  esquery.query(block, "[id.name=__apm$wrapped]")[0].init = node;
+  return block;
+}
+function wrapSuper(node) {
+  const members = /* @__PURE__ */ new Set();
+  esquery.traverse(
+    node.body,
+    esquery.parse("[object.type=Super]"),
+    (node2, parent) => {
+      const { name } = node2.property;
+      let child;
+      if (parent.callee) {
+        const { expression } = parse(`__apm$super['${name}'].call(this)`).body[0];
+        parent.callee = child = expression.callee;
+        parent.arguments.unshift(...expression.arguments);
+      } else {
+        parent.expression = child = parse(`__apm$super['${name}']`).body[0];
+      }
+      child.computed = parent.callee.computed;
+      child.optional = parent.callee.optional;
+      members.add(name);
+    }
+  );
+  for (const name of members) {
+    const member = parse(`
+      class Wrapper {
+        wrapper () {
+          __apm$super['${name}'] = super['${name}']
+        }
+      }
+    `).body[0].body.body[0].value.body.body[0];
+    node.body.body.unshift(member);
+  }
+  if (members.size > 0) {
+    node.body.body.unshift(parse("const __apm$super = {}").body[0]);
+  }
+}
+function wrapCallback(state2) {
+  const {
+    channelName,
+    functionQuery: { callbackIndex = -1 }
+  } = state2;
+  const channelVariable = formatChannelVariable(channelName);
+  return parse(`
+    function wrapper () {
+      const __apm$cb = Array.prototype.at.call(arguments, ${callbackIndex});
+
+      if (!${channelVariable}.start.hasSubscribers) return __apm$traced();
+
+      function __apm$wrappedCb(err, res) {
+        if (err) {
+          __apm$ctx.error = err;
+          ${channelVariable}.error.publish(__apm$ctx);
+        } else {
+          __apm$ctx.result = res;
+        }
+
+        ${channelVariable}.asyncStart.runStores(__apm$ctx, () => {
+          try {
+            if (__apm$cb) {
+              return __apm$cb.apply(this, arguments);
+            }
+          } finally {
+            ${channelVariable}.asyncEnd.publish(__apm$ctx);
+          }
+        });
+      }
+
+      if (typeof __apm$cb !== 'function') {
+        return __apm$traced();
+      }
+      Array.prototype.splice.call(arguments, ${callbackIndex}, 1, __apm$wrappedCb);
+
+      return ${channelVariable}.start.runStores(__apm$ctx, () => {
+        try {
+          return __apm$traced();
+        } catch (err) {
+          __apm$ctx.error = err;
+          ${channelVariable}.error.publish(__apm$ctx);
+          throw err;
+        } finally {
+          __apm$ctx.self ??= this;
+          ${channelVariable}.end.publish(__apm$ctx);
+        }
+      });
+    }
+  `);
+}
+function wrapPromise(state2) {
+  const { channelName } = state2;
+  const channelVariable = formatChannelVariable(channelName);
+  return parse(`
+    function wrapper () {
+      if (!tr_ch_apm_hasSubscribers(${channelVariable})) return __apm$traced();
+
+      return ${channelVariable}.start.runStores(__apm$ctx, () => {
+        try {
+          let promise = __apm$traced();
+          if (typeof promise?.then !== 'function') {
+            __apm$ctx.result = promise;
+            return promise;
+          }
+          // Mirror Node.js core diagnostics_channel behaviour: for native Promise
+          // instances, chain normally (safe since there is no subclass API to
+          // preserve). For Promise subclasses and other thenables, side-chain the
+          // callbacks for event publishing and return the original so that any
+          // subclass-specific methods (e.g. APIPromise.withResponse()) remain
+          // accessible to the caller.
+          if (promise instanceof Promise && promise.constructor === Promise) {
+            return promise.then(
+              result => {
+                __apm$ctx.result = result;
+                ${channelVariable}.asyncStart.publish(__apm$ctx);
+                ${channelVariable}.asyncEnd.publish(__apm$ctx);
+                return result;
+              },
+              err => {
+                __apm$ctx.error = err;
+                ${channelVariable}.error.publish(__apm$ctx);
+                ${channelVariable}.asyncStart.publish(__apm$ctx);
+                ${channelVariable}.asyncEnd.publish(__apm$ctx);
+                throw err;
+              }
+            );
+          }
+          promise.then(
+            result => {
+              __apm$ctx.result = result;
+              ${channelVariable}.asyncStart.publish(__apm$ctx);
+              ${channelVariable}.asyncEnd.publish(__apm$ctx);
+            },
+            err => {
+              __apm$ctx.error = err;
+              ${channelVariable}.error.publish(__apm$ctx);
+              ${channelVariable}.asyncStart.publish(__apm$ctx);
+              ${channelVariable}.asyncEnd.publish(__apm$ctx);
+            }
+          );
+          return promise;
+        } catch (err) {
+          __apm$ctx.error = err;
+          ${channelVariable}.error.publish(__apm$ctx);
+          throw err;
+        } finally {
+          __apm$ctx.self ??= this;
+          ${channelVariable}.end.publish(__apm$ctx);
+        }
+      });
+    }
+  `);
+}
+function wrapSync(state2) {
+  const { channelName } = state2;
+  const channelVariable = formatChannelVariable(channelName);
+  return parse(`
+    function wrapper () {
+      if (!tr_ch_apm_hasSubscribers(${channelVariable})) return __apm$traced();
+
+      return ${channelVariable}.start.runStores(__apm$ctx, () => {
+        try {
+          const result = __apm$traced();
+          __apm$ctx.result = result;
+          return result;
+        } catch (err) {
+          __apm$ctx.error = err;
+          ${channelVariable}.error.publish(__apm$ctx);
+          throw err;
+        } finally {
+         __apm$ctx.self ??= this;
+          ${channelVariable}.end.publish(__apm$ctx);
+        }
+      });
+    }
+  `);
+}
+
+// src/auto-instrumentations/orchestrion-js/transformer.ts
+var Transformer = class {
+  moduleName;
+  version;
+  filePath;
+  configs = [];
+  dcModule;
+  constructor(moduleName, version, filePath, configs, dcModule) {
+    this.moduleName = moduleName;
+    this.version = version;
+    this.filePath = filePath;
+    this.configs = configs;
+    this.dcModule = dcModule;
+  }
+  /**
+   * Instruments `code` by injecting diagnostics_channel tracing around the
+   * target functions defined by this transformer's configs.
+   */
+  transform(code, moduleType) {
+    if (Buffer.isBuffer(code)) {
+      code = code.toString();
+    }
+    if (!code) {
+      return { code };
+    }
+    let ast;
+    let aliases = {};
+    let injectionCount = 0;
+    for (const config of this.configs) {
+      const { astQuery, functionQuery } = config;
+      if (!ast) {
+        const options = {
+          loc: true,
+          ranges: true,
+          raw: true,
+          module: moduleType === "esm"
+        };
+        try {
+          ast = parse2(code, options);
+        } catch {
+          ast = parse2(code, { ...options, module: !options.module });
+        }
+        if (moduleType === "esm") {
+          aliases = this.collectExportAliases(ast);
+        }
+      }
+      const resolvedFunctionQuery = this.resolveExportAlias(
+        functionQuery,
+        aliases
+      );
+      const query = astQuery || this.fromFunctionQuery(resolvedFunctionQuery);
+      const state2 = {
+        ...config,
+        dcModule: this.dcModule,
+        moduleType,
+        moduleVersion: this.version,
+        functionQuery: resolvedFunctionQuery,
+        operator: this.getOperator(resolvedFunctionQuery.kind)
+      };
+      esquery2.traverse(ast, esquery2.parse(query), (...args) => {
+        injectionCount++;
+        this.visit(state2, ...args);
+      });
+    }
+    if (injectionCount === 0 && this.configs.length > 0) {
+      const names = this.configs.map(({ astQuery, functionQuery }) => {
+        const resolvedQuery = this.resolveExportAlias(functionQuery, aliases);
+        const queryName = (q) => q.methodName || q.privateMethodName || q.functionName || q.propertyName || astQuery || "unknown";
+        const originalName = queryName(functionQuery);
+        const originalAlias = functionQuery.className || functionQuery.functionName;
+        const resolvedAlias = resolvedQuery.className || resolvedQuery.functionName;
+        if (originalAlias && originalAlias !== resolvedAlias) {
+          return `${originalAlias} (local name: ${resolvedAlias})`;
+        }
+        return originalName;
+      });
+      throw new Error(
+        `Failed to find injection points for: ${JSON.stringify(names)}`
+      );
+    }
+    if (ast) {
+      const file = `${this.moduleName}/${this.filePath}`;
+      const sourceMap = new SourceMapGenerator({ file });
+      const transformedCode = generate(ast, { sourceMap });
+      const map = sourceMap.toString();
+      return { code: transformedCode, map };
+    }
+    return { code };
+  }
+  free() {
+  }
+  visit(state2, ...args) {
+    const transform = transforms[state2.operator];
+    const { index = 0 } = state2.functionQuery;
+    const [node] = args;
+    const type = node.init?.type || node.type;
+    if (type !== "ClassDeclaration" && type !== "ClassExpression") {
+      if (node.type === "VariableDeclarator") {
+        return;
+      }
+      state2.functionIndex = state2.functionIndex === void 0 ? 0 : state2.functionIndex + 1;
+      if (index !== null && index !== state2.functionIndex) {
+        return;
+      }
+    }
+    transform(state2, ...args);
+  }
+  getOperator(kind) {
+    switch (kind) {
+      case "Async":
+        return "tracePromise";
+      case "Callback":
+        return "traceCallback";
+      case "Sync":
+        return "traceSync";
+    }
+  }
+  collectExportAliases(ast) {
+    const aliases = {};
+    for (const node of ast.body) {
+      if (node.type === "ExportNamedDeclaration" && !node.source) {
+        for (const spec of node.specifiers) {
+          if (spec.exported && spec.local) {
+            const exportedName = spec.exported.name ?? spec.exported.value;
+            const localName = spec.local.name ?? spec.local.value;
+            if (exportedName && localName) {
+              aliases[exportedName] = localName;
+            }
+          }
+        }
+      }
+    }
+    return aliases;
+  }
+  resolveExportAlias(functionQuery, aliases) {
+    if (!("isExportAlias" in functionQuery) || !functionQuery.isExportAlias) {
+      return functionQuery;
+    }
+    if ("className" in functionQuery && aliases[functionQuery.className]) {
+      return {
+        ...functionQuery,
+        className: aliases[functionQuery.className]
+      };
+    }
+    if ("functionName" in functionQuery && aliases[functionQuery.functionName]) {
+      return {
+        ...functionQuery,
+        functionName: aliases[functionQuery.functionName]
+      };
+    }
+    return functionQuery;
+  }
+  functionQueryLabel(functionQuery) {
+    if ("methodName" in functionQuery) {
+      return functionQuery.methodName;
+    }
+    if ("privateMethodName" in functionQuery) {
+      return functionQuery.privateMethodName;
+    }
+    if ("functionName" in functionQuery) {
+      return functionQuery.functionName;
+    }
+    if ("propertyName" in functionQuery) {
+      return functionQuery.propertyName;
+    }
+    return "unknown";
+  }
+  fromFunctionQuery(functionQuery) {
+    const queries = [];
+    if ("className" in functionQuery) {
+      const { className } = functionQuery;
+      const methodName = this.functionQueryLabel(functionQuery);
+      const keyType = "privateMethodName" in functionQuery ? "PrivateIdentifier" : "Identifier";
+      queries.push(
+        `[id.name="${className}"]`,
+        `[id.name="${className}"] > ClassExpression`,
+        `[id.name="${className}"] > ClassBody > [key.name="${methodName}"][key.type=${keyType}] > [async]`,
+        `[id.name="${className}"] > ClassExpression > ClassBody > [key.name="${methodName}"][key.type=${keyType}] > [async]`
+      );
+    } else if ("methodName" in functionQuery) {
+      const { methodName } = functionQuery;
+      queries.push(
+        `ClassBody > [key.name="${methodName}"][key.type=Identifier] > [async]`,
+        `Property[key.name="${methodName}"][key.type=Identifier] > [async]`
+      );
+    }
+    if ("functionName" in functionQuery) {
+      const { functionName } = functionQuery;
+      queries.push(`FunctionDeclaration[id.name="${functionName}"][async]`);
+    }
+    if ("objectName" in functionQuery) {
+      const { objectName, propertyName } = functionQuery;
+      const objectSelector = objectName === "this" ? "left.object.type=ThisExpression" : `left.object.name="${objectName}"`;
+      queries.push(
+        `AssignmentExpression[${objectSelector}][left.property.name="${propertyName}"] > [async]`
+      );
+    }
+    return queries.join(", ");
+  }
+};
+
+// src/auto-instrumentations/orchestrion-js/matcher.ts
+var InstrumentationMatcher = class {
+  configs = [];
+  dcModule;
+  transformers = {};
+  constructor(configs, dcModule) {
+    this.configs = configs;
+    this.dcModule = dcModule || "diagnostics_channel";
+  }
+  /**
+   * Returns a Transformer for the given module/file/version, or undefined if no
+   * registered config matches.
+   */
+  getTransformer(moduleName, version, filePath) {
+    filePath = filePath.replace(/\\/g, "/");
+    const id = `${moduleName}/${filePath}@${version}`;
+    if (this.transformers[id]) {
+      return this.transformers[id];
+    }
+    const configs = this.configs.filter(
+      ({ module: mod }) => mod.name === moduleName && mod.filePath === filePath && semifies(version, mod.versionRange)
+    );
+    if (configs.length === 0) {
+      return void 0;
+    }
+    this.transformers[id] = new Transformer(
+      moduleName,
+      version,
+      filePath,
+      configs,
+      this.dcModule
+    );
+    return this.transformers[id];
+  }
+  free() {
+    this.transformers = {};
+  }
+};
+
+// src/auto-instrumentations/orchestrion-js/index.ts
+function create(configs, dcModule) {
+  return new InstrumentationMatcher(configs, dcModule);
+}
+
 // src/auto-instrumentations/loader/cjs-patch.ts
-import {
-  create
-} from "@apm-js-collab/code-transformer";
 import * as NodeModule from "node:module";
 import { sep } from "node:path";
 import moduleDetailsFromPath from "module-details-from-path";
@@ -2701,8 +3285,6 @@ var ModulePatch = class {
             args[0] = transformedCode?.code;
           } catch (error) {
             console.warn(`Error transforming module ${filename}:`, error);
-          } finally {
-            transformer.free();
           }
         }
       }