fix(google-genai): keep the finish reason of a streamed response that has no content

.changeset/keep-google-genai-stream-finish-reason.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix(google-genai): Keep the finish reason of a streamed response that has no content

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

@@ -542,6 +542,98 @@ describe("GoogleGenAIPlugin", () => {
       expect(JSON.stringify(span.log.mock.calls)).not.toContain("aGVsbG8=");
       expect(span.end).toHaveBeenCalledTimes(1);
     });
+
+    it("keeps the finish reason of a stream that ends without content", async () => {
+      plugin.enable();
+      const handlers = subscribeSpy.mock.calls[1][0];
+      const event: any = {
+        arguments: [
+          { contents: "something blocked", model: "gemini-3.8-flash" },
+        ],
+      };
+
+      async function* stream() {
+        yield {
+          candidates: [
+            {
+              finishReason: "SAFETY",
+              safetyRatings: [
+                {
+                  category: "HARM_CATEGORY_DANGEROUS_CONTENT",
+                  probability: "HIGH",
+                },
+              ],
+            },
+          ],
+        };
+        yield { usageMetadata: { promptTokenCount: 3, totalTokenCount: 3 } };
+      }
+
+      handlers.begin(event);
+      event.result = stream();
+      handlers.resolve(event);
+      for await (const _chunk of event.result) {
+        // Consume the provider stream so the instrumentation finalizes it.
+      }
+
+      const span = mockStartSpan.mock.results.at(-1)?.value as {
+        log: ReturnType<typeof vi.fn>;
+      };
+      expect(span.log).toHaveBeenCalledWith(
+        expect.objectContaining({
+          output: {
+            candidates: [
+              {
+                finishReason: "SAFETY",
+                safetyRatings: [
+                  {
+                    category: "HARM_CATEGORY_DANGEROUS_CONTENT",
+                    probability: "HIGH",
+                  },
+                ],
+              },
+            ],
+            usageMetadata: { promptTokenCount: 3, totalTokenCount: 3 },
+          },
+        }),
+      );
+    });
+
+    it("keeps the prompt feedback of a stream whose prompt was blocked", async () => {
+      plugin.enable();
+      const handlers = subscribeSpy.mock.calls[1][0];
+      const event: any = {
+        arguments: [
+          { contents: "something blocked", model: "gemini-3.8-flash" },
+        ],
+      };
+
+      async function* stream() {
+        yield {
+          promptFeedback: { blockReason: "SAFETY" },
+          usageMetadata: { promptTokenCount: 3, totalTokenCount: 3 },
+        };
+      }
+
+      handlers.begin(event);
+      event.result = stream();
+      handlers.resolve(event);
+      for await (const _chunk of event.result) {
+        // Consume the provider stream so the instrumentation finalizes it.
+      }
+
+      const span = mockStartSpan.mock.results.at(-1)?.value as {
+        log: ReturnType<typeof vi.fn>;
+      };
+      expect(span.log).toHaveBeenCalledWith(
+        expect.objectContaining({
+          output: {
+            promptFeedback: { blockReason: "SAFETY" },
+            usageMetadata: { promptTokenCount: 3, totalTokenCount: 3 },
+          },
+        }),
+      );
+    });
   });
 
   describe("generateImages channel interception", () => {

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

@@ -2001,19 +2001,22 @@ function aggregateGenerateContentChunks(
   const otherParts: Record<string, unknown>[] = [];
   let groundingMetadata: unknown = undefined;
   let usageMetadata: GoogleGenAIUsageMetadata | null = null;
-  let lastResponse: GoogleGenAIGenerateContentResponse | null = null;
+  let promptFeedback: GoogleGenAIGenerateContentResponse["promptFeedback"];
+  let lastCandidates: GoogleGenAIGenerateContentResponse["candidates"];
 
   for (const chunk of chunks) {
-    lastResponse = chunk;
-
     if (chunk.usageMetadata) {
       usageMetadata = chunk.usageMetadata;
     }
+    if (chunk.promptFeedback) {
+      promptFeedback = chunk.promptFeedback;
+    }
     if (chunk.groundingMetadata !== undefined) {
       groundingMetadata = chunk.groundingMetadata;
     }
 
     if (chunk.candidates && Array.isArray(chunk.candidates)) {
+      lastCandidates = chunk.candidates;
       for (const candidate of chunk.candidates) {
         if (candidate.content?.parts) {
           for (const part of candidate.content.parts) {
@@ -2054,15 +2057,15 @@ function aggregateGenerateContentChunks(
   }
   parts.push(...otherParts);
 
-  if (parts.length > 0 && lastResponse?.candidates) {
+  if (lastCandidates) {
     const candidates: Record<string, unknown>[] = [];
-    for (const candidate of lastResponse.candidates) {
-      const candidateDict: Record<string, unknown> = {
-        content: {
-          parts,
-          role: "model",
-        },
-      };
+    for (const candidate of lastCandidates) {
+      // A stream can end without content (for example a SAFETY block), but its
+      // finish reason and ratings still belong in the output.
+      const candidateDict: Record<string, unknown> = {};
+      if (parts.length > 0) {
+        candidateDict.content = { parts, role: "model" };
+      }
 
       if (candidate.finishReason !== undefined) {
         candidateDict.finishReason = candidate.finishReason;
@@ -2077,15 +2080,22 @@ function aggregateGenerateContentChunks(
         candidateDict.safetyRatings = candidate.safetyRatings;
       }
 
-      candidates.push(candidateDict);
+      if (Object.keys(candidateDict).length > 0) {
+        candidates.push(candidateDict);
+      }
+    }
+    if (candidates.length > 0) {
+      aggregated.candidates = candidates;
     }
-    aggregated.candidates = candidates;
   }
 
   if (usageMetadata) {
     aggregated.usageMetadata = usageMetadata;
     populateUsageMetrics(metrics, usageMetadata);
   }
+  if (promptFeedback) {
+    aggregated.promptFeedback = promptFeedback;
+  }
   if (groundingMetadata !== undefined) {
     aggregated.groundingMetadata = groundingMetadata;
   }

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

@@ -237,6 +237,7 @@ export interface GoogleGenAIGenerateContentResponse {
     safetyRatings?: Record<string, unknown>[];
   }[];
   groundingMetadata?: GoogleGenAIGroundingMetadata;
+  promptFeedback?: Record<string, unknown>;
   usageMetadata?: GoogleGenAIUsageMetadata;
   text?: string;
   [key: string]: unknown;