fix: preserve non-text output in streaming responses

Mateusz Krawiec · copybara-github · commit b8be90d24ef8 · 2026-07-03T08:12:50.000-07:00
The streaming response aggregator only accumulated text and function-call
parts, so any other part (an image returned as inline data, server-side tool
calls/responses, etc.) was dropped from the final aggregated response when text
accompanied it. Mirror ADK Python and preserve any part that is not text or a
function call, keeping the existing exclusion of standalone thought-signature
parts.

PiperOrigin-RevId: 942147787
diff --git a/core/src/main/java/com/google/adk/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java
@@ -356,18 +356,24 @@ private static String generateClientFunctionCallId() {
     }
 
     /**
-     * Accumulates text and function calls from incoming parts. Function-call parts passed to this
-     * method are expected to already have IDs (see {@link #ensureFunctionCallIds}).
+     * Accumulates content from incoming parts: text, function calls, and any other content part
+     * (inline image/audio data, file data, code execution, server-side tool calls/responses, and
+     * future part types). Standalone thought-signature/thought parts are the one exception: their
+     * signature is captured and re-attached to the last real part in {@link #processFinalResponse},
+     * so they are not emitted on their own. Function-call parts passed to this method are expected
+     * to already have IDs (see {@link #ensureFunctionCallIds}).
      *
-     * @return true if any text or function call was present, false otherwise.
+     * @return true if any content part was present, false otherwise.
      */
     private boolean accumulateParts(List<Part> parts) {
-      boolean hasTextOrFc = false;
+      boolean hasContent = false;
       for (Part part : parts) {
-        part.thoughtSignature().ifPresent(sig -> currentThoughtSignature = sig);
         String text = part.text().orElse("");
         if (!text.isEmpty()) {
-          hasTextOrFc = true;
+          hasContent = true;
+          // The signature belongs to this text; capture it so flushTextBufferToSequence attaches
+          // it.
+          part.thoughtSignature().ifPresent(sig -> currentThoughtSignature = sig);
           boolean isThought = part.thought().orElse(false);
           // Immediately flush the active text buffer to preserve the exact interleaved blocks of
           // text/thoughts.
@@ -378,13 +384,28 @@ private boolean accumulateParts(List<Part> parts) {
             currentTextIsThought = isThought;
           }
           currentTextBuffer.append(text);
-        }
-        if (part.functionCall().isPresent()) {
-          hasTextOrFc = true;
+        } else if (part.functionCall().isPresent()) {
+          hasContent = true;
           processFunctionCallPart(part);
+        } else if (part.text().isEmpty() && !part.thought().orElse(false)) {
+          // Mirror ADK Python's catch-all: preserve any part that is not text or a function call
+          // (inline image/audio data, file data, code execution, server-side tool calls/responses,
+          // future part types) rather than an allowlist that silently drops unlisted types. Flush
+          // buffered text first so parts keep their order, then append the part verbatim keeping
+          // any
+          // thoughtSignature it carries. The signature is intentionally not captured into
+          // currentThoughtSignature, which would leak it onto the preceding part.
+          hasContent = true;
+          flushTextBufferToSequence();
+          accumulatedSequence.add(part);
+        } else {
+          // Standalone thought/thought-signature part with no renderable content: not emitted on
+          // its
+          // own; capture its signature to re-attach to the last real part in processFinalResponse.
+          part.thoughtSignature().ifPresent(sig -> currentThoughtSignature = sig);
         }
       }
-      return hasTextOrFc;
+      return hasContent;
     }
 
     /**
diff --git a/core/src/test/java/com/google/adk/models/GeminiTest.java b/core/src/test/java/com/google/adk/models/GeminiTest.java
@@ -21,6 +21,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
+import com.google.genai.types.Blob;
 import com.google.genai.types.Candidate;
 import com.google.genai.types.Content;
 import com.google.genai.types.FinishReason;
@@ -29,6 +30,8 @@
 import com.google.genai.types.GenerateContentResponseUsageMetadata;
 import com.google.genai.types.Part;
 import com.google.genai.types.PartialArg;
+import com.google.genai.types.ToolCall;
+import com.google.genai.types.ToolResponse;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.functions.Predicate;
 import io.reactivex.rxjava3.subscribers.TestSubscriber;
@@ -368,6 +371,126 @@ public void processRawResponses_twoStreamingFunctionCalls_keepArgsSeparate() {
     assertThat(second.args().get()).containsExactly("b", "2");
   }
 
+  @Test
+  public void processRawResponses_imageOnlyWithStop_emitsFinalImagePart() {
+    Part imagePart = Part.fromBytes(new byte[] {1, 2, 3}, "image/png");
+    GenerateContentResponse imageWithStop =
+        toResponse(
+            Candidate.builder()
+                .content(Content.builder().role("model").parts(imagePart).build())
+                .finishReason(new FinishReason(FinishReason.Known.STOP))
+                .build());
+
+    ImmutableList<LlmResponse> responses =
+        ImmutableList.copyOf(
+            Gemini.processRawResponses(Flowable.just(imageWithStop)).blockingIterable());
+
+    LlmResponse finalResponse = Iterables.getLast(responses);
+    assertThat(finalResponse.content().get().parts().get()).hasSize(1);
+    assertThat(finalResponse.content().get().parts().get().get(0).inlineData()).isPresent();
+  }
+
+  // Image-generation models return image bytes as an inline-data part, often alongside text.
+  // Regression test: the aggregated response must retain the image, not just the text.
+  @Test
+  public void processRawResponses_textThenImageWithStop_finalKeepsTextAndImage() {
+    Part imagePart = Part.fromBytes(new byte[] {1, 2, 3}, "image/png");
+    GenerateContentResponse textChunk = toResponseWithText("Here is your image:");
+    GenerateContentResponse imageWithStop =
+        toResponse(
+            Candidate.builder()
+                .content(Content.builder().role("model").parts(imagePart).build())
+                .finishReason(new FinishReason(FinishReason.Known.STOP))
+                .build());
+
+    ImmutableList<LlmResponse> responses =
+        ImmutableList.copyOf(
+            Gemini.processRawResponses(Flowable.just(textChunk, imageWithStop)).blockingIterable());
+
+    LlmResponse finalResponse = Iterables.getLast(responses);
+    assertThat(finalResponse.content().get().parts().get()).hasSize(2);
+    assertThat(finalResponse.content().get().parts().get().get(0).text())
+        .hasValue("Here is your image:");
+    assertThat(finalResponse.content().get().parts().get().get(1).inlineData()).isPresent();
+  }
+
+  // The aggregator must pass through any non-text, non-function-call part, not just an allowlist.
+  // These guard the part types that were being silently dropped: server-side tool calls/responses
+  // and function responses. Each uses a text-then-part sequence so the part must survive the final
+  // aggregation (a lone part would otherwise slip through via the empty-sequence fallback).
+  @Test
+  public void processRawResponses_textThenToolCall_finalKeepsBoth() {
+    Part toolCallPart =
+        Part.builder()
+            .toolCall(ToolCall.builder().id("tc-1").args(ImmutableMap.of("q", "weather")).build())
+            .build();
+
+    LlmResponse finalResponse = aggregateTextThenPart(toolCallPart);
+
+    assertThat(finalResponse.content().get().parts().get()).hasSize(2);
+    assertThat(finalResponse.content().get().parts().get().get(1).toolCall()).isPresent();
+  }
+
+  @Test
+  public void processRawResponses_textThenToolResponse_finalKeepsBoth() {
+    Part toolResponsePart =
+        Part.builder()
+            .toolResponse(
+                ToolResponse.builder().id("tc-1").response(ImmutableMap.of("ok", true)).build())
+            .build();
+
+    LlmResponse finalResponse = aggregateTextThenPart(toolResponsePart);
+
+    assertThat(finalResponse.content().get().parts().get()).hasSize(2);
+    assertThat(finalResponse.content().get().parts().get().get(1).toolResponse()).isPresent();
+  }
+
+  @Test
+  public void processRawResponses_textThenFunctionResponse_finalKeepsBoth() {
+    Part functionResponsePart = Part.fromFunctionResponse("my_tool", ImmutableMap.of("result", 42));
+
+    LlmResponse finalResponse = aggregateTextThenPart(functionResponsePart);
+
+    assertThat(finalResponse.content().get().parts().get()).hasSize(2);
+    assertThat(finalResponse.content().get().parts().get().get(1).functionResponse()).isPresent();
+  }
+
+  // Per the Gemini docs, a data part (e.g. inlineData) can carry a thoughtSignature with the
+  // thought
+  // flag unset (multi-turn image editing). The part must be kept verbatim, and its signature must
+  // not leak onto the preceding text part (the docs forbid putting a signature on a part that did
+  // not originally carry one).
+  @Test
+  public void processRawResponses_textThenDataPartWithSignature_keepsSignatureOnDataPartOnly() {
+    Part imageWithSignature =
+        Part.builder()
+            .inlineData(Blob.builder().mimeType("image/png").data(new byte[] {1, 2, 3}).build())
+            .thoughtSignature("sig".getBytes(UTF_8))
+            .build();
+
+    LlmResponse finalResponse = aggregateTextThenPart(imageWithSignature);
+
+    assertThat(finalResponse.content().get().parts().get()).hasSize(2);
+    assertThat(finalResponse.content().get().parts().get().get(0).text())
+        .hasValue("Working on it:");
+    assertThat(finalResponse.content().get().parts().get().get(0).thoughtSignature()).isEmpty();
+    assertThat(finalResponse.content().get().parts().get().get(1).inlineData()).isPresent();
+    assertThat(finalResponse.content().get().parts().get().get(1).thoughtSignature()).isPresent();
+  }
+
+  private LlmResponse aggregateTextThenPart(Part part) {
+    GenerateContentResponse textChunk = toResponseWithText("Working on it:");
+    GenerateContentResponse partWithStop =
+        toResponse(
+            Candidate.builder()
+                .content(Content.builder().role("model").parts(part).build())
+                .finishReason(new FinishReason(FinishReason.Known.STOP))
+                .build());
+    return Iterables.getLast(
+        ImmutableList.copyOf(
+            Gemini.processRawResponses(Flowable.just(textChunk, partWithStop)).blockingIterable()));
+  }
+
   @Test
   public void processRawResponses_textAndStopReason_emitsPartialThenFinalText() {
     Flowable<GenerateContentResponse> rawResponses =
