fix: make internally-invoked built-in transforms overridable

README.md

@@ -252,6 +252,18 @@ matcher.addTransform("myCustomTransform", (state, node) => {
 });
 ```
 
+Registering a name matching a built-in transform (`traceSync`,
+`tracePromise`, `traceCallback`, `traceAuto`,
+`tracingChannelImport`, `tracingChannelDeclaration`) overrides it
+everywhere it is dispatched — including when one built-in invokes
+another internally (e.g. the trace operators calling
+`tracingChannelDeclaration`, which in turn calls
+`tracingChannelImport`). `state.transforms` holds the merged map
+(built-ins plus registered overrides) used for this dispatch; to
+delegate to the original built-in from an override, call it via
+the `@apm-js-collab/code-transformer/lib/transforms` module
+directly.
+
 (The CLI accepts the same transforms via the `customTransforms`
 field of a configuration module. See: [CLI Tool](#cli-tool).)
 

index.d.ts

@@ -29,11 +29,32 @@ export type FunctionKind = "Sync" | "Async" | "Callback" | "Auto";
  */
 export type FunctionQuery = { className: string; methodName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean } | { className: string; privateMethodName: string; kind: FunctionKind; index?: number | null } | { className: string; index?: number | null; isExportAlias?: boolean } | { methodName: string; kind: FunctionKind; index?: number | null } | { functionName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean } | { expressionName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean };
 
+/**
+ * The merged instrumentation state passed to a transform function: the fields
+ * of the matched {@link InstrumentationConfig} (with `functionQuery`'s export
+ * aliases resolved to local names) plus runtime fields added by the
+ * transformer for the current file.
+ */
+export type KnownState = InstrumentationConfig & {
+    /** The diagnostics_channel module specifier injected into instrumented code */
+    dcModule: string;
+    /** Whether the file being transformed is ESM or CJS */
+    moduleType: ModuleType;
+    /** The version of the module being instrumented */
+    moduleVersion: string;
+    /** The resolved operator name: a built-in (e.g. `'traceSync'`) or a custom transform name */
+    operator: string;
+    /** The merged transform map (built-ins plus `addTransform` overrides) used for dispatch */
+    transforms: Record<string, CustomTransform>;
+    /** Counter of function nodes matched so far, used for `index`-based selection */
+    functionIndex?: number;
+};
+
 /**
  * A custom transform function registered via `addTransform`.
  * Receives the instrumentation state and the matched AST node.
  */
-export type CustomTransform = (state: unknown, node: Node, parent: Node, ancestry: Node[]) => void;
+export type CustomTransform<ExtraState = Record<string, unknown>> = (state: KnownState & ExtraState, node: Node, parent: Node, ancestry: Node[]) => void;
 
 /**
  * The behaviour-only fields of a `FunctionQuery`. Used together with `astQuery`,

lib/matcher.js

@@ -31,8 +31,11 @@ class InstrumentationMatcher {
   /**
    * Registers a custom transform function under the given operator name.
    *
-   * Custom transforms override built-in ones when an instrumentation config
-   * specifies the same `transform` value.
+   * Custom transforms override built-in ones of the same name everywhere they
+   * are dispatched: when an instrumentation config references the name via
+   * `transform`, when it is selected via `functionQuery.kind`, and when one
+   * transform invokes another internally (e.g. `tracingChannelDeclaration`
+   * calling `tracingChannelImport`).
    *
    * @param {string} name - Operator name (e.g. `'traceSync'`).
    * @param {Function} fn - Transform function `(state, node, parent, ancestry) => void`.

lib/transformer.js

@@ -78,6 +78,8 @@ class Transformer {
     let aliases = {}
     let injectionCount = 0
 
+    const mergedTransforms = { ...transforms, ...this.#customTransforms }
+
     for (const config of this.#configs) {
       const { astQuery, functionQuery = {} } = config
 
@@ -112,7 +114,8 @@ class Transformer {
         dcModule: this.#dcModule,
         moduleType,
         moduleVersion: this.#version,
-        functionQuery: resolvedFunctionQuery
+        functionQuery: resolvedFunctionQuery,
+        transforms: mergedTransforms
       }
 
       state.operator = this.#getOperator(state)
@@ -170,7 +173,7 @@ class Transformer {
    * @param {...unknown} args - `(node, parent, ancestry)` from esquery traverse.
    */
   #visit (state, ...args) {
-    const transform = this.#customTransforms[state.operator] ?? transforms[state.operator]
+    const transform = state.transforms[state.operator]
     const { index = 0 } = state.functionQuery
     const [node] = args
     const type = node.init?.type || node.type

lib/transforms.js

@@ -22,7 +22,7 @@ const CHANNEL_REGEX = /[^\w]/g
  */
 const formatChannelVariable = (channelName) => `tr_ch_apm$${channelName.replace(CHANNEL_REGEX, '_')}`
 
-const transforms = module.exports = {
+module.exports = {
   /**
    * Injects a `tracingChannel` import/require into the program body if one is not
    * already present.
@@ -58,7 +58,7 @@ const transforms = module.exports = {
    * Injects a `tracingChannel(...)` variable declaration for the config's channel
    * into the program body, also ensuring the import is present.
    *
-   * @param {{ channelName: string, module: { name: string }, dcModule: string, sourceType: 'module'|'script' }} state
+   * @param {{ channelName: string, module: { name: string }, dcModule: string, sourceType: 'module'|'script', transforms: Record<string, Function> }} state
    * @param {import('estree').Program} node - The program root node.
    */
   tracingChannelDeclaration (state, node) {
@@ -67,7 +67,7 @@ const transforms = module.exports = {
 
     if (node.body.some(child => child.declarations?.[0]?.id?.name === channelVariable)) return
 
-    transforms.tracingChannelImport(state, node)
+    state.transforms.tracingChannelImport(state, node)
 
     const index = node.body.findIndex(tracingChannelPredicate)
     const code = `
@@ -135,7 +135,7 @@ function traceAny (state, node, _parent, ancestry) {
  * @param {import('estree').Program} program
  */
 function traceFunction (state, node, program) {
-  transforms.tracingChannelDeclaration(state, program)
+  state.transforms.tracingChannelDeclaration(state, program)
 
   const { functionQuery: { methodName, privateMethodName, functionName, expressionName, propertyName } } = state
   const isConstructor = methodName === 'constructor' ||
@@ -194,7 +194,7 @@ function traceInstanceMethod (state, node, program) {
   // wrap it in the constructor instead.
   let ctor = classBody.body.find(({ kind }) => kind === 'constructor')
 
-  transforms.tracingChannelDeclaration(state, program)
+  state.transforms.tracingChannelDeclaration(state, program)
 
   if (!ctor) {
     ctor = parse(

tests/custom_transform_override_cjs/mod.js

@@ -0,0 +1,5 @@
+function fetch (url) {
+  return 42
+}
+
+module.exports = { fetch }

tests/custom_transform_override_cjs/test.js

@@ -0,0 +1,5 @@
+const { fetch } = require('./instrumented.js')
+const assert = require('node:assert')
+
+assert.strictEqual(global.__importOverridden, true)
+assert.strictEqual(fetch('https://example.com'), 42)

tests/tests.test.mjs

@@ -1,4 +1,5 @@
 import { create } from '../lib/index.js'
+import builtinTransforms from '../lib/transforms.js'
 import { describe, test } from 'node:test'
 import assert from 'node:assert'
 import { readFileSync, writeFileSync, rmSync, existsSync } from 'node:fs'
@@ -549,6 +550,41 @@ describe('custom_transform_cjs', () => {
   })
 })
 
+describe('custom_transform_override_cjs', () => {
+  test('overrides a built-in transform invoked internally by another transform', () => {
+    runTest('custom_transform_override_cjs', [
+      {
+        channelName: 'fetch_override',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        functionQuery: { functionName: 'fetch', kind: 'Sync' },
+      },
+    ], {
+      customTransforms: {
+        // Not referenced by any config's `transform` field: it is only reached
+        // through the built-in traceSync -> tracingChannelDeclaration chain.
+        tracingChannelImport (state, node) {
+          builtinTransforms.tracingChannelImport(state, node)
+          node.body.unshift({
+            type: 'ExpressionStatement',
+            expression: {
+              type: 'AssignmentExpression',
+              operator: '=',
+              left: {
+                type: 'MemberExpression',
+                object: { type: 'Identifier', name: 'global' },
+                property: { type: 'Identifier', name: '__importOverridden' },
+                computed: false,
+                optional: false,
+              },
+              right: { type: 'Literal', value: true, raw: 'true' },
+            },
+          })
+        },
+      },
+    })
+  })
+})
+
 describe('buffer_input', () => {
   test('accepts a Buffer and produces the same output as a string', () => {
     const code = [