diagnostics_channel: return original thenable

doc/api/diagnostics_channel.md

@@ -948,20 +948,23 @@ added:
  - v19.9.0
  - v18.19.0
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/62407
+    description: Non-native-Promise thenables are now returned as-is,
+                 preserving their original type and methods.
   - version: v26.0.0
     pr-url: https://github.com/nodejs/node/pull/61766
-    description: Custom thenables will no longer be wrapped in native Promises.
-                 Non-thenables will be returned with a warning.
+    description: Non-thenables will be returned with a warning.
 -->
 
 * `fn` {Function} Function to wrap a trace around
 * `context` {Object} Shared object to correlate trace events through
 * `thisArg` {any} The receiver to be used for the function call
 * `...args` {any} Optional arguments to pass to the function
-* Returns: {any} The return value of the given function, or the result of
-  calling `.then(...)` on the return value if the tracing channel has active
-  subscribers. If the return value is not a Promise or thenable, then
-  it is returned as-is and a warning is emitted.
+* Returns: {any} The return value of the given function. If the return value
+  is a Promise or thenable, tracing events will be published when it settles.
+  If the return value is not a Promise or thenable, it is returned as-is and
+  a warning is emitted.
 
 Trace an asynchronous function call which returns a {Promise} or
 [thenable object][]. This will always produce a [`start` event][] and

lib/diagnostics_channel.js

@@ -10,8 +10,8 @@ const {
   ObjectDefineProperty,
   ObjectGetPrototypeOf,
   ObjectSetPrototypeOf,
+  PromisePrototype,
   PromisePrototypeThen,
-  PromiseReject,
   ReflectApply,
   SafeFinalizationRegistry,
   SafeMap,
@@ -546,17 +546,21 @@ class TracingChannel {
     const { error } = this;
     const continuationWindow = this.#continuationWindow;
 
-    function reject(err) {
+    function onReject(err) {
       context.error = err;
       error.publish(context);
       // Use continuation window for asyncStart/asyncEnd
       // eslint-disable-next-line no-unused-vars
       using scope = continuationWindow.withScope(context);
       // TODO: Is there a way to have asyncEnd _after_ the continuation?
-      return PromiseReject(err);
     }
 
-    function resolve(result) {
+    function onRejectWithRethrow(err) {
+      onReject(err);
+      throw err;
+    }
+
+    function onResolve(result) {
       context.result = result;
       // Use continuation window for asyncStart/asyncEnd
       // eslint-disable-next-line no-unused-vars
@@ -576,12 +580,17 @@ class TracingChannel {
         context.result = result;
         return result;
       }
-      // For native Promises use PromisePrototypeThen to avoid user overrides.
-      if (isPromise(result)) {
-        return PromisePrototypeThen(result, resolve, reject);
+      // isPromise() matches sub-classes, but we need to match only direct
+      // instances of the native Promise type to safely use PromisePrototypeThen.
+      if (isPromise(result) && ObjectGetPrototypeOf(result) === PromisePrototype) {
+        return PromisePrototypeThen(result, onResolve, onRejectWithRethrow);
       }
-      // For custom thenables, call .then() directly to preserve the thenable type.
-      return result.then(resolve, reject);
+      // For non-native thenables, subscribe to the result but return the
+      // original thenable so the consumer can continue handling it directly.
+      // Non-native thenables don't have unhandledRejection tracking, so
+      // swallowing the rejection here doesn't change existing behaviour.
+      result.then(onResolve, onReject);
+      return result;
     } catch (err) {
       context.error = err;
       error.publish(context);

test/parallel/test-diagnostics-channel-tracing-channel-promise-spoofed-constructor.js

@@ -0,0 +1,52 @@
+'use strict';
+
+const common = require('../common');
+const dc = require('diagnostics_channel');
+const assert = require('assert');
+
+class SpoofedPromise extends Promise {
+  customMethod() {
+    return 'works';
+  }
+}
+
+const channel = dc.tracingChannel('test');
+
+const expectedResult = { foo: 'bar' };
+const input = { foo: 'bar' };
+const thisArg = { baz: 'buz' };
+
+function check(found) {
+  assert.strictEqual(found, input);
+}
+
+function checkAsync(found) {
+  check(found);
+  assert.strictEqual(found.error, undefined);
+  assert.deepStrictEqual(found.result, expectedResult);
+}
+
+const handlers = {
+  start: common.mustCall(check),
+  end: common.mustCall(check),
+  asyncStart: common.mustCall(checkAsync),
+  asyncEnd: common.mustCall(checkAsync),
+  error: common.mustNotCall()
+};
+
+channel.subscribe(handlers);
+
+let innerPromise;
+
+const result = channel.tracePromise(common.mustCall(function() {
+  innerPromise = SpoofedPromise.resolve(expectedResult);
+  // Spoof the constructor to try to trick the brand check
+  innerPromise.constructor = Promise;
+  return innerPromise;
+}), input, thisArg);
+
+// Despite the spoofed constructor, the original subclass instance should be
+// returned directly so that custom methods remain accessible.
+assert(result instanceof SpoofedPromise);
+assert.strictEqual(result, innerPromise);
+assert.strictEqual(result.customMethod(), 'works');

test/parallel/test-diagnostics-channel-tracing-channel-promise-thenable.js

@@ -12,6 +12,9 @@ class ResolvedThenable {
   then(resolve) {
     return new ResolvedThenable(resolve(this.#result));
   }
+  customMethod() {
+    return this.#result;
+  }
 }
 
 const channel = dc.tracingChannel('test');
@@ -49,7 +52,10 @@ const result = channel.tracePromise(common.mustCall(function(value) {
 }), input, thisArg, expectedResult);
 
 assert(result instanceof ResolvedThenable);
-assert.notStrictEqual(result, innerThenable);
+// With branching then, the original thenable is returned directly so that
+// extra methods defined on it remain accessible to the caller.
+assert.strictEqual(result, innerThenable);
+assert.deepStrictEqual(result.customMethod(), expectedResult);
 result.then(common.mustCall((value) => {
   assert.deepStrictEqual(value, expectedResult);
 }));