feat: Capture grounding metadata for Google GenAI

e2e/scenarios/google-genai-instrumentation/assertions.ts

@@ -6,7 +6,11 @@ import {
   resolveFileSnapshotPath,
 } from "../../helpers/file-snapshot";
 import { withScenarioHarness } from "../../helpers/scenario-harness";
-import { findChildSpans, findLatestSpan } from "../../helpers/trace-selectors";
+import {
+  findChildSpans,
+  findLatestChildSpan,
+  findLatestSpan,
+} from "../../helpers/trace-selectors";
 import { summarizeWrapperContract } from "../../helpers/wrapper-contract";
 
 import { GOOGLE_MODEL, ROOT_NAME, SCENARIO_NAME } from "./scenario.impl.mjs";
@@ -33,7 +37,9 @@ function findGoogleSpan(
   names: string[],
 ) {
   for (const name of names) {
-    const span = findChildSpans(events, name, parentId)[0];
+    const span =
+      findLatestChildSpan(events, name, parentId) ??
+      findChildSpans(events, name, parentId)[0];
     if (span) {
       return span;
     }
@@ -63,6 +69,35 @@ function isRecord(value: Json | undefined): value is Record<string, Json> {
   return typeof value === "object" && value !== null && !Array.isArray(value);
 }
 
+function extractGroundingMetadataFromOutput(
+  output: Record<string, unknown> | undefined,
+): Record<string, unknown> | undefined {
+  if (!output) {
+    return undefined;
+  }
+
+  if (isRecord(output.groundingMetadata as Json)) {
+    return output.groundingMetadata as Record<string, unknown>;
+  }
+
+  const candidates = output.candidates;
+  if (!Array.isArray(candidates)) {
+    return undefined;
+  }
+
+  for (const candidate of candidates) {
+    if (!isRecord(candidate as Json)) {
+      continue;
+    }
+
+    if (isRecord(candidate.groundingMetadata as Json)) {
+      return candidate.groundingMetadata as Record<string, unknown>;
+    }
+  }
+
+  return undefined;
+}
+
 function normalizeGoogleVariableTokenCounts(value: Json): Json {
   if (Array.isArray(value)) {
     return value.map((entry) =>
@@ -425,6 +460,62 @@ export function defineGoogleGenAIInstrumentationAssertions(options: {
       },
     );
 
+    test("captures grounding metadata for generateContent", testConfig, () => {
+      const operation = findLatestSpan(
+        events,
+        "google-grounded-generate-operation",
+      );
+      const span = findGoogleSpan(events, operation?.span.id, [
+        "generate_content",
+        "google-genai.generateContent",
+      ]);
+      const metadata = span?.row.metadata as
+        | Record<string, unknown>
+        | undefined;
+      const output = span?.output as Record<string, unknown> | undefined;
+      const metadataGrounding = metadata?.groundingMetadata as
+        | Record<string, unknown>
+        | undefined;
+      const outputGrounding = extractGroundingMetadataFromOutput(output);
+
+      expect(operation).toBeDefined();
+      expect(span).toBeDefined();
+      expect(metadataGrounding).toBeDefined();
+      expect(outputGrounding).toBeDefined();
+      expect(Array.isArray(metadataGrounding?.webSearchQueries)).toBe(true);
+      expect(Array.isArray(outputGrounding?.webSearchQueries)).toBe(true);
+    });
+
+    test(
+      "captures grounding metadata for generateContentStream",
+      testConfig,
+      () => {
+        const operation = findLatestSpan(
+          events,
+          "google-grounded-stream-operation",
+        );
+        const span = findGoogleSpan(events, operation?.span.id, [
+          "generate_content_stream",
+          "google-genai.generateContentStream",
+        ]);
+        const metadata = span?.row.metadata as
+          | Record<string, unknown>
+          | undefined;
+        const output = span?.output as Record<string, unknown> | undefined;
+        const metadataGrounding = metadata?.groundingMetadata as
+          | Record<string, unknown>
+          | undefined;
+        const outputGrounding = extractGroundingMetadataFromOutput(output);
+
+        expect(operation).toBeDefined();
+        expect(span).toBeDefined();
+        expect(metadataGrounding).toBeDefined();
+        expect(outputGrounding).toBeDefined();
+        expect(Array.isArray(metadataGrounding?.webSearchQueries)).toBe(true);
+        expect(Array.isArray(outputGrounding?.webSearchQueries)).toBe(true);
+      },
+    );
+
     test("captures trace for tool calling", testConfig, () => {
       const root = findLatestSpan(events, ROOT_NAME);
       const operation = findLatestSpan(events, "google-tool-operation");

e2e/scenarios/google-genai-instrumentation/scenario.impl.mjs

@@ -7,12 +7,13 @@ import {
 } from "../../helpers/provider-runtime.mjs";
 
 const GOOGLE_MODEL = "gemini-2.5-flash-lite";
+const GOOGLE_GROUNDING_MODEL = "gemini-2.0-flash";
 const ROOT_NAME = "google-genai-instrumentation-root";
 const SCENARIO_NAME = "google-genai-instrumentation";
 const GOOGLE_GENAI_RETRY_OPTIONS = {
-  attempts: 3,
+  attempts: 4,
   delayMs: 1_000,
-  maxDelayMs: 5_000,
+  maxDelayMs: 8_000,
   shouldRetry: isRetriableGoogleGenAIError,
 };
 const WEATHER_TOOL = {
@@ -33,6 +34,9 @@ const WEATHER_TOOL = {
     },
   ],
 };
+const GOOGLE_SEARCH_TOOL = {
+  googleSearch: {},
+};
 
 function isObject(value) {
   return value !== null && typeof value === "object" && !Array.isArray(value);
@@ -62,15 +66,24 @@ function getRetryStatus(error) {
 
 function isRetriableGoogleGenAIError(error) {
   const status = getRetryStatus(error);
-  if (status === 429 || status === 500 || status === 503 || status === 504) {
+  if (
+    status === 408 ||
+    status === 429 ||
+    status === 500 ||
+    status === 502 ||
+    status === 503 ||
+    status === 504
+  ) {
     return true;
   }
 
   const message = error instanceof Error ? error.message : String(error ?? "");
+  const normalizedMessage = message.toLowerCase();
   return (
-    message.includes("request timed out") ||
-    message.includes("UNAVAILABLE") ||
-    message.includes("temporarily unavailable")
+    normalizedMessage.includes("request timed out") ||
+    normalizedMessage.includes("timed out") ||
+    normalizedMessage.includes("unavailable") ||
+    normalizedMessage.includes("high demand")
   );
 }
 
@@ -192,6 +205,45 @@ async function runGoogleGenAIInstrumentationScenario(sdk, options = {}) {
         },
       );
 
+      await runOperation(
+        "google-grounded-generate-operation",
+        "grounded-generate",
+        async () => {
+          await withRetry(async () => {
+            await client.models.generateContent({
+              model: GOOGLE_GROUNDING_MODEL,
+              contents:
+                "Use Google Search grounding and answer in one sentence: What is the current population of Paris, France?",
+              config: {
+                maxOutputTokens: 256,
+                temperature: 0,
+                tools: [GOOGLE_SEARCH_TOOL],
+              },
+            });
+          }, GOOGLE_GENAI_RETRY_OPTIONS);
+        },
+      );
+
+      await runOperation(
+        "google-grounded-stream-operation",
+        "grounded-stream",
+        async () => {
+          await withRetry(async () => {
+            const stream = await client.models.generateContentStream({
+              model: GOOGLE_GROUNDING_MODEL,
+              contents:
+                "Use Google Search grounding and answer in one sentence: What is the current weather in Paris?",
+              config: {
+                maxOutputTokens: 256,
+                temperature: 0,
+                tools: [GOOGLE_SEARCH_TOOL],
+              },
+            });
+            await collectAsync(stream);
+          }, GOOGLE_GENAI_RETRY_OPTIONS);
+        },
+      );
+
       await runOperation("google-tool-operation", "tool", async () => {
         await withRetry(async () => {
           await client.models.generateContent({

js/src/instrumentation/plugins/google-genai-plugin.ts

@@ -142,7 +142,9 @@ export class GoogleGenAIPlugin extends BasePlugin {
           }
 
           try {
+            const responseMetadata = extractResponseMetadata(event.result);
             spanState.span.log({
+              ...(responseMetadata ? { metadata: responseMetadata } : {}),
               metrics: cleanMetrics(
                 extractGenerateContentMetrics(
                   event.result,
@@ -331,7 +333,11 @@ function patchGoogleGenAIStreamingResult(args: {
 
     if (options.result) {
       const { end, ...metricsWithoutEnd } = options.result.metrics;
+      const responseMetadata = extractResponseMetadata(
+        options.result.aggregated,
+      );
       span.log({
+        ...(responseMetadata ? { metadata: responseMetadata } : {}),
         metrics: cleanMetrics(metricsWithoutEnd),
         output: options.result.aggregated,
       });
@@ -740,6 +746,7 @@ function aggregateGenerateContentChunks(
   let text = "";
   let thoughtText = "";
   const otherParts: Record<string, unknown>[] = [];
+  let groundingMetadata: unknown = undefined;
   let usageMetadata: GoogleGenAIUsageMetadata | null = null;
   let lastResponse: GoogleGenAIGenerateContentResponse | null = null;
 
@@ -749,6 +756,9 @@ function aggregateGenerateContentChunks(
     if (chunk.usageMetadata) {
       usageMetadata = chunk.usageMetadata;
     }
+    if (chunk.groundingMetadata !== undefined) {
+      groundingMetadata = chunk.groundingMetadata;
+    }
 
     if (chunk.candidates && Array.isArray(chunk.candidates)) {
       for (const candidate of chunk.candidates) {
@@ -799,6 +809,12 @@ function aggregateGenerateContentChunks(
       if (candidate.finishReason !== undefined) {
         candidateDict.finishReason = candidate.finishReason;
       }
+      if (candidate.groundingMetadata !== undefined) {
+        candidateDict.groundingMetadata = candidate.groundingMetadata;
+        if (groundingMetadata === undefined) {
+          groundingMetadata = candidate.groundingMetadata;
+        }
+      }
       if (candidate.safetyRatings) {
         candidateDict.safetyRatings = candidate.safetyRatings;
       }
@@ -812,6 +828,9 @@ function aggregateGenerateContentChunks(
     aggregated.usageMetadata = usageMetadata;
     populateUsageMetrics(metrics, usageMetadata);
   }
+  if (groundingMetadata !== undefined) {
+    aggregated.groundingMetadata = groundingMetadata;
+  }
 
   if (text) {
     aggregated.text = text;
@@ -830,6 +849,38 @@ function cleanMetrics(metrics: Record<string, number>): Record<string, number> {
   return cleaned;
 }
 
+function extractResponseMetadata(
+  response: unknown,
+): Record<string, unknown> | undefined {
+  const responseDict = tryToDict(response);
+  if (!responseDict) {
+    return undefined;
+  }
+
+  const metadata: Record<string, unknown> = {};
+  const responseGroundingMetadata = responseDict.groundingMetadata;
+  const candidateGroundingMetadata: unknown[] = [];
+
+  if (Array.isArray(responseDict.candidates)) {
+    for (const candidate of responseDict.candidates) {
+      const candidateDict = tryToDict(candidate);
+      if (candidateDict?.groundingMetadata !== undefined) {
+        candidateGroundingMetadata.push(candidateDict.groundingMetadata);
+      }
+    }
+  }
+
+  if (responseGroundingMetadata !== undefined) {
+    metadata.groundingMetadata = responseGroundingMetadata;
+  } else if (candidateGroundingMetadata.length === 1) {
+    [metadata.groundingMetadata] = candidateGroundingMetadata;
+  } else if (candidateGroundingMetadata.length > 1) {
+    metadata.groundingMetadata = candidateGroundingMetadata;
+  }
+
+  return Object.keys(metadata).length > 0 ? metadata : undefined;
+}
+
 /**
  * Helper to convert objects to dictionaries.
  */

js/src/vendor-sdk-types/google-genai.ts

@@ -69,13 +69,30 @@ export interface GoogleGenAIGenerateContentResponse {
       role?: string;
     };
     finishReason?: string;
+    groundingMetadata?: GoogleGenAIGroundingMetadata;
     safetyRatings?: Record<string, unknown>[];
   }[];
+  groundingMetadata?: GoogleGenAIGroundingMetadata;
   usageMetadata?: GoogleGenAIUsageMetadata;
   text?: string;
   [key: string]: unknown;
 }
 
+export interface GoogleGenAIGroundingMetadata {
+  groundingChunks?: Array<{
+    web?: {
+      title?: string;
+      uri?: string;
+      [key: string]: unknown;
+    };
+    [key: string]: unknown;
+  }>;
+  groundingSupports?: Record<string, unknown>[];
+  searchEntryPoint?: Record<string, unknown>;
+  webSearchQueries?: string[];
+  [key: string]: unknown;
+}
+
 export interface GoogleGenAIUsageMetadata {
   promptTokenCount?: number;
   candidatesTokenCount?: number;