feat: allow message.result to be changed

README.md

@@ -87,6 +87,80 @@ const instrumentation = {
 
 This also works for class exports (e.g., `export { MyClass as PublicClass }`).
 
+### Mutating the Return Value
+
+A subscriber can both observe or mutate a function's return value
+via `message.result`. Reassigning `message.result` is useful when
+a function returns another function (or object) that you need to
+wrap, like a factory that returns a per-request handler.
+
+For **synchronous** functions, reassign `message.result` in the
+`end` handler:
+
+```js
+const instrumentation = {
+    channelName: "create-handler",
+    module: { name: "my-framework", versionRange: ">=1.0.0", filePath: "lib/router.js" },
+    functionQuery: { methodName: "create", kind: "Sync" },
+};
+```
+
+```js
+const { tracingChannel } = require("node:diagnostics_channel");
+
+tracingChannel("orchestrion:my-framework:create-handler").subscribe({
+    end(message) {
+        const original = message.result;
+        // Replace the returned handler with a wrapped version.
+        message.result = function wrapped(...args) {
+            // ...start a span, etc.
+            return original.apply(this, args);
+        };
+    },
+});
+```
+
+For **asynchronous** functions, reassign `message.result` in the
+`asyncEnd` handler to substitute the value the returned promise
+resolves to — for example, to wrap a function the promise resolves
+to:
+
+```js
+const instrumentation = {
+    channelName: "load-handler",
+    module: { name: "my-framework", versionRange: ">=1.0.0", filePath: "lib/router.js" },
+    functionQuery: { methodName: "load", kind: "Async" },
+};
+```
+
+```js
+tracingChannel("orchestrion:my-framework:load-handler").subscribe({
+    asyncEnd(message) {
+        const original = message.result;
+        message.result = function wrapped(...args) {
+            // ...start a span, etc.
+            return original.apply(this, args);
+        };
+    },
+});
+```
+
+> [!INFO]
+>
+> Mutating the resolved value of an async function only works
+> when it returns a native `Promise` (`value instanceof
+> Promise`). Promise subclasses and other userland thenables are
+> side-chained and returned to the caller unchanged so that their
+> subclass-specific methods (e.g. `APIPromise.withResponse()`)
+> remain accessible; their resolved value therefore **cannot** be
+> mutated, and reassigning `message.result` for them has no
+> effect.
+
+On the throw path the original error still propagates; the
+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.
+
 ### API Reference
 
 ```ts

lib/transforms.js

@@ -210,7 +210,7 @@ function traceInstanceMethod (state, node, program) {
     const __apm$${methodName} = this["${methodName}"]
     this["${methodName}"] = function () {}
     if (typeof __apm$${methodName} === 'function') {
-      Object.defineProperty(this["${methodName}"], 'length', { 
+      Object.defineProperty(this["${methodName}"], 'length', {
         value: __apm$${methodName}.length,
         configurable: true
       })
@@ -468,6 +468,13 @@ function wrapPromise (state, node, iterPatch = '') {
   const { channelName } = state
   const channelVariable = formatChannelVariable(channelName)
 
+  // A subscriber can substitute the resolved value by reassigning
+  // `message.result` in `asyncEnd` (the wrapper returns `__apm$ctx.result`),
+  // but ONLY for native `Promise` instances. For Promise subclasses and other
+  // thenables we side-chain the callbacks and return the original object so
+  // that any subclass-specific methods (e.g. `APIPromise.withResponse()`)
+  // remain accessible to the caller; that means their resolved value cannot be
+  // mutated.
   return parse(`
     function wrapper () {
       if (!tr_ch_apm_hasSubscribers(${channelVariable})) return __apm$traced();
@@ -478,7 +485,7 @@ function wrapPromise (state, node, iterPatch = '') {
           if (typeof promise?.then !== 'function') {
             __apm$ctx.result = promise;
             ${iterPatch}
-            return promise;
+            return __apm$ctx.result;
           }
           // Mirror Node.js core diagnostics_channel behaviour: for native Promise
           // instances, chain normally (safe since there is no subclass API to
@@ -493,7 +500,7 @@ function wrapPromise (state, node, iterPatch = '') {
                 ${iterPatch}
                 ${channelVariable}.asyncStart.publish(__apm$ctx);
                 ${channelVariable}.asyncEnd.publish(__apm$ctx);
-                return result;
+                return __apm$ctx.result;
               },
               err => {
                 __apm$ctx.error = err;
@@ -536,7 +543,9 @@ function wrapPromise (state, node, iterPatch = '') {
  * Builds the wrapper AST for a synchronous function.
  *
  * Uses `runStores` for context propagation and publishes `error` channel events
- * on throw. The result is stored in `__apm$ctx.result` on success.
+ * on throw. The result is stored in `__apm$ctx.result` on success, and the
+ * wrapper returns whatever `__apm$ctx.result` holds after `end` publishes, so a
+ * subscriber's `end` handler can substitute the return value via `message.result`.
  *
  * @param {{ channelName: string }} state
  * @param {import('estree').Node} node
@@ -546,16 +555,19 @@ function wrapSync (state, node, iterPatch = '') {
   const { channelName } = state
   const channelVariable = formatChannelVariable(channelName)
 
+  // The return value is whatever `__apm$ctx.result` holds *after* `end` has
+  // published, so a subscriber's `end` handler can reassign `message.result`
+  // to substitute it (e.g. to wrap a returned handler). The `return`
+  // therefore sits below the `try`/`finally`: on the throw path the `catch`
+  // re-throws, so the trailing `return` is unreachable.
   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;
+          __apm$ctx.result = __apm$traced();
           ${iterPatch}
-          return result;
         } catch (err) {
           __apm$ctx.error = err;
           ${channelVariable}.error.publish(__apm$ctx);
@@ -564,6 +576,7 @@ function wrapSync (state, node, iterPatch = '') {
          __apm$ctx.self ??= this;
           ${channelVariable}.end.publish(__apm$ctx);
         }
+        return __apm$ctx.result;
       });
     }
   `)

tests/mutable_result_async_cjs/mod.js

@@ -0,0 +1,24 @@
+// Async factory: returns a native Promise that resolves to a per-request
+// handler. A subscriber can replace the resolved handler in `asyncEnd`.
+async function load (instance, callback) {
+  return function handler (req) {
+    return callback.call(instance, req)
+  }
+}
+
+// Returns a userland thenable (NOT a native Promise) that resolves to a
+// handler. Because the returned value is not `instanceof Promise`, the wrapper
+// side-chains it and returns the original object, so a subscriber cannot mutate
+// its resolved value.
+function loadThenable (instance, callback) {
+  const handler = function handler (req) {
+    return callback.call(instance, req)
+  }
+  return {
+    then (onResolve) {
+      onResolve(handler)
+    }
+  }
+}
+
+module.exports = { load, loadThenable }

tests/mutable_result_async_cjs/test.js

@@ -0,0 +1,56 @@
+const { load, loadThenable } = require('./instrumented.js')
+const assert = require('node:assert')
+const { tracingChannel } = require('node:diagnostics_channel')
+
+const instance = { name: 'Ctrl' }
+const callback = function (req) {
+  return `${this.name}:${req}`
+}
+
+;(async () => {
+  // Scenario 1: a native Promise. A subscriber substitutes the resolved handler
+  // by reassigning `message.result` in `asyncEnd`. The caller awaits and
+  // receives the substituted wrapper.
+  const events = []
+  tracingChannel('orchestrion:undici:load_mutable').subscribe({
+    asyncEnd (message) {
+      events.push('asyncEnd')
+      const original = message.result
+      assert.strictEqual(typeof original, 'function')
+      message.result = function wrapped (...args) {
+        events.push('wrapped')
+        return `[${original.apply(this, args)}]`
+      }
+    }
+  })
+
+  const handler = await load(instance, callback)
+  assert.strictEqual(typeof handler, 'function')
+  // The handler returned to the caller is the substituted wrapper.
+  assert.strictEqual(handler('GET'), '[Ctrl:GET]')
+  assert.deepStrictEqual(events, ['asyncEnd', 'wrapped'])
+
+  // Scenario 2: a userland thenable (not a native Promise). The subscriber can
+  // observe `message.result` but reassigning it has no effect. The original
+  // resolved handler is preserved, since the wrapper returns the thenable
+  // itself rather than `__apm$ctx.result`.
+  let observed
+  tracingChannel('orchestrion:undici:load_thenable_mutable').subscribe({
+    asyncEnd (message) {
+      observed = message.result
+      message.result = function wrapped () {
+        return 'SHOULD_NOT_APPLY'
+      }
+    }
+  })
+
+  const thenableHandler = await loadThenable(instance, callback)
+  assert.strictEqual(typeof thenableHandler, 'function')
+  // Subscriber observed the resolved value...
+  assert.strictEqual(typeof observed, 'function')
+  // ...but its mutation was ignored: the caller gets the original handler.
+  assert.strictEqual(thenableHandler('GET'), 'Ctrl:GET')
+})().catch((err) => {
+  console.error(err)
+  process.exit(1)
+})

tests/mutable_result_cjs/mod.js

@@ -0,0 +1,18 @@
+function create (instance, callback) {
+  // Returns a per-request handler that closes over `instance`/`callback`,
+  // mirroring NestJS `RouterExecutionContext.create`. A subscriber needs to
+  // replace this returned function with a span-wrapping version.
+  return function handler (req) {
+    return callback.call(instance, req)
+  }
+}
+
+function compute (x) {
+  return x + 1
+}
+
+function boom () {
+  throw new Error('boom')
+}
+
+module.exports = { create, compute, boom }

tests/mutable_result_cjs/test.js

@@ -0,0 +1,54 @@
+const { create, compute, boom } = require('./instrumented.js')
+const assert = require('node:assert')
+const { tracingChannel } = require('node:diagnostics_channel')
+
+// Scenario 1: a subscriber substitutes the returned handler by reassigning
+// `message.result` in `end`. The wrapper must return the substituted function.
+const events = []
+tracingChannel('orchestrion:undici:create_mutable').subscribe({
+  start () {
+    events.push('start')
+  },
+  end (message) {
+    events.push('end')
+    const original = message.result
+    assert.strictEqual(typeof original, 'function')
+    message.result = function wrapped (...args) {
+      events.push('wrapped')
+      return `[${original.apply(this, args)}]`
+    }
+  }
+})
+
+const instance = { name: 'Ctrl' }
+const callback = function (req) {
+  return `${this.name}:${req}`
+}
+const handler = create(instance, callback)
+assert.strictEqual(typeof handler, 'function')
+// The handler returned to the caller is the substituted wrapper.
+assert.strictEqual(handler('GET'), '[Ctrl:GET]')
+assert.deepStrictEqual(events, ['start', 'end', 'wrapped'])
+
+// Scenario 2: a subscriber observes `message.result` but does NOT reassign it.
+// The original return value must be preserved unchanged.
+let observed
+tracingChannel('orchestrion:undici:compute_mutable').subscribe({
+  end (message) {
+    observed = message.result
+  }
+})
+assert.strictEqual(compute(41), 42)
+assert.strictEqual(observed, 42)
+
+// Scenario 3 (throw path): the restructured try/finally must still propagate
+// the error and publish it; the trailing `return message.result` is unreachable.
+let errored
+tracingChannel('orchestrion:undici:boom_mutable').subscribe({
+  error (message) {
+    errored = message.error
+  }
+})
+assert.throws(() => boom(), /boom/)
+assert.ok(errored instanceof Error)
+assert.strictEqual(errored.message, 'boom')

tests/tests.test.mjs

@@ -394,6 +394,45 @@ describe('ast_query_cjs', () => {
   })
 })
 
+describe('mutable_result_cjs', () => {
+  test('lets a subscriber substitute the synchronous return value via message.result', () => {
+    runTest('mutable_result_cjs', [
+      {
+        channelName: 'create_mutable',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        functionQuery: { functionName: 'create', kind: 'Sync' },
+      },
+      {
+        channelName: 'compute_mutable',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        functionQuery: { functionName: 'compute', kind: 'Sync' },
+      },
+      {
+        channelName: 'boom_mutable',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        functionQuery: { functionName: 'boom', kind: 'Sync' },
+      },
+    ])
+  })
+})
+
+describe('mutable_result_async_cjs', () => {
+  test('lets a subscriber substitute a native Promise result, but not a thenable result', () => {
+    runTest('mutable_result_async_cjs', [
+      {
+        channelName: 'load_mutable',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        functionQuery: { functionName: 'load', kind: 'Async' },
+      },
+      {
+        channelName: 'load_thenable_mutable',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        functionQuery: { functionName: 'loadThenable', kind: 'Async' },
+      },
+    ])
+  })
+})
+
 describe('polyfill_cjs', () => {
   test('instruments with a custom dc module (cjs)', () => {
     runTest('polyfill_cjs', [