fix: Preserve Promise subclass return values in wrapPromise (#57)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Qard

lib/transforms.js

@@ -391,24 +391,46 @@ function wrapPromise (state, node) {
         try {
           let promise = __apm$traced();
           if (typeof promise?.then !== 'function') {
-            __apm$ctx.result = promise; 
+            __apm$ctx.result = promise;
             return promise;
           }
-          return promise.then(
+          // Mirror Node.js core diagnostics_channel behaviour: for native Promise
+          // instances, chain normally (safe since there is no subclass API to
+          // preserve). For Promise subclasses and other thenables, side-chain the
+          // callbacks for event publishing and return the original so that any
+          // subclass-specific methods (e.g. APIPromise.withResponse()) remain
+          // accessible to the caller.
+          if (promise instanceof Promise && promise.constructor === Promise) {
+            return promise.then(
+              result => {
+                __apm$ctx.result = result;
+                ${channelVariable}.asyncStart.publish(__apm$ctx);
+                ${channelVariable}.asyncEnd.publish(__apm$ctx);
+                return result;
+              },
+              err => {
+                __apm$ctx.error = err;
+                ${channelVariable}.error.publish(__apm$ctx);
+                ${channelVariable}.asyncStart.publish(__apm$ctx);
+                ${channelVariable}.asyncEnd.publish(__apm$ctx);
+                throw err;
+              }
+            );
+          }
+          promise.then(
             result => {
               __apm$ctx.result = result;
               ${channelVariable}.asyncStart.publish(__apm$ctx);
               ${channelVariable}.asyncEnd.publish(__apm$ctx);
-              return result;
             },
             err => {
               __apm$ctx.error = err;
               ${channelVariable}.error.publish(__apm$ctx);
               ${channelVariable}.asyncStart.publish(__apm$ctx);
               ${channelVariable}.asyncEnd.publish(__apm$ctx);
-              throw err;
             }
           );
+          return promise;
         } catch (err) {
           __apm$ctx.error = err;
           ${channelVariable}.error.publish(__apm$ctx);

tests/promise_subclass/mod.js

@@ -0,0 +1,27 @@
+/**
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc.
+ **/
+
+/**
+ * A Promise subclass with an extra method, similar to Anthropic SDK's APIPromise.
+ * The instrumented wrapper must preserve the original return type so callers
+ * can still access subclass-specific methods like `.withResponse()`.
+ */
+class ExtendedPromise extends Promise {
+  withResponse () {
+    return this.then(result => ({ data: result, response: { status: 200 } }))
+  }
+}
+
+/**
+ * NOTE: Intentionally NOT async. Async functions always wrap their return value
+ * in a native Promise, losing the subclass type. The real-world scenario
+ * (e.g. Anthropic SDK's APIPromise) uses a non-async function that explicitly
+ * constructs and returns a Promise subclass instance.
+ */
+function fetch (url) {
+  return new ExtendedPromise(resolve => resolve(42))
+}
+
+module.exports = { fetch }

tests/promise_subclass/test.js

@@ -0,0 +1,21 @@
+/**
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc.
+ **/
+const { fetch } = require('./instrumented.js')
+const { assert, getContext } = require('../common/preamble.js')
+const context = getContext('orchestrion:undici:fetch_subclass');
+(async () => {
+  const promise = fetch('https://example.com')
+  // The instrumented wrapper must return the original promise, not promise.then(),
+  // so subclass methods like withResponse() remain accessible.
+  assert.strictEqual(typeof promise.withResponse, 'function', 'withResponse should be available on the returned promise')
+  const { data } = await promise.withResponse()
+  assert.strictEqual(data, 42)
+  assert.deepStrictEqual(context, {
+    start: true,
+    end: true,
+    asyncStart: 42,
+    asyncEnd: 42
+  })
+})()

tests/tests.test.mjs

@@ -518,6 +518,18 @@ describe('wrap_promise_non_promise', () => {
   })
 })
 
+describe('promise_subclass', () => {
+  test('instruments async function and preserves original Promise subclass return type', () => {
+    runTest('promise_subclass', [
+      {
+        channelName: 'fetch_subclass',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        functionQuery: { functionName: 'fetch', kind: 'Async' }
+      }
+    ])
+  })
+})
+
 describe('IIFE with class', () => {
   test('instruments a class within a IIFE, variable same name as class', () => {
     runTest('iife_nested_class', [