From d7355a712345864682134762df890bf7b713b8c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Tue, 11 Aug 2026 05:01:37 -0700 Subject: [PATCH] fix: keep thought signature and tool call parts through streaming and history PiperOrigin-RevId: 962714340 --- .../google/adk/flows/llmflows/Contents.java | 63 +- .../java/com/google/adk/models/Gemini.java | 83 ++- .../adk/flows/llmflows/ContentsTest.java | 301 ++++++++ .../com/google/adk/models/GeminiTest.java | 682 +++++++++++++++++- 4 files changed, 1072 insertions(+), 57 deletions(-) diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Contents.java b/core/src/main/java/com/google/adk/flows/llmflows/Contents.java index 1f81bfa3e..64bb793d7 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Contents.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Contents.java @@ -155,7 +155,10 @@ private ImmutableList getContents( // TODO: Skip auth events. if (isOtherAgentReply(agentName, event)) { - filteredEvents.add(convertForeignEvent(event)); + Event foreignEvent = convertForeignEvent(event); + if (foreignEvent != null) { + filteredEvents.add(foreignEvent); + } } else { filteredEvents.add(event); } @@ -180,8 +183,9 @@ private ImmutableList getContents( * *

This can happen to the events that only changed session state. When both content and * transcriptions are empty, the event will be considered as empty. The content is considered - * empty if none of its parts contain text, inline data, file data, function call, or function - * response. Parts with only thoughts are also considered empty. + * empty if none of its parts contain text, inline data, file data, function call, function + * response, server-side tool call, or server-side tool response. Parts with only thoughts are + * also considered empty. * * @param event the event to check. * @return {@code true} if the event is considered to have empty content, {@code false} otherwise. @@ -205,12 +209,16 @@ private boolean isEmptyContent(Event event) { * *

    *
  • It has no meaningful content (text, inline_data, file_data, function_call, - * function_response, executable_code, or code_execution_result), OR - *
  • It is marked as a thought AND does not contain function_call or function_response + * function_response, tool_call, tool_response, executable_code, or code_execution_result) + * and no thought_signature, OR + *
  • It is marked as a thought AND does not contain function_call, function_response, + * tool_call, tool_response or thought_signature *
* *

Function calls and responses are never invisible, even if marked as thought, because they - * represent actions that need to be executed or results that need to be processed. + * represent actions that need to be executed or results that need to be processed. Parts carrying + * a thought signature, and server-side tool calls and their responses, are never invisible + * either, because the caller is required to echo them back on the next request. * * @param part the part to check. * @return {@code true} if the part is invisible, {@code false} otherwise. @@ -219,6 +227,18 @@ private boolean isPartInvisible(Part part) { if (part.functionCall().isPresent() || part.functionResponse().isPresent()) { return false; } + + // A thought signature is opaque state to hand back verbatim, and it routinely arrives on a part + // with nothing else in it, so it has to be checked before the emptiness test below. + if (part.thoughtSignature().map(signature -> signature.length > 0).orElse(false)) { + return false; + } + + // Server-side tool calls/responses must be echoed back to the model. + if (part.toolCall().isPresent() || part.toolResponse().isPresent()) { + return false; + } + return part.thought().orElse(false) || !(part.text().isPresent() || part.inlineData().isPresent() @@ -387,8 +407,13 @@ private static boolean isOtherAgentReply(String agentName, Event event) { && !event.author().equals("user"); } - /** Converts an {@code event} authored by another agent to a 'contextual-only' event. */ - private static Event convertForeignEvent(Event event) { + /** + * Converts an {@code event} authored by another agent to a 'contextual-only' event. + * + *

Returns {@code null} when nothing but the "For context:" preamble survives the conversion, + * so the caller drops the event instead of sending a preamble with no context after it. + */ + private static @Nullable Event convertForeignEvent(Event event) { if (event.content().isEmpty() || event.content().get().parts().isEmpty() || event.content().get().parts().get().isEmpty()) { @@ -401,9 +426,14 @@ private static Event convertForeignEvent(Event event) { String originalAuthor = event.author(); for (Part part : event.content().get().parts().get()) { - if (part.text().isPresent() - && !part.text().get().isEmpty() - && !part.thought().orElse(false)) { + // Thoughts belong to the agent that produced them and are never narrated, whatever else the + // part carries. ADK Python and ADK Kotlin both skip them before the branches below. + if (part.thought().orElse(false)) { + continue; + } + // Blank text is not narrated: such a part is a signature carrier, and a bare "said:" would + // both pollute the prompt and keep the event alive on nothing. + if (part.text().map(text -> !text.isBlank()).orElse(false)) { parts.add(Part.fromText(String.format("[%s] said: %s", originalAuthor, part.text().get()))); } else if (part.functionCall().isPresent()) { FunctionCall functionCall = part.functionCall().get(); @@ -423,9 +453,18 @@ private static Event convertForeignEvent(Event event) { originalAuthor, functionResponse.name().orElse("unknown_tool"), functionResponse.response().map(Contents::convertMapToJson).orElse("{}")))); - } else { + } else if (part.inlineData().isPresent() + || part.fileData().isPresent() + || part.executableCode().isPresent() + || part.codeExecutionResult().isPresent()) { parts.add(part); } + // Anything else - a bare signature, a server-side call - belongs to the model instance that + // produced it, so claiming it for another agent would be wrong. + } + + if (parts.size() == 1) { + return null; } Content content = Content.builder().role("user").parts(parts).build(); diff --git a/core/src/main/java/com/google/adk/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java index 36267551c..8b4d95298 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -321,6 +321,24 @@ private static final class StreamingResponseAggregator { private final StringBuilder currentTextBuffer = new StringBuilder(); // Always reassigned in accumulateParts() before it is read; the initializer is never observed. private boolean currentTextIsThought = false; + + /** + * Returns whether the part is the empty-text terminator Gemini 3 ends a stream with: empty text + * and nothing else worth keeping. Compared by rebuilding rather than against a single literal, + * so a terminator that also carries an explicit {@code thought=false} is still recognised. + */ + private static boolean isStreamTerminator(Part part) { + if (!part.text().map(String::isEmpty).orElse(false)) { + return false; + } + Part.Builder terminator = Part.builder().text(""); + part.thought().ifPresent(terminator::thought); + return terminator.build().equals(part); + } + + // Signature of the buffered text run, kept apart from the call's slot below so an interleaved + // chunk cannot flush one part carrying the other's signature. + private byte[] currentTextThoughtSignature = null; private byte[] currentThoughtSignature = null; private GenerateContentResponse lastRawResponse = null; @@ -407,11 +425,11 @@ private static String generateClientFunctionCallId() { /** * 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}). + * (inline image/audio data, file data, code execution, server-side tool calls/responses, + * standalone thought signatures, and future part types), which are appended verbatim as ADK + * Python does. The empty-text part that ends a Gemini 3 stream is the one thing dropped. + * Function-call parts passed to this method are expected to already have IDs (see {@link + * #ensureFunctionCallIds}). * * @return true if any content part was present, false otherwise. */ @@ -421,38 +439,33 @@ private boolean accumulateParts(List parts) { String text = part.text().orElse(""); if (!text.isEmpty()) { 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. + // Flush before capturing this chunk's signature below, or the signature of the run + // starting here lands on the run being flushed. if (!currentTextBuffer.isEmpty() && isThought != currentTextIsThought) { flushTextBufferToSequence(); } if (currentTextBuffer.isEmpty()) { currentTextIsThought = isThought; } + // Keep the first signature of the run, as ADK Python does; the merged part takes it in + // flushTextBufferToSequence. + if (currentTextThoughtSignature == null + && part.thoughtSignature().map(sig -> sig.length > 0).orElse(false)) { + currentTextThoughtSignature = part.thoughtSignature().get(); + } currentTextBuffer.append(text); } 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. + } else if (isStreamTerminator(part)) { + // Gemini 3 ends a stream with a bare empty text part; it carries nothing to keep. + } else { + // Everything else is appended as the model sent it, signature included. Relocating a + // signature onto a neighbouring part would hand it back on a part the model never signed. 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 hasContent; @@ -476,7 +489,8 @@ private void processFunctionCallPart(Part part) { || (currentFcName != null && !hasName); if (streamedPart) { // Capture the thought signature from the first chunk that carries one. - if (part.thoughtSignature().isPresent() && currentThoughtSignature == null) { + if (currentThoughtSignature == null + && part.thoughtSignature().map(sig -> sig.length > 0).orElse(false)) { currentThoughtSignature = part.thoughtSignature().get(); } processStreamingFunctionCall(fc); @@ -605,9 +619,9 @@ private void flushTextBufferToSequence() { if (!currentTextBuffer.isEmpty()) { Part.Builder partBuilder = Part.builder().text(currentTextBuffer.toString()).thought(currentTextIsThought); - if (currentThoughtSignature != null) { - partBuilder.thoughtSignature(currentThoughtSignature); - currentThoughtSignature = null; + if (currentTextThoughtSignature != null) { + partBuilder.thoughtSignature(currentTextThoughtSignature); + currentTextThoughtSignature = null; } accumulatedSequence.add(partBuilder.build()); currentTextBuffer.setLength(0); @@ -652,18 +666,9 @@ private Flowable processFinalResponse() { return Flowable.just(finalResponseBuilder.build()); } - // If the final chunk carries a thoughtSignature (e.g. from a preceding function call or - // thought), attach it to the last accumulated part in the sequence. - GeminiUtil.getPart0FromLlmResponse(currentResponse) - .flatMap(Part::thoughtSignature) - .ifPresent( - signature -> { - int targetIndex = accumulatedSequence.size() - 1; - Part targetPart = accumulatedSequence.get(targetIndex); - accumulatedSequence.set( - targetIndex, targetPart.toBuilder().thoughtSignature(signature).build()); - }); - + // No re-attach of the final chunk's signature: every part now keeps the signature the model + // put on it, so reading part 0 and stamping the last part could only mis-attribute one. ADK + // Python and the ADK Kotlin sibling have no equivalent either. return Flowable.just( finalResponseBuilder .content(Content.builder().role("model").parts(accumulatedSequence).build()) diff --git a/core/src/test/java/com/google/adk/flows/llmflows/ContentsTest.java b/core/src/test/java/com/google/adk/flows/llmflows/ContentsTest.java index 5d2c3d5fc..146a4e92f 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/ContentsTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/ContentsTest.java @@ -19,6 +19,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Correspondence.transforming; import static com.google.common.truth.Truth.assertThat; +import static java.nio.charset.StandardCharsets.UTF_8; import static org.junit.Assert.assertThrows; import com.google.adk.agents.InvocationContext; @@ -33,10 +34,13 @@ import com.google.adk.sessions.Session; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Blob; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; +import com.google.genai.types.ToolCall; +import com.google.genai.types.ToolResponse; import java.util.ArrayList; import java.util.ConcurrentModificationException; import java.util.Iterator; @@ -946,6 +950,294 @@ public void processRequest_notEmptyContent() { assertThat(contents).containsExactly(e.content().get()); } + // On models that return a signature for every part, it arrives on parts holding nothing else. + // Dropping those as "empty" loses the reasoning the model expects back on the next turn. + @Test + public void processRequest_contentFreeThoughtSignatureEvent_notSkipped() { + Event signatureEvent = + createModelEvent( + "e2", Part.builder().thoughtSignature("call-context".getBytes(UTF_8)).build()); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "Summarize the video."), signatureEvent); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(2); + assertThat(contents.get(1).parts().get().get(0).thoughtSignature()) + .hasValue("call-context".getBytes(UTF_8)); + } + + // A thought part carrying a signature is kept for the same reason, even though a bare thought + // part is dropped. + @Test + public void processRequest_thoughtWithSignatureEvent_notSkipped() { + Event thoughtEvent = + createModelEvent( + "e2", + Part.builder() + .thought(true) + .text("Let me check the frame at 0:05.") + .thoughtSignature("thought-sig".getBytes(UTF_8)) + .build()); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "Summarize the video."), thoughtEvent); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(2); + assertThat(contents.get(1).parts().get().get(0).thoughtSignature()) + .hasValue("thought-sig".getBytes(UTF_8)); + } + + // The caller must echo server-side tool parts back, so dropping them as "empty" makes the model + // redo the work or fail on a call with no matching response. + @Test + public void processRequest_serverSideToolCallAndResponseEvents_notSkipped() { + Event toolCallEvent = + createModelEvent( + "e2", + Part.builder() + .toolCall( + ToolCall.builder() + .id("tc1") + .args(ImmutableMap.of("url", "https://example.com")) + .build()) + .build()); + Event toolResponseEvent = + createModelEvent( + "e3", + Part.builder() + .toolResponse( + ToolResponse.builder() + .id("tc1") + .response(ImmutableMap.of("content", "page text")) + .build()) + .build()); + ImmutableList events = + ImmutableList.of( + createUserEvent("e1", "Summarize the linked page."), toolCallEvent, toolResponseEvent); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(3); + assertThat(contents.get(1).parts().get().get(0).toolCall().get().id()).hasValue("tc1"); + ToolResponse toolResponse = contents.get(2).parts().get().get(0).toolResponse().get(); + assertThat(toolResponse.id()).hasValue("tc1"); + assertThat(toolResponse.response()).hasValue(ImmutableMap.of("content", "page text")); + } + + // The echo-back contract holds regardless of how the model labels the part, so a thought marking + // must not drop it. + @Test + public void processRequest_serverSideToolCallMarkedAsThought_notSkipped() { + Event toolCallEvent = + createModelEvent( + "e2", + Part.builder() + .thought(true) + .toolCall( + ToolCall.builder() + .id("tc1") + .args(ImmutableMap.of("url", "https://example.com")) + .build()) + .build()); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "Summarize the linked page."), toolCallEvent); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(2); + assertThat(contents.get(1).parts().get().get(0).toolCall().get().id()).hasValue("tc1"); + } + + // A server-side call belongs to the model instance that made it, so the other-agent path must + // keep dropping it rather than claiming the call on this agent's behalf. + @Test + public void processRequest_serverSideToolCallFromOtherAgent_isDropped() { + Event otherAgentToolCall = + Event.builder() + .id("e2") + .author(OTHER_AGENT) + .content( + Content.builder() + .role("model") + .parts( + ImmutableList.of( + Part.builder() + .toolCall( + ToolCall.builder() + .id("tc1") + .args(ImmutableMap.of("url", "https://example.com")) + .build()) + .build())) + .build()) + .invocationId("invocationId") + .build(); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "Summarize the linked page."), otherAgentToolCall); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(1); + assertThat(contents.get(0).parts().get().get(0).text()).hasValue("Summarize the linked page."); + } + + @Test + public void processRequest_serverSideToolCallWithThoughtFromOtherAgent_isDropped() { + Event otherAgentToolCall = + Event.builder() + .id("e2") + .author(OTHER_AGENT) + .content( + Content.builder() + .role("model") + .parts( + ImmutableList.of( + Part.builder().thought(true).text("Let me look it up.").build(), + Part.builder() + .toolCall( + ToolCall.builder() + .id("tc1") + .args(ImmutableMap.of("url", "https://example.com")) + .build()) + .build())) + .build()) + .invocationId("invocationId") + .build(); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "Summarize the linked page."), otherAgentToolCall); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(1); + assertThat(contents.get(0).parts().get().get(0).text()).hasValue("Summarize the linked page."); + } + + // A thought-marked function call from another agent must not be narrated: the thought guard runs + // before the branches that would turn it into "[agent] called tool ...". + @Test + public void processRequest_thoughtMarkedFunctionCallFromOtherAgent_isDropped() { + Event otherAgentEvent = + Event.builder() + .id("e2") + .author(OTHER_AGENT) + .content( + Content.builder() + .role("model") + .parts( + ImmutableList.of( + Part.builder() + .thought(true) + .functionCall( + FunctionCall.builder() + .name("lookup") + .args(ImmutableMap.of("q", "x")) + .build()) + .build())) + .build()) + .invocationId("invocationId") + .build(); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "What is in the picture?"), otherAgentEvent); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(1); + } + + // Whitespace-only text is not content either, so the carrier that carries it must not be narrated + // as a bare "said:". Matches what the emptiness rule already treats as blank. + @Test + public void processRequest_blankTextSignaturePartFromOtherAgent_isNotNarrated() { + Event otherAgentEvent = + Event.builder() + .id("e2") + .author(OTHER_AGENT) + .content( + Content.builder() + .role("model") + .parts( + ImmutableList.of( + Part.builder() + .text(" ") + .thoughtSignature(new byte[] {7, 7, 7}) + .build())) + .build()) + .invocationId("invocationId") + .build(); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "Summarize the video."), otherAgentEvent); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(1); + } + + // Another agent's reasoning belongs to that agent and is never narrated: only the answer text + // beside it may be attributed. + @Test + public void processRequest_thoughtTextFromOtherAgent_isNotNarrated() { + Event otherAgentEvent = + Event.builder() + .id("e2") + .author(OTHER_AGENT) + .content( + Content.builder() + .role("model") + .parts( + ImmutableList.of( + Part.builder().thought(true).text("Let me check the map.").build(), + Part.fromText("It is in Paris."))) + .build()) + .invocationId("invocationId") + .build(); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "Where is it?"), otherAgentEvent); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(2); + assertThat( + contents.get(1).parts().get().stream() + .map(part -> part.text().orElse("")) + .collect(toImmutableList())) + .containsExactly("For context:", "[" + OTHER_AGENT + "] said: It is in Paris."); + } + + // The other-agent path still narrates what it can: media parts pass through unchanged, so the + // drop above is about attribution rather than a blanket filter. + @Test + public void processRequest_mediaPartFromOtherAgent_isKept() { + Event otherAgentImage = + Event.builder() + .id("e2") + .author(OTHER_AGENT) + .content( + Content.builder() + .role("model") + .parts( + ImmutableList.of( + Part.builder() + .inlineData( + Blob.builder() + .mimeType("image/png") + .data(new byte[] {1, 2, 3}) + .build()) + .build())) + .build()) + .invocationId("invocationId") + .build(); + ImmutableList events = + ImmutableList.of(createUserEvent("e1", "What is in the picture?"), otherAgentImage); + + List contents = runContentsProcessor(events); + + assertThat(contents).hasSize(2); + assertThat(contents.get(1).parts().get()).hasSize(2); + assertThat(contents.get(1).parts().get().get(0).text()).hasValue("For context:"); + assertThat(contents.get(1).parts().get().get(1).inlineData()).isPresent(); + } + @Test public void processRequest_concurrentReadAndWrite_noException() throws Exception { LlmAgent agent = @@ -1028,6 +1320,15 @@ private static Event createUserEvent( .build(); } + private static Event createModelEvent(String id, Part part) { + return Event.builder() + .id(id) + .author(AGENT) + .content(Content.builder().role("model").parts(ImmutableList.of(part)).build()) + .invocationId("invocationId") + .build(); + } + private static Event createAgentEvent(String id, String text) { return createAgentEvent(AGENT, id, text); } diff --git a/core/src/test/java/com/google/adk/models/GeminiTest.java b/core/src/test/java/com/google/adk/models/GeminiTest.java index a965e5b68..a56628493 100644 --- a/core/src/test/java/com/google/adk/models/GeminiTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiTest.java @@ -920,8 +920,7 @@ public void processRawResponses_thoughtAndTextWithStop_onlyFinalTextIncludesUsag } @Test - public void - processRawResponses_thoughtThenEmptyWithSignatureAndStop_flushesThoughtWithSignature() { + public void processRawResponses_thoughtThenSignatureAndStop_keepsSignatureOnItsOwnPart() { GenerateContentResponseUsageMetadata metadata1 = createUsageMetadata(5, 10, 15); GenerateContentResponseUsageMetadata metadata2 = createUsageMetadata(5, 20, 25); GenerateContentResponse chunk1 = toResponseWithThoughtText("Thinking", metadata1); @@ -948,7 +947,16 @@ public void processRawResponses_thoughtAndTextWithStop_onlyFinalTextIncludesUsag assertLlmResponses( llmResponses, isPartialThoughtResponseWithUsageMetadata("Thinking", metadata1), - isFinalThoughtResponseWithUsageMetadataAndSignature("Thinking", metadata2, "sig")); + isPartialSignatureResponse("sig"), + response -> { + ImmutableList parts = ImmutableList.copyOf(response.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).text()).hasValue("Thinking"); + assertThat(parts.get(0).thoughtSignature()).isEmpty(); + assertThat(parts.get(1).thoughtSignature()).hasValue("sig".getBytes(UTF_8)); + assertThat(response.usageMetadata()).hasValue(metadata2); + return true; + }); } @Test @@ -998,7 +1006,7 @@ public void processRawResponses_thoughtAndTextWithStop_onlyFinalTextIncludesUsag @Test public void - processRawResponses_thoughtThenFunctionCallWithSignatureAndStop_attachesSignatureToFunctionCall() { + processRawResponses_thoughtThenFunctionCallThenSignature_keepsSignatureOnItsOwnPart() { GenerateContentResponseUsageMetadata metadata1 = createUsageMetadata(5, 10, 15); GenerateContentResponseUsageMetadata metadata2 = createUsageMetadata(5, 20, 25); GenerateContentResponse chunk1 = toResponseWithThoughtText("Thinking", metadata1); @@ -1028,8 +1036,17 @@ public void processRawResponses_thoughtAndTextWithStop_onlyFinalTextIncludesUsag llmResponses, isPartialThoughtResponseWithUsageMetadata("Thinking", metadata1), isPartialFunctionCallResponse("my_tool"), - isFinalThoughtAndFunctionCallResponseWithUsageMetadataAndSignature( - "Thinking", metadata2, "sig", "my_tool")); + isPartialSignatureResponse("sig"), + response -> { + ImmutableList parts = ImmutableList.copyOf(response.content().get().parts().get()); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).text()).hasValue("Thinking"); + assertThat(parts.get(1).functionCall().get().name()).hasValue("my_tool"); + assertThat(parts.get(1).thoughtSignature()).isEmpty(); + assertThat(parts.get(2).thoughtSignature()).hasValue("sig".getBytes(UTF_8)); + assertThat(response.usageMetadata()).hasValue(metadata2); + return true; + }); } @Test @@ -1065,9 +1082,630 @@ public void processRawResponses_emptyPartsThenSignature_doesNotThrowException() assertLlmResponses( llmResponses, isEmptyResponse(), + isPartialSignatureResponse("sig"), isFinalThoughtResponseWithUsageMetadataAndSignature("", metadata, "sig")); } + // Consecutive text chunks are merged into a single part the aggregator builds from scratch, so a + // thought signature the chunks carried is lost unless it is copied across. The model expects its + // signature back verbatim; without it, it redoes the reasoning the signature stood for. Mirrors + // ADK Python's TestStreamingThoughtSignature. + @Test + public void processRawResponses_signatureOnMergedText_isPreserved() { + GenerateContentResponse chunk1 = toResponseWithTextAndSignature("At minute 5 ", "text-sig"); + GenerateContentResponse chunk2 = + toResponseWithText("the presenter speaks.", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).text()).hasValue("At minute 5 the presenter speaks."); + assertThat(parts.get(0).thoughtSignature()).hasValue("text-sig".getBytes(UTF_8)); + } + + // The signature can land on any chunk of the run, not just the first. + @Test + public void processRawResponses_signatureOnLaterTextChunk_isPreserved() { + GenerateContentResponse chunk1 = toResponseWithText("At minute 5 "); + GenerateContentResponse chunk2 = toResponseWithTextAndSignature("the presenter ", "late-sig"); + GenerateContentResponse chunk3 = toResponseWithText("speaks.", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).text()).hasValue("At minute 5 the presenter speaks."); + assertThat(parts.get(0).thoughtSignature()).hasValue("late-sig".getBytes(UTF_8)); + } + + // A merged part carries one signature; the run keeps the first it saw, as ADK Python does. + @Test + public void processRawResponses_multipleSignaturesInOneRun_keepsTheFirst() { + GenerateContentResponse chunk1 = toResponseWithTextAndSignature("At minute 5 ", "first-sig"); + GenerateContentResponse chunk2 = toResponseWithTextAndSignature("the presenter ", "second-sig"); + GenerateContentResponse chunk3 = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts(Part.fromFunctionCall("done", ImmutableMap.of())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thoughtSignature()).hasValue("first-sig".getBytes(UTF_8)); + } + + // A thought run and an answer run flush separately and must not swap signatures: the answer's + // signature arrives on the chunk that triggers the flush of the thought. + @Test + public void processRawResponses_thoughtAndAnswerRuns_keepTheirOwnSignatures() { + GenerateContentResponse chunk1 = + toResponse( + Part.builder() + .text("Let me check.") + .thought(true) + .thoughtSignature("thought-sig".getBytes(UTF_8)) + .build()); + GenerateContentResponse chunk2 = + toResponseWithTextAndSignature("It is a dog.", "answer-sig", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thought()).hasValue(true); + assertThat(parts.get(0).thoughtSignature()).hasValue("thought-sig".getBytes(UTF_8)); + assertThat(parts.get(1).thoughtSignature()).hasValue("answer-sig".getBytes(UTF_8)); + } + + // A signature-only thought part keeps its signature on itself, as ADK Python does, rather than + // having it relocated onto the text around it. + @Test + public void processRawResponses_standaloneSignatureMidTextRun_keepsItsOwnSignature() { + GenerateContentResponse chunk1 = toResponseWithText("At minute 5 "); + GenerateContentResponse chunk2 = + toResponse( + Part.builder().thought(true).thoughtSignature("carried-sig".getBytes(UTF_8)).build()); + GenerateContentResponse chunk3 = + toResponseWithText("the presenter speaks.", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).text()).hasValue("At minute 5 "); + assertThat(parts.get(0).thoughtSignature()).isEmpty(); + assertThat(parts.get(1).thoughtSignature()).hasValue("carried-sig".getBytes(UTF_8)); + assertThat(parts.get(2).text()).hasValue("the presenter speaks."); + assertThat(parts.get(2).thoughtSignature()).isEmpty(); + } + + // A text chunk arriving mid-stream of a function call must not take the call's signature with it: + // the two runs flush together and each keeps its own. + @Test + public void processRawResponses_textInterleavedWithStreamedCall_keepsBothSignatures() { + GenerateContentResponse chunk1 = + toResponse( + Part.builder() + .functionCall( + FunctionCall.builder() + .name("search") + .partialArgs( + PartialArg.builder().jsonPath("$.q").stringValue("hel").build()) + .willContinue(true) + .build()) + .thoughtSignature("fc-sig".getBytes(UTF_8)) + .build()); + GenerateContentResponse chunk2 = toResponseWithTextAndSignature("Working on it.", "text-sig"); + GenerateContentResponse chunk3 = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder() + .jsonPath("$.q") + .stringValue("lo") + .build()) + .willContinue(false) + .build())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).text()).hasValue("Working on it."); + assertThat(parts.get(0).thoughtSignature()).hasValue("text-sig".getBytes(UTF_8)); + assertThat(parts.get(1).functionCall().get().name()).hasValue("search"); + assertThat(parts.get(1).thoughtSignature()).hasValue("fc-sig".getBytes(UTF_8)); + } + + // A signature-only part with no text run open must not be dropped, and the streamed call that + // follows must not inherit its signature. + @Test + public void processRawResponses_standaloneSignatureThenStreamedCall_keepsItOnItsOwnPart() { + GenerateContentResponse chunk1 = + toResponse(Part.builder().thought(true).thoughtSignature("sig-A".getBytes(UTF_8)).build()); + GenerateContentResponse chunk2 = + toResponse( + functionCallPart( + FunctionCall.builder() + .name("search") + .partialArgs(PartialArg.builder().jsonPath("$.q").stringValue("hel").build()) + .willContinue(true) + .build())); + GenerateContentResponse chunk3 = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder() + .jsonPath("$.q") + .stringValue("lo") + .build()) + .willContinue(false) + .build())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thoughtSignature()).hasValue("sig-A".getBytes(UTF_8)); + assertThat(parts.get(1).functionCall().get().name()).hasValue("search"); + assertThat(parts.get(1).thoughtSignature()).isEmpty(); + } + + // Two runs inside the final chunk each keep their own signature: the final-chunk re-attach must + // not stamp the first run's signature over the second's. + @Test + public void processRawResponses_thoughtAndAnswerInFinalChunk_keepTheirOwnSignatures() { + Part thought = + Part.builder() + .text("Let me check.") + .thought(true) + .thoughtSignature("thought-sig".getBytes(UTF_8)) + .build(); + Part answer = + Part.builder().text("It is a dog.").thoughtSignature("answer-sig".getBytes(UTF_8)).build(); + GenerateContentResponse chunk = + toResponse( + Candidate.builder() + .content(Content.builder().parts(thought, answer).build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thoughtSignature()).hasValue("thought-sig".getBytes(UTF_8)); + assertThat(parts.get(1).thoughtSignature()).hasValue("answer-sig".getBytes(UTF_8)); + } + + // A carrier between two text runs ends the first and is emitted on its own; neither run's + // signature moves, so nothing is attributed to a part the model did not sign. + @Test + public void processRawResponses_carrierBetweenTwoRuns_isEmittedOnItsOwn() { + GenerateContentResponse chunk1 = toResponseWithTextAndSignature("Hello", "sig-A"); + GenerateContentResponse chunk2 = + toResponse(Part.builder().thought(true).thoughtSignature("sig-B".getBytes(UTF_8)).build()); + GenerateContentResponse chunk3 = toResponseWithText(" world", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).text()).hasValue("Hello"); + assertThat(parts.get(0).thoughtSignature()).hasValue("sig-A".getBytes(UTF_8)); + assertThat(parts.get(1).thoughtSignature()).hasValue("sig-B".getBytes(UTF_8)); + assertThat(parts.get(2).text()).hasValue(" world"); + assertThat(parts.get(2).thoughtSignature()).isEmpty(); + } + + // An empty signature must not occupy the run's slot and block the real one behind it. + @Test + public void processRawResponses_emptySignatureThenRealOne_keepsTheRealOne() { + GenerateContentResponse chunk1 = toResponseWithTextAndSignature("Hel", ""); + GenerateContentResponse chunk2 = + toResponseWithTextAndSignature("lo", "real-sig", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).thoughtSignature()).hasValue("real-sig".getBytes(UTF_8)); + } + + // A signature the aggregator already placed must not be handed out again by the final-chunk + // re-attach when the last part happens to be unsigned. + @Test + public void processRawResponses_signedThenUnsignedRunInFinalChunk_doesNotDuplicate() { + Part signed = Part.builder().text("A").thoughtSignature("sig-1".getBytes(UTF_8)).build(); + Part unsigned = Part.builder().text("B").thought(true).build(); + GenerateContentResponse chunk = + toResponse( + Candidate.builder() + .content(Content.builder().parts(signed, unsigned).build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thoughtSignature()).hasValue("sig-1".getBytes(UTF_8)); + assertThat(parts.get(1).thoughtSignature()).isEmpty(); + } + + // A carrier's signature stays on the carrier: neither the call after it nor the text after that + // may end up carrying the same bytes. + @Test + public void processRawResponses_carrierThenCallThenText_doesNotDuplicate() { + GenerateContentResponse chunk1 = + toResponse(Part.builder().thought(true).thoughtSignature("sig-A".getBytes(UTF_8)).build()); + GenerateContentResponse chunk2 = + toResponse( + functionCallPart( + FunctionCall.builder() + .name("search") + .partialArgs(PartialArg.builder().jsonPath("$.q").stringValue("hel").build()) + .willContinue(true) + .build())); + GenerateContentResponse chunk3 = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs(PartialArg.builder().jsonPath("$.q").stringValue("lo").build()) + .willContinue(false) + .build())); + GenerateContentResponse chunk4 = toResponseWithText("Done.", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3, chunk4); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).thoughtSignature()).hasValue("sig-A".getBytes(UTF_8)); + assertThat(parts.get(1).functionCall()).isPresent(); + assertThat(parts.get(1).thoughtSignature()).isEmpty(); + assertThat(parts.get(2).text()).hasValue("Done."); + assertThat(parts.get(2).thoughtSignature()).isEmpty(); + } + + // An empty text part that also carries payload must survive: Optional.isEmpty() is false for + // text="", so such a part misses the catch-all unless the emptiness is tested on the value. + @Test + public void processRawResponses_emptyTextPartWithInlineData_isKept() { + GenerateContentResponse chunk1 = toResponseWithText("Here."); + GenerateContentResponse chunk2 = + toResponse( + Part.builder() + .text("") + .inlineData(Blob.builder().mimeType("image/png").data(new byte[] {1, 2}).build()) + .build()); + GenerateContentResponse chunk3 = toResponseWithText("", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(1).inlineData()).isPresent(); + } + + // A signature appears on exactly one part: the one the model put it on. Neither the text run + // after the carrier nor the call after that may emit the same bytes. + @Test + public void processRawResponses_carriedSignature_isNotEmittedOnTwoParts() { + GenerateContentResponse chunk1 = + toResponse(Part.builder().thought(true).thoughtSignature("sig-A".getBytes(UTF_8)).build()); + GenerateContentResponse chunk2 = toResponseWithText("Working on it."); + GenerateContentResponse chunk3 = + toResponse( + functionCallPart( + FunctionCall.builder() + .name("search") + .partialArgs(PartialArg.builder().jsonPath("$.q").stringValue("hel").build()) + .willContinue(true) + .build())); + GenerateContentResponse chunk4 = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder() + .jsonPath("$.q") + .stringValue("lo") + .build()) + .willContinue(false) + .build())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3, chunk4); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).thoughtSignature()).hasValue("sig-A".getBytes(UTF_8)); + assertThat(parts.get(1).text()).hasValue("Working on it."); + assertThat(parts.get(1).thoughtSignature()).isEmpty(); + assertThat(parts.get(2).functionCall()).isPresent(); + assertThat(parts.get(2).thoughtSignature()).isEmpty(); + } + + // A streamed call that carries its own signature keeps it, and the carrier before it keeps its + // own: two signatures in, two signatures out, neither displaced. + @Test + public void processRawResponses_streamedCallKeepsItsOwnSignatureAfterACarrier() { + GenerateContentResponse chunk1 = + toResponse(Part.builder().thought(true).thoughtSignature("sig-A".getBytes(UTF_8)).build()); + GenerateContentResponse chunk2 = + toResponse( + Part.builder() + .functionCall( + FunctionCall.builder() + .name("search") + .partialArgs( + PartialArg.builder().jsonPath("$.q").stringValue("hel").build()) + .willContinue(true) + .build()) + .thoughtSignature("fc-B".getBytes(UTF_8)) + .build()); + GenerateContentResponse chunk3 = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder() + .jsonPath("$.q") + .stringValue("lo") + .build()) + .willContinue(false) + .build())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thoughtSignature()).hasValue("sig-A".getBytes(UTF_8)); + assertThat(parts.get(1).thoughtSignature()).hasValue("fc-B".getBytes(UTF_8)); + } + + // Server-side media tools return signatures on parts holding nothing else. Such a part must + // survive as its own part rather than being folded into the surrounding text. + @Test + public void processRawResponses_contentFreeSignaturePart_isKept() { + GenerateContentResponse chunk1 = toResponseWithText("At minute 5 the presenter speaks."); + GenerateContentResponse chunk2 = + toResponse(Part.builder().thoughtSignature("call-context".getBytes(UTF_8)).build()); + GenerateContentResponse chunk3 = toResponseWithText("", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thoughtSignature()).isEmpty(); + assertThat(parts.get(1).thoughtSignature()).hasValue("call-context".getBytes(UTF_8)); + } + + // The other half of the rule above: an empty text part carrying nothing at all only marks the end + // of a Gemini 3 stream, so it must not reach the caller as a part of its own. + @Test + public void processRawResponses_bareEmptyTextPart_isDropped() { + GenerateContentResponse chunk1 = toResponseWithText("Let me check."); + GenerateContentResponse chunk2 = toResponseWithText("", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).text()).hasValue("Let me check."); + } + + // The wire shape the standard Gemini API actually sends for the same thing: the signature rides + // on a part whose text is present but empty, which Optional.isEmpty() does not recognise. + @Test + public void processRawResponses_emptyTextSignaturePart_isKept() { + GenerateContentResponse chunk1 = toResponseWithText("The answer is 42."); + GenerateContentResponse chunk2 = + toResponse( + Part.builder().text("").thoughtSignature("trailing-sig".getBytes(UTF_8)).build()); + GenerateContentResponse chunk3 = toResponseWithText("", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thoughtSignature()).isEmpty(); + assertThat(parts.get(1).thoughtSignature()).hasValue("trailing-sig".getBytes(UTF_8)); + } + + // Three parts, two signatures, and no relocation: the carrier keeps its own and the signed text + // run behind the call keeps its own. + @Test + public void processRawResponses_carrierThenCallThenSignedText_keepsEachSignatureInPlace() { + GenerateContentResponse chunk1 = + toResponse(Part.builder().thought(true).thoughtSignature("carry".getBytes(UTF_8)).build()); + GenerateContentResponse chunk2 = + toResponse( + functionCallPart( + FunctionCall.builder() + .name("search") + .partialArgs(PartialArg.builder().jsonPath("$.q").stringValue("hel").build()) + .willContinue(true) + .build())); + GenerateContentResponse chunk3 = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs(PartialArg.builder().jsonPath("$.q").stringValue("lo").build()) + .willContinue(false) + .build())); + GenerateContentResponse chunk4 = + toResponseWithTextAndSignature("Here you go.", "text-B", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3, chunk4); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).thoughtSignature()).hasValue("carry".getBytes(UTF_8)); + assertThat(parts.get(1).functionCall()).isPresent(); + assertThat(parts.get(1).thoughtSignature()).isEmpty(); + assertThat(parts.get(2).thoughtSignature()).hasValue("text-B".getBytes(UTF_8)); + } + + // Same invariant with a complete rather than a streamed call. The model did not sign the call, + // so the call goes out unsigned rather than inheriting the thought's signature. + @Test + public void processRawResponses_carrierThenCompleteCall_leavesTheCallUnsigned() { + GenerateContentResponse chunk1 = + toResponse(Part.builder().thought(true).thoughtSignature("carry".getBytes(UTF_8)).build()); + GenerateContentResponse chunk2 = + toResponse(functionCallPart(FunctionCall.builder().name("search").id("fc-1").build())); + GenerateContentResponse chunk3 = + toResponseWithTextAndSignature("Here you go.", "text-B", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).thoughtSignature()).hasValue("carry".getBytes(UTF_8)); + assertThat(parts.get(1).functionCall()).isPresent(); + assertThat(parts.get(1).thoughtSignature()).isEmpty(); + assertThat(parts.get(2).thoughtSignature()).hasValue("text-B".getBytes(UTF_8)); + } + + // A thought-marked server-side tool call carries payload the model has to see again, so only the + // marker and the signature may be folded away - the part itself has to reach the session intact. + @Test + public void processRawResponses_thoughtMarkedServerSideToolCall_survivesTheStream() { + Part toolCallPart = + Part.builder() + .thought(true) + .toolCall(ToolCall.builder().id("tc1").build()) + .thoughtSignature("tool-sig".getBytes(UTF_8)) + .build(); + GenerateContentResponse chunk1 = toResponse(toolCallPart); + GenerateContentResponse chunk2 = toResponseWithText("Found it.", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0)).isEqualTo(toolCallPart); + assertThat(parts.get(1).text()).hasValue("Found it."); + } + + // A zero-length signature must not occupy the streamed call's own slot and block the real one + // behind it, the same rule the text run's slot follows. + @Test + public void processRawResponses_emptySignatureThenRealOneOnAStreamedCall_keepsTheRealOne() { + GenerateContentResponse chunk1 = + toResponse( + Part.builder() + .functionCall( + FunctionCall.builder() + .name("search") + .partialArgs( + PartialArg.builder().jsonPath("$.q").stringValue("hel").build()) + .willContinue(true) + .build()) + .thoughtSignature(new byte[0]) + .build()); + GenerateContentResponse chunk2 = + toResponse( + Part.builder() + .functionCall( + FunctionCall.builder() + .partialArgs(PartialArg.builder().jsonPath("$.q").stringValue("lo").build()) + .willContinue(false) + .build()) + .thoughtSignature("real-sig".getBytes(UTF_8)) + .build()); + // The stream ends unsigned, so the final-chunk re-attach cannot supply the signature and the + // assertion is about the call's own slot rather than a fallback filling the gap. + GenerateContentResponse chunk3 = toResponseWithText("", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).thoughtSignature()).hasValue("real-sig".getBytes(UTF_8)); + } + + // The stream terminator is recognised by shape, not by identity: one that also carries an + // explicit thought marker is still nothing to keep. + @Test + public void processRawResponses_emptyTextPartWithExplicitThoughtFalse_isDropped() { + GenerateContentResponse chunk1 = toResponseWithText("Let me check."); + GenerateContentResponse chunk2 = toResponse(Part.builder().text("").thought(false).build()); + GenerateContentResponse chunk3 = toResponseWithText("", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).text()).hasValue("Let me check."); + } + + // A multi-part final chunk already carries each call's own signature. The re-attach reads part0 + // only, so it must not stamp the first call's signature onto the last one. + @Test + public void processRawResponses_multiCallFinalChunkSignedOnPart0_doesNotStampTheLastCall() { + Part signedCall = + Part.builder() + .functionCall(FunctionCall.builder().name("get_weather").id("fc-0").build()) + .thoughtSignature("call-sig".getBytes(UTF_8)) + .build(); + Part secondCall = + functionCallPart(FunctionCall.builder().name("get_weather").id("fc-1").build()); + Part thirdCall = + functionCallPart(FunctionCall.builder().name("get_weather").id("fc-2").build()); + GenerateContentResponse chunk = + toResponse( + Candidate.builder() + .content(Content.builder().parts(signedCall, secondCall, thirdCall).build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).thoughtSignature()).hasValue("call-sig".getBytes(UTF_8)); + assertThat(parts.get(1).thoughtSignature()).isEmpty(); + assertThat(parts.get(2).thoughtSignature()).isEmpty(); + } + @Test public void functionCallThenEmptyTextWithStop_emitsPartialThenFinalAggregatedFunctionCall() { Flowable rawResponses = @@ -1303,6 +1941,16 @@ private static Predicate isFinalThoughtResponseWithUsageMetadata( }; } + /** A partial chunk holding nothing but a thought marker and a signature. */ + private static Predicate isPartialSignatureResponse(String expectedSignature) { + return response -> { + assertThat(response.partial()).hasValue(true); + assertThat(GeminiUtil.getPart0FromLlmResponse(response).flatMap(Part::thoughtSignature)) + .hasValue(expectedSignature.getBytes(UTF_8)); + return true; + }; + } + private static Predicate isFinalThoughtResponseWithUsageMetadataAndSignature( String expectedText, GenerateContentResponseUsageMetadata expectedMetadata, @@ -1477,6 +2125,28 @@ private GenerateContentResponse toResponseWithText( .build(); } + private GenerateContentResponse toResponseWithTextAndSignature(String text, String signature) { + return toResponse( + Part.builder().text(text).thoughtSignature(signature.getBytes(UTF_8)).build()); + } + + private GenerateContentResponse toResponseWithTextAndSignature( + String text, String signature, FinishReason.Known finishReason) { + Part part = Part.builder().text(text).thoughtSignature(signature.getBytes(UTF_8)).build(); + return toResponse( + Candidate.builder() + .content(Content.builder().parts(part).build()) + .finishReason(new FinishReason(finishReason)) + .build()); + } + + /** Runs the chunks through the aggregator and returns the final (non-partial) response. */ + private static LlmResponse aggregateFinalResponse(GenerateContentResponse... chunks) { + return Iterables.getLast( + ImmutableList.copyOf( + Gemini.processRawResponses(Flowable.fromArray(chunks)).blockingIterable())); + } + private static Part functionCallPart(FunctionCall functionCall) { return Part.builder().functionCall(functionCall).build(); }