From 4385b969afacfc49d42bbe99ca2f8d28e84f362c Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Sat, 22 Aug 2026 22:31:37 +0100
Subject: [PATCH] fix(flows): drop stale async tool replies

---
 src/google/adk/flows/llm_flows/contents.py    | 34 ++++++++++++++++---
 .../flows/llm_flows/test_contents.py          | 33 ++++++++++++++++++
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/src/google/adk/flows/llm_flows/contents.py b/src/google/adk/flows/llm_flows/contents.py
index e2ef2f9292..5f276382af 100644
--- a/src/google/adk/flows/llm_flows/contents.py
+++ b/src/google/adk/flows/llm_flows/contents.py
@@ -148,7 +148,7 @@ async def run_async(
 def _rearrange_events_for_async_function_responses_in_history(
     events: list[Event],
 ) -> list[Event]:
-  """Rearrange the async function_response events in the history."""
+  """Rearrange async function responses and their model replies in history."""
   # 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
@@ -164,6 +164,7 @@ def _rearrange_events_for_async_function_responses_in_history(
       call_event_indices_by_id.setdefault(function_call.id, []).append(i)
 
   response_event_index_by_call: dict[tuple[str | None, int], int] = {}
+  response_event_indices_by_call: dict[tuple[str | None, int], list[int]] = {}
   history_has_function_responses = False
   for i, event in enumerate(events):
     for function_response in event.get_function_responses():
@@ -176,15 +177,40 @@ def _rearrange_events_for_async_function_responses_in_history(
       # 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
+      call_key = (function_response.id, owning_call_event_index)
+      response_event_index_by_call[call_key] = i
+      response_event_indices_by_call.setdefault(call_key, []).append(i)
 
   if not history_has_function_responses:
     return events
 
+  # A text-only model reply immediately after an old tool update was generated
+  # from the result that is about to be discarded. Remove that reply as well,
+  # while stopping at any user input, function call, or other event boundary.
+  latest_response_event_indices = set(response_event_index_by_call.values())
+  superseded_response_event_indices = {
+      response_event_index
+      for response_event_indices in response_event_indices_by_call.values()
+      for response_event_index in response_event_indices
+      if response_event_index not in latest_response_event_indices
+  }
+  stale_model_event_indices: set[int] = set()
+  for response_event_index in superseded_response_event_indices:
+    for i in range(response_event_index + 1, len(events)):
+      event = events[i]
+      if (
+          not event.content
+          or event.content.role != 'model'
+          or event.get_function_calls()
+          or not any(part.text for part in event.content.parts or [])
+      ):
+        break
+      stale_model_event_indices.add(i)
+
   result_events: list[Event] = []
   for i, event in enumerate(events):
+    if i in stale_model_event_indices:
+      continue
     if event.get_function_responses():
       # function_response should be handled together with function_call below.
       continue
diff --git a/tests/unittests/flows/llm_flows/test_contents.py b/tests/unittests/flows/llm_flows/test_contents.py
index 309b82af06..af6f9375b5 100644
--- a/tests/unittests/flows/llm_flows/test_contents.py
+++ b/tests/unittests/flows/llm_flows/test_contents.py
@@ -1971,6 +1971,39 @@ def _function_response_event(call_id: str, name: str, result: str) -> Event:
   )
 
 
+def _model_text_event(text: str) -> Event:
+  return Event(
+      invocation_id="inv1",
+      author="test_agent",
+      content=types.ModelContent(text),
+  )
+
+
+def test_rearrange_async_function_responses_drops_stale_model_reply():
+  """A model reply to a superseded tool update must not remain in history."""
+  events = [
+      _function_call_event("call_1", "watch"),
+      _function_response_event("call_1", "watch", "progress"),
+      _model_text_event("Still working."),
+      _function_response_event("call_1", "watch", "done"),
+      _model_text_event("Finished."),
+      Event(
+          invocation_id="inv2",
+          author="user",
+          content=types.UserContent("What happened?"),
+      ),
+  ]
+
+  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[2].content.parts[0].text == "Finished."
+  assert result[3].content.parts[0].text == "What happened?"
+
+
 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 = [