fix(workflow): avoid O(n^2) deep event comparison during rehydration

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

@@ -144,16 +144,24 @@ def get_events_for_rehydration(
     # Top-level user text prompts live under root key ("").
     # Merge them so multi-turn turn inputs remain visible during state reconstruction.
     root_events = self._events_by_parent.get("", [])
+    # Membership is by identity, matching the `id()` set used below: both lists
+    # are views onto `session.events`, so an event is already indexed under this
+    # node only if it is the same object. `e not in node_events` would instead
+    # invoke `Event.__eq__` per probe, deep-comparing whole events and making
+    # this loop quadratic in the session size.
+    node_event_ids = {id(e) for e in node_events}
     user_prompts = [
-        e for e in root_events if e.author == "user" and e not in node_events
+        e
+        for e in root_events
+        if e.author == "user" and id(e) not in node_event_ids
     ]
 
     if not user_prompts:
       return node_events
 
     # Retain exact chronological ordering of session events.
     session_events = ctx._invocation_context.session.events
-    event_ids = {id(e) for e in node_events}.union(id(e) for e in user_prompts)
+    event_ids = node_event_ids.union(id(e) for e in user_prompts)
     return [e for e in session_events if id(e) in event_ids]
 
   def _add_event_to_index(self, parent_path: str, event: Event) -> None:

tests/unittests/workflow/utils/test_replay_manager.py

@@ -170,6 +170,110 @@ def test_get_events_for_rehydration_lazily_builds_event_index():
   assert events == [e_a]
 
 
+def _duplicate_user_response_event():
+  """A user function-response event with fixed `id`/`timestamp`.
+
+  Two of these are distinct objects that nonetheless compare equal, which is how
+  a value-based membership test can confuse one for the other.
+  """
+  from google.genai import types
+
+  return Event(
+      author="user",
+      invocation_id="inv-1",
+      id="duplicate-event-id",
+      timestamp=1000.0,
+      content=types.Content(
+          role="user",
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      name="RequestInput", id="fc-1", response={"result": "ok"}
+                  )
+              )
+          ],
+      ),
+  )
+
+
+def test_get_events_for_rehydration_merges_user_prompts_by_identity():
+  """A user prompt is not dropped just because an equal-valued event is indexed.
+
+  `e_user_root` is indexed under root because it arrives before the interrupt id
+  that would route it, while `e_user_routed` arrives after and is indexed under
+  the node path. The two compare equal, so testing membership with `in` treats
+  the root prompt as already present and silently drops it from rehydration.
+  """
+  mgr = ReplayManager()
+  e_user_root = _duplicate_user_response_event()
+  e_node = Event(
+      author="node",
+      node_info=NodeInfo(path="wf@1/child_a@1"),
+      invocation_id="inv-1",
+      long_running_tool_ids=["fc-1"],
+  )
+  e_user_routed = _duplicate_user_response_event()
+  assert e_user_root == e_user_routed
+  assert e_user_root is not e_user_routed
+
+  session_events = [e_user_root, e_node, e_user_routed]
+  ctx = MagicMock()
+  ctx._invocation_context = MagicMock()
+  ctx._invocation_context.invocation_id = "inv-1"
+  ctx._invocation_context.session = MagicMock()
+  ctx._invocation_context.session.events = session_events
+
+  events = mgr.get_events_for_rehydration(ctx, "wf@1/child_a@1")
+
+  # Every event is preserved, in session order.
+  assert [id(e) for e in events] == [id(e) for e in session_events]
+
+
+def test_get_events_for_rehydration_does_not_deep_compare_events():
+  """Merging user prompts must not invoke `Event.__eq__`.
+
+  `Event.__eq__` is a recursive deep comparison, so using it for a membership
+  test over a list makes this path quadratic in the number of session events.
+  """
+  eq_calls = 0
+
+  class CountingEvent(Event):
+
+    def __eq__(self, other):
+      nonlocal eq_calls
+      eq_calls += 1
+      return super().__eq__(other)
+
+    __hash__ = None
+
+  mgr = ReplayManager()
+  # A root-indexed user prompt is required for the merge path to run at all.
+  e_user = CountingEvent(**_duplicate_user_response_event().model_dump())
+  session_events = [
+      e_user,
+      *[
+          CountingEvent(
+              author="node",
+              node_info=NodeInfo(path="wf@1/child_a@1"),
+              invocation_id="inv-1",
+          )
+          for _ in range(20)
+      ],
+  ]
+  ctx = MagicMock()
+  ctx._invocation_context = MagicMock()
+  ctx._invocation_context.invocation_id = "inv-1"
+  ctx._invocation_context.session = MagicMock()
+  ctx._invocation_context.session.events = session_events
+
+  events = mgr.get_events_for_rehydration(ctx, "wf@1/child_a@1")
+
+  # Sanity-check that the merge path actually ran, so `eq_calls == 0` below
+  # means "no deep comparison" rather than "nothing was compared".
+  assert len(events) == len(session_events)
+  assert eq_calls == 0
+
+
 def test_scan_workflow_events_recovers_children_from_transitive_descendant_events():
   """Scanning workflow events recovers child nodes when events are emitted deep in child subtrees."""
   mgr = ReplayManager()