From 00ac382cf2f474c942ad10fd5e1c50194626d8db Mon Sep 17 00:00:00 2001 From: rahulmansharamani14 Date: Tue, 24 Feb 2026 03:37:33 -0600 Subject: [PATCH 1/4] fix: skip rewound invocations when resuming agent in _find_agent_to_run --- src/google/adk/flows/llm_flows/contents.py | 47 +++++++---- src/google/adk/runners.py | 5 +- tests/unittests/test_runners.py | 91 ++++++++++++++++++++++ 3 files changed, 124 insertions(+), 19 deletions(-) diff --git a/src/google/adk/flows/llm_flows/contents.py b/src/google/adk/flows/llm_flows/contents.py index 9b7ef9e121..cd9fdc43c5 100644 --- a/src/google/adk/flows/llm_flows/contents.py +++ b/src/google/adk/flows/llm_flows/contents.py @@ -406,6 +406,35 @@ def _is_timestamp_compacted(ts: float) -> bool: return [event for _, _, event in processed_items] +def _filter_rewound_events(events: list[Event]) -> list[Event]: + """Returns events with those annulled by a rewind removed. + + Iterates backward; when a rewind marker is found, skips all events + back to the rewind_before_invocation_id. + + Args: + events: The full event list from the session. + + Returns: + A new list with rewound events removed, in the original order. + """ + filtered = [] + i = len(events) - 1 + while i >= 0: + event = events[i] + if event.actions and event.actions.rewind_before_invocation_id: + rewind_id = event.actions.rewind_before_invocation_id + for j in range(0, i): + if events[j].invocation_id == rewind_id: + i = j + break + else: + filtered.append(event) + i -= 1 + filtered.reverse() + return filtered + + def _get_contents( current_branch: Optional[str], events: list[Event], @@ -430,23 +459,7 @@ def _get_contents( accumulated_output_transcription = '' # Filter out events that are annulled by a rewind. - # By iterating backward, when a rewind event is found, we skip all events - # from that point back to the `rewind_before_invocation_id`, thus removing - # them from the history used for the LLM request. - rewind_filtered_events = [] - i = len(events) - 1 - while i >= 0: - event = events[i] - if event.actions and event.actions.rewind_before_invocation_id: - rewind_invocation_id = event.actions.rewind_before_invocation_id - for j in range(0, i, 1): - if events[j].invocation_id == rewind_invocation_id: - i = j - break - else: - rewind_filtered_events.append(event) - i -= 1 - rewind_filtered_events.reverse() + rewind_filtered_events = _filter_rewound_events(events) # Parse the events, leaving the contents and the function calls and # responses from the current agent. diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index cdb878cf24..d362981c0d 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -1055,7 +1055,8 @@ def _find_agent_to_run( # the agent that returned the corresponding function call regardless the # type of the agent. e.g. a remote a2a agent may surface a credential # request as a special long-running function tool call. - event = find_matching_function_call(session.events) + filtered_events = contents._filter_rewound_events(session.events) + event = find_matching_function_call(filtered_events) if event and event.author: return root_agent.find_agent(event.author) @@ -1067,7 +1068,7 @@ def _event_filter(event: Event) -> bool: return False return True - for event in filter(_event_filter, reversed(session.events)): + for event in filter(_event_filter, reversed(filtered_events)): if event.author == root_agent.name: # Found root agent. return root_agent diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index ca7eb37533..f53bdacb23 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -30,6 +30,8 @@ from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService from google.adk.cli.utils.agent_loader import AgentLoader from google.adk.events.event import Event +from google.adk.events.event import EventActions +from google.adk.flows.llm_flows.contents import _filter_rewound_events from google.adk.plugins.base_plugin import BasePlugin from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService @@ -641,6 +643,95 @@ def test_is_transferable_across_agent_tree_with_non_llm_agent(self): assert result is False +def test_find_agent_to_run_ignores_rewound_sub_agent_event(): + """After a rewind, events from the rewound invocation are ignored.""" + root_agent = MockLlmAgent("root_agent") + sub_agent1 = MockLlmAgent("sub_agent1", parent_agent=root_agent) + root_agent.sub_agents = [sub_agent1] + + runner = Runner( + app_name="test_app", + agent=root_agent, + session_service=InMemorySessionService(), + artifact_service=InMemoryArtifactService(), + ) + + # sub_agent1 was the last active agent during inv1 + sub_agent_event = Event( + invocation_id="inv1", + author="sub_agent1", + content=types.Content( + role="model", parts=[types.Part(text="Sub agent response")] + ), + ) + # Rewind event that annuls inv1 and everything after it + rewind_event = Event( + invocation_id="inv2", + author="user", + actions=EventActions(rewind_before_invocation_id="inv1"), + ) + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[sub_agent_event, rewind_event], + ) + + result = runner._find_agent_to_run(session, root_agent) + assert result == root_agent + + +def test_find_agent_to_run_ignores_rewound_function_call(): + """After a rewind, a function call from the rewound invocation is not matched.""" + root_agent = MockLlmAgent("root_agent") + sub_agent2 = MockLlmAgent("sub_agent2", parent_agent=root_agent) + root_agent.sub_agents = [sub_agent2] + + runner = Runner( + app_name="test_app", + agent=root_agent, + session_service=InMemorySessionService(), + artifact_service=InMemoryArtifactService(), + ) + + function_call = types.FunctionCall(id="func_789", name="test_func", args={}) + function_response = types.FunctionResponse( + id="func_789", name="test_func", response={} + ) + + # sub_agent2 issued a function call in inv1 + call_event = Event( + invocation_id="inv1", + author="sub_agent2", + content=types.Content( + role="model", parts=[types.Part(function_call=function_call)] + ), + ) + # User provides the function response, also in inv1 + response_event = Event( + invocation_id="inv1", + author="user", + content=types.Content( + role="user", parts=[types.Part(function_response=function_response)] + ), + ) + # Rewind event that annuls inv1 + rewind_event = Event( + invocation_id="inv2", + author="user", + actions=EventActions(rewind_before_invocation_id="inv1"), + ) + session = Session( + id="test_session", + user_id="test_user", + app_name="test_app", + events=[call_event, response_event, rewind_event], + ) + + # The rewound function call should not be matched; root_agent is returned + result = runner._find_agent_to_run(session, root_agent) + assert result == root_agent + @pytest.mark.asyncio async def test_run_config_custom_metadata_propagates_to_events(): session_service = InMemorySessionService() From cba37cdb72dccdb06ae380a2684f717d664049ab Mon Sep 17 00:00:00 2001 From: rahulmansharamani14 Date: Tue, 24 Feb 2026 04:44:19 -0600 Subject: [PATCH 2/4] refactor: make filter_rewound_events a public function --- src/google/adk/flows/llm_flows/contents.py | 4 ++-- src/google/adk/runners.py | 2 +- tests/unittests/test_runners.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/google/adk/flows/llm_flows/contents.py b/src/google/adk/flows/llm_flows/contents.py index cd9fdc43c5..ae6665ac6a 100644 --- a/src/google/adk/flows/llm_flows/contents.py +++ b/src/google/adk/flows/llm_flows/contents.py @@ -406,7 +406,7 @@ def _is_timestamp_compacted(ts: float) -> bool: return [event for _, _, event in processed_items] -def _filter_rewound_events(events: list[Event]) -> list[Event]: +def filter_rewound_events(events: list[Event]) -> list[Event]: """Returns events with those annulled by a rewind removed. Iterates backward; when a rewind marker is found, skips all events @@ -459,7 +459,7 @@ def _get_contents( accumulated_output_transcription = '' # Filter out events that are annulled by a rewind. - rewind_filtered_events = _filter_rewound_events(events) + rewind_filtered_events = filter_rewound_events(events) # Parse the events, leaving the contents and the function calls and # responses from the current agent. diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index d362981c0d..a4212f3b38 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -1055,7 +1055,7 @@ def _find_agent_to_run( # the agent that returned the corresponding function call regardless the # type of the agent. e.g. a remote a2a agent may surface a credential # request as a special long-running function tool call. - filtered_events = contents._filter_rewound_events(session.events) + filtered_events = contents.filter_rewound_events(session.events) event = find_matching_function_call(filtered_events) if event and event.author: return root_agent.find_agent(event.author) diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index f53bdacb23..4ef3297b2e 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -31,7 +31,7 @@ from google.adk.cli.utils.agent_loader import AgentLoader from google.adk.events.event import Event from google.adk.events.event import EventActions -from google.adk.flows.llm_flows.contents import _filter_rewound_events +from google.adk.flows.llm_flows.contents import filter_rewound_events from google.adk.plugins.base_plugin import BasePlugin from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService From a5d97ec16db8b74eb45c196d071ea1d9684c2011 Mon Sep 17 00:00:00 2001 From: rahulmansharamani14 Date: Tue, 24 Feb 2026 04:49:16 -0600 Subject: [PATCH 3/4] =?UTF-8?q?perf:=20optimize=20filter=5Frewound=5Fevent?= =?UTF-8?q?s=20from=20O(N=C2=B2)=20to=20O(N)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/google/adk/flows/llm_flows/contents.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/google/adk/flows/llm_flows/contents.py b/src/google/adk/flows/llm_flows/contents.py index ae6665ac6a..be1714321b 100644 --- a/src/google/adk/flows/llm_flows/contents.py +++ b/src/google/adk/flows/llm_flows/contents.py @@ -418,16 +418,20 @@ def filter_rewound_events(events: list[Event]) -> list[Event]: Returns: A new list with rewound events removed, in the original order. """ + # Pre-compute the first occurrence index of each invocation_id for O(1) lookup. + first_occurrence: dict[str, int] = {} + for idx, event in enumerate(events): + if event.invocation_id not in first_occurrence: + first_occurrence[event.invocation_id] = idx + filtered = [] i = len(events) - 1 while i >= 0: event = events[i] if event.actions and event.actions.rewind_before_invocation_id: rewind_id = event.actions.rewind_before_invocation_id - for j in range(0, i): - if events[j].invocation_id == rewind_id: - i = j - break + if rewind_id in first_occurrence and first_occurrence[rewind_id] < i: + i = first_occurrence[rewind_id] else: filtered.append(event) i -= 1 From d853bb62720c0ceba8f0e4b755da8854b0f72461 Mon Sep 17 00:00:00 2001 From: rahulmansharamani14 Date: Sat, 28 Feb 2026 01:38:37 -0600 Subject: [PATCH 4/4] style: apply autoformat --- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - tests/unittests/test_runners.py | 1 + 3 files changed, 1 insertion(+), 2 deletions(-) diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index f3751206a8..2710c3894c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index d857da9635..e31db15788 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index daa6de0e2d..071c3b3ff2 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -733,6 +733,7 @@ def test_find_agent_to_run_ignores_rewound_function_call(): result = runner._find_agent_to_run(session, root_agent) assert result == root_agent + @pytest.mark.asyncio async def test_run_config_custom_metadata_propagates_to_events(): session_service = InMemorySessionService()