fix: recover compacted function calls and pair reused call IDs in LLM content assembly

PiperOrigin-RevId: 995900828

Author: google-genai-bot

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

@@ -29,6 +29,8 @@
 import com.google.adk.models.LlmRequest;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.genai.types.Content;
@@ -41,6 +43,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -176,7 +179,8 @@ private ImmutableList<Content> getContents(
     }
 
     if (hasCompactEvent) {
-      filteredEvents = processCompactionEvent(filteredEvents, agentName);
+      List<Event> compactedEvents = processCompactionEvent(filteredEvents, agentName);
+      filteredEvents = recoverCompactedFunctionCalls(compactedEvents, filteredEvents);
     }
 
     // Merge after compaction, which may end mid-run, and before other agents are narrated.
@@ -529,6 +533,99 @@ private static Event createCompactionEvent(Event event, String agentName) {
         .build();
   }
 
+  /**
+   * Re-injects function-call events from {@code sourceEvents} when compaction summarized away a
+   * call whose response arrived later. Restores the full call event to preserve first-part thought
+   * signatures on parallel calls, along with the latest compacted response for each sibling call so
+   * siblings do not appear pending. Resolves call ownership in {@code sourceEvents} order by event
+   * identity so reused call IDs bind to the latest preceding call.
+   */
+  private static List<Event> recoverCompactedFunctionCalls(
+      List<Event> events, List<Event> sourceEvents) {
+    if (events.stream().allMatch(event -> event.functionResponses().isEmpty())) {
+      return events;
+    }
+
+    Set<Event> presentEvents = Collections.newSetFromMap(new IdentityHashMap<>());
+    presentEvents.addAll(events);
+
+    Map<String, Event> latestCallById = new HashMap<>();
+    IdentityHashMap<Event, Map<String, Event>> callByResponse = new IdentityHashMap<>();
+    IdentityHashMap<Event, Map<String, Event>> responseByCall = new IdentityHashMap<>();
+    IdentityHashMap<Event, Set<String>> survivingResponses = new IdentityHashMap<>();
+
+    for (Event event : sourceEvents) {
+      for (FunctionCall functionCall : event.functionCalls()) {
+        functionCall
+            .id()
+            .filter(id -> !id.isEmpty())
+            .ifPresent(id -> latestCallById.put(id, event));
+      }
+      for (FunctionResponse functionResponse : event.functionResponses()) {
+        String responseId = functionResponse.id().orElse("");
+        if (responseId.isEmpty()) {
+          continue;
+        }
+        Event callEvent = latestCallById.get(responseId);
+        if (callEvent == null) {
+          continue;
+        }
+        callByResponse.computeIfAbsent(event, k -> new HashMap<>()).put(responseId, callEvent);
+        if (presentEvents.contains(event)) {
+          survivingResponses.computeIfAbsent(callEvent, k -> new HashSet<>()).add(responseId);
+        }
+        // Keep the latest update for this call; equal timestamps retain source order.
+        Map<String, Event> responsesForCall =
+            responseByCall.computeIfAbsent(callEvent, k -> new HashMap<>());
+        Event existing = responsesForCall.get(responseId);
+        if (existing == null || event.timestamp() >= existing.timestamp()) {
+          responsesForCall.put(responseId, event);
+        }
+      }
+    }
+
+    List<Event> result = new ArrayList<>();
+    Set<Event> reinjectedCalls = Collections.newSetFromMap(new IdentityHashMap<>());
+    for (Event event : events) {
+      Map<String, Event> callsForEvent = callByResponse.get(event);
+      if (callsForEvent != null) {
+        for (FunctionResponse functionResponse : event.functionResponses()) {
+          String responseId = functionResponse.id().orElse("");
+          Event callEvent = callsForEvent.get(responseId);
+          if (callEvent == null
+              || presentEvents.contains(callEvent)
+              || !reinjectedCalls.add(callEvent)) {
+            continue;
+          }
+          result.add(callEvent);
+          Set<String> survivingForCall =
+              survivingResponses.getOrDefault(callEvent, ImmutableSet.of());
+          Map<String, Event> responsesForCall =
+              responseByCall.getOrDefault(callEvent, ImmutableMap.of());
+          Set<Event> siblingResponsesToAdd = Collections.newSetFromMap(new IdentityHashMap<>());
+          for (FunctionCall siblingCall : callEvent.functionCalls()) {
+            String siblingId = siblingCall.id().orElse("");
+            if (!siblingId.isEmpty() && !survivingForCall.contains(siblingId)) {
+              Event siblingResponse = responsesForCall.get(siblingId);
+              if (siblingResponse != null) {
+                siblingResponsesToAdd.add(siblingResponse);
+              }
+            }
+          }
+          if (!siblingResponsesToAdd.isEmpty()) {
+            for (Event sourceEvent : sourceEvents) {
+              if (siblingResponsesToAdd.contains(sourceEvent)) {
+                result.add(sourceEvent);
+              }
+            }
+          }
+        }
+      }
+      result.add(event);
+    }
+    return reinjectedCalls.isEmpty() ? events : result;
+  }
+
   /** Whether the event is a reply from another agent. */
   private static boolean isOtherAgentReply(String agentName, Event event) {
     return !agentName.isEmpty()
@@ -770,27 +867,50 @@ private static List<Event> rearrangeEventsForLatestFunctionResponse(List<Event>
 
   private static List<Event> rearrangeEventsForAsyncFunctionResponsesInHistory(
       List<Event> events, boolean groupFunctionResponses) {
-    Map<String, Integer> functionCallIdToResponseEventIndex = new HashMap<>();
+    // Pair each response with the latest preceding call of that ID, keeping the last update.
+    Map<String, List<Integer>> callEventIndicesById = new HashMap<>();
     for (int i = 0; i < events.size(); i++) {
-      final int index = i;
-      Event event = events.get(index);
-      event
-          .content()
-          .flatMap(Content::parts)
-          .ifPresent(
-              parts -> {
-                for (Part part : parts) {
-                  part.functionResponse()
-                      .ifPresent(
-                          response ->
-                              response
-                                  .id()
-                                  .ifPresent(
-                                      functionCallId ->
-                                          functionCallIdToResponseEventIndex.put(
-                                              functionCallId, index)));
-                }
-              });
+      Event event = events.get(i);
+      if (!event.functionResponses().isEmpty()) {
+        continue;
+      }
+      for (FunctionCall call : event.functionCalls()) {
+        final int callIndex = i;
+        call.id()
+            .ifPresent(
+                id ->
+                    callEventIndicesById
+                        .computeIfAbsent(id, k -> new ArrayList<>())
+                        .add(callIndex));
+      }
+    }
+
+    Map<Integer, Map<String, Integer>> responseEventIndexByCall = new HashMap<>();
+    for (int i = 0; i < events.size(); i++) {
+      final int responseIndex = i;
+      Event event = events.get(responseIndex);
+      for (FunctionResponse response : event.functionResponses()) {
+        response
+            .id()
+            .ifPresent(
+                functionCallId -> {
+                  List<Integer> callIndices = callEventIndicesById.get(functionCallId);
+                  if (callIndices == null || callIndices.isEmpty()) {
+                    return;
+                  }
+                  int owningCallIndex = callIndices.get(0);
+                  for (int callIdx : callIndices) {
+                    if (callIdx < responseIndex) {
+                      owningCallIndex = callIdx;
+                    } else {
+                      break;
+                    }
+                  }
+                  responseEventIndexByCall
+                      .computeIfAbsent(owningCallIndex, k -> new HashMap<>())
+                      .put(functionCallId, responseIndex);
+                });
+      }
     }
 
     List<Event> resultEvents = new ArrayList<>();
@@ -818,6 +938,8 @@ private static List<Event> rearrangeEventsForAsyncFunctionResponsesInHistory(
               .orElse(false);
 
       if (hasFunctionCalls) {
+        Map<String, Integer> responsesForCall =
+            responseEventIndexByCall.getOrDefault(i, ImmutableMap.of());
         Set<Integer> responseEventIndices = new HashSet<>();
         // Iterate through parts again to get function call IDs
         partsOptional
@@ -830,11 +952,9 @@ private static List<Event> rearrangeEventsForAsyncFunctionResponsesInHistory(
                                 call.id()
                                     .ifPresent(
                                         functionCallId -> {
-                                          if (functionCallIdToResponseEventIndex.containsKey(
-                                              functionCallId)) {
-                                            responseEventIndices.add(
-                                                functionCallIdToResponseEventIndex.get(
-                                                    functionCallId));
+                                          Integer respIdx = responsesForCall.get(functionCallId);
+                                          if (respIdx != null) {
+                                            responseEventIndices.add(respIdx);
                                           }
                                         })));