From 9c7191af590c9ae318ff6e1df737994c8ae018ba Mon Sep 17 00:00:00 2001 From: isaacs Date: Wed, 24 Jun 2026 08:03:31 -0700 Subject: [PATCH 1/2] feat: make astQuery a first-class feature This is the cleanest way to do certain things that cannot be reached by name alone. However, it's brittle to rely on a feature, even a stable one, which is not documented or tested. So, treat it as a first-class feature. --- README.md | 72 ++++++++++++++++++++-- index.d.ts | 68 ++++++++++++++------ lib/transformer.js | 8 ++- tests/ast_query_returned_arrow_cjs/mod.js | 12 ++++ tests/ast_query_returned_arrow_cjs/test.js | 26 ++++++++ tests/tests.test.mjs | 29 +++++++++ 6 files changed, 191 insertions(+), 24 deletions(-) create mode 100644 tests/ast_query_returned_arrow_cjs/mod.js create mode 100644 tests/ast_query_returned_arrow_cjs/test.js diff --git a/README.md b/README.md index cb2e75b..cc7cb88 100644 --- a/README.md +++ b/README.md @@ -161,6 +161,52 @@ substituted return value only applies when the function returns (or resolves) normally. If no subscriber reassigns `message.result`, the original return value is preserved unchanged. +### AST Query + +The name-based `FunctionQuery` variants cover the common cases +(named functions, class/object methods, expressions). When you +need to target a node they can't express, such as an **anonymous +function returned by a factory** (a decorator factory, a +per-request handler), you can set `astQuery` to a raw +[esquery](https://github.com/estools/esquery) selector instead. + +When present, `astQuery` chooses the nodes to instrument and +takes precedence over `functionQuery`'s matching fields; +`functionQuery` then only supplies behaviour (`kind`, `index`, +`callbackIndex`, `mutableResult`) and may be omitted (it defaults +to `kind: "Sync"`). + +For example, to instrument the decorator returned by a factory: + +```js +function Injectable(options) { + return (target) => { /* applied to the decorated class */ }; +} +``` + +```js +const instrumentation = { + channelName: "injectable-apply", + module: { name: "@nestjs/common", versionRange: ">=8.0.0", filePath: "decorators/core/injectable.decorator.js" }, + // Match the arrow returned from `Injectable`. There is no name to target! + astQuery: 'FunctionDeclaration[id.name="Injectable"] ReturnStatement > ArrowFunctionExpression', + functionQuery: { kind: "Sync" }, +}; +``` + +The channel then fires each time the decorator is applied, with +the decorated target available as `message.arguments[0]`, which a +subscriber can mutate, for example to wrap prototype methods. + +An `astQuery` is used verbatim, so it can match any node, +including ones the name-based variants don't expose, such as +anonymous or deeply nested functions. (Both name-based and +`astQuery` matching work on synchronous and async functions +alike.) + +If an `astQuery` matches no nodes, the "failed to find injection +points" error includes the selector so it can be debugged. + ### API Reference ```ts @@ -206,11 +252,29 @@ type ModuleMatcher = { #### **`InstrumentationConfig`** ```ts -type InstrumentationConfig = { - channelName: string; // Name of the diagnostics channel - module: ModuleMatcher; - functionQuery: FunctionQuery; +// Behaviour-only fields, used when `astQuery` does the matching. +type FunctionBehavior = { + kind?: FunctionKind; + index?: number | null; + callbackIndex?: number; + mutableResult?: boolean; }; + +type InstrumentationConfig = + | { + channelName: string; // Name of the diagnostics channel + module: ModuleMatcher; + functionQuery: FunctionQuery; // Name-based matching + astQuery?: string; // Raw esquery selector; takes precedence over functionQuery matching + transform?: string; // Name of a custom transform registered via addTransform + } + | { + channelName: string; + module: ModuleMatcher; + astQuery: string; // Raw esquery selector chooses the node(s) + functionQuery?: FunctionBehavior; // Behaviour only; matching fields ignored + transform?: string; + }; ``` ### Functions diff --git a/index.d.ts b/index.d.ts index 380e404..1e3c8a4 100644 --- a/index.d.ts +++ b/index.d.ts @@ -36,28 +36,58 @@ export type FunctionQuery = { className: string; methodName: string; kind: Funct export type CustomTransform = (state: unknown, node: Node, parent: Node, ancestry: Node[]) => void; /** - * Configuration for injecting instrumentation code + * The behaviour-only fields of a `FunctionQuery`. Used together with `astQuery`, + * where the raw selector chooses the node and these fields drive how it is + * wrapped (the name-based matching fields are ignored). */ -export interface InstrumentationConfig { - /** - * The name of the diagnostics channel to publish to - */ - channelName: string; - /** - * The module matcher to identify the module and file to instrument - */ - module: ModuleMatcher; - /** - * The function query to identify the function to instrument - */ - functionQuery: FunctionQuery; - /** - * The name of a custom transform registered via `addTransform`. - * When set, takes precedence over `functionQuery.kind`. - */ - transform?: string; +export interface FunctionBehavior { + kind?: FunctionKind; + index?: number | null; + callbackIndex?: number; + mutableResult?: boolean; } +/** + * Configuration for injecting instrumentation code. + * + * Either `functionQuery` (name-based matching) or `astQuery` (a raw esquery + * selector) must identify the node(s) to instrument. When `astQuery` is set it + * takes precedence over `functionQuery`'s matching fields, and `functionQuery` + * becomes an optional bag of behaviour options ({@link FunctionBehavior}). + */ +export type InstrumentationConfig = + | { + /** The name of the diagnostics channel to publish to */ + channelName: string; + /** The module matcher to identify the module and file to instrument */ + module: ModuleMatcher; + /** The function query to identify the function to instrument */ + functionQuery: FunctionQuery; + /** + * A raw esquery selector that chooses the node(s) to instrument. When + * set, it takes precedence over `functionQuery`'s matching fields. + */ + astQuery?: string; + /** + * The name of a custom transform registered via `addTransform`. + * When set, takes precedence over `functionQuery.kind`. + */ + transform?: string; + } + | { + channelName: string; + module: ModuleMatcher; + /** + * A raw esquery selector that chooses the node(s) to instrument. This is + * the escape hatch for shapes the name-based `functionQuery` can't + * express, e.g. an anonymous arrow returned by a factory function. + */ + astQuery: string; + /** Behaviour options for the matched node(s); matching fields are ignored. */ + functionQuery?: FunctionBehavior; + transform?: string; + }; + /** * Describes the module and file path you would like to match */ diff --git a/lib/transformer.js b/lib/transformer.js index 76d6c4d..3f7a5c7 100644 --- a/lib/transformer.js +++ b/lib/transformer.js @@ -101,6 +101,11 @@ class Transformer { } const resolvedFunctionQuery = this.#resolveExportAlias(functionQuery, aliases) + // A raw `astQuery` selector takes precedence over the name-based + // `functionQuery`: it chooses which node(s) to instrument, while + // `functionQuery` still supplies behaviour (kind/index/callbackIndex/ + // mutableResult). This is the escape hatch for shapes the name-based + // queries can't express (e.g. an anonymous arrow returned by a factory). const query = astQuery || this.#fromFunctionQuery(resolvedFunctionQuery) const state = { ...config, @@ -119,7 +124,8 @@ class Transformer { } if (injectionCount === 0 && this.#configs.length > 0) { - const names = this.#configs.map(({ functionQuery = {} }) => { + const names = this.#configs.map(({ astQuery, functionQuery = {} }) => { + if (astQuery) return astQuery const resolvedQuery = this.#resolveExportAlias(functionQuery, aliases) const queryName = (q) => q.methodName || q.privateMethodName || q.functionName || q.expressionName || 'constructor' const originalName = queryName(functionQuery) diff --git a/tests/ast_query_returned_arrow_cjs/mod.js b/tests/ast_query_returned_arrow_cjs/mod.js new file mode 100644 index 0000000..4834abb --- /dev/null +++ b/tests/ast_query_returned_arrow_cjs/mod.js @@ -0,0 +1,12 @@ +// Mirrors a NestJS-style decorator factory: `Decorator(options)` returns the +// actual decorator `(target) => {...}` that is applied to a class. The returned +// arrow is anonymous and synchronous, so the name-based FunctionQuery variants +// cannot target it. Only a raw `astQuery` can do that. +function Decorator (options) { + return (target) => { + target.decorated = options + return target + } +} + +module.exports = { Decorator } diff --git a/tests/ast_query_returned_arrow_cjs/test.js b/tests/ast_query_returned_arrow_cjs/test.js new file mode 100644 index 0000000..2992473 --- /dev/null +++ b/tests/ast_query_returned_arrow_cjs/test.js @@ -0,0 +1,26 @@ +const { Decorator } = require('./instrumented.js') +const assert = require('node:assert') +const { tracingChannel } = require('node:diagnostics_channel') + +// The channel fires when the (anonymous, sync) returned decorator is applied. +// A subscriber can mutate the decorated target (arg 0): building block for +// instrumenting decorated classes. +let started = false +tracingChannel('orchestrion:undici:decorator_apply').subscribe({ + start (message) { + started = true + const target = message.arguments[0] + target.instrumented = true + } +}) + +class MyService {} +const decorate = Decorator({ scope: 'request' }) +const result = decorate(MyService) + +assert.strictEqual(started, true) +// Original decorator still ran. +assert.strictEqual(result, MyService) +assert.deepStrictEqual(MyService.decorated, { scope: 'request' }) +// Subscriber's mutation of the target took effect. +assert.strictEqual(MyService.instrumented, true) diff --git a/tests/tests.test.mjs b/tests/tests.test.mjs index 27b8c95..41da70e 100644 --- a/tests/tests.test.mjs +++ b/tests/tests.test.mjs @@ -433,6 +433,35 @@ describe('mutable_result_async_cjs', () => { }) }) +describe('ast_query_returned_arrow_cjs', () => { + test('instruments an anonymous arrow returned by a factory via astQuery', () => { + runTest('ast_query_returned_arrow_cjs', [ + { + channelName: 'decorator_apply', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + astQuery: 'FunctionDeclaration[id.name="Decorator"] ReturnStatement > ArrowFunctionExpression', + functionQuery: { kind: 'Sync' }, + }, + ]) + }) + + test('reports the astQuery selector when no injection points are found', () => { + const instrumentor = create([ + { + channelName: 'no_match', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + astQuery: 'FunctionDeclaration[id.name="doesNotExist"]', + functionQuery: { kind: 'Sync' }, + }, + ]) + const transformer = instrumentor.getTransformer(TEST_MODULE_NAME, TEST_MODULE_VERSION, TEST_MODULE_PATH) + assert.throws( + () => transformer.transform('function fetch () { return 42 }', 'cjs'), + /doesNotExist/ + ) + }) +}) + describe('polyfill_cjs', () => { test('instruments with a custom dc module (cjs)', () => { runTest('polyfill_cjs', [ From 215af67601f35dfc6767bec052bd74cc1571bc32 Mon Sep 17 00:00:00 2001 From: isaacs Date: Wed, 24 Jun 2026 11:01:15 -0700 Subject: [PATCH 2/2] doc: document addTransform method --- README.md | 61 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/README.md b/README.md index cc7cb88..1ae493f 100644 --- a/README.md +++ b/README.md @@ -207,6 +207,54 @@ alike.) If an `astQuery` matches no nodes, the "failed to find injection points" error includes the selector so it can be debugged. +### Custom Transforms + +The built-in operators (`traceSync`, `tracePromise`, +`traceCallback`, `traceAuto`, selected by `functionQuery.kind`) +cover diagnostics-channel tracing. When you need to rewrite a +matched node differently, register a custom operator with +`matcher.addTransform(name, fn)` and reference it from a config's +`transform` field. When `transform` is set it takes precedence +over `functionQuery.kind`. + +A custom transform receives the same arguments an esquery +traversal yields, plus the merged instrumentation state, and +mutates the matched node in place: + +```ts +type CustomTransform = (state, node, parent, ancestry) => void; +``` + +- `state` - the config merged with runtime fields + (`moduleVersion`, `moduleType`, `operator`, the resolved + `functionQuery`, …). +- `node` - the matched AST node (the function/expression selected + by `functionQuery` or `astQuery`). +- `parent`, `ancestry` - the parent node and full ancestor chain + (root last); `ancestry[ancestry.length - 1]` is the `Program`. + +A custom transform fully owns how `node` is rewritten. It is not +required to inject a tracing channel at all. + +```js +const matcher = codeTransformer.create([ + { + channelName: "fetch", + module: { name: "my-module", versionRange: ">=1.0.0", filePath: "index.js" }, + functionQuery: { functionName: "fetch", kind: "Sync" }, + transform: "myCustomTransform", + }, +]); + +matcher.addTransform("myCustomTransform", (state, node) => { + // Mutate the matched node's AST however you need. + node.body.body.unshift(/* an ESTree statement */); +}); +``` + +(The CLI accepts the same transforms via the `customTransforms` +field of a configuration module. See: [CLI Tool](#cli-tool).) + ### API Reference ```ts @@ -303,6 +351,19 @@ matching instrumentation configurations. - `version` - Version of the module. - `filePath` - Relative Unix-style path to the file from the module root (e.g. `"lib/index.js"`). Windows-style backslash paths are also accepted and will be normalized automatically. +```ts +addTransform(name: string, fn: CustomTransform): void; +``` + +Registers a custom transform operator under `name`, referenced +from a config's `transform` field. See [Custom +Transforms](#custom-transforms). + +- `name` - Operator name (also used as the config's `transform` + value). +- `fn` - `(state, node, parent, ancestry) => void`; mutates the + matched node. + #### **`Transformer`** ```ts