fix: scope replay sequence to the current invocation

Merge #6498

Closes #6497

PiperOrigin-RevId: 955036653

Author: chelsealong21

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

@@ -63,7 +63,12 @@ def _ensure_index(self, ctx: Context) -> None:
       self._indexed_event_count = len(events)
 
   def _build_event_index(self, events: list[Event], invocation_id: str) -> None:
-    """Builds index of events grouped by parent path (both direct and transitive)."""
+    """Builds index of events grouped by parent path (both direct and transitive).
+
+    The index intentionally spans every invocation in the session so multi-turn
+    conversation context stays visible during rehydration. Consumers that need
+    a single invocation must therefore filter by `invocation_id` themselves.
+    """
     self._events_by_parent = {}
     self._transitive_events_by_parent = {}
     fc_to_parent: dict[str, str] = {}
@@ -176,8 +181,17 @@ def _scan_sequence(
     """Extract chronological child completion sequence under base_path."""
     base_path_builder = _NodePathBuilder.from_string(base_path)
     sequence: list[str] = []
+    invocation_id = ctx._invocation_context.invocation_id
 
     for event in events:
+      # The event index spans the whole session so multi-turn context stays
+      # visible during rehydration. The replay sequence, however, must describe
+      # only the current invocation: a terminal event from an earlier completed
+      # invocation would otherwise block the barrier on a node that never runs
+      # during this resume.
+      if invocation_id and event.invocation_id != invocation_id:
+        continue
+
       event_node_path = event.node_info.path or ""
       event_path_builder = _NodePathBuilder.from_string(event_node_path)
 

tests/unittests/workflow/test_workflow_hitl.py

@@ -2091,3 +2091,111 @@ def source_b():
   # We assert that the FIRST trigger (for from_a) was processed FIRST!
   assert interrupt_id_2_first == 'req_d_from_a'
   assert interrupt_id_2_second == 'req_e_from_b'
+
+
+class _BranchRouterNode(BaseNode):
+  """Routes the first invocation and later invocations down different branches.
+
+  Stands in for the LlmAgent classification in the original report: the model
+  is incidental, what matters is that two invocations in one session take
+  different branches, so the first invocation ends on a node the second never
+  runs.
+  """
+
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+  seen_invocations: list[str] = Field(default_factory=list)
+
+  @override
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    self.seen_invocations.append(ctx.invocation_id)
+    distinct = list(dict.fromkeys(self.seen_invocations))
+    route = 'DONE' if len(distinct) == 1 else 'NEEDS_INPUT'
+    yield Event(output=node_input, route=route)
+
+
+class _ClarifyNode(BaseNode):
+  """Pauses for input on first execution, emits the answer once resumed."""
+
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+  rerun_on_resume: bool = Field(default=True)
+
+  @override
+  async def _run_impl(
+      self, *, ctx: Context, node_input: Any
+  ) -> AsyncGenerator[Any, None]:
+    interrupt_id = f'clarify:{ctx.run_id}'
+    response = ctx.resume_inputs.get(interrupt_id)
+    if response is None:
+      yield RequestInput(interrupt_id=interrupt_id, message='Which city?')
+      return
+    yield Event(output=f'resumed:{response}')
+
+
+@pytest.mark.asyncio
+async def test_request_input_resume_after_earlier_invocation_completed(
+    request: pytest.FixtureRequest,
+):
+  """A completed earlier invocation must not block a later HITL resume.
+
+  Regression test for #6497. The first invocation finishes on the `finish`
+  branch. The second invocation takes the `clarify` branch and pauses for
+  input. When the replay sequence was built from every event in the session,
+  the terminal `finish` event of the first invocation entered the sequence
+  and, being chronologically first, was the only key the sequence barrier
+  unblocked. `finish` never runs during the resume, so the pending node waited
+  out the barrier timeout and raised "Replay divergence detected".
+  """
+  prepare = _TestingNode(name='prepare_intent_text', message='prepped')
+  router = _BranchRouterNode(name='route')
+  finish = _TestingNode(name='finish', message='done')
+  clarify = _ClarifyNode(name='clarify')
+
+  agent = Workflow(
+      name='test_workflow_replay_across_invocations',
+      edges=[
+          Edge(from_node=START, to_node=prepare),
+          Edge(from_node=prepare, to_node=router),
+          Edge(from_node=router, to_node=finish, route='DONE'),
+          Edge(from_node=router, to_node=clarify, route='NEEDS_INPUT'),
+      ],
+  )
+  app = App(
+      name=request.function.__name__,
+      root_agent=agent,
+      resumability_config=ResumabilityConfig(is_resumable=True),
+  )
+  runner = testing_utils.InMemoryRunner(app=app)
+
+  # Invocation 1: runs to completion down the `finish` branch.
+  events1 = await runner.run_async(
+      testing_utils.get_user_content('who are you')
+  )
+  outputs1 = [e.output for e in events1 if e.output is not None]
+  assert 'done' in outputs1
+
+  # Invocation 2, same session: takes the `clarify` branch and pauses.
+  events2 = await runner.run_async(
+      testing_utils.get_user_content('weather please')
+  )
+  request_input_event = workflow_testing_utils.find_function_call_event(
+      events2, REQUEST_INPUT_FUNCTION_CALL_NAME
+  )
+  assert request_input_event is not None
+  interrupt_id = get_request_input_interrupt_ids(request_input_event)[0]
+  invocation_id = request_input_event.invocation_id
+
+  # Resuming must reach the pending node instead of timing out on `finish`.
+  events3 = await runner.run_async(
+      new_message=testing_utils.UserContent(
+          create_request_input_response(interrupt_id, {'result': 'Berlin'})
+      ),
+      invocation_id=invocation_id,
+  )
+
+  outputs3 = [e.output for e in events3 if e.output is not None]
+  assert any(
+      isinstance(output, str) and output.startswith('resumed:')
+      for output in outputs3
+  ), outputs3

tests/unittests/workflow/utils/test_replay_manager.py

@@ -14,6 +14,7 @@
 
 """Tests for ReplayManager utility."""
 
+import asyncio
 from unittest.mock import MagicMock
 
 from google.adk.events.event import Event
@@ -185,3 +186,110 @@ def test_scan_workflow_events_recovers_children_from_transitive_descendant_event
   recovered, _ = mgr.scan_workflow_events(ctx)
 
   assert "child_a@1" in recovered
+
+
+def test_scan_workflow_events_sequence_excludes_prior_invocation_events():
+  """Replay sequence covers only the current invocation.
+
+  A session may hold a completed earlier invocation followed by a second
+  invocation that pauses for human input. Terminal events from the earlier
+  invocation must not enter the replay sequence, otherwise the sequence
+  barrier blocks on a node that never runs during the resume.
+  """
+  mgr = ReplayManager()
+  # Completed earlier invocation in the same session.
+  prior = Event(
+      author="node",
+      node_info=NodeInfo(path="wf@1/finish@1", run_id="1"),
+      invocation_id="inv-1",
+      output="prior_out",
+  )
+  # Current invocation, ending on an unresolved RequestInput interrupt.
+  current_first = Event(
+      author="node",
+      node_info=NodeInfo(path="wf@1/alpha@1", run_id="1"),
+      invocation_id="inv-2",
+      output="alpha_out",
+  )
+  current_pending = Event(
+      author="node",
+      node_info=NodeInfo(path="wf@1/beta@1", run_id="1"),
+      invocation_id="inv-2",
+      long_running_tool_ids=["clarify:1"],
+  )
+
+  ctx = MagicMock()
+  ctx._invocation_context = MagicMock()
+  ctx._invocation_context.invocation_id = "inv-2"
+  ctx._invocation_context.session = MagicMock()
+  ctx._invocation_context.session.events = [
+      prior,
+      current_first,
+      current_pending,
+  ]
+  ctx.node_path = "wf@1"
+
+  recovered, sequence = mgr.scan_workflow_events(ctx)
+
+  assert sequence == ["alpha@1", "beta@1"]
+  # Sequence and recovered state must agree; disagreement was the defect.
+  assert "finish@1" not in recovered
+  # The fix belongs in _scan_sequence, NOT in the event index: the index
+  # deliberately spans the whole session so multi-turn context stays visible
+  # during rehydration. Filtering there instead would pass the assertions
+  # above while silently breaking cross-turn context.
+  assert prior in mgr._transitive_events_by_parent["wf@1"]
+
+
+def test_prepare_parent_sequence_barrier_excludes_prior_invocation_events():
+  """Dynamic-node sequence barriers are also scoped to the current invocation."""
+  mgr = ReplayManager()
+  prior = Event(
+      author="node",
+      node_info=NodeInfo(path="wf@1/finish@1", run_id="1"),
+      invocation_id="inv-1",
+      output="prior_out",
+  )
+  current = Event(
+      author="node",
+      node_info=NodeInfo(path="wf@1/alpha@1", run_id="1"),
+      invocation_id="inv-2",
+      output="alpha_out",
+  )
+
+  ctx = MagicMock()
+  ctx._invocation_context = MagicMock()
+  ctx._invocation_context.invocation_id = "inv-2"
+  ctx._invocation_context.session = MagicMock()
+  ctx._invocation_context.session.events = [prior, current]
+  ctx.node_path = "wf@1"
+
+  barrier = mgr.prepare_parent_sequence_barrier(ctx, "wf@1")
+
+  assert barrier.sequence == ["alpha@1"]
+  assert prior in mgr._events_by_parent["wf@1"]
+
+
+@pytest.mark.asyncio
+async def test_scan_workflow_events_sequence_empty_when_all_events_are_prior():
+  """A session holding only prior-invocation events yields a non-blocking barrier."""
+  mgr = ReplayManager()
+  prior = Event(
+      author="node",
+      node_info=NodeInfo(path="wf@1/finish@1", run_id="1"),
+      invocation_id="inv-1",
+      output="prior_out",
+  )
+
+  ctx = MagicMock()
+  ctx._invocation_context = MagicMock()
+  ctx._invocation_context.invocation_id = "inv-2"
+  ctx._invocation_context.session = MagicMock()
+  ctx._invocation_context.session.events = [prior]
+  ctx.node_path = "wf@1"
+
+  _, sequence = mgr.scan_workflow_events(ctx)
+
+  assert sequence == []
+  # An empty sequence must fast-forward rather than deadlock.
+  await asyncio.wait_for(mgr.sequence_barrier.wait("anything"), timeout=1)