From d6e734608bc9f88ea59c9b107ba2841930797138 Mon Sep 17 00:00:00 2001 From: zkasuran Date: Fri, 14 Aug 2026 14:58:57 +0530 Subject: [PATCH] fix(server-utils): Deduplicate Google GenAI streaming tool calls The streaming handler recorded every tool call twice: once from `chunk.functionCalls` and again from the `functionCall` parts of each candidate. `functionCalls` is an SDK getter over those same parts, so one real tool call produced two span entries with mismatched shapes (one keyed by the non-spec `args`, one by `arguments`). Take tool calls only from `chunk.functionCalls`, the same source the non-streaming path uses, so each call is recorded once in one shape. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/ai/google-genai/streaming.ts | 12 +- .../tracing/google-genai-streaming.test.ts | 128 ++++++++++++++++++ 2 files changed, 132 insertions(+), 8 deletions(-) create mode 100644 packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts diff --git a/packages/server-utils/src/ai/google-genai/streaming.ts b/packages/server-utils/src/ai/google-genai/streaming.ts index 646884b62e4b..22aa1032daea 100644 --- a/packages/server-utils/src/ai/google-genai/streaming.ts +++ b/packages/server-utils/src/ai/google-genai/streaming.ts @@ -68,6 +68,10 @@ function handleResponseMetadata(chunk: GoogleGenAIResponse, state: StreamingStat * @param recordOutputs - Whether to record outputs */ function handleCandidateContent(chunk: GoogleGenAIResponse, state: StreamingState, recordOutputs: boolean): void { + // `chunk.functionCalls` is the SDK accessor over the candidate's function-call parts, so it is the + // single source of truth for tool calls. Also reading `part.functionCall` from those same parts + // would record every call twice, with two different shapes. This mirrors the non-streaming path, + // which likewise takes tool calls from `response.functionCalls`. if (Array.isArray(chunk.functionCalls)) { state.toolCalls.push(...chunk.functionCalls); } @@ -79,14 +83,6 @@ function handleCandidateContent(chunk: GoogleGenAIResponse, state: StreamingStat for (const part of candidate?.content?.parts ?? []) { if (recordOutputs && part.text) state.responseTexts.push(part.text); - if (part.functionCall) { - state.toolCalls.push({ - type: 'function', - id: part.functionCall.id, - name: part.functionCall.name, - arguments: part.functionCall.args, - }); - } } } } diff --git a/packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts b/packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts new file mode 100644 index 000000000000..924dc7b8df18 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts @@ -0,0 +1,128 @@ +import { describe, expect, it } from 'vitest'; +import type { Span } from '@sentry/core'; +import { GEN_AI_RESPONSE_TOOL_CALLS } from '@sentry/conventions/attributes'; +import { addResponseAttributes } from '../../../../src/ai/google-genai/index'; +import { instrumentStream } from '../../../../src/ai/google-genai/streaming'; +import type { ContentPart, GoogleGenAIResponse } from '../../../../src/ai/google-genai/types'; + +function createMockSpan(): { span: Span; attributes: Record } { + const attributes: Record = {}; + let isEnded = false; + const span = { + isRecording: () => !isEnded, + setAttribute: (key: string, value: unknown) => { + attributes[key] = value; + }, + setAttributes: (attrs: Record) => { + Object.assign(attributes, attrs); + }, + setStatus: () => {}, + end: () => { + isEnded = true; + }, + } as unknown as Span; + return { span, attributes }; +} + +// Mirrors the real `@google/genai` response object: `functionCalls` is a getter that reads the +// function-call parts of the first candidate, i.e. the very same parts exposed under +// `candidates[].content.parts`. A chunk therefore surfaces each tool call through both accessors. +function chunkWithParts(parts: ContentPart[], extra: Record = {}): GoogleGenAIResponse { + const chunk = { + candidates: [{ content: { parts, role: 'model' }, index: 0 }], + ...extra, + } as Record; + Object.defineProperty(chunk, 'functionCalls', { + enumerable: false, + get() { + const calls = parts.map(part => part.functionCall).filter(fc => fc !== undefined); + return calls.length ? calls : undefined; + }, + }); + return chunk as GoogleGenAIResponse; +} + +async function* streamOf(chunks: GoogleGenAIResponse[]): AsyncGenerator { + for (const chunk of chunks) { + yield chunk; + } +} + +async function drain(stream: AsyncGenerator): Promise { + for await (const _ of stream) { + void _; + } +} + +describe('instrumentStream (google-genai tool calls)', () => { + it('records exactly one entry per streamed tool call, in the native function-call shape', async () => { + const { span, attributes } = createMockSpan(); + + const functionCall = { + id: 'call_light_stream_1', + name: 'controlLight', + args: { brightness: 0.5, colorTemperature: 'cool' }, + }; + + const chunks = [ + chunkWithParts([{ text: 'Let me control the lights for you.' }], { responseId: 'resp-1' }), + chunkWithParts([{ functionCall }]), + chunkWithParts([{ text: ' Done!' }], { + candidates: [{ content: { parts: [{ text: ' Done!' }], role: 'model' }, finishReason: 'STOP', index: 0 }], + usageMetadata: { promptTokenCount: 12, candidatesTokenCount: 10, totalTokenCount: 22 }, + }), + ]; + + await drain(instrumentStream(streamOf(chunks), span, true)); + + const raw = attributes[GEN_AI_RESPONSE_TOOL_CALLS]; + expect(raw).toBeDefined(); + const toolCalls = JSON.parse(raw as string) as Array>; + + // A single real tool call must yield a single entry (the pre-fix code double-dipped and produced two). + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0]).toEqual(functionCall); + // No entry should carry the divergent `arguments` shape the duplicate push used to emit. + expect(toolCalls.every(call => !('arguments' in call))).toBe(true); + }); + + it('records one entry per call when several tool calls arrive across chunks', async () => { + const { span, attributes } = createMockSpan(); + + const first = { id: 'c1', name: 'controlLight', args: { brightness: 0.2 } }; + const second = { id: 'c2', name: 'setColor', args: { color: 'red' } }; + + const chunks = [chunkWithParts([{ functionCall: first }]), chunkWithParts([{ functionCall: second }])]; + + await drain(instrumentStream(streamOf(chunks), span, true)); + + const toolCalls = JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string) as Array>; + expect(toolCalls).toEqual([first, second]); + }); + + it('does not record tool calls when recordOutputs is false', async () => { + const { span, attributes } = createMockSpan(); + + const chunks = [chunkWithParts([{ functionCall: { id: 'c1', name: 'controlLight', args: { brightness: 1 } } }])]; + + await drain(instrumentStream(streamOf(chunks), span, false)); + + expect(attributes[GEN_AI_RESPONSE_TOOL_CALLS]).toBeUndefined(); + }); + + it('non-streaming addResponseAttributes still records one entry per call in the same shape', () => { + const { span, attributes } = createMockSpan(); + + const functionCall = { + id: 'call_light_control_1', + name: 'controlLight', + args: { brightness: 0.3, colorTemperature: 'warm' }, + }; + const response = chunkWithParts([{ text: 'I need to check the light status first.' }, { functionCall }]); + + addResponseAttributes(span, response, true); + + const toolCalls = JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string) as Array>; + expect(toolCalls).toEqual([functionCall]); + }); +});