Fix more centrally

Author: lforst

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,
   });
 }
 

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,28 +200,28 @@ 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();
         },
       );
 
       await runOperation(
         "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);
         },
       );

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<string> & {
+      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();

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

@@ -10,6 +10,10 @@
  * and in configureNode/configureBrowser for the bundler plugin path.
  */
 
+function hasOwnPromiseAugmentations(promise: Promise<unknown>): boolean {
+  return;
+}
+
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export function patchTracingChannel(
   tracingChannelFn: (name: string) => any,
@@ -77,7 +81,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 +93,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);