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 0d34514c84441a29d1fc488fc93b1bdf05663ec7 Mon Sep 17 00:00:00 2001 From: isaacs Date: Wed, 24 Jun 2026 10:40:45 -0700 Subject: [PATCH 2/2] test: verify [async] marker is structural, not a filter The 'functionQuery' generated variants include a '[async]' token in the generated selector. This was somewhat confusing, leading me to think that it was filtering for async functions, though this is not the case. The added test does nothing but verify the result of this investigation, and prevents our functionQuery selectors from regressing to only affecting async functions. At least, the presence of a test makes it less confusing. --- tests/sync_name_match_cjs/mod.js | 16 ++++++++++++++++ tests/sync_name_match_cjs/test.js | 16 ++++++++++++++++ tests/tests.test.mjs | 17 +++++++++++++++++ 3 files changed, 49 insertions(+) create mode 100644 tests/sync_name_match_cjs/mod.js create mode 100644 tests/sync_name_match_cjs/test.js diff --git a/tests/sync_name_match_cjs/mod.js b/tests/sync_name_match_cjs/mod.js new file mode 100644 index 0000000..2550bbe --- /dev/null +++ b/tests/sync_name_match_cjs/mod.js @@ -0,0 +1,16 @@ +// A synchronous function expression and a synchronous object method. +// This fixture guards that the name-based `functionQuery` variants match +// synchronous functions: the `[async]` token in the generated selectors +// is a structural "is-a-function-node" existence check (every function +// node has an `async` field), NOT an async-only filter. +const fetchExpr = function fetchExpr (x) { + return x + 1 +} + +const client = { + query (x) { + return x * 2 + } +} + +module.exports = { fetchExpr, client } diff --git a/tests/sync_name_match_cjs/test.js b/tests/sync_name_match_cjs/test.js new file mode 100644 index 0000000..c5c35e8 --- /dev/null +++ b/tests/sync_name_match_cjs/test.js @@ -0,0 +1,16 @@ +// This test guards that the name-based `functionQuery` variants match +// synchronous functions: the `[async]` token in the generated selectors +// is a structural "is-a-function-node" existence check (every function +// node has an `async` field), NOT an async-only filter. +const { fetchExpr, client } = require('./instrumented.js') +const { assert, getContext } = require('../common/preamble.js') + +// expressionName, kind: 'Sync' matched, fetchExpr is not async. +const exprCtx = getContext('orchestrion:undici:fetch_expr_sync') +assert.strictEqual(fetchExpr(1), 2) +assert.deepStrictEqual(exprCtx, { start: true, end: 2 }) + +// methodName (object method), kind: 'Sync' matched, query is not async. +const methodCtx = getContext('orchestrion:undici:client_query_sync') +assert.strictEqual(client.query(3), 6) +assert.deepStrictEqual(methodCtx, { start: true, end: 6 }) diff --git a/tests/tests.test.mjs b/tests/tests.test.mjs index 41da70e..bc25441 100644 --- a/tests/tests.test.mjs +++ b/tests/tests.test.mjs @@ -462,6 +462,23 @@ describe('ast_query_returned_arrow_cjs', () => { }) }) +describe('sync_name_match_cjs', () => { + test('name-based queries match synchronous functions (the [async] token is structural)', () => { + runTest('sync_name_match_cjs', [ + { + channelName: 'fetch_expr_sync', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + functionQuery: { expressionName: 'fetchExpr', kind: 'Sync' }, + }, + { + channelName: 'client_query_sync', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + functionQuery: { methodName: 'query', kind: 'Sync' }, + }, + ]) + }) +}) + describe('polyfill_cjs', () => { test('instruments with a custom dc module (cjs)', () => { runTest('polyfill_cjs', [