From 0335659788daa1960ca17977d1315cfafd35fef7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Fri, 7 Aug 2026 06:59:57 -0700 Subject: [PATCH] fix: keep each streamed text run's own thought signature PiperOrigin-RevId: 960907051 --- .../google/adk/flows/llmflows/Contents.java | 47 ++++- .../java/com/google/adk/models/Gemini.java | 28 +-- .../adk/flows/llmflows/ContentsTest.java | 170 ++++++++++++++++ .../com/google/adk/models/GeminiTest.java | 188 ++++++++++++++++++ 4 files changed, 411 insertions(+), 22 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..b6ddcd5ea 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), + * OR + *
  • It is marked as a thought AND does not contain function_call, function_response, + * tool_call or tool_response *
* *

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. 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,12 @@ private boolean isPartInvisible(Part part) { if (part.functionCall().isPresent() || part.functionResponse().isPresent()) { 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 +401,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()) { @@ -423,9 +442,19 @@ 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 thought-only part, a bare signature, a server-side tool call - carries no + // narratable content, and a server-side call in particular belongs to the model instance that + // made 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..fe4ce99a5 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,9 @@ 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; + // Signature of the buffered text run, kept apart from the streamed function 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; @@ -421,18 +424,17 @@ 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; } + // The signature rides on the merged part that flushTextBufferToSequence builds. + part.thoughtSignature().ifPresent(sig -> currentTextThoughtSignature = sig); currentTextBuffer.append(text); } else if (part.functionCall().isPresent()) { hasContent = true; @@ -443,16 +445,16 @@ private boolean accumulateParts(List parts) { // 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. + // thoughtSignature it carries. The signature is intentionally not captured, 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); + // its own; its signature rides on the text run it sits in, and overrides what that run's + // own chunks carried, because a signature-only part is an explicit carrier. + part.thoughtSignature().ifPresent(sig -> currentTextThoughtSignature = sig); } } return hasContent; @@ -605,9 +607,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); 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..55462c7fc 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 @@ -33,10 +33,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 +949,164 @@ public void processRequest_notEmptyContent() { assertThat(contents).containsExactly(e.content().get()); } + // 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."); + } + + // 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 +1189,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..80efe3a24 100644 --- a/core/src/test/java/com/google/adk/models/GeminiTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiTest.java @@ -15,6 +15,7 @@ */ package com.google.adk.models; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import static java.nio.charset.StandardCharsets.UTF_8; @@ -35,6 +36,7 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.functions.Predicate; import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.Optional; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -1068,6 +1070,170 @@ public void processRawResponses_emptyPartsThenSignature_doesNotThrowException() 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. This class keeps the last of the run rather than the + // first, because the part it flushes can still be rewritten by the final chunk's signature. + @Test + public void processRawResponses_multipleSignaturesInOneRun_keepsTheLast() { + 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("second-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 is not emitted on its own, so its signature has to ride on the + // text run it sits in. It is dropped if it lands in the function call's slot instead. + @Test + public void processRawResponses_standaloneSignatureMidTextRun_ridesOnTheMergedText() { + 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(1); + assertThat(parts.get(0).text()).hasValue("At minute 5 the presenter speaks."); + assertThat(parts.get(0).thoughtSignature()).hasValue("carried-sig".getBytes(UTF_8)); + } + + // 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)); + } + + // 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 signatures = + finalResponse.content().get().parts().get().stream() + .map(Part::thoughtSignature) + .flatMap(Optional::stream) + .map(signature -> new String(signature, UTF_8)) + .collect(toImmutableList()); + assertThat(signatures).containsExactly("call-context"); + } + @Test public void functionCallThenEmptyTextWithStop_emitsPartialThenFinalAggregatedFunctionCall() { Flowable rawResponses = @@ -1477,6 +1643,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(); }