From 60c04aebd2e565aca565f5fdfe0b2f503013fa39 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Thu, 26 Mar 2026 21:39:03 +0100 Subject: [PATCH 1/3] Fix more centrally --- .../scenario.impl.mjs | 47 ++++++------------- .../openai-instrumentation/scenario.impl.mjs | 42 +++++++---------- .../patch-tracing-channel.test.ts | 21 +++++++++ .../patch-tracing-channel.ts | 10 +++- 4 files changed, 59 insertions(+), 61 deletions(-) diff --git a/e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs b/e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs index a5563296..d6c6a7c8 100644 --- a/e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs @@ -26,11 +26,7 @@ const WEATHER_TOOL = { async function runAnthropicInstrumentationScenario( Anthropic, - { - decorateClient, - useBetaMessages = true, - useMessagesStreamHelper = true, - } = {}, + { decorateClient, useBetaMessages = true } = {}, ) { const imageBase64 = ( await readFile(new URL("./test-image.png", import.meta.url)) @@ -123,33 +119,19 @@ async function runAnthropicInstrumentationScenario( "anthropic-stream-with-response-operation", "stream-with-response", async () => { - const stream = - useMessagesStreamHelper === false - ? await client.messages.create({ - model: ANTHROPIC_MODEL, - max_tokens: 32, - temperature: 0, - stream: true, - messages: [ - { - role: "user", - content: - "Count from 1 to 3 and include the words one two three.", - }, - ], - }) - : client.messages.stream({ - model: ANTHROPIC_MODEL, - max_tokens: 32, - temperature: 0, - messages: [ - { - role: "user", - content: - "Count from 1 to 3 and include the words one two three.", - }, - ], - }); + const stream = client.messages.stream({ + model: ANTHROPIC_MODEL, + max_tokens: 32, + temperature: 0, + messages: [ + { + role: "user", + content: + "Count from 1 to 3 and include the words one two three.", + }, + ], + }); + await stream.withResponse(); await collectAsync(stream); }, ); @@ -251,7 +233,6 @@ export async function runWrappedAnthropicInstrumentation(Anthropic, options) { export async function runAutoAnthropicInstrumentation(Anthropic, options) { await runAnthropicInstrumentationScenario(Anthropic, { ...options, - useMessagesStreamHelper: false, }); } diff --git a/e2e/scenarios/openai-instrumentation/scenario.impl.mjs b/e2e/scenarios/openai-instrumentation/scenario.impl.mjs index 5383462a..58fb0577 100644 --- a/e2e/scenarios/openai-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/openai-instrumentation/scenario.impl.mjs @@ -34,16 +34,6 @@ async function collectOneAndReturn(stream) { } } -async function awaitMaybeWithResponse(request) { - if (typeof request?.withResponse === "function") { - return await request.withResponse(); - } - - return { - data: await request, - }; -} - export async function runOpenAIInstrumentationScenario(options) { const baseClient = new options.OpenAI({ apiKey: process.env.OPENAI_API_KEY, @@ -68,14 +58,14 @@ export async function runOpenAIInstrumentationScenario(options) { "openai-chat-with-response-operation", "chat-with-response", async () => { - await awaitMaybeWithResponse( - client.chat.completions.create({ + await client.chat.completions + .create({ model: OPENAI_MODEL, messages: [{ role: "user", content: "Reply with exactly FOUR." }], max_tokens: 8, temperature: 0, - }), - ); + }) + .withResponse(); }, ); @@ -97,8 +87,8 @@ export async function runOpenAIInstrumentationScenario(options) { "openai-stream-with-response-operation", "stream-with-response", async () => { - const { data: chatStream } = await awaitMaybeWithResponse( - client.chat.completions.create({ + const { data: chatStream } = await client.chat.completions + .create({ model: OPENAI_MODEL, messages: [ { @@ -112,8 +102,8 @@ export async function runOpenAIInstrumentationScenario(options) { stream_options: { include_usage: true, }, - }), - ); + }) + .withResponse(); await collectAsync(chatStream); }, ); @@ -210,13 +200,13 @@ export async function runOpenAIInstrumentationScenario(options) { "openai-responses-with-response-operation", "responses-with-response", async () => { - await awaitMaybeWithResponse( - client.responses.create({ + await client.responses + .create({ model: OPENAI_MODEL, input: "What is 2 + 2? Reply with just the number.", max_output_tokens: 16, - }), - ); + }) + .withResponse(); }, ); @@ -224,14 +214,14 @@ export async function runOpenAIInstrumentationScenario(options) { "openai-responses-create-stream-operation", "responses-create-stream", async () => { - const { data: responseStream } = await awaitMaybeWithResponse( - client.responses.create({ + const { data: responseStream } = await client.responses + .create({ model: OPENAI_MODEL, input: "Reply with exactly RESPONSE STREAM.", max_output_tokens: 16, stream: true, - }), - ); + }) + .withResponse(); await collectAsync(responseStream); }, ); diff --git a/js/src/auto-instrumentations/patch-tracing-channel.test.ts b/js/src/auto-instrumentations/patch-tracing-channel.test.ts index a78b9e08..f1dde61c 100644 --- a/js/src/auto-instrumentations/patch-tracing-channel.test.ts +++ b/js/src/auto-instrumentations/patch-tracing-channel.test.ts @@ -201,6 +201,27 @@ describe("patchTracingChannel", () => { expect(withResponse.response.ok).toBe(true); }); + it("patched tracePromise preserves helper methods on augmented native Promise instances", async () => { + const FakeTCClass = makeUnpatchedTracingChannel(); + const channel = new FakeTCClass(); + patchTracingChannel(() => channel); + + const nativePromise = Promise.resolve("hello") as Promise & { + withResponse: () => Promise<{ data: string; response: { ok: boolean } }>; + }; + nativePromise.withResponse = async () => ({ + data: await nativePromise, + response: { ok: true }, + }); + + const traced = channel.tracePromise(() => nativePromise, {}, null); + const withResponse = await traced.withResponse(); + + expect(traced).toBe(nativePromise); + 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(); diff --git a/js/src/auto-instrumentations/patch-tracing-channel.ts b/js/src/auto-instrumentations/patch-tracing-channel.ts index b7f71893..e1bfb2d5 100644 --- a/js/src/auto-instrumentations/patch-tracing-channel.ts +++ b/js/src/auto-instrumentations/patch-tracing-channel.ts @@ -77,7 +77,7 @@ export function patchTracingChannel( // established by bindStore — required for span context to propagate across awaits. // PATCHED: inside the callback, use duck-type thenable check instead of // PromisePrototypeThen, which triggers Symbol.species and breaks Promise subclasses - // like Anthropic's APIPromise that have non-standard constructors. + // like Anthropic's and Openai's APIPromise that have non-standard constructors. return start.runStores(context, () => { try { // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -89,7 +89,13 @@ export function patchTracingChannel( (typeof result === "object" || typeof result === "function") && typeof result.then === "function" ) { - if (result.constructor === Promise) { + if ( + // We only want to return the Promise chain when it's an actual + // promise and also doesn't have any additional fields + result.constructor === Promise && + Object.getOwnPropertyNames(result).length === 0 && + Object.getOwnPropertySymbols(result).length === 0 + ) { return result.then( (res) => { publishResolved(res); From 91298eefc3467713968c5cd08dd8dd5853d17666 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Fri, 27 Mar 2026 14:31:43 +0100 Subject: [PATCH 2/3] Also check prototype and fix assertion for older versions --- .../scenario.anthropic-v0273.mjs | 5 ++- .../scenario.anthropic-v0273.ts | 5 ++- .../scenario.impl.mjs | 21 +++++++++++-- .../patch-tracing-channel.test.ts | 31 +++++++++++++++++++ .../patch-tracing-channel.ts | 18 ++++++++--- 5 files changed, 71 insertions(+), 9 deletions(-) diff --git a/e2e/scenarios/anthropic-instrumentation/scenario.anthropic-v0273.mjs b/e2e/scenarios/anthropic-instrumentation/scenario.anthropic-v0273.mjs index f33a4f98..65f94765 100644 --- a/e2e/scenarios/anthropic-instrumentation/scenario.anthropic-v0273.mjs +++ b/e2e/scenarios/anthropic-instrumentation/scenario.anthropic-v0273.mjs @@ -3,5 +3,8 @@ import { runMain } from "../../helpers/provider-runtime.mjs"; import { runAutoAnthropicInstrumentation } from "./scenario.impl.mjs"; runMain(async () => - runAutoAnthropicInstrumentation(Anthropic, { useBetaMessages: false }), + runAutoAnthropicInstrumentation(Anthropic, { + expectStreamWithResponse: false, + useBetaMessages: false, + }), ); diff --git a/e2e/scenarios/anthropic-instrumentation/scenario.anthropic-v0273.ts b/e2e/scenarios/anthropic-instrumentation/scenario.anthropic-v0273.ts index f6d3267e..4c75349d 100644 --- a/e2e/scenarios/anthropic-instrumentation/scenario.anthropic-v0273.ts +++ b/e2e/scenarios/anthropic-instrumentation/scenario.anthropic-v0273.ts @@ -3,5 +3,8 @@ import { runMain } from "../../helpers/scenario-runtime"; import { runWrappedAnthropicInstrumentation } from "./scenario.impl.mjs"; runMain(async () => - runWrappedAnthropicInstrumentation(Anthropic, { useBetaMessages: false }), + runWrappedAnthropicInstrumentation(Anthropic, { + expectStreamWithResponse: false, + useBetaMessages: false, + }), ); diff --git a/e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs b/e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs index d6c6a7c8..6d3eb714 100644 --- a/e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/anthropic-instrumentation/scenario.impl.mjs @@ -26,7 +26,11 @@ const WEATHER_TOOL = { async function runAnthropicInstrumentationScenario( Anthropic, - { decorateClient, useBetaMessages = true } = {}, + { + decorateClient, + expectStreamWithResponse = true, + useBetaMessages = true, + } = {}, ) { const imageBase64 = ( await readFile(new URL("./test-image.png", import.meta.url)) @@ -131,7 +135,20 @@ async function runAnthropicInstrumentationScenario( }, ], }); - await stream.withResponse(); + + if (expectStreamWithResponse) { + if (typeof stream.withResponse !== "function") { + throw new Error( + "Expected messages.stream() to expose withResponse()", + ); + } + await stream.withResponse(); + } else if (typeof stream.withResponse === "function") { + throw new Error( + "Expected messages.stream() to not expose withResponse()", + ); + } + await collectAsync(stream); }, ); diff --git a/js/src/auto-instrumentations/patch-tracing-channel.test.ts b/js/src/auto-instrumentations/patch-tracing-channel.test.ts index f1dde61c..897f8a76 100644 --- a/js/src/auto-instrumentations/patch-tracing-channel.test.ts +++ b/js/src/auto-instrumentations/patch-tracing-channel.test.ts @@ -222,6 +222,37 @@ describe("patchTracingChannel", () => { expect(withResponse.response.ok).toBe(true); }); + it("patched tracePromise preserves helper methods on prototype-augmented native Promise instances", async () => { + const FakeTCClass = makeUnpatchedTracingChannel(); + const channel = new FakeTCClass(); + patchTracingChannel(() => channel); + + const nativePromise = Promise.resolve("hello"); + const augmentedProto = Object.create( + Promise.prototype, + ) as Promise & { + withResponse: () => Promise<{ data: string; response: { ok: boolean } }>; + }; + + augmentedProto.withResponse = async function () { + const data = await this; + return { data, response: { ok: true } }; + }; + + Object.setPrototypeOf(nativePromise, augmentedProto); + + const traced = channel.tracePromise( + () => nativePromise, + {}, + null, + ) as typeof nativePromise & typeof augmentedProto; + const withResponse = await traced.withResponse(); + + expect(traced).toBe(nativePromise); + 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(); diff --git a/js/src/auto-instrumentations/patch-tracing-channel.ts b/js/src/auto-instrumentations/patch-tracing-channel.ts index e1bfb2d5..660a7f69 100644 --- a/js/src/auto-instrumentations/patch-tracing-channel.ts +++ b/js/src/auto-instrumentations/patch-tracing-channel.ts @@ -10,6 +10,15 @@ * and in configureNode/configureBrowser for the bundler plugin path. */ +function isPlainNativePromiseWithoutHelpers(result: Promise): boolean { + return ( + result.constructor === Promise && + Object.getPrototypeOf(result) === Promise.prototype && + Object.getOwnPropertyNames(result).length === 0 && + Object.getOwnPropertySymbols(result).length === 0 + ); +} + // eslint-disable-next-line @typescript-eslint/no-explicit-any export function patchTracingChannel( tracingChannelFn: (name: string) => any, @@ -90,11 +99,10 @@ export function patchTracingChannel( typeof result.then === "function" ) { if ( - // We only want to return the Promise chain when it's an actual - // promise and also doesn't have any additional fields - result.constructor === Promise && - Object.getOwnPropertyNames(result).length === 0 && - Object.getOwnPropertySymbols(result).length === 0 + // Return the Promise chain only for plain native Promises. + // Promise subclasses and prototype-augmented Promises must be + // returned as-is so SDK helper methods stay intact. + isPlainNativePromiseWithoutHelpers(result) ) { return result.then( (res) => { From 6eb8bd40ac320b62e3f9c57514361a71e38d5848 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Fri, 27 Mar 2026 14:57:07 +0100 Subject: [PATCH 3/3] Fix type --- js/src/auto-instrumentations/patch-tracing-channel.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/js/src/auto-instrumentations/patch-tracing-channel.ts b/js/src/auto-instrumentations/patch-tracing-channel.ts index 660a7f69..df07d40c 100644 --- a/js/src/auto-instrumentations/patch-tracing-channel.ts +++ b/js/src/auto-instrumentations/patch-tracing-channel.ts @@ -105,11 +105,11 @@ export function patchTracingChannel( isPlainNativePromiseWithoutHelpers(result) ) { return result.then( - (res) => { + (res: unknown) => { publishResolved(res); return res; }, - (err) => { + (err: unknown) => { publishRejected(err); return Promise.reject(err); },