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

Fix the auto-instrumentation node hook so Anthropic `messages.stream()`
keeps working under `node --import braintrust/hook.mjs`.

- Preserve the original promise-like return value in the hook instead of
returning a chained promise
- Keep Anthropic `APIPromise` helpers like `withResponse()` intact
- Remove the e2e workaround so the Anthropic node-hook scenario uses
`messages.stream()` again
- Add regression coverage for promise subclass helper preservation

Fixes
#1591

Author: lforst

js/src/auto-instrumentations/patch-tracing-channel.test.ts

@@ -135,6 +135,13 @@ class MockAPIPromise<T> extends Promise<T> {
     return this.#innerPromise.then(onfulfilled, onrejected);
   }
 
+  async withResponse() {
+    return {
+      data: await this.#innerPromise,
+      response: { ok: true },
+    };
+  }
+
   // Symbol.species returns a class that throws, reproducing the exact error seen
   // when Node.js's PromisePrototypeThen tries to construct the result promise.
   static override get [Symbol.species](): PromiseConstructor {
@@ -180,6 +187,20 @@ describe("patchTracingChannel", () => {
     expect(context.result).toBe("hello");
   });
 
+  it("patched tracePromise preserves helper methods on promise subclasses", async () => {
+    const FakeTCClass = makeUnpatchedTracingChannel();
+    const channel = new FakeTCClass();
+    patchTracingChannel(() => channel);
+
+    const apiPromise = new MockAPIPromise(Promise.resolve("hello"));
+    const traced = channel.tracePromise(() => apiPromise, {}, null);
+    const withResponse = await traced.withResponse();
+
+    expect(traced).toBe(apiPromise);
+    expect(withResponse.data).toBe("hello");
+    expect(withResponse.response.ok).toBe(true);
+  });
+
   it("patched tracePromise correctly handles plain async functions", async () => {
     const FakeTCClass = makeUnpatchedTracingChannel();
     const channel = new FakeTCClass();

js/src/auto-instrumentations/patch-tracing-channel.ts

@@ -59,20 +59,18 @@ export function patchTracingChannel(
       const { start, end, asyncStart, asyncEnd, error } = this;
 
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      function reject(err: any) {
+      function publishRejected(err: any) {
         context.error = err;
         error?.publish(context);
         asyncStart?.publish(context);
         asyncEnd?.publish(context);
-        return Promise.reject(err);
       }
 
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      function resolve(result: any) {
+      function publishResolved(result: any) {
         context.result = result;
         asyncStart?.publish(context);
         asyncEnd?.publish(context);
-        return result;
       }
 
       // Use runStores (not just publish) so fn runs inside the ALS context
@@ -91,7 +89,41 @@ export function patchTracingChannel(
             (typeof result === "object" || typeof result === "function") &&
             typeof result.then === "function"
           ) {
-            return result.then(resolve, reject);
+            if (result.constructor === Promise) {
+              return result.then(
+                (res) => {
+                  publishResolved(res);
+                  return res;
+                },
+                (err) => {
+                  publishRejected(err);
+                  return Promise.reject(err);
+                },
+              );
+            }
+
+            // Preserve the original promise-like object so SDK helper methods
+            // like Anthropic APIPromise.withResponse() remain available.
+            void result.then(
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
+              (resolved: any) => {
+                try {
+                  publishResolved(resolved);
+                } catch {
+                  // Preserve wrapped promise semantics even if instrumentation fails.
+                }
+              },
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
+              (err: any) => {
+                try {
+                  publishRejected(err);
+                } catch {
+                  // Preserve wrapped promise semantics even if instrumentation fails.
+                }
+              },
+            );
+
+            return result;
           }
 
           context.result = result;

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

@@ -0,0 +1,45 @@
+import { tracingChannel } from "node:diagnostics_channel";
+import { parentPort } from "node:worker_threads";
+
+class HelperPromise extends Promise {
+  #innerPromise;
+
+  constructor(responsePromise) {
+    let resolveOuter;
+    super((resolve) => {
+      resolveOuter = resolve;
+    });
+    this.#innerPromise = Promise.resolve(responsePromise);
+    this.#innerPromise.then(resolveOuter);
+  }
+
+  then(onfulfilled, onrejected) {
+    return this.#innerPromise.then(onfulfilled, onrejected);
+  }
+
+  async withResponse() {
+    return {
+      data: await this.#innerPromise,
+      response: { ok: true },
+    };
+  }
+}
+
+const channel = tracingChannel("braintrust:test:helper-promise");
+const traced = channel.tracePromise(
+  () => new HelperPromise(Promise.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;
       }
     });