fix(openai): Fix wrapOpenAI breaking native private fields of client (#1704)

Aims to address
#1693

- Adds regression tests
- Adds a more robust method of returning the accessed functions

Please see comments in code for further explanation.

We didn't chose the original approach the referenced issue takes
(`Reflect.get(..., target)`) because it doesn't work since `this` is
still the proxied instance.

Author: lforst

e2e/scenarios/openai-instrumentation/assertions.ts

@@ -427,6 +427,7 @@ function buildPayloadSummary(events: CapturedLogEvent[]): Json {
 }
 
 export function defineOpenAIInstrumentationAssertions(options: {
+  assertPrivateFieldMethodsOperation?: boolean;
   name: string;
   runScenario: RunOpenAIScenario;
   snapshotName: string;
@@ -466,6 +467,26 @@ export function defineOpenAIInstrumentationAssertions(options: {
       });
     });
 
+    if (options.assertPrivateFieldMethodsOperation) {
+      test(
+        "keeps wrapped v6 client private-field methods callable",
+        testConfig,
+        () => {
+          const root = findLatestSpan(events, ROOT_NAME);
+          const operation = findLatestSpan(
+            events,
+            "openai-client-private-fields-operation",
+          );
+
+          expect(operation).toBeDefined();
+          expect(operation?.row.metadata).toMatchObject({
+            operation: "client-private-fields",
+          });
+          expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]);
+        },
+      );
+    }
+
     for (const spec of OPERATION_SPECS) {
       test(spec.testName, testConfig, () => {
         const root = findLatestSpan(events, ROOT_NAME);

e2e/scenarios/openai-instrumentation/scenario.impl.mjs

@@ -44,6 +44,15 @@ async function awaitMaybeWithResponse(request) {
   };
 }
 
+function parseMajorVersion(version) {
+  if (typeof version !== "string") {
+    return null;
+  }
+
+  const major = Number.parseInt(version.split(".")[0], 10);
+  return Number.isNaN(major) ? null : major;
+}
+
 export async function runOpenAIInstrumentationScenario(options) {
   const baseClient = new options.OpenAI({
     apiKey: process.env.OPENAI_API_KEY,
@@ -52,9 +61,51 @@ export async function runOpenAIInstrumentationScenario(options) {
   const client = options.decorateClient
     ? options.decorateClient(baseClient)
     : baseClient;
+  const openAIMajorVersion = parseMajorVersion(options.openaiSdkVersion);
+  const shouldCheckPrivateFieldMethods =
+    typeof options.decorateClient === "function" &&
+    openAIMajorVersion !== null &&
+    openAIMajorVersion >= 6;
 
   await runTracedScenario({
     callback: async () => {
+      if (shouldCheckPrivateFieldMethods) {
+        await runOperation(
+          "openai-client-private-fields-operation",
+          "client-private-fields",
+          async () => {
+            if (
+              typeof client.buildURL !== "function" ||
+              typeof client.buildRequest !== "function"
+            ) {
+              throw new Error(
+                "Expected wrapped OpenAI v6 client to expose buildURL and buildRequest",
+              );
+            }
+
+            const builtUrl = client.buildURL("/files", null);
+            if (typeof builtUrl !== "string" || !builtUrl.includes("/files")) {
+              throw new Error(
+                `Unexpected buildURL result: ${String(builtUrl)}`,
+              );
+            }
+
+            const builtRequest = await client.buildRequest(
+              { method: "post", path: "/files" },
+              { retryCount: 0 },
+            );
+            if (
+              typeof builtRequest?.url !== "string" ||
+              !builtRequest.url.includes("/files")
+            ) {
+              throw new Error(
+                `Unexpected buildRequest result: ${String(builtRequest?.url)}`,
+              );
+            }
+          },
+        );
+      }
+
       await runOperation("openai-chat-operation", "chat", async () => {
         await client.chat.completions.create({
           model: OPENAI_MODEL,

e2e/scenarios/openai-instrumentation/scenario.test.ts

@@ -14,12 +14,14 @@ const openaiScenarios = await Promise.all(
   [
     {
       autoEntry: "scenario.openai-v4.mjs",
+      disablePrivateFieldMethodsAssertion: true,
       dependencyName: "openai-v4",
       snapshotName: "openai-v4",
       wrapperEntry: "scenario.openai-v4.ts",
     },
     {
       autoEntry: "scenario.openai-v5.mjs",
+      disablePrivateFieldMethodsAssertion: true,
       dependencyName: "openai-v5",
       snapshotName: "openai-v5",
       wrapperEntry: "scenario.openai-v5.ts",
@@ -40,8 +42,12 @@ const openaiScenarios = await Promise.all(
 );
 
 for (const scenario of openaiScenarios) {
+  const assertPrivateFieldMethodsOperation =
+    !scenario.disablePrivateFieldMethodsAssertion;
+
   describe(`openai sdk ${scenario.version}`, () => {
     defineOpenAIInstrumentationAssertions({
+      assertPrivateFieldMethodsOperation,
       name: "wrapped instrumentation",
       runScenario: async ({ runScenarioDir }) => {
         await runScenarioDir({

js/src/wrappers/oai.test.ts

@@ -35,6 +35,20 @@ test("openai is installed", () => {
   assert.ok(OpenAI);
 });
 
+test("wrapOpenAI keeps private-field helper methods callable", async () => {
+  const wrapped = wrapOpenAI(new OpenAI({ apiKey: "test-key" }));
+
+  expect(wrapped.buildURL).toBe(wrapped.buildURL);
+  expect(wrapped.buildRequest).toBe(wrapped.buildRequest);
+  expect(() => wrapped.buildURL("/files", null)).not.toThrow();
+
+  const request = await wrapped.buildRequest(
+    { method: "post", path: "/files" },
+    { retryCount: 0 },
+  );
+  expect(request.url).toContain("/files");
+});
+
 describe("openai client unit tests", TEST_SUITE_OPTIONS, () => {
   let oai: OpenAI;
   let client: OpenAI;

js/src/wrappers/oai.ts

@@ -74,6 +74,12 @@ type OpenAILike = OpenAIV4Client;
 export function wrapOpenAIv4<T extends OpenAILike>(openai: T): T {
   // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
   const typedOpenai = openai as OpenAIV4Client;
+  // Recover `this` for fallback methods so private fields and internal slots
+  // keep seeing the original OpenAI instance instead of the proxy.
+  const privateMethodWorkaroundCache = new WeakMap<
+    (...args: unknown[]) => unknown,
+    (...args: unknown[]) => unknown
+  >();
 
   const completionProxy = new Proxy(typedOpenai.chat.completions, {
     get(target, name, receiver) {
@@ -142,8 +148,8 @@ export function wrapOpenAIv4<T extends OpenAILike>(openai: T): T {
   }
 
   // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
-  return new Proxy(typedOpenai, {
-    get(target, name, receiver) {
+  const topLevelProxy = new Proxy(typedOpenai, {
+    get(target, name) {
       switch (name) {
         case "chat":
           return chatProxy;
@@ -158,9 +164,44 @@ export function wrapOpenAIv4<T extends OpenAILike>(openai: T): T {
       if (name === "beta" && betaProxy) {
         return betaProxy;
       }
-      return Reflect.get(target, name, receiver);
+
+      // The following rather convoluted code is a workaround for https://github.com/braintrustdata/braintrust-sdk-javascript/issues/1693
+      // The problem is that Proxies are inherently difficult to work with native private class fields because when a
+      // class function accesses a private field, JS checks whether `this` is equal to the actual instance with the
+      // private field and if that's not the case, it throws a `TypeError`.
+      // We could have also done `if (typeof value === "function") return value.bind(target);`, but it would have
+      // created a new function on each function access, so we are caching, and also it would have always stomped on
+      // someone passing another `this` which may clash with different instrumentations.
+
+      // Use the real client as receiver when reading fallback members.
+      const value = Reflect.get(target, name, target);
+      if (typeof value !== "function") {
+        return value;
+      }
+
+      const cachedValue = privateMethodWorkaroundCache.get(value);
+      if (cachedValue) {
+        return cachedValue;
+      }
+
+      const thisBoundValue = function (
+        this: unknown,
+        ...args: unknown[]
+      ): unknown {
+        // Calling through the proxy would set `this` to the proxy and break
+        // native private-field methods, so recover the original target.
+        const thisArg = this === topLevelProxy ? target : this;
+        const output = Reflect.apply(value, thisArg, args);
+        // Preserve chaining on wrapped clients (method returns `this`).
+        return output === target ? topLevelProxy : output;
+      };
+
+      privateMethodWorkaroundCache.set(value, thisBoundValue);
+      return thisBoundValue;
     },
-  }) as T;
+  });
+
+  return topLevelProxy as T;
 }
 
 type SpanInfo = {