fix(auto-instrumentation): Fully preserve returned promise in tracing channel hook

e2e/scenarios/anthropic-auto-instrumentation-node-hook/scenario.mjs

@@ -9,6 +9,5 @@ runMain(async () =>
     rootName: "anthropic-auto-hook-root",
     scenarioName: "anthropic-auto-instrumentation-node-hook",
     testImageUrl: new URL("./test-image.png", import.meta.url),
-    useMessagesStreamHelper: false,
   }),
 );

js/src/auto-instrumentations/hook.mts

@@ -65,19 +65,17 @@ if (TracingChannel && TracingChannel.prototype.tracePromise) {
   ) {
     const { start, end, asyncStart, asyncEnd, error } = this;
 
-    function reject(err: any) {
+    function publishRejected(err: any) {
       context.error = err;
       error?.publish(context);
       asyncStart?.publish(context);
       asyncEnd?.publish(context);
-      return Promise.reject(err);
     }
 
-    function resolve(result: any) {
+    function publishResolved(result: any) {
       context.result = result;
       asyncStart?.publish(context);
       asyncEnd?.publish(context);
-      return result;
     }
 
     start?.publish(context);
@@ -93,16 +91,31 @@ if (TracingChannel && TracingChannel.prototype.tracePromise) {
         (typeof result === "object" || typeof result === "function") &&
         typeof result.then === "function"
       ) {
-        return result.then(resolve, reject);
+        // Preserve the original promise-like object so SDK helper methods
+        // like Anthropic APIPromise.withResponse() remain available.
+        void result.then(
+          (resolved: any) => {
+            try {
+              publishResolved(resolved);
+            } catch {
+              // Preserve wrapped promise semantics even if instrumentation fails.
+            }
+          },
+          (err: any) => {
+            try {
+              publishRejected(err);
+            } catch {
+              // Preserve wrapped promise semantics even if instrumentation fails.
+            }
+          },
+        );
+        return result;
       }
 
-      context.result = result;
-      asyncStart?.publish(context);
-      asyncEnd?.publish(context);
+      publishResolved(result);
       return result;
     } catch (err) {
-      context.error = err;
-      error?.publish(context);
+      publishRejected(err);
       throw err;
     } finally {
       end?.publish(context);

js/tests/auto-instrumentations/fixtures/test-api-promise-preservation.mjs

@@ -0,0 +1,30 @@
+import { tracingChannel } from "node:diagnostics_channel";
+import { parentPort } from "node:worker_threads";
+
+class HelperPromise extends Promise {
+  async withResponse() {
+    return {
+      data: await this,
+      response: { ok: true },
+    };
+  }
+}
+
+const channel = tracingChannel("braintrust:test:helper-promise");
+const traced = channel.tracePromise(
+  () => new HelperPromise((resolve) => resolve("ok")),
+  {},
+);
+
+const withResponse = await traced.withResponse();
+
+parentPort?.postMessage({
+  type: "helper-result",
+  result: {
+    awaitedValue: await traced,
+    constructorName: traced.constructor.name,
+    hasWithResponse: typeof traced.withResponse === "function",
+    withResponseData: withResponse.data,
+    withResponseOk: withResponse.response.ok,
+  },
+});

js/tests/auto-instrumentations/loader-hook.test.ts

@@ -1,12 +1,10 @@
 import { describe, it, expect, beforeAll, afterAll } from "vitest";
 import { Worker } from "node:worker_threads";
-import * as fs from "node:fs";
 import * as path from "node:path";
 import { fileURLToPath, pathToFileURL } from "node:url";
 
 const __dirname = path.dirname(fileURLToPath(import.meta.url));
 const fixturesDir = path.join(__dirname, "fixtures");
-const nodeModulesDir = path.join(fixturesDir, "node_modules");
 
 // Path to unified loader hook (built dist file)
 const hookPath = path.join(
@@ -18,6 +16,10 @@ const hookPath = path.join(
 const listenerPath = path.join(fixturesDir, "listener-esm.mjs");
 const testAppEsmPath = path.join(fixturesDir, "test-app-esm.mjs");
 const testAppCjsPath = path.join(fixturesDir, "test-app-cjs.cjs");
+const helperPromisePath = path.join(
+  fixturesDir,
+  "test-api-promise-preservation.mjs",
+);
 
 interface TestResult {
   events: { start: any[]; end: any[]; error: any[] };
@@ -53,15 +55,46 @@ describe("Unified Loader Hook Integration Tests", () => {
       expect(result.events.start.length).toBeGreaterThan(0);
       expect(result.events.end.length).toBeGreaterThan(0);
     });
+
+    it("should preserve helper methods on promise subclasses", async () => {
+      const result = await runWithWorkerMessage<{
+        awaitedValue: string;
+        constructorName: string;
+        hasWithResponse: boolean;
+        withResponseData: string;
+        withResponseOk: boolean;
+      }>({
+        execArgv: ["--import", hookPath],
+        messageType: "helper-result",
+        script: helperPromisePath,
+      });
+
+      expect(result.hasWithResponse).toBe(true);
+      expect(result.awaitedValue).toBe("ok");
+      expect(result.withResponseData).toBe("ok");
+      expect(result.withResponseOk).toBe(true);
+      expect(result.constructorName).toBe("HelperPromise");
+    });
   });
 });
 
 async function runWithWorker(options: {
   execArgv: string[];
   script: string;
 }): Promise<TestResult> {
+  return runWithWorkerMessage({
+    ...options,
+    messageType: "events",
+  });
+}
+
+async function runWithWorkerMessage<T>(options: {
+  execArgv: string[];
+  messageType: string;
+  script: string;
+}): Promise<T> {
   return new Promise((resolve, reject) => {
-    let result: TestResult | null = null;
+    let result: T | null = null;
 
     // Convert execArgv paths to file URLs on Windows
     // On Windows, Node.js --import requires file:// URLs
@@ -94,8 +127,8 @@ async function runWithWorker(options: {
     });
 
     worker.on("message", (msg) => {
-      if (msg.type === "events") {
-        result = { events: msg.events };
+      if (msg.type === options.messageType) {
+        result = (msg.result ?? { events: msg.events }) as T;
       }
     });