feat(workflow): filter rewound invocations during workflow rehydration

Runner.rewind_async records rewinds by appending an event with
rewind_before_invocation_id rather than deleting historical events from
session.events. ReplayManager._ensure_index, ReplayManager.scan_workflow_events,
NodeRunner._create_child_context, _find_unresolved_task_delegations,
_find_active_task_scope, restore_branch_from_history,
_find_user_message_for_invocation, and _resolve_invocation_id_from_fr
previously read raw session.events directly, so re-running or resuming a
workflow after rewind_async replayed rewound node outputs, branches, user
inputs, and task scopes instead of executing fresh.

Filter session.events through _apply_rewinds before indexing or scanning
workflow history.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 987118607

Author: DeanChensj

src/google/adk/agents/_agent_router.py

@@ -194,11 +194,13 @@ def restore_branch_from_history(
   (a fresh direct-node turn, or a new invocation continuing a sub-agent), the
   most recent matching event across the session is used.
   """
+  from ..events._rewind_events import _apply_rewinds
   from ..workflow._base_node import find_static_node_path
 
+  live_events = _apply_rewinds(invocation_context.session.events)
   expected_static_path = find_static_node_path(root, node)
-  tool_call_ids = _collect_function_call_ids(invocation_context.session.events)
-  for event in reversed(invocation_context.session.events):
+  tool_call_ids = _collect_function_call_ids(live_events)
+  for event in reversed(live_events):
     if invocation_id is not None and event.invocation_id != invocation_id:
       continue
     if not event.branch:

src/google/adk/runners.py

@@ -49,6 +49,7 @@
 from .auth.credential_service.base_credential_service import BaseCredentialService
 from .errors._stale_session_error import StaleSessionError
 from .errors.session_not_found_error import SessionNotFoundError
+from .events._rewind_events import _apply_rewinds
 from .events.event import Event
 from .events.event_actions import EventActions
 from .flows.llm_flows.context import _contents as contents
@@ -120,8 +121,9 @@ def _find_active_task_scope(session: Session) -> Optional[tuple[str, str]]:
   # We must do this in a separate pass because walking backward directly would
   # hit post-finish events (like status updates or duplicate FRs) before hitting
   # the older success FR, falsely indicating the scope is still active.
+  live_events = _apply_rewinds(session.events)
   finished_scopes: set[str] = set()
-  for event in session.events:
+  for event in live_events:
     scope = event.isolation_scope
     if not scope:
       continue
@@ -138,7 +140,7 @@ def _find_active_task_scope(session: Session) -> Optional[tuple[str, str]]:
           break
 
   # Pass 2: Walk backward to find the latest active scope that is not finished.
-  for event in reversed(session.events):
+  for event in reversed(live_events):
     scope = event.isolation_scope
     if not scope:
       continue
@@ -637,7 +639,7 @@ def _resolve_invocation_id_from_fr(
 
     # Find invocation_id for each FR by matching its FC in session
     invocation_ids = set()
-    for event in reversed(session.events):
+    for event in reversed(_apply_rewinds(session.events)):
       for fc in event.get_function_calls():
         if fc.id in fr_ids:
           invocation_ids.add(event.invocation_id)
@@ -736,7 +738,7 @@ def _find_user_message_for_invocation(
     invocation used to fail outright, because the caller treats "not found" as
     an error.
     """
-    for event in events:
+    for event in _apply_rewinds(events):
       if (
           event.invocation_id == invocation_id
           and event.author == 'user'

src/google/adk/workflow/_llm_agent_wrapper.py

@@ -174,11 +174,12 @@ def _find_unresolved_task_delegations(
   current turn's scope would hide the coordinator's own FC from a
   prior turn.  Author + tool-name filtering is sufficient.
   """
+  from ..events._rewind_events import _apply_rewinds
   from ..tools.agent_tool import _TaskAgentTool
 
   fc_by_id: dict[str, types.FunctionCall] = {}
   fr_ids: set[str] = set()
-  for event in session.events:
+  for event in _apply_rewinds(session.events):
     if event.author != owner and event.author != 'user':
       continue
     if not event.content or not event.content.parts:

src/google/adk/workflow/_node_runner.py

@@ -239,11 +239,14 @@ def _create_child_context(
     )
 
     if ic.session and ic.session.events:
+      from ..events._rewind_events import _apply_rewinds
+
+      live_events = _apply_rewinds(ic.session.events)
       node_path = ctx.node_path
       node_path_builder = _NodePathBuilder.from_string(node_path)
       has_prior_node_events = bool(self._prior_interrupt_ids)
       if not has_prior_node_events:
-        for ev in ic.session.events:
+        for ev in live_events:
           if ic.invocation_id and ev.invocation_id != ic.invocation_id:
             continue
           if ev.node_info is not None and ev.node_info.path:
@@ -257,7 +260,7 @@ def _create_child_context(
         from .utils._rehydration_utils import _reconstruct_node_states
 
         states = _reconstruct_node_states(
-            events=ic.session.events,
+            events=live_events,
             base_path=node_path,
             invocation_id=ic.invocation_id,
         )

src/google/adk/workflow/utils/_replay_manager.py

@@ -21,6 +21,7 @@
 from ...agents.context import Context
 from ...events._branch_path import _BranchPath
 from ...events._node_path_builder import _NodePathBuilder
+from ...events._rewind_events import _apply_rewinds
 from ...events.event import Event
 from ._rehydration_utils import _ChildScanState
 from ._rehydration_utils import _reconstruct_node_states
@@ -50,6 +51,7 @@ def __init__(self) -> None:
     self._events_by_parent: dict[str, list[Event]] = {}
     self._transitive_events_by_parent: dict[str, list[Event]] = {}
     self._fc_to_parent: dict[str, str] = {}
+    self._live_events: list[Event] = []
     self._indexed_event_count: int = 0
     self._indexed_last_event: Event | None = None
 
@@ -76,14 +78,25 @@ def _ensure_index(self, ctx: Context) -> list[Event]:
     existing buckets then hold events the session no longer has.
     """
     ic = ctx._invocation_context
-    events = ic.session.events
-    if self._indexed_prefix_is_intact(events):
-      if len(events) > self._indexed_event_count:
-        self._index_events(events[self._indexed_event_count :])
-        self._record_indexed_through(events)
-    else:
-      self._build_event_index(events)
-    return events
+    raw_events = ic.session.events
+    if self._indexed_event_count > 0 and self._indexed_prefix_is_intact(
+        raw_events
+    ):
+      delta = raw_events[self._indexed_event_count :]
+      has_new_rewind = any(
+          ev.actions and ev.actions.rewind_before_invocation_id for ev in delta
+      )
+      if not has_new_rewind:
+        if delta:
+          self._live_events.extend(delta)
+          self._index_events(delta)
+          self._record_indexed_through(raw_events)
+        return self._live_events
+
+    self._live_events = _apply_rewinds(raw_events)
+    self._build_event_index(self._live_events)
+    self._record_indexed_through(raw_events)
+    return self._live_events
 
   def _indexed_prefix_is_intact(self, events: list[Event]) -> bool:
     """Whether the already-indexed events are still a prefix of `events`.
@@ -116,6 +129,7 @@ def _build_event_index(self, events: list[Event]) -> None:
     self._events_by_parent = {}
     self._transitive_events_by_parent = {}
     self._fc_to_parent = {}
+    self._live_events = list(events)
     self._index_events(events)
     self._record_indexed_through(events)
 
@@ -317,8 +331,7 @@ def scan_workflow_events(
     """Scan session events for direct child workflow nodes and initialize sequence barrier."""
     ic = ctx._invocation_context
 
-    # Build the index
-    self._build_event_index(ic.session.events)
+    self._ensure_index(ctx)
 
     # Use transitive parent events for static child nodes so deeper descendant events (e.g. delegated outputs/interrupts) are recovered
     filtered_events = self._transitive_events_by_parent.get(ctx.node_path, [])

tests/unittests/agents/test_agent_router.py

@@ -25,6 +25,7 @@
 from google.adk.agents.run_config import RunConfig
 from google.adk.apps.app import ResumabilityConfig
 from google.adk.events.event import Event
+from google.adk.events.event_actions import EventActions
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.sessions.session import Session
 from google.genai import types
@@ -495,3 +496,78 @@ def test_restore_branch_from_history():
 
   _agent_router.restore_branch_from_history(ic, sub1, root=root)
   assert ic.branch == "root@1.sub_agent1@1"
+
+
+def test_restore_branch_from_history_skips_rewound_events():
+  """restore_branch_from_history ignores branches authored in rewound invocations."""
+  session_service = InMemorySessionService()
+  session = Session(
+      id="s1",
+      app_name="app",
+      user_id="u1",
+      events=[
+          Event(
+              author="sub_agent1",
+              branch="root@1.sub_agent1@1",
+              invocation_id="inv_1",
+          ),
+          Event(
+              author="sub_agent1",
+              branch="root@1.sub_agent1@2",
+              invocation_id="inv_2",
+          ),
+          Event(
+              author="user",
+              invocation_id="inv_3",
+              actions=EventActions(rewind_before_invocation_id="inv_2"),
+          ),
+      ],
+  )
+  root, sub1, _, _ = _make_agent_tree()
+
+  ic = InvocationContext(
+      session_service=session_service,
+      invocation_id="inv_3",
+      agent=sub1,
+      session=session,
+      run_config=RunConfig(),
+  )
+  ic.branch = None
+
+  _agent_router.restore_branch_from_history(ic, sub1, root=root)
+  assert ic.branch == "root@1.sub_agent1@1"
+
+
+def test_restore_branch_from_history_all_rewound_leaves_branch_none():
+  """restore_branch_from_history leaves branch as None if all matches are rewound."""
+  session_service = InMemorySessionService()
+  session = Session(
+      id="s1",
+      app_name="app",
+      user_id="u1",
+      events=[
+          Event(
+              author="sub_agent1",
+              branch="root@1.sub_agent1@1",
+              invocation_id="inv_1",
+          ),
+          Event(
+              author="user",
+              invocation_id="inv_2",
+              actions=EventActions(rewind_before_invocation_id="inv_1"),
+          ),
+      ],
+  )
+  root, sub1, _, _ = _make_agent_tree()
+
+  ic = InvocationContext(
+      session_service=session_service,
+      invocation_id="inv_2",
+      agent=sub1,
+      session=session,
+      run_config=RunConfig(),
+  )
+  ic.branch = None
+
+  _agent_router.restore_branch_from_history(ic, sub1, root=root)
+  assert ic.branch is None

tests/unittests/workflow/test_dynamic_node_scheduler.py

@@ -118,6 +118,7 @@ def _make_fr_event(fc_id, response, invocation_id='inv-1'):
   event.branch = None
   event.isolation_scope = None
   event.long_running_tool_ids = None
+  event.actions = None
 
   fr = MagicMock()
   fr.id = fc_id