From 5ec1251dd66da8c24df66d02420fb1a31c7ed030 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 23 Mar 2026 12:22:54 +0100 Subject: [PATCH 1/4] fix(js): Align Google GenAI auto traces with wrappers --- e2e/helpers/google-genai-trace-contract.ts | 104 +- .../__snapshots__/log-payloads.json | 915 ------------------ .../__snapshots__/span-events.json | 191 ---- .../scenario.test.ts | 54 -- .../__snapshots__/log-payloads.json | 4 +- .../__snapshots__/span-events.json | 4 +- .../scenario.test.ts | 157 ++- .../plugins/google-genai-plugin.ts | 541 +++++++++-- 8 files changed, 713 insertions(+), 1257 deletions(-) delete mode 100644 e2e/scenarios/google-genai-auto-instrumentation-node-hook/__snapshots__/log-payloads.json delete mode 100644 e2e/scenarios/google-genai-auto-instrumentation-node-hook/__snapshots__/span-events.json delete mode 100644 e2e/scenarios/google-genai-auto-instrumentation-node-hook/scenario.test.ts diff --git a/e2e/helpers/google-genai-trace-contract.ts b/e2e/helpers/google-genai-trace-contract.ts index 10eb3a30d..e1b284272 100644 --- a/e2e/helpers/google-genai-trace-contract.ts +++ b/e2e/helpers/google-genai-trace-contract.ts @@ -3,6 +3,7 @@ import { normalizeForSnapshot, type Json } from "./normalize"; import type { CapturedLogEvent, CapturedLogPayload, + CapturedLogRow, } from "./mock-braintrust-server"; import { findChildSpans, findLatestSpan } from "./trace-selectors"; import { @@ -48,11 +49,98 @@ function normalizeGooglePayloads(payloadRows: unknown[]): unknown[] { }); } +function normalizeGoogleSnapshotSummary(options: { + event: CapturedLogEvent; + metadataKeys: string[]; + rootSpanId: string | undefined; + rootName: string; + snapshotRootName?: string; + snapshotScenarioName?: string; +}): Json { + const summary = summarizeWrapperContract( + options.event, + options.metadataKeys, + ) as + | { + metadata?: Record | null; + name?: string | null; + } + | Json; + + if ( + options.event.span.id === options.rootSpanId && + summary && + typeof summary === "object" && + !Array.isArray(summary) + ) { + if (options.snapshotRootName && summary.name === options.rootName) { + summary.name = options.snapshotRootName; + } + + if ( + options.snapshotScenarioName && + summary.metadata && + typeof summary.metadata === "object" && + !Array.isArray(summary.metadata) + ) { + summary.metadata = { + ...summary.metadata, + scenario: options.snapshotScenarioName, + }; + } + } + + return summary as Json; +} + +function normalizeGoogleSnapshotPayloadRows(options: { + payloadRows: CapturedLogRow[]; + rootSpanId: string | undefined; + snapshotRootName?: string; + snapshotScenarioName?: string; +}): CapturedLogRow[] { + return options.payloadRows.map((row) => { + if (row.span_id !== options.rootSpanId) { + return row; + } + + const normalizedRow = structuredClone(row); + + if ( + options.snapshotScenarioName && + normalizedRow.metadata && + typeof normalizedRow.metadata === "object" && + !Array.isArray(normalizedRow.metadata) + ) { + normalizedRow.metadata = { + ...normalizedRow.metadata, + scenario: options.snapshotScenarioName, + }; + } + + if ( + options.snapshotRootName && + normalizedRow.span_attributes && + typeof normalizedRow.span_attributes === "object" && + !Array.isArray(normalizedRow.span_attributes) + ) { + normalizedRow.span_attributes = { + ...normalizedRow.span_attributes, + name: options.snapshotRootName, + }; + } + + return normalizedRow; + }); +} + export function assertGoogleGenAITraceContract(options: { capturedEvents: CapturedLogEvent[]; payloads: CapturedLogPayload[]; rootName: string; scenarioName: string; + snapshotRootName?: string; + snapshotScenarioName?: string; }): { payloadSummary: Json; spanSummary: Json } { const root = findLatestSpan(options.capturedEvents, options.rootName); const generateOperation = findLatestSpan( @@ -206,12 +294,24 @@ export function assertGoogleGenAITraceContract(options: { toolOperation, toolSpan, ].map((event) => - summarizeWrapperContract(event!, ["model", "operation", "scenario"]), + normalizeGoogleSnapshotSummary({ + event: event!, + metadataKeys: ["model", "operation", "scenario"], + rootName: options.rootName, + rootSpanId: root?.span.id, + snapshotRootName: options.snapshotRootName, + snapshotScenarioName: options.snapshotScenarioName, + }), ) as Json, ), payloadSummary: normalizeForSnapshot( normalizeGooglePayloads( - payloadRowsForRootSpan(options.payloads, root?.span.id), + normalizeGoogleSnapshotPayloadRows({ + payloadRows: payloadRowsForRootSpan(options.payloads, root?.span.id), + rootSpanId: root?.span.id, + snapshotRootName: options.snapshotRootName, + snapshotScenarioName: options.snapshotScenarioName, + }), ) as Json, ), }; diff --git a/e2e/scenarios/google-genai-auto-instrumentation-node-hook/__snapshots__/log-payloads.json b/e2e/scenarios/google-genai-auto-instrumentation-node-hook/__snapshots__/log-payloads.json deleted file mode 100644 index 082a95b8d..000000000 --- a/e2e/scenarios/google-genai-auto-instrumentation-node-hook/__snapshots__/log-payloads.json +++ /dev/null @@ -1,915 +0,0 @@ -[ - { - "_is_merge": false, - "context": { - "caller_filename": "/e2e/helpers/provider-runtime.mjs", - "caller_functionname": "runTracedScenario", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "log_id": "g", - "metadata": { - "scenario": "google-genai-auto-instrumentation-node-hook", - "testRunId": "" - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 0, - "name": "google-genai-auto-hook-root", - "type": "task" - }, - "span_id": "" - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "" - }, - { - "_is_merge": false, - "context": { - "caller_filename": "/e2e/helpers/provider-runtime.mjs", - "caller_functionname": "runOperation", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "log_id": "g", - "metadata": { - "operation": "generate", - "testRunId": "" - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 1, - "name": "google-generate-operation" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "context": {}, - "created": "", - "id": "", - "input": { - "config": { - "maxOutputTokens": 16, - "temperature": 0 - }, - "contents": { - "text": "Reply with exactly PARIS." - }, - "model": "gemini-2.0-flash-001" - }, - "log_id": "g", - "metadata": { - "maxOutputTokens": 16, - "model": "gemini-2.0-flash-001", - "provider": "google-genai", - "temperature": 0 - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 2, - "name": "google-genai.generateContent", - "type": "llm" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "completion_tokens": 3, - "duration": 0, - "prompt_tokens": 5, - "tokens": 8 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "output": { - "candidates": [ - { - "avgLogprobs": 0, - "content": { - "parts": [ - { - "text": "PARIS\n" - } - ], - "role": "model" - }, - "finishReason": "STOP" - } - ], - "modelVersion": "gemini-2.0-flash-001", - "responseId": "", - "sdkHttpResponse": { - "headers": { - "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", - "content-encoding": "gzip", - "content-type": "application/json; charset=UTF-8", - "date": "", - "server": "scaffolding on HTTPServer2", - "server-timing": "", - "transfer-encoding": "chunked", - "vary": "Origin, X-Origin, Referer", - "x-content-type-options": "nosniff", - "x-frame-options": "SAMEORIGIN", - "x-gemini-service-tier": "", - "x-xss-protection": "0" - } - }, - "usageMetadata": { - "candidatesTokenCount": 3, - "candidatesTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 3 - } - ], - "promptTokenCount": 5, - "promptTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 5 - } - ], - "totalTokenCount": 8 - } - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": false, - "context": { - "caller_filename": "/e2e/helpers/provider-runtime.mjs", - "caller_functionname": "runOperation", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "log_id": "g", - "metadata": { - "operation": "attachment", - "testRunId": "" - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 3, - "name": "google-attachment-operation" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "context": {}, - "created": "", - "id": "", - "input": { - "config": { - "maxOutputTokens": 24, - "temperature": 0 - }, - "contents": [ - { - "parts": [ - { - "image_url": { - "url": { - "content_type": "image/png", - "filename": "file.png", - "key": "", - "type": "braintrust_attachment" - } - } - }, - { - "text": "Describe the attached image in one short sentence." - } - ], - "role": "user" - } - ], - "model": "gemini-2.0-flash-001" - }, - "log_id": "g", - "metadata": { - "maxOutputTokens": 24, - "model": "gemini-2.0-flash-001", - "provider": "google-genai", - "temperature": 0 - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 4, - "name": "google-genai.generateContent", - "type": "llm" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "completion_tokens": 16, - "duration": 0, - "prompt_tokens": 1299, - "tokens": 1315 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "output": { - "candidates": [ - { - "avgLogprobs": 0, - "content": { - "parts": [ - { - "text": "A tall ship braves a stormy sea with lightning flashing in the background.\n" - } - ], - "role": "model" - }, - "finishReason": "STOP" - } - ], - "modelVersion": "gemini-2.0-flash-001", - "responseId": "", - "sdkHttpResponse": { - "headers": { - "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", - "content-encoding": "gzip", - "content-type": "application/json; charset=UTF-8", - "date": "", - "server": "scaffolding on HTTPServer2", - "server-timing": "", - "transfer-encoding": "chunked", - "vary": "Origin, X-Origin, Referer", - "x-content-type-options": "nosniff", - "x-frame-options": "SAMEORIGIN", - "x-gemini-service-tier": "", - "x-xss-protection": "0" - } - }, - "usageMetadata": { - "candidatesTokenCount": 16, - "candidatesTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 16 - } - ], - "promptTokenCount": 1299, - "promptTokensDetails": [ - { - "modality": "IMAGE", - "tokenCount": 1290 - }, - { - "modality": "TEXT", - "tokenCount": 9 - } - ], - "totalTokenCount": 1315 - } - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": false, - "context": { - "caller_filename": "/e2e/helpers/provider-runtime.mjs", - "caller_functionname": "runOperation", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "log_id": "g", - "metadata": { - "operation": "stream", - "testRunId": "" - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 5, - "name": "google-stream-operation" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "context": {}, - "created": "", - "id": "", - "input": { - "config": { - "maxOutputTokens": 32, - "temperature": 0 - }, - "contents": { - "text": "Count from 1 to 3 and include the words one two three." - }, - "model": "gemini-2.0-flash-001" - }, - "log_id": "g", - "metadata": { - "maxOutputTokens": 32, - "model": "gemini-2.0-flash-001", - "provider": "google-genai", - "temperature": 0 - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 6, - "name": "google-genai.generateContentStream", - "type": "llm" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "completion_tokens": 7, - "duration": 0, - "prompt_tokens": 15, - "time_to_first_token": 0, - "tokens": 22 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "output": { - "candidates": [ - { - "content": { - "parts": [ - { - "text": "One, two, three.\n" - } - ], - "role": "model" - }, - "finishReason": "STOP" - } - ], - "text": "One, two, three.\n", - "usageMetadata": { - "candidatesTokenCount": 7, - "candidatesTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 7 - } - ], - "promptTokenCount": 15, - "promptTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 15 - } - ], - "totalTokenCount": 22 - } - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": false, - "context": { - "caller_filename": "/e2e/helpers/provider-runtime.mjs", - "caller_functionname": "runOperation", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "log_id": "g", - "metadata": { - "operation": "stream-return", - "testRunId": "" - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 7, - "name": "google-stream-return-operation" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "context": {}, - "created": "", - "id": "", - "input": { - "config": { - "maxOutputTokens": 48, - "temperature": 0 - }, - "contents": { - "text": "Write a short poem about Paris." - }, - "model": "gemini-2.0-flash-001" - }, - "log_id": "g", - "metadata": { - "maxOutputTokens": 48, - "model": "gemini-2.0-flash-001", - "provider": "google-genai", - "temperature": 0 - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 8, - "name": "google-genai.generateContentStream", - "type": "llm" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "duration": 0, - "prompt_tokens": 8, - "time_to_first_token": 0, - "tokens": 8 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "output": { - "candidates": [ - { - "content": { - "parts": [ - { - "text": "City" - } - ], - "role": "model" - } - } - ], - "text": "City", - "usageMetadata": { - "promptTokenCount": 8, - "promptTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 8 - } - ], - "totalTokenCount": 8 - } - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": false, - "context": { - "caller_filename": "/e2e/helpers/provider-runtime.mjs", - "caller_functionname": "runOperation", - "caller_lineno": 0 - }, - "created": "", - "id": "", - "log_id": "g", - "metadata": { - "operation": "tool", - "testRunId": "" - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 9, - "name": "google-tool-operation" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "context": {}, - "created": "", - "id": "", - "input": { - "config": { - "maxOutputTokens": 128, - "temperature": 0, - "toolConfig": { - "functionCallingConfig": { - "allowedFunctionNames": [ - "get_weather" - ], - "mode": "ANY" - } - }, - "tools": [ - { - "functionDeclarations": [ - { - "description": "Get the current weather in a given location", - "name": "get_weather", - "parametersJsonSchema": { - "properties": { - "location": { - "description": "The city and state or city and country", - "type": "string" - } - }, - "required": [ - "location" - ], - "type": "object" - } - } - ] - } - ] - }, - "contents": { - "text": "Use the get_weather function for Paris, France. Do not answer from memory." - }, - "model": "gemini-2.0-flash-001" - }, - "log_id": "g", - "metadata": { - "maxOutputTokens": 128, - "model": "gemini-2.0-flash-001", - "provider": "google-genai", - "temperature": 0, - "toolConfig": { - "functionCallingConfig": { - "allowedFunctionNames": [ - "get_weather" - ], - "mode": "ANY" - } - } - }, - "metrics": { - "start": 0 - }, - "project_id": "", - "root_span_id": "", - "span_attributes": { - "exec_counter": 10, - "name": "google-genai.generateContent", - "type": "llm" - }, - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "completion_tokens": 7, - "duration": 0, - "prompt_tokens": 28, - "tokens": 35 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "output": { - "candidates": [ - { - "avgLogprobs": 0, - "content": { - "parts": [ - { - "functionCall": { - "args": { - "location": "Paris, France" - }, - "name": "get_weather" - } - } - ], - "role": "model" - }, - "finishReason": "STOP" - } - ], - "modelVersion": "gemini-2.0-flash-001", - "responseId": "", - "sdkHttpResponse": { - "headers": { - "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", - "content-encoding": "gzip", - "content-type": "application/json; charset=UTF-8", - "date": "", - "server": "scaffolding on HTTPServer2", - "server-timing": "", - "transfer-encoding": "chunked", - "vary": "Origin, X-Origin, Referer", - "x-content-type-options": "nosniff", - "x-frame-options": "SAMEORIGIN", - "x-gemini-service-tier": "", - "x-xss-protection": "0" - } - }, - "usageMetadata": { - "candidatesTokenCount": 7, - "candidatesTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 7 - } - ], - "promptTokenCount": 28, - "promptTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 28 - } - ], - "totalTokenCount": 35 - } - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - }, - { - "_is_merge": true, - "id": "", - "log_id": "g", - "metrics": { - "end": 0 - }, - "project_id": "", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ] - } -] diff --git a/e2e/scenarios/google-genai-auto-instrumentation-node-hook/__snapshots__/span-events.json b/e2e/scenarios/google-genai-auto-instrumentation-node-hook/__snapshots__/span-events.json deleted file mode 100644 index 562d46fc4..000000000 --- a/e2e/scenarios/google-genai-auto-instrumentation-node-hook/__snapshots__/span-events.json +++ /dev/null @@ -1,191 +0,0 @@ -[ - { - "has_input": false, - "has_output": false, - "metadata": { - "scenario": "google-genai-auto-instrumentation-node-hook" - }, - "metric_keys": [], - "name": "google-genai-auto-hook-root", - "root_span_id": "", - "span_id": "", - "span_parents": [], - "type": "task" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "generate" - }, - "metric_keys": [], - "name": "google-generate-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "gemini-2.0-flash-001" - }, - "metric_keys": [ - "completion_tokens", - "duration", - "prompt_tokens", - "tokens" - ], - "name": "google-genai.generateContent", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "attachment" - }, - "metric_keys": [], - "name": "google-attachment-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "gemini-2.0-flash-001" - }, - "metric_keys": [ - "completion_tokens", - "duration", - "prompt_tokens", - "tokens" - ], - "name": "google-genai.generateContent", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "stream" - }, - "metric_keys": [], - "name": "google-stream-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "gemini-2.0-flash-001" - }, - "metric_keys": [ - "completion_tokens", - "duration", - "prompt_tokens", - "time_to_first_token", - "tokens" - ], - "name": "google-genai.generateContentStream", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "stream-return" - }, - "metric_keys": [], - "name": "google-stream-return-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "gemini-2.0-flash-001" - }, - "metric_keys": [ - "duration", - "prompt_tokens", - "time_to_first_token", - "tokens" - ], - "name": "google-genai.generateContentStream", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "tool" - }, - "metric_keys": [], - "name": "google-tool-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "gemini-2.0-flash-001" - }, - "metric_keys": [ - "completion_tokens", - "duration", - "prompt_tokens", - "tokens" - ], - "name": "google-genai.generateContent", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - } -] diff --git a/e2e/scenarios/google-genai-auto-instrumentation-node-hook/scenario.test.ts b/e2e/scenarios/google-genai-auto-instrumentation-node-hook/scenario.test.ts deleted file mode 100644 index d01bf4eec..000000000 --- a/e2e/scenarios/google-genai-auto-instrumentation-node-hook/scenario.test.ts +++ /dev/null @@ -1,54 +0,0 @@ -import { expect, test } from "vitest"; -import { assertGoogleGenAITraceContract } from "../../helpers/google-genai-trace-contract"; -import { - formatJsonFileSnapshot, - resolveFileSnapshotPath, -} from "../../helpers/file-snapshot"; -import { - prepareScenarioDir, - resolveScenarioDir, - withScenarioHarness, -} from "../../helpers/scenario-harness"; -import { E2E_TAGS } from "../../helpers/tags"; - -const scenarioDir = await prepareScenarioDir({ - scenarioDir: resolveScenarioDir(import.meta.url), -}); -const TIMEOUT_MS = 90_000; - -test( - "google genai auto-instrumentation via node hook collects the shared google genai trace contract", - { - tags: [E2E_TAGS.externalApi], - timeout: TIMEOUT_MS, - }, - async () => { - await withScenarioHarness( - async ({ events, payloads, runNodeScenarioDir }) => { - await runNodeScenarioDir({ - nodeArgs: ["--import", "braintrust/hook.mjs"], - scenarioDir, - timeoutMs: TIMEOUT_MS, - }); - - const contract = assertGoogleGenAITraceContract({ - capturedEvents: events(), - payloads: payloads(), - rootName: "google-genai-auto-hook-root", - scenarioName: "google-genai-auto-instrumentation-node-hook", - }); - - await expect( - formatJsonFileSnapshot(contract.spanSummary), - ).toMatchFileSnapshot( - resolveFileSnapshotPath(import.meta.url, "span-events.json"), - ); - await expect( - formatJsonFileSnapshot(contract.payloadSummary), - ).toMatchFileSnapshot( - resolveFileSnapshotPath(import.meta.url, "log-payloads.json"), - ); - }, - ); - }, -); diff --git a/e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/log-payloads.json b/e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/log-payloads.json index e712905b6..3458ff0c1 100644 --- a/e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/log-payloads.json +++ b/e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/log-payloads.json @@ -10,7 +10,7 @@ "id": "", "log_id": "g", "metadata": { - "scenario": "wrap-google-genai-content-traces", + "scenario": "google-genai-traces", "testRunId": "" }, "metrics": { @@ -20,7 +20,7 @@ "root_span_id": "", "span_attributes": { "exec_counter": 0, - "name": "google-genai-wrapper-root", + "name": "google-genai-root", "type": "task" }, "span_id": "" diff --git a/e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/span-events.json b/e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/span-events.json index e250f11d9..0eae373d5 100644 --- a/e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/span-events.json +++ b/e2e/scenarios/wrap-google-genai-content-traces/__snapshots__/span-events.json @@ -3,10 +3,10 @@ "has_input": false, "has_output": false, "metadata": { - "scenario": "wrap-google-genai-content-traces" + "scenario": "google-genai-traces" }, "metric_keys": [], - "name": "google-genai-wrapper-root", + "name": "google-genai-root", "root_span_id": "", "span_id": "", "span_parents": [], diff --git a/e2e/scenarios/wrap-google-genai-content-traces/scenario.test.ts b/e2e/scenarios/wrap-google-genai-content-traces/scenario.test.ts index 1b1916187..44736f435 100644 --- a/e2e/scenarios/wrap-google-genai-content-traces/scenario.test.ts +++ b/e2e/scenarios/wrap-google-genai-content-traces/scenario.test.ts @@ -11,38 +11,151 @@ import { } from "../../helpers/scenario-harness"; import { E2E_TAGS } from "../../helpers/tags"; -const scenarioDir = await prepareScenarioDir({ - scenarioDir: resolveScenarioDir(import.meta.url), -}); const TIMEOUT_MS = 90_000; +const SHARED_ROOT_NAME = "google-genai-root"; +const SHARED_SCENARIO_NAME = "google-genai-traces"; +const sharedSpanSnapshotPath = resolveFileSnapshotPath( + import.meta.url, + "span-events.json", +); +const sharedPayloadSnapshotPath = resolveFileSnapshotPath( + import.meta.url, + "log-payloads.json", +); + +type GoogleGenAIContract = ReturnType; + +let wrapperScenarioDirPromise: Promise | undefined; +let autoScenarioDirPromise: Promise | undefined; +let wrapperContractPromise: Promise | undefined; +let autoContractPromise: Promise | undefined; + +function getWrapperScenarioDir(): Promise { + wrapperScenarioDirPromise ??= prepareScenarioDir({ + scenarioDir: resolveScenarioDir(import.meta.url), + }); + + return wrapperScenarioDirPromise; +} + +function getAutoScenarioDir(): Promise { + autoScenarioDirPromise ??= prepareScenarioDir({ + scenarioDir: resolveScenarioDir( + new URL( + "../google-genai-auto-instrumentation-node-hook/scenario.mjs", + import.meta.url, + ).href, + ), + }); + + return autoScenarioDirPromise; +} + +function getWrapperContract(): Promise { + wrapperContractPromise ??= (async () => { + let contract: GoogleGenAIContract | undefined; + const wrapperScenarioDir = await getWrapperScenarioDir(); -test( - "wrap-google-genai-content-traces captures generate, attachment, stream, early-return, and tool spans", - { - tags: [E2E_TAGS.externalApi], - timeout: TIMEOUT_MS, - }, - async () => { await withScenarioHarness(async ({ events, payloads, runScenarioDir }) => { - await runScenarioDir({ scenarioDir, timeoutMs: TIMEOUT_MS }); + await runScenarioDir({ + scenarioDir: wrapperScenarioDir, + timeoutMs: TIMEOUT_MS, + }); - const contract = assertGoogleGenAITraceContract({ + contract = assertGoogleGenAITraceContract({ capturedEvents: events(), payloads: payloads(), rootName: "google-genai-wrapper-root", scenarioName: "wrap-google-genai-content-traces", + snapshotRootName: SHARED_ROOT_NAME, + snapshotScenarioName: SHARED_SCENARIO_NAME, }); + }); - await expect( - formatJsonFileSnapshot(contract.spanSummary), - ).toMatchFileSnapshot( - resolveFileSnapshotPath(import.meta.url, "span-events.json"), - ); - await expect( - formatJsonFileSnapshot(contract.payloadSummary), - ).toMatchFileSnapshot( - resolveFileSnapshotPath(import.meta.url, "log-payloads.json"), + if (!contract) { + throw new Error("Failed to capture Google GenAI wrapper contract"); + } + + return contract; + })(); + + return wrapperContractPromise; +} + +function getAutoContract(): Promise { + autoContractPromise ??= (async () => { + let contract: GoogleGenAIContract | undefined; + const autoScenarioDir = await getAutoScenarioDir(); + + await withScenarioHarness( + async ({ events, payloads, runNodeScenarioDir }) => { + await runNodeScenarioDir({ + nodeArgs: ["--import", "braintrust/hook.mjs"], + scenarioDir: autoScenarioDir, + timeoutMs: TIMEOUT_MS, + }); + + contract = assertGoogleGenAITraceContract({ + capturedEvents: events(), + payloads: payloads(), + rootName: "google-genai-auto-hook-root", + scenarioName: "google-genai-auto-instrumentation-node-hook", + snapshotRootName: SHARED_ROOT_NAME, + snapshotScenarioName: SHARED_SCENARIO_NAME, + }); + }, + ); + + if (!contract) { + throw new Error( + "Failed to capture Google GenAI auto-instrumentation contract", ); - }); + } + + return contract; + })(); + + return autoContractPromise; +} + +test( + "wrap-google-genai-content-traces captures the shared Google GenAI trace contract", + { + tags: [E2E_TAGS.externalApi], + timeout: TIMEOUT_MS, + }, + async () => { + const contract = await getWrapperContract(); + + await expect( + formatJsonFileSnapshot(contract.spanSummary), + ).toMatchFileSnapshot(sharedSpanSnapshotPath); + await expect( + formatJsonFileSnapshot(contract.payloadSummary), + ).toMatchFileSnapshot(sharedPayloadSnapshotPath); + }, +); + +test( + "google genai auto-instrumentation via node hook matches wrapper instrumentation", + { + tags: [E2E_TAGS.externalApi], + timeout: TIMEOUT_MS, + }, + async () => { + const [wrapperContract, autoContract] = await Promise.all([ + getWrapperContract(), + getAutoContract(), + ]); + + expect(autoContract.spanSummary).toEqual(wrapperContract.spanSummary); + expect(autoContract.payloadSummary).toEqual(wrapperContract.payloadSummary); + + await expect( + formatJsonFileSnapshot(autoContract.spanSummary), + ).toMatchFileSnapshot(sharedSpanSnapshotPath); + await expect( + formatJsonFileSnapshot(autoContract.payloadSummary), + ).toMatchFileSnapshot(sharedPayloadSnapshotPath); }, ); diff --git a/js/src/instrumentation/plugins/google-genai-plugin.ts b/js/src/instrumentation/plugins/google-genai-plugin.ts index 21a65525e..2fe17ae95 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.ts @@ -1,10 +1,22 @@ import { BasePlugin } from "../core"; +import { unsubscribeAll } from "../core/channel-tracing"; +import type { + ChannelMessage, + ErrorOf, + StartOf, +} from "../core/channel-definitions"; +import type { + IsoAsyncLocalStorage, + IsoChannelHandlers, + IsoTracingChannel, +} from "../../isomorph"; import { - traceAsyncChannel, - traceStreamingChannel, - unsubscribeAll, -} from "../core/channel-tracing"; -import { Attachment } from "../../logger"; + _internalGetGlobalState, + Attachment, + BRAINTRUST_CURRENT_SPAN_STORE, + startSpan, + type Span, +} from "../../logger"; import { SpanTypeAttribute } from "../../../util/index"; import { getCurrentUnixTimestamp } from "../../util"; import { googleGenAIChannels } from "./google-genai-channels"; @@ -16,6 +28,26 @@ import type { GoogleGenAIUsageMetadata, } from "../../vendor-sdk-types/google-genai"; +type GenerateContentChannel = typeof googleGenAIChannels.generateContent; +type GenerateContentStreamChannel = + typeof googleGenAIChannels.generateContentStream; +type GenerateContentStreamEvent = + ChannelMessage & { + googleGenAIInput?: Record; + googleGenAIMetadata?: Record; + }; + +type SpanState = { + span: Span; + startTime: number; +}; + +const GOOGLE_GENAI_INTERNAL_CONTEXT = { + caller_filename: "", + caller_functionname: "", + caller_lineno: 0, +}; + /** * Auto-instrumentation plugin for the Google GenAI SDK. * @@ -40,55 +72,409 @@ export class GoogleGenAIPlugin extends BasePlugin { } private subscribeToGoogleGenAIChannels(): void { - // GenerativeModel.generateContent (non-streaming) - this.unsubscribers.push( - traceAsyncChannel(googleGenAIChannels.generateContent, { - name: "google-genai.generateContent", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const input = serializeInput(params); - const metadata = extractMetadata(params); - return { + this.subscribeToGenerateContentChannel(); + this.subscribeToGenerateContentStreamChannel(); + } + + private subscribeToGenerateContentChannel(): void { + const tracingChannel = + googleGenAIChannels.generateContent.tracingChannel() as IsoTracingChannel< + ChannelMessage + >; + const states = new WeakMap(); + const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( + tracingChannel, + states, + (event) => { + const params = event.arguments[0]; + const input = serializeInput(params); + const metadata = extractMetadata(params); + const span = startSpan({ + name: "generate_content", + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + event: { + context: GOOGLE_GENAI_INTERNAL_CONTEXT, input, - metadata: { ...metadata, provider: "google-genai" }, - }; - }, - extractOutput: (result) => { - return result; - }, - extractMetrics: (result, startTime) => { - return extractGenerateContentMetrics(result, startTime); - }, - }), + metadata, + }, + }); + + return { + span, + startTime: getCurrentUnixTimestamp(), + }; + }, ); - // GenerativeModel.generateContentStream (streaming) - this.unsubscribers.push( - traceStreamingChannel(googleGenAIChannels.generateContentStream, { - name: "google-genai.generateContentStream", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const input = serializeInput(params); - const metadata = extractMetadata(params); - return { - input, - metadata: { ...metadata, provider: "google-genai" }, - }; - }, - extractOutput: (result) => { - return result; + const handlers: IsoChannelHandlers> = + { + start: (event) => { + ensureSpanState(states, event, () => { + const params = event.arguments[0]; + const input = serializeInput(params); + const metadata = extractMetadata(params); + const span = startSpan({ + name: "generate_content", + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + event: { + context: GOOGLE_GENAI_INTERNAL_CONTEXT, + input, + metadata, + }, + }); + + return { + span, + startTime: getCurrentUnixTimestamp(), + }; + }); }, - extractMetrics: () => { - return {}; + asyncEnd: (event) => { + const spanState = states.get(event as object); + if (!spanState) { + return; + } + + try { + spanState.span.log({ + metrics: cleanMetrics( + extractGenerateContentMetrics( + event.result, + spanState.startTime, + ), + ), + output: event.result, + }); + } finally { + spanState.span.end(); + states.delete(event as object); + } }, - aggregateChunks: (chunks, _result, _endEvent, startTime) => { - return aggregateGenerateContentChunks(chunks, startTime); + error: (event) => { + logErrorAndEndSpan(states, event as ErrorOf); }, - }), - ); + }; + + tracingChannel.subscribe(handlers); + this.unsubscribers.push(() => { + unbindCurrentSpanStore?.(); + tracingChannel.unsubscribe(handlers); + }); + } + + private subscribeToGenerateContentStreamChannel(): void { + const tracingChannel = + googleGenAIChannels.generateContentStream.tracingChannel() as IsoTracingChannel< + ChannelMessage + >; + + const handlers: IsoChannelHandlers< + ChannelMessage + > = { + start: (event) => { + const streamEvent = event as GenerateContentStreamEvent; + const params = event.arguments[0]; + streamEvent.googleGenAIInput = serializeInput(params); + streamEvent.googleGenAIMetadata = extractMetadata(params); + }, + asyncEnd: (event) => { + const streamEvent = event as GenerateContentStreamEvent; + patchGoogleGenAIStreamingResult({ + input: streamEvent.googleGenAIInput, + metadata: streamEvent.googleGenAIMetadata, + result: streamEvent.result, + }); + }, + error: () => {}, + }; + + tracingChannel.subscribe(handlers); + this.unsubscribers.push(() => { + tracingChannel.unsubscribe(handlers); + }); } } +function ensureSpanState( + states: WeakMap, + event: TEvent, + create: () => SpanState, +): SpanState { + const existing = states.get(event); + if (existing) { + return existing; + } + + const created = create(); + states.set(event, created); + return created; +} + +function bindCurrentSpanStoreToStart( + tracingChannel: IsoTracingChannel>, + states: WeakMap, + create: (event: StartOf) => SpanState, +): (() => void) | undefined { + const state = _internalGetGlobalState(); + const startChannel = tracingChannel.start as + | ({ + bindStore?: ( + store: IsoAsyncLocalStorage, + callback: (event: ChannelMessage) => Span, + ) => void; + unbindStore?: (store: IsoAsyncLocalStorage) => void; + } & object) + | undefined; + const currentSpanStore = state?.contextManager + ? ( + state.contextManager as { + [BRAINTRUST_CURRENT_SPAN_STORE]?: IsoAsyncLocalStorage; + } + )[BRAINTRUST_CURRENT_SPAN_STORE] + : undefined; + + if (!startChannel?.bindStore || !currentSpanStore) { + return undefined; + } + + startChannel.bindStore( + currentSpanStore, + (event) => + ensureSpanState(states, event as object, () => + create(event as StartOf), + ).span, + ); + + return () => { + startChannel.unbindStore?.(currentSpanStore); + }; +} + +function logErrorAndEndSpan( + states: WeakMap, + event: ErrorOf, +): void { + const spanState = states.get(event as object); + if (!spanState) { + return; + } + + spanState.span.log({ + error: event.error.message, + }); + spanState.span.end(); + states.delete(event as object); +} + +function patchGoogleGenAIStreamingResult(args: { + input: Record | undefined; + metadata: Record | undefined; + result: unknown; +}): boolean { + const { input, metadata, result } = args; + + if ( + !input || + !metadata || + !result || + typeof result !== "object" || + typeof (result as AsyncIterator) + .next !== "function" + ) { + return false; + } + + const chunks: GoogleGenAIGenerateContentResponse[] = []; + let firstTokenTime: number | null = null; + let finalized = false; + let span: Span | null = null; + let startTime: number | null = null; + + const ensureSpan = () => { + if (!span) { + span = startSpan({ + name: "generate_content_stream", + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + event: { + input, + metadata, + }, + }); + startTime = getCurrentUnixTimestamp(); + } + + return span; + }; + + const finalize = (options: { + error?: unknown; + result?: { + aggregated: Record; + metrics: Record; + }; + }) => { + if (finalized || !span) { + return; + } + + finalized = true; + + if (options.result) { + const { end, ...metricsWithoutEnd } = options.result.metrics; + span.log({ + metrics: cleanMetrics(metricsWithoutEnd), + output: options.result.aggregated, + }); + span.end(typeof end === "number" ? { endTime: end } : undefined); + return; + } + + if (options.error !== undefined) { + span.log({ + error: + options.error instanceof Error + ? options.error.message + : String(options.error), + }); + } + + span.end(); + }; + + const patchIterator = ( + iterator: AsyncIterator, + ): AsyncIterator => { + if ( + typeof iterator !== "object" || + iterator === null || + "__braintrustGoogleGenAIPatched" in (iterator as object) + ) { + return iterator; + } + + const iteratorRecord = + iterator as AsyncIterator & + Record; + const originalNext = + typeof iteratorRecord.next === "function" + ? iteratorRecord.next.bind(iterator) + : undefined; + const originalReturn = + typeof iteratorRecord.return === "function" + ? iteratorRecord.return.bind(iterator) + : undefined; + const originalThrow = + typeof iteratorRecord.throw === "function" + ? iteratorRecord.throw.bind(iterator) + : undefined; + const originalAsyncIterator = + typeof iteratorRecord[Symbol.asyncIterator] === "function" + ? iteratorRecord[Symbol.asyncIterator].bind(iterator) + : undefined; + + Object.defineProperty(iteratorRecord, "__braintrustGoogleGenAIPatched", { + configurable: true, + enumerable: false, + value: true, + writable: false, + }); + + if (originalNext) { + iteratorRecord.next = async (...nextArgs: [] | [unknown]) => { + ensureSpan(); + + try { + const nextResult = (await originalNext( + ...nextArgs, + )) as IteratorResult; + + if (!nextResult.done && nextResult.value) { + if (firstTokenTime === null) { + firstTokenTime = getCurrentUnixTimestamp(); + } + chunks.push(nextResult.value); + } + + if (nextResult.done && startTime !== null) { + finalize({ + result: aggregateGenerateContentChunks( + chunks, + startTime, + firstTokenTime, + ), + }); + } + + return nextResult; + } catch (error) { + finalize({ error }); + throw error; + } + }; + } + + if (originalReturn) { + iteratorRecord.return = async (...returnArgs: [] | [unknown]) => { + ensureSpan(); + + try { + return (await originalReturn( + ...returnArgs, + )) as IteratorResult; + } finally { + if (startTime !== null) { + finalize({ + result: + chunks.length > 0 + ? aggregateGenerateContentChunks( + chunks, + startTime, + firstTokenTime, + ) + : undefined, + }); + } else { + finalize({}); + } + } + }; + } + + if (originalThrow) { + iteratorRecord.throw = async (...throwArgs: [] | [unknown]) => { + ensureSpan(); + + try { + return (await originalThrow( + ...throwArgs, + )) as IteratorResult; + } catch (error) { + finalize({ error }); + throw error; + } + }; + } + + iteratorRecord[Symbol.asyncIterator] = () => { + const asyncIterator = originalAsyncIterator + ? (originalAsyncIterator() as AsyncIterator) + : iterator; + return patchIterator(asyncIterator); + }; + + return iterator; + }; + + patchIterator(result as AsyncIterator); + return true; +} + /** * Serialize input parameters for Google GenAI API calls. */ @@ -103,11 +489,13 @@ function serializeInput( if (params.config) { const config = tryToDict(params.config); if (config) { - const tools = serializeTools(params); - if (tools) { - config.tools = tools; - } - input.config = config; + const filteredConfig: Record = {}; + Object.keys(config).forEach((key) => { + if (key !== "tools") { + filteredConfig[key] = config[key]; + } + }); + input.config = filteredConfig; } } @@ -255,6 +643,11 @@ function extractMetadata( } } + const tools = serializeTools(params); + if (tools) { + metadata.tools = tools; + } + return metadata; } @@ -267,8 +660,10 @@ function extractGenerateContentMetrics( ): Record { const metrics: Record = {}; - if (startTime) { + if (startTime !== undefined) { const end = getCurrentUnixTimestamp(); + metrics.start = startTime; + metrics.end = end; metrics.duration = end - startTime; } @@ -305,27 +700,25 @@ function populateUsageMetrics( */ function aggregateGenerateContentChunks( chunks: GoogleGenAIGenerateContentResponse[], - startTime?: number, + startTime: number, + firstTokenTime: number | null, ): { - output: Record; + aggregated: Record; metrics: Record; } { - const metrics: Record = {}; - - if (startTime !== undefined) { - const end = getCurrentUnixTimestamp(); - metrics.duration = end - startTime; - } - - let firstTokenTime: number | null = null; + const end = getCurrentUnixTimestamp(); + const metrics: Record = { + start: startTime, + end, + duration: end - startTime, + }; - if (chunks.length > 0 && firstTokenTime === null && startTime !== undefined) { - firstTokenTime = getCurrentUnixTimestamp(); + if (firstTokenTime !== null) { metrics.time_to_first_token = firstTokenTime - startTime; } if (chunks.length === 0) { - return { output: {}, metrics }; + return { aggregated: {}, metrics }; } let text = ""; @@ -366,7 +759,7 @@ function aggregateGenerateContentChunks( } } - const output: Record = {}; + const aggregated: Record = {}; const parts: Record[] = []; if (thoughtText) { @@ -396,19 +789,29 @@ function aggregateGenerateContentChunks( candidates.push(candidateDict); } - output.candidates = candidates; + aggregated.candidates = candidates; } if (usageMetadata) { - output.usageMetadata = usageMetadata; + aggregated.usageMetadata = usageMetadata; populateUsageMetrics(metrics, usageMetadata); } if (text) { - output.text = text; + aggregated.text = text; } - return { output, metrics }; + return { aggregated, metrics }; +} + +function cleanMetrics(metrics: Record): Record { + const cleaned: Record = {}; + for (const [key, value] of Object.entries(metrics)) { + if (value !== null && value !== undefined) { + cleaned[key] = value; + } + } + return cleaned; } /** From d4bcaf1c38513fe2234d7a2c133910cc552704cb Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 23 Mar 2026 12:47:22 +0100 Subject: [PATCH 2/4] . --- .../plugins/google-genai-plugin.ts | 51 +++++++++++++------ 1 file changed, 35 insertions(+), 16 deletions(-) diff --git a/js/src/instrumentation/plugins/google-genai-plugin.ts b/js/src/instrumentation/plugins/google-genai-plugin.ts index 2fe17ae95..55cb70f6c 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.ts @@ -15,6 +15,7 @@ import { Attachment, BRAINTRUST_CURRENT_SPAN_STORE, startSpan, + type StartSpanArgs, type Span, } from "../../logger"; import { SpanTypeAttribute } from "../../../util/index"; @@ -48,6 +49,17 @@ const GOOGLE_GENAI_INTERNAL_CONTEXT = { caller_lineno: 0, }; +function createWrapperParityEvent(args: { + input: Record; + metadata: Record; +}): StartSpanArgs["event"] { + return { + context: GOOGLE_GENAI_INTERNAL_CONTEXT, + input: args.input, + metadata: args.metadata, + } as StartSpanArgs["event"]; +} + /** * Auto-instrumentation plugin for the Google GenAI SDK. * @@ -94,11 +106,7 @@ export class GoogleGenAIPlugin extends BasePlugin { spanAttributes: { type: SpanTypeAttribute.LLM, }, - event: { - context: GOOGLE_GENAI_INTERNAL_CONTEXT, - input, - metadata, - }, + event: createWrapperParityEvent({ input, metadata }), }); return { @@ -120,11 +128,7 @@ export class GoogleGenAIPlugin extends BasePlugin { spanAttributes: { type: SpanTypeAttribute.LLM, }, - event: { - context: GOOGLE_GENAI_INTERNAL_CONTEXT, - input, - metadata, - }, + event: createWrapperParityEvent({ input, metadata }), }); return { @@ -363,19 +367,34 @@ function patchGoogleGenAIStreamingResult(args: { Record; const originalNext = typeof iteratorRecord.next === "function" - ? iteratorRecord.next.bind(iterator) + ? ( + iteratorRecord.next as ( + ...args: [] | [undefined] + ) => Promise> + ).bind(iterator) : undefined; const originalReturn = typeof iteratorRecord.return === "function" - ? iteratorRecord.return.bind(iterator) + ? ( + iteratorRecord.return as ( + ...args: [] | [unknown] + ) => Promise> + ).bind(iterator) : undefined; const originalThrow = typeof iteratorRecord.throw === "function" - ? iteratorRecord.throw.bind(iterator) + ? ( + iteratorRecord.throw as ( + ...args: [] | [unknown] + ) => Promise> + ).bind(iterator) : undefined; + const asyncIteratorMethod = iteratorRecord[Symbol.asyncIterator]; const originalAsyncIterator = - typeof iteratorRecord[Symbol.asyncIterator] === "function" - ? iteratorRecord[Symbol.asyncIterator].bind(iterator) + typeof asyncIteratorMethod === "function" + ? ( + asyncIteratorMethod as () => AsyncIterator + ).bind(iterator) : undefined; Object.defineProperty(iteratorRecord, "__braintrustGoogleGenAIPatched", { @@ -386,7 +405,7 @@ function patchGoogleGenAIStreamingResult(args: { }); if (originalNext) { - iteratorRecord.next = async (...nextArgs: [] | [unknown]) => { + iteratorRecord.next = async (...nextArgs: [] | [undefined]) => { ensureSpan(); try { From 2d4e8d6617b489cacaf6aadcc68fa6492edb484f Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 23 Mar 2026 15:47:38 +0100 Subject: [PATCH 3/4] test(e2e): normalize google cache metric --- e2e/helpers/google-genai-trace-contract.ts | 44 ++++++++++++++++++---- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/e2e/helpers/google-genai-trace-contract.ts b/e2e/helpers/google-genai-trace-contract.ts index d8bd8f295..cea9d3ddd 100644 --- a/e2e/helpers/google-genai-trace-contract.ts +++ b/e2e/helpers/google-genai-trace-contract.ts @@ -33,12 +33,20 @@ function normalizeGooglePayloads(payloadRows: unknown[]): unknown[] { } const row = structuredClone(payload) as { + metrics?: Record; output?: { usageMetadata?: { + cachedContentTokenCount?: number; promptTokensDetails?: Array<{ modality?: string }>; }; }; }; + if (row.metrics) { + delete row.metrics.prompt_cached_tokens; + } + if (row.output?.usageMetadata) { + delete row.output.usageMetadata.cachedContentTokenCount; + } const promptTokensDetails = row.output?.usageMetadata?.promptTokensDetails; if (promptTokensDetails) { promptTokensDetails.sort((left, right) => @@ -93,6 +101,24 @@ function normalizeGoogleSnapshotSummary(options: { return summary as Json; } +function normalizeGoogleSummaryMetrics(summary: Json): Json { + if ( + !summary || + typeof summary !== "object" || + Array.isArray(summary) || + !Array.isArray(summary.metric_keys) + ) { + return summary; + } + + return { + ...summary, + metric_keys: summary.metric_keys.filter( + (metric): metric is string => metric !== "prompt_cached_tokens", + ), + } satisfies Json; +} + function normalizeGoogleSnapshotPayloadRows(options: { payloadRows: CapturedLogRow[]; rootSpanId: string | undefined; @@ -294,14 +320,16 @@ export function assertGoogleGenAITraceContract(options: { toolOperation, toolSpan, ].map((event) => - normalizeGoogleSnapshotSummary({ - event: event!, - metadataKeys: ["model", "operation", "scenario"], - rootName: options.rootName, - rootSpanId: root?.span.id, - snapshotRootName: options.snapshotRootName, - snapshotScenarioName: options.snapshotScenarioName, - }), + normalizeGoogleSummaryMetrics( + normalizeGoogleSnapshotSummary({ + event: event!, + metadataKeys: ["model", "operation", "scenario"], + rootName: options.rootName, + rootSpanId: root?.span.id, + snapshotRootName: options.snapshotRootName, + snapshotScenarioName: options.snapshotScenarioName, + }), + ), ) as Json, ), payloadSummary: normalizeForSnapshot( From 15e6c6d346140077ecb3576f624d4cdbfe8a9608 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 23 Mar 2026 16:22:22 +0100 Subject: [PATCH 4/4] fix --- .../package.json | 1 + e2e/scripts/run-canary-tests.mjs | 17 ++++++++++++----- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/e2e/scenarios/google-genai-auto-instrumentation-node-hook/package.json b/e2e/scenarios/google-genai-auto-instrumentation-node-hook/package.json index 512358ec5..6b9a8d4f8 100644 --- a/e2e/scenarios/google-genai-auto-instrumentation-node-hook/package.json +++ b/e2e/scenarios/google-genai-auto-instrumentation-node-hook/package.json @@ -3,6 +3,7 @@ "private": true, "braintrustScenario": { "canary": { + "testFile": "../wrap-google-genai-content-traces/scenario.test.ts", "dependencies": { "@google/genai": "latest" } diff --git a/e2e/scripts/run-canary-tests.mjs b/e2e/scripts/run-canary-tests.mjs index 419838ced..2c346b413 100644 --- a/e2e/scripts/run-canary-tests.mjs +++ b/e2e/scripts/run-canary-tests.mjs @@ -19,7 +19,7 @@ async function fileExists(filePath) { async function getCanaryTestFiles() { const entries = await readdir(SCENARIOS_DIR, { withFileTypes: true }); - const testFiles = []; + const testFiles = new Set(); for (const entry of entries) { if (!entry.isDirectory()) { @@ -44,17 +44,24 @@ async function getCanaryTestFiles() { continue; } - const testPath = path.join(scenarioDir, "scenario.test.ts"); + const configuredTestFile = manifest?.braintrustScenario?.canary?.testFile; + const testPath = + typeof configuredTestFile === "string" && configuredTestFile.length > 0 + ? path.resolve(scenarioDir, configuredTestFile) + : path.join(scenarioDir, "scenario.test.ts"); if (!(await fileExists(testPath))) { throw new Error( - `Canary scenario ${entry.name} is missing scenario.test.ts`, + `Canary scenario ${entry.name} is missing test file ${path.relative( + E2E_ROOT, + testPath, + )}`, ); } - testFiles.push(path.relative(E2E_ROOT, testPath)); + testFiles.add(path.relative(E2E_ROOT, testPath)); } - return testFiles.sort(); + return [...testFiles].sort(); } async function runVitest(testFiles) {