fix: propagate thought signatures in BaseLlmFlow

PiperOrigin-RevId: 836673368

Author: google-genai-bot

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -40,7 +40,9 @@
 import com.google.adk.tools.ToolContext;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.genai.types.Content;
 import com.google.genai.types.FunctionResponse;
+import com.google.genai.types.Part;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.context.Context;
@@ -56,6 +58,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -138,7 +141,8 @@ protected Flowable<Event> postprocess(
       InvocationContext context,
       Event baseEventForLlmResponse,
       LlmRequest llmRequest,
-      LlmResponse llmResponse) {
+      LlmResponse llmResponse,
+      AtomicReference<Optional<byte[]>> thoughtSignatureState) {
 
     List<Iterable<Event>> eventIterables = new ArrayList<>();
     Single<LlmResponse> currentLlmResponse = Single.just(llmResponse);
@@ -167,7 +171,8 @@ protected Flowable<Event> postprocess(
           }
 
           Event modelResponseEvent =
-              buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse);
+              buildModelResponseEvent(
+                  baseEventForLlmResponse, llmRequest, updatedResponse, thoughtSignatureState);
 
           Flowable<Event> modelEventStream = Flowable.just(modelResponseEvent);
 
@@ -335,6 +340,8 @@ private Single<LlmResponse> handleAfterModelCallback(
    */
   private Flowable<Event> runOneStep(InvocationContext context) {
     LlmRequest initialLlmRequest = LlmRequest.builder().build();
+    AtomicReference<Optional<byte[]>> thoughtSignatureState =
+        new AtomicReference<>(Optional.empty());
 
     return preprocess(context, initialLlmRequest)
         .flatMapPublisher(
@@ -373,7 +380,8 @@ private Flowable<Event> runOneStep(InvocationContext context) {
                                       context,
                                       mutableEventTemplate,
                                       llmRequestAfterPreprocess,
-                                      llmResponse)
+                                      llmResponse,
+                                      thoughtSignatureState)
                                   .doFinally(
                                       () -> {
                                         String oldId = mutableEventTemplate.id();
@@ -453,6 +461,8 @@ public Flowable<Event> run(InvocationContext invocationContext) {
   @Override
   public Flowable<Event> runLive(InvocationContext invocationContext) {
     LlmRequest llmRequest = LlmRequest.builder().build();
+    AtomicReference<Optional<byte[]>> thoughtSignatureState =
+        new AtomicReference<>(Optional.empty());
 
     return preprocess(invocationContext, llmRequest)
         .flatMapPublisher(
@@ -568,7 +578,8 @@ public void onError(Throwable e) {
                                 invocationContext,
                                 baseEventForThisLlmResponse,
                                 llmRequestAfterPreprocess,
-                                llmResponse);
+                                llmResponse,
+                                thoughtSignatureState);
                           })
                       .flatMap(
                           event -> {
@@ -624,7 +635,10 @@ public void onError(Throwable e) {
    * @return A fully constructed {@link Event} representing the LLM response.
    */
   private Event buildModelResponseEvent(
-      Event baseEventForLlmResponse, LlmRequest llmRequest, LlmResponse llmResponse) {
+      Event baseEventForLlmResponse,
+      LlmRequest llmRequest,
+      LlmResponse llmResponse,
+      AtomicReference<Optional<byte[]>> thoughtSignatureState) {
     Event.Builder eventBuilder =
         baseEventForLlmResponse.toBuilder()
             .content(llmResponse.content())
@@ -638,6 +652,33 @@ private Event buildModelResponseEvent(
             .finishReason(llmResponse.finishReason())
             .usageMetadata(llmResponse.usageMetadata());
 
+    List<Part> parts = llmResponse.content().flatMap(Content::parts).orElse(ImmutableList.of());
+    if (!parts.isEmpty()) {
+      boolean signaturePresentInResponse = false;
+      for (Part part : parts) {
+        if (part.thoughtSignature().isPresent()) {
+          signaturePresentInResponse = true;
+          thoughtSignatureState.set(part.thoughtSignature());
+        }
+      }
+
+      if (!signaturePresentInResponse
+          && parts.stream().anyMatch(p -> p.functionCall().isPresent())) {
+        ArrayList<Part> updatedParts = new ArrayList<>();
+        for (Part part : parts) {
+          Part partToAdd = part;
+          if (part.functionCall().isPresent() && part.thoughtSignature().isEmpty()) {
+            Optional<byte[]> signatureToApply = thoughtSignatureState.get();
+            if (signatureToApply.isPresent()) {
+              partToAdd = part.toBuilder().thoughtSignature(signatureToApply.get()).build();
+            }
+          }
+          updatedParts.add(partToAdd);
+        }
+        eventBuilder.content(llmResponse.content().get().toBuilder().parts(updatedParts).build());
+      }
+    }
+
     Event event = eventBuilder.build();
 
     if (!event.functionCalls().isEmpty()) {