diff --git a/packages/core/src/tracing/google-genai/index.ts b/packages/core/src/tracing/google-genai/index.ts index 68e4d414586a..9390a82e0e14 100644 --- a/packages/core/src/tracing/google-genai/index.ts +++ b/packages/core/src/tracing/google-genai/index.ts @@ -360,9 +360,16 @@ function instrumentMethod( * Recursively instruments methods and handles special cases like chats.create */ function createDeepProxy(target: T, currentPath = '', options: GoogleGenAIOptions): T { + const propertyCache = new Map(); + return new Proxy(target, { get: (t, prop, receiver) => { const value = Reflect.get(t, prop, receiver); + const cachedEntry = propertyCache.get(prop); + if (cachedEntry && cachedEntry.original === value) { + return cachedEntry.wrapped; + } + const methodPath = buildMethodPath(currentPath, String(prop)); const instrumentedMethod: InstrumentedMethodEntry | undefined = @@ -374,6 +381,7 @@ function createDeepProxy(target: T, currentPath = '', options: : value.bind(t); if (!instrumentedMethod.proxyResultPath) { + propertyCache.set(prop, { original: value, wrapped: wrappedMethod }); return wrappedMethod; } @@ -381,22 +389,28 @@ function createDeepProxy(target: T, currentPath = '', options: // 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; diff --git a/packages/core/src/tracing/openai/index.ts b/packages/core/src/tracing/openai/index.ts index 821e9c68e0ff..f5335253f14a 100644 --- a/packages/core/src/tracing/openai/index.ts +++ b/packages/core/src/tracing/openai/index.ts @@ -238,30 +238,43 @@ function instrumentMethod( * Create a deep proxy for OpenAI client instrumentation */ function createDeepProxy(target: T, currentPath = '', options: OpenAiOptions): T { + const propertyCache = new Map(); + return new Proxy(target, { - get(obj: object, prop: string): unknown { - const value = (obj as Record)[prop]; + get(obj: object, prop: string | symbol): unknown { + const value = Reflect.get(obj, prop) as unknown; + const cachedEntry = propertyCache.get(prop); + if (cachedEntry && cachedEntry.original === value) { + return cachedEntry.wrapped; + } + 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, 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; diff --git a/packages/core/test/tracing/ai-integration-proxy.test.ts b/packages/core/test/tracing/ai-integration-proxy.test.ts new file mode 100644 index 000000000000..5d36884aa757 --- /dev/null +++ b/packages/core/test/tracing/ai-integration-proxy.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from 'vitest'; +import { instrumentGoogleGenAIClient } from '../../src/tracing/google-genai'; +import { instrumentOpenAiClient } from '../../src/tracing/openai'; + +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 completions = instrumentedClient.chat.completions; + const create = completions.create; + const getBaseURL = instrumentedClient.getBaseURL; + + expect(instrumentedClient.chat.completions).toBe(completions); + expect(instrumentedClient.chat.completions.create).toBe(create); + expect(instrumentedClient.getBaseURL).toBe(getBaseURL); + }); + + it('returns stable references for Google GenAI client properties', () => { + const client = { + chats: { create: () => ({ sendMessage: () => Promise.resolve({}) }) }, + getVersion: () => '1.0.0', + }; + + const instrumentedClient = instrumentGoogleGenAIClient(client); + const chats = instrumentedClient.chats; + const create = chats.create; + const getVersion = instrumentedClient.getVersion; + + expect(instrumentedClient.chats).toBe(chats); + expect(instrumentedClient.chats.create).toBe(create); + expect(instrumentedClient.getVersion).toBe(getVersion); + }); +});