From f1a2e3754924efd8909c9b0b1181d79ad595ff82 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Fri, 20 Mar 2026 09:43:29 +0100 Subject: [PATCH] fix(auto-instrumentation): Fully preserve returned promise in tracing channel hook --- .../scenario.mjs | 1 - js/src/auto-instrumentations/hook.mts | 33 +++++++++----- .../test-api-promise-preservation.mjs | 30 +++++++++++++ .../auto-instrumentations/loader-hook.test.ts | 43 ++++++++++++++++--- 4 files changed, 91 insertions(+), 16 deletions(-) create mode 100644 js/tests/auto-instrumentations/fixtures/test-api-promise-preservation.mjs diff --git a/e2e/scenarios/anthropic-auto-instrumentation-node-hook/scenario.mjs b/e2e/scenarios/anthropic-auto-instrumentation-node-hook/scenario.mjs index 43e01c4fd..de8cdc5b3 100644 --- a/e2e/scenarios/anthropic-auto-instrumentation-node-hook/scenario.mjs +++ b/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, }), ); diff --git a/js/src/auto-instrumentations/hook.mts b/js/src/auto-instrumentations/hook.mts index af4987533..578b42fd9 100644 --- a/js/src/auto-instrumentations/hook.mts +++ b/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); diff --git a/js/tests/auto-instrumentations/fixtures/test-api-promise-preservation.mjs b/js/tests/auto-instrumentations/fixtures/test-api-promise-preservation.mjs new file mode 100644 index 000000000..c0487667d --- /dev/null +++ b/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, + }, +}); diff --git a/js/tests/auto-instrumentations/loader-hook.test.ts b/js/tests/auto-instrumentations/loader-hook.test.ts index 8c4e8a332..54809fef0 100644 --- a/js/tests/auto-instrumentations/loader-hook.test.ts +++ b/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,6 +55,26 @@ 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"); + }); }); }); @@ -60,8 +82,19 @@ async function runWithWorker(options: { execArgv: string[]; script: string; }): Promise { + return runWithWorkerMessage({ + ...options, + messageType: "events", + }); +} + +async function runWithWorkerMessage(options: { + execArgv: string[]; + messageType: string; + script: string; +}): Promise { 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; } });