fix(core): Preserve AI proxy property references

packages/server-utils/src/ai/google-genai/index.ts

@@ -378,9 +378,19 @@ function instrumentMethod<T extends unknown[], R>(
  * Recursively instruments methods and handles special cases like chats.create
  */
 function createDeepProxy<T extends object>(target: T, currentPath = '', options: GoogleGenAIOptions): T {
+  const propertyCache = new Map<PropertyKey, { original: unknown; wrapped: unknown }>();
+
   return new Proxy(target, {
     get: (t, prop, receiver) => {
       const value = Reflect.get(t, prop, receiver);
+      const cachedEntry = propertyCache.get(prop);
+      if (cachedEntry) {
+        if (cachedEntry.original === value) {
+          return cachedEntry.wrapped;
+        }
+        propertyCache.delete(prop);
+      }
+
       const methodPath = buildMethodPath(currentPath, String(prop));
 
       const instrumentedMethod: InstrumentedMethodEntry | undefined =
@@ -392,29 +402,36 @@ function createDeepProxy<T extends object>(target: T, currentPath = '', options:
           : value.bind(t);
 
         if (!instrumentedMethod.proxyResultPath) {
+          propertyCache.set(prop, { original: value, wrapped: wrappedMethod });
           return wrappedMethod;
         }
 
         // If a proxyResultPath is specified, we need to proxy the result of the method.
         // Note: This currently only properly handles synchronous methods. For async methods,
         // the Promise itself would be proxied instead of the resolved value. Currently we
         // don't have a case where this is needed, so I'll keep it simple for now.
-        return function (...args: unknown[]): unknown {
+        const methodWithProxiedResult = function (...args: unknown[]): unknown {
           const result = wrappedMethod(...args);
           if (result && typeof result === 'object') {
             return createDeepProxy(result as object, instrumentedMethod.proxyResultPath, options);
           }
           return result;
         };
+        propertyCache.set(prop, { original: value, wrapped: methodWithProxiedResult });
+        return methodWithProxiedResult;
       }
 
       if (typeof value === 'function') {
         // Bind non-instrumented functions to preserve the original `this` context
-        return value.bind(t);
+        const boundMethod = value.bind(t);
+        propertyCache.set(prop, { original: value, wrapped: boundMethod });
+        return boundMethod;
       }
 
       if (value && typeof value === 'object') {
-        return createDeepProxy(value, methodPath, options);
+        const nestedProxy = createDeepProxy(value, methodPath, options);
+        propertyCache.set(prop, { original: value, wrapped: nestedProxy });
+        return nestedProxy;
       }
 
       return value;

packages/server-utils/src/ai/openai/index.ts

@@ -225,30 +225,46 @@ function instrumentMethod<T extends unknown[], R>(
  * Create a deep proxy for OpenAI client instrumentation
  */
 function createDeepProxy<T extends object>(target: T, currentPath = '', options: OpenAiOptions): T {
+  const propertyCache = new Map<PropertyKey, { original: unknown; wrapped: unknown }>();
+
   return new Proxy(target, {
-    get(obj: object, prop: string): unknown {
-      const value = (obj as Record<string, unknown>)[prop];
+    get(obj: object, prop: string | symbol): unknown {
+      const value = Reflect.get(obj, prop) as unknown;
+      const cachedEntry = propertyCache.get(prop);
+      if (cachedEntry) {
+        if (cachedEntry.original === value) {
+          return cachedEntry.wrapped;
+        }
+        propertyCache.delete(prop);
+      }
+
       const methodPath = buildMethodPath(currentPath, String(prop));
 
       const instrumentedMethod = OPENAI_METHOD_REGISTRY[methodPath as keyof typeof OPENAI_METHOD_REGISTRY];
       if (typeof value === 'function' && instrumentedMethod) {
-        return instrumentMethod(
+        const wrappedMethod = instrumentMethod(
           value as (...args: unknown[]) => Promise<unknown>,
           methodPath,
           instrumentedMethod,
           obj,
           options,
         );
+        propertyCache.set(prop, { original: value, wrapped: wrappedMethod });
+        return wrappedMethod;
       }
 
       if (typeof value === 'function') {
         // Bind non-instrumented functions to preserve the original `this` context,
         // which is required for accessing private class fields (e.g. #baseURL) in OpenAI SDK v5.
-        return value.bind(obj);
+        const boundMethod = value.bind(obj);
+        propertyCache.set(prop, { original: value, wrapped: boundMethod });
+        return boundMethod;
       }
 
       if (value && typeof value === 'object') {
-        return createDeepProxy(value, methodPath, options);
+        const nestedProxy = createDeepProxy(value, methodPath, options);
+        propertyCache.set(prop, { original: value, wrapped: nestedProxy });
+        return nestedProxy;
       }
 
       return value;

packages/server-utils/test/ai/tracing/ai-integration-proxy.test.ts

@@ -0,0 +1,82 @@
+import { describe, expect, it } from 'vitest';
+import { instrumentGoogleGenAIClient } from '../../../src/ai/google-genai';
+import { instrumentOpenAiClient } from '../../../src/ai/openai';
+
+interface MutableClient {
+  service: { execute: () => Promise<void> } | null;
+  getVersion: (() => string) | null;
+}
+
+const clientInstrumenters: Array<{
+  provider: string;
+  instrument: (client: MutableClient) => MutableClient;
+}> = [
+  { provider: 'OpenAI', instrument: client => instrumentOpenAiClient(client) },
+  { provider: 'Google GenAI', instrument: client => instrumentGoogleGenAIClient(client) },
+];
+
+describe('AI integration deep proxies', () => {
+  it('returns stable references for OpenAI client properties', () => {
+    const client = {
+      chat: { completions: { create: () => Promise.resolve({}) } },
+      getBaseURL: () => 'https://api.openai.com',
+    };
+
+    const instrumentedClient = instrumentOpenAiClient(client);
+    const chat = instrumentedClient.chat;
+    const completions = chat.completions;
+    const create = completions.create;
+    const getBaseURL = instrumentedClient.getBaseURL;
+
+    expect(instrumentedClient.chat).toBe(chat);
+    expect(chat.completions).toBe(completions);
+    expect(completions.create).toBe(create);
+    expect(instrumentedClient.chat.completions.create).toBe(create);
+    expect(instrumentedClient.getBaseURL).toBe(getBaseURL);
+  });
+
+  it('returns stable references for Google GenAI client properties', () => {
+    const client = {
+      models: { generateContent: () => Promise.resolve({}) },
+      chats: { create: () => ({ sendMessage: () => Promise.resolve({}) }) },
+      getVersion: () => '1.0.0',
+    };
+
+    const instrumentedClient = instrumentGoogleGenAIClient(client);
+    const models = instrumentedClient.models;
+    const generateContent = models.generateContent;
+    const chats = instrumentedClient.chats;
+    const create = chats.create;
+    const chat = chats.create();
+    const sendMessage = chat.sendMessage;
+    const getVersion = instrumentedClient.getVersion;
+
+    expect(instrumentedClient.models).toBe(models);
+    expect(models.generateContent).toBe(generateContent);
+    expect(instrumentedClient.models.generateContent).toBe(generateContent);
+    expect(instrumentedClient.chats).toBe(chats);
+    expect(chats.create).toBe(create);
+    expect(instrumentedClient.chats.create).toBe(create);
+    expect(chat.sendMessage).toBe(sendMessage);
+    expect(instrumentedClient.getVersion).toBe(getVersion);
+  });
+
+  it.each(clientInstrumenters)('drops cached $provider wrappers after properties are cleared', ({ instrument }) => {
+    const service = { execute: () => Promise.resolve() };
+    const getVersion = () => '1.0.0';
+    const client: MutableClient = { service, getVersion };
+    const instrumentedClient = instrument(client);
+    const cachedService = instrumentedClient.service;
+    const cachedGetVersion = instrumentedClient.getVersion;
+
+    client.service = null;
+    client.getVersion = null;
+    expect(instrumentedClient.service).toBeNull();
+    expect(instrumentedClient.getVersion).toBeNull();
+
+    client.service = service;
+    client.getVersion = getVersion;
+    expect(instrumentedClient.service).not.toBe(cachedService);
+    expect(instrumentedClient.getVersion).not.toBe(cachedGetVersion);
+  });
+});