diff --git a/packages/server-utils/src/ai/google-genai/index.ts b/packages/server-utils/src/ai/google-genai/index.ts index c7ec34b5b70d..2e1e0d2ea5a9 100644 --- a/packages/server-utils/src/ai/google-genai/index.ts +++ b/packages/server-utils/src/ai/google-genai/index.ts @@ -250,6 +250,46 @@ export function addResponseAttributes(span: Span, response: GoogleGenAIResponse, } } +function asConfigObject(config: unknown): Record | undefined { + return config && typeof config === 'object' ? (config as Record) : undefined; +} + +/** + * Merge the parameters captured at `chats.create()` time onto a `chat.sendMessage()` / + * `chat.sendMessageStream()` call, so the config that is defined once on the chat lands on every + * message span. The SDK resolves the request config as `params.config ?? chat.config`, so a + * per-message config replaces the create-time config wholesale and the create-time config only + * applies when the message omits one. Only `model` and `config` are inherited; the create `history` + * is intentionally left off the message spans. See issue #20086. + */ +function mergeChatCreateParams( + chatCreateParams: Record | undefined, + callParams: Record | undefined, +): Record | undefined { + if (!chatCreateParams) { + return callParams; + } + + const merged: Record = { ...callParams }; + + if (!('model' in merged) && 'model' in chatCreateParams) { + merged.model = chatCreateParams.model; + } + + // @google/genai sends `params.config ?? chat.config`, so a per-message config replaces the + // create-time config wholesale rather than merging into it. Fall back to the create-time config + // only when the message did not carry one, otherwise the span reports fields that were not sent. + const callConfig = asConfigObject(callParams?.config); + if (!callConfig) { + const createConfig = asConfigObject(chatCreateParams.config); + if (createConfig) { + merged.config = createConfig; + } + } + + return merged; +} + /** * Instrument any async or synchronous genai method with Sentry spans * Handles operations like models.generateContent and chat.sendMessage and chats.create @@ -261,6 +301,7 @@ function instrumentMethod( instrumentedMethod: InstrumentedMethodEntry, context: unknown, options: GoogleGenAIOptions, + chatCreateParams?: Record, ): (...args: T) => R | Promise { const isEmbeddings = instrumentedMethod.operation === 'embeddings'; @@ -268,7 +309,9 @@ function instrumentMethod( apply(target, _, args: T): R | Promise { const operationName = instrumentedMethod.operation || 'unknown'; const params = args[0] as Record | undefined; - const requestAttributes = extractRequestAttributes(operationName, params, context); + // The chat config is set once on chats.create() and reused for every message, so weld it on. + const attributeParams = mergeChatCreateParams(chatCreateParams, params); + const requestAttributes = extractRequestAttributes(operationName, attributeParams, context); const model = requestAttributes[GEN_AI_REQUEST_MODEL] ?? 'unknown'; // Check if this is a streaming method @@ -282,8 +325,8 @@ function instrumentMethod( }, async (span: Span) => { try { - if (options.recordInputs && params) { - addPrivateRequestAttributes(span, params, operationName); + if (options.recordInputs && attributeParams) { + addPrivateRequestAttributes(span, attributeParams, operationName); } const stream = await target.apply(context, args); return instrumentStream(stream, span, Boolean(options.recordOutputs)) as R; @@ -310,8 +353,8 @@ function instrumentMethod( attributes: requestAttributes, }, (span: Span) => { - if (options.recordInputs && params) { - addPrivateRequestAttributes(span, params, operationName); + if (options.recordInputs && attributeParams) { + addPrivateRequestAttributes(span, attributeParams, operationName); } return handleCallbackErrors( @@ -339,7 +382,12 @@ function instrumentMethod( * Create a deep proxy for Google GenAI client instrumentation * Recursively instruments methods and handles special cases like chats.create */ -function createDeepProxy(target: T, currentPath = '', options: GoogleGenAIOptions): T { +function createDeepProxy( + target: T, + currentPath = '', + options: GoogleGenAIOptions, + chatCreateParams?: Record, +): T { return new Proxy(target, { get: (t, prop, receiver) => { const value = Reflect.get(t, prop, receiver); @@ -350,7 +398,14 @@ function createDeepProxy(target: T, currentPath = '', options: if (typeof value === 'function' && instrumentedMethod) { // If an operation is specified, we need to instrument the method itself const wrappedMethod = instrumentedMethod.operation - ? instrumentMethod(value as (...args: unknown[]) => unknown, methodPath, instrumentedMethod, t, options) + ? instrumentMethod( + value as (...args: unknown[]) => unknown, + methodPath, + instrumentedMethod, + t, + options, + chatCreateParams, + ) : value.bind(t); if (!instrumentedMethod.proxyResultPath) { @@ -364,7 +419,14 @@ function createDeepProxy(target: T, currentPath = '', options: return function (...args: unknown[]): unknown { const result = wrappedMethod(...args); if (result && typeof result === 'object') { - return createDeepProxy(result as object, instrumentedMethod.proxyResultPath, options); + // The result (e.g. a chat object from chats.create()) carries the config passed here, so + // hand those params down to instrument its message methods with them (issue #20086). + return createDeepProxy( + result as object, + instrumentedMethod.proxyResultPath, + options, + args[0] as Record | undefined, + ); } return result; }; @@ -376,7 +438,7 @@ function createDeepProxy(target: T, currentPath = '', options: } if (value && typeof value === 'object') { - return createDeepProxy(value, methodPath, options); + return createDeepProxy(value, methodPath, options, chatCreateParams); } return value; diff --git a/packages/server-utils/test/ai/lib/tracing/google-genai.test.ts b/packages/server-utils/test/ai/lib/tracing/google-genai.test.ts new file mode 100644 index 000000000000..9493ca9a40f1 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/google-genai.test.ts @@ -0,0 +1,180 @@ +import { + GEN_AI_INPUT_MESSAGES, + GEN_AI_OPERATION_NAME, + GEN_AI_PROVIDER_NAME, + GEN_AI_REQUEST_FREQUENCY_PENALTY, + GEN_AI_REQUEST_MAX_TOKENS, + GEN_AI_REQUEST_MODEL, + GEN_AI_REQUEST_PRESENCE_PENALTY, + GEN_AI_REQUEST_TEMPERATURE, + GEN_AI_REQUEST_TOP_K, + GEN_AI_REQUEST_TOP_P, + GEN_AI_SYSTEM_INSTRUCTIONS, + GEN_AI_TOOL_DEFINITIONS, +} from '@sentry/conventions/attributes'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { getMainCarrier, setCurrentClient, spanToJSON } from '@sentry/core'; +import type { Span } from '@sentry/core'; +import { instrumentGoogleGenAIClient } from '../../../../src/ai/google-genai'; +import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; + +const MODEL = 'gemini-1.5-pro'; + +const CHAT_CONFIG = { + temperature: 0.8, + topP: 0.9, + topK: 40, + maxOutputTokens: 150, + frequencyPenalty: 0.5, + presencePenalty: 0.3, + tools: [{ functionDeclarations: [{ name: 'getWeather' }] }], + systemInstruction: 'You are a friendly robot.', +}; + +const MOCK_RESPONSE = { + modelVersion: MODEL, + candidates: [{ content: { parts: [{ text: 'Hi there!' }], role: 'model' } }], + usageMetadata: { promptTokenCount: 8, candidatesTokenCount: 12, totalTokenCount: 20 }, +}; + +/** + * Minimal stand-in for a `@google/genai` client. `chats.create()` returns a chat object that keeps + * the model (as the real SDK does) and exposes the two message-sending methods. The config passed to + * `create()` is stored on the SDK internally and is not repeated on each `sendMessage()` call. + */ +function createFakeClient(): { chats: { create: (params: Record) => unknown } } { + return { + chats: { + create: (params: Record) => ({ + model: params.model, + sendMessage: async (_params: Record) => MOCK_RESPONSE, + sendMessageStream: async (_params: Record) => + (async function* () { + yield MOCK_RESPONSE; + })(), + }), + }, + }; +} + +function setupClient(): Span[] { + const client = new TestClient( + getDefaultTestClientOptions({ + dsn: 'https://public@dsn.ingest.sentry.io/1337', + tracesSampleRate: 1, + }), + ); + setCurrentClient(client); + client.init(); + + const endedSpans: Span[] = []; + client.on('spanEnd', span => endedSpans.push(span)); + return endedSpans; +} + +async function drain(stream: AsyncIterable): Promise { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + for await (const _ of stream) { + // consume so the streaming span is ended + } +} + +describe('instrumentGoogleGenAIClient chat config propagation (#20086)', () => { + beforeEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + afterEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + it('welds chats.create() config onto chat.sendMessage() spans', async () => { + const endedSpans = setupClient(); + const instrumented = instrumentGoogleGenAIClient(createFakeClient()); + + const chat = instrumented.chats.create({ model: MODEL, config: CHAT_CONFIG }) as { + sendMessage: (params: Record) => Promise; + }; + await chat.sendMessage({ message: 'Tell me a joke' }); + + expect(endedSpans).toHaveLength(1); + const data = spanToJSON(endedSpans[0]!).data; + + expect(data[GEN_AI_OPERATION_NAME]).toBe('chat'); + expect(data[GEN_AI_PROVIDER_NAME]).toBe('google_genai'); + expect(data[GEN_AI_REQUEST_MODEL]).toBe(MODEL); + expect(data[GEN_AI_REQUEST_TEMPERATURE]).toBe(0.8); + expect(data[GEN_AI_REQUEST_TOP_P]).toBe(0.9); + expect(data[GEN_AI_REQUEST_TOP_K]).toBe(40); + expect(data[GEN_AI_REQUEST_MAX_TOKENS]).toBe(150); + expect(data[GEN_AI_REQUEST_FREQUENCY_PENALTY]).toBe(0.5); + expect(data[GEN_AI_REQUEST_PRESENCE_PENALTY]).toBe(0.3); + expect(data[GEN_AI_TOOL_DEFINITIONS]).toBe('[{"name":"getWeather"}]'); + expect(data[GEN_AI_SYSTEM_INSTRUCTIONS]).toBe('[{"type":"text","content":"You are a friendly robot."}]'); + // The chat message stays as the only input message; the system instruction is split out above. + expect(data[GEN_AI_INPUT_MESSAGES]).toBe('[{"role":"user","content":"Tell me a joke"}]'); + }); + + it('welds chats.create() config onto chat.sendMessageStream() spans', async () => { + const endedSpans = setupClient(); + const instrumented = instrumentGoogleGenAIClient(createFakeClient()); + + const chat = instrumented.chats.create({ model: MODEL, config: CHAT_CONFIG }) as { + sendMessageStream: (params: Record) => Promise>; + }; + await drain(await chat.sendMessageStream({ message: 'Tell me a joke' })); + + expect(endedSpans).toHaveLength(1); + const data = spanToJSON(endedSpans[0]!).data; + + expect(data[GEN_AI_OPERATION_NAME]).toBe('chat'); + expect(data[GEN_AI_REQUEST_MODEL]).toBe(MODEL); + expect(data[GEN_AI_REQUEST_TEMPERATURE]).toBe(0.8); + expect(data[GEN_AI_REQUEST_TOP_P]).toBe(0.9); + expect(data[GEN_AI_REQUEST_TOP_K]).toBe(40); + expect(data[GEN_AI_REQUEST_MAX_TOKENS]).toBe(150); + expect(data[GEN_AI_REQUEST_FREQUENCY_PENALTY]).toBe(0.5); + expect(data[GEN_AI_REQUEST_PRESENCE_PENALTY]).toBe(0.3); + expect(data[GEN_AI_TOOL_DEFINITIONS]).toBe('[{"name":"getWeather"}]'); + expect(data[GEN_AI_SYSTEM_INSTRUCTIONS]).toBe('[{"type":"text","content":"You are a friendly robot."}]'); + }); + + it('replaces the chats.create() config when a message provides its own', async () => { + const endedSpans = setupClient(); + const instrumented = instrumentGoogleGenAIClient(createFakeClient()); + + const chat = instrumented.chats.create({ model: MODEL, config: CHAT_CONFIG }) as { + sendMessage: (params: Record) => Promise; + }; + await chat.sendMessage({ message: 'Tell me a joke', config: { temperature: 0.1 } }); + + const data = spanToJSON(endedSpans[0]!).data; + // @google/genai resolves the request config as `params.config ?? chat.config`, so the + // per-message config is sent on its own. The create-time fields it omits are not part of the + // request, so they must not appear on the span. + expect(data[GEN_AI_REQUEST_TEMPERATURE]).toBe(0.1); + expect(data[GEN_AI_REQUEST_TOP_P]).toBeUndefined(); + expect(data[GEN_AI_REQUEST_MAX_TOKENS]).toBeUndefined(); + expect(data[GEN_AI_TOOL_DEFINITIONS]).toBeUndefined(); + expect(data[GEN_AI_SYSTEM_INSTRUCTIONS]).toBeUndefined(); + }); + + it('does not leak chat config onto models.generateContent spans', async () => { + const endedSpans = setupClient(); + const client = { + chats: { create: createFakeClient().chats.create }, + models: { generateContent: async (_params: Record) => MOCK_RESPONSE }, + }; + const instrumented = instrumentGoogleGenAIClient(client); + + // Create a chat (with config) first, then make an unrelated generateContent call. + instrumented.chats.create({ model: MODEL, config: CHAT_CONFIG }); + await instrumented.models.generateContent({ model: 'gemini-1.5-flash' }); + + const genContentSpan = endedSpans.find(span => spanToJSON(span).data[GEN_AI_OPERATION_NAME] === 'generate_content'); + const data = spanToJSON(genContentSpan!).data; + expect(data[GEN_AI_REQUEST_MODEL]).toBe('gemini-1.5-flash'); + expect(data[GEN_AI_REQUEST_TEMPERATURE]).toBeUndefined(); + expect(data[GEN_AI_SYSTEM_INSTRUCTIONS]).toBeUndefined(); + }); +});