fix(flows): pair a function response with the call it answers

Close #6761

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 967478519

src/google/adk/flows/llm_flows/contents.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from bisect import bisect_left
 import copy
 import logging
 from typing import AsyncGenerator
@@ -147,19 +148,42 @@ def _rearrange_events_for_async_function_responses_in_history(
     events: list[Event],
 ) -> list[Event]:
   """Rearrange the async function_response events in the history."""
-  function_call_id_to_response_events_index: dict[str | None, int] = {}
+  # A model may hand out the same function call id more than once in a session,
+  # so an id on its own does not identify a single call. Each response is
+  # attributed to the newest call that precedes it and carries the same id, and
+  # a call then takes the last response attributed to it. Taking the last one
+  # keeps the closing update of a long-running tool, which reports progress
+  # several times under one id, while attributing first stops a reused id from
+  # handing a call the response that belongs to a different call.
+  call_event_indices_by_id: dict[str | None, list[int]] = {}
   for i, event in enumerate(events):
-    function_responses = event.get_function_responses()
-    if function_responses:
-      for function_response in function_responses:
-        function_call_id = function_response.id
-        function_call_id_to_response_events_index[function_call_id] = i
+    if event.get_function_responses():
+      continue
+    for function_call in event.get_function_calls():
+      call_event_indices_by_id.setdefault(function_call.id, []).append(i)
 
-  if not function_call_id_to_response_events_index:
+  response_event_index_by_call: dict[tuple[str | None, int], int] = {}
+  history_has_function_responses = False
+  for i, event in enumerate(events):
+    for function_response in event.get_function_responses():
+      history_has_function_responses = True
+      call_event_indices = call_event_indices_by_id.get(function_response.id)
+      if not call_event_indices:
+        continue
+      # Indices are collected in ascending order, so the call that owns this
+      # response is the one just before it. A response preceding every call
+      # that carries its id keeps the first, as it did before ids could repeat.
+      preceding_calls = bisect_left(call_event_indices, i)
+      owning_call_event_index = call_event_indices[max(preceding_calls - 1, 0)]
+      response_event_index_by_call[
+          (function_response.id, owning_call_event_index)
+      ] = i
+
+  if not history_has_function_responses:
     return events
 
   result_events: list[Event] = []
-  for event in events:
+  for i, event in enumerate(events):
     if event.get_function_responses():
       # function_response should be handled together with function_call below.
       continue
@@ -167,11 +191,11 @@ def _rearrange_events_for_async_function_responses_in_history(
 
       function_response_events_indices = set()
       for function_call in event.get_function_calls():
-        function_call_id = function_call.id
-        if function_call_id in function_call_id_to_response_events_index:
-          function_response_events_indices.add(
-              function_call_id_to_response_events_index[function_call_id]
-          )
+        response_event_index = response_event_index_by_call.get(
+            (function_call.id, i)
+        )
+        if response_event_index is not None:
+          function_response_events_indices.add(response_event_index)
       result_events.append(event)
       if not function_response_events_indices:
         continue

tests/unittests/flows/llm_flows/test_contents.py

@@ -1964,6 +1964,99 @@ def test_rearrange_async_function_responses_early_returns_when_no_responses():
   assert result is events
 
 
+def _function_call_event(call_id: str, name: str) -> Event:
+  return Event(
+      invocation_id="inv1",
+      author="test_agent",
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part(
+                  function_call=types.FunctionCall(
+                      id=call_id, name=name, args={}
+                  )
+              )
+          ],
+      ),
+  )
+
+
+def _function_response_event(call_id: str, name: str, result: str) -> Event:
+  return Event(
+      invocation_id="inv1",
+      author="user",
+      content=types.Content(
+          role="user",
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      id=call_id, name=name, response={"result": result}
+                  )
+              )
+          ],
+      ),
+  )
+
+
+def test_rearrange_async_function_responses_reused_id_across_tools():
+  """A reused call id must not pair a call with a different tool's response."""
+  events = [
+      _function_call_event("call_807", "site_posture"),
+      _function_response_event("call_807", "site_posture", "site"),
+      _function_call_event("call_807", "fleet_summary"),
+      _function_response_event("call_807", "fleet_summary", "fleet"),
+  ]
+
+  result = contents._rearrange_events_for_async_function_responses_in_history(  # pylint: disable=protected-access
+      events
+  )
+
+  assert len(result) == 4
+  assert result[0].get_function_calls()[0].name == "site_posture"
+  assert result[1].get_function_responses()[0].name == "site_posture"
+  assert result[2].get_function_calls()[0].name == "fleet_summary"
+  assert result[3].get_function_responses()[0].name == "fleet_summary"
+
+
+def test_rearrange_async_function_responses_reused_id_same_tool():
+  """A call id reused by one tool must keep each call's own response."""
+  events = [
+      _function_call_event("call_42", "lookup"),
+      _function_response_event("call_42", "lookup", "first"),
+      _function_call_event("call_42", "lookup"),
+      _function_response_event("call_42", "lookup", "second"),
+  ]
+
+  result = contents._rearrange_events_for_async_function_responses_in_history(  # pylint: disable=protected-access
+      events
+  )
+
+  assert len(result) == 4
+  assert result[1].get_function_responses()[0].response == {"result": "first"}
+  assert result[3].get_function_responses()[0].response == {"result": "second"}
+
+
+def test_rearrange_async_function_responses_reused_id_keeps_last_update():
+  """A call reporting progress twice keeps its last update, not a later call's."""
+  events = [
+      _function_call_event("call_7", "watch"),
+      _function_response_event("call_7", "watch", "progress"),
+      _function_response_event("call_7", "watch", "done"),
+      _function_call_event("call_7", "watch"),
+      _function_response_event("call_7", "watch", "second_call"),
+  ]
+
+  result = contents._rearrange_events_for_async_function_responses_in_history(  # pylint: disable=protected-access
+      events
+  )
+
+  assert len(result) == 4
+  assert result[1].get_function_responses()[0].response == {"result": "done"}
+  assert result[3].get_function_responses()[0].response == {
+      "result": "second_call"
+  }
+
+
 def _long_running_call_event() -> Event:
   return Event(
       invocation_id="inv2",