fix: Support multi-turn nested HITL pause and resumption across workflows and agent tools

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

Author: DeanChensj

src/google/adk/flows/llm_flows/_resume_utils.py

@@ -0,0 +1,259 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Deciding how a resumable LLM flow continues from the events it already has.
+
+A resumed invocation replays the branch's events and has to answer one
+question before it may call the LLM again: is this branch still waiting on a
+tool, does it owe a tool call that was never executed, or is it free to carry
+on? The matching that answers it is fiddly -- ids, names, long-running calls
+and HITL answers that come back on a sub-branch rather than against the
+original call -- so it lives here rather than inline in the flow.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import enum
+from typing import Any
+from typing import TYPE_CHECKING
+
+from google.genai import types
+
+from ...events._branch_path import _BranchPath
+from ...events.event import Event
+
+if TYPE_CHECKING:
+  from ...agents.invocation_context import InvocationContext
+
+
+class ResumeAction(enum.Enum):
+  """What the flow should do with the events it resumed from."""
+
+  CONTINUE = 'continue'
+  """Nothing outstanding; proceed to the LLM call."""
+
+  PAUSE = 'pause'
+  """A tool call is still unanswered; stop without emitting anything."""
+
+  REPLAY_CALLS = 'replay_calls'
+  """A tool call was never executed; run the calls on `ResumeDecision.event`."""
+
+
+@dataclasses.dataclass(frozen=True)
+class ResumeDecision:
+  """The action to take, and the event it applies to."""
+
+  action: ResumeAction
+  event: Event | None = None
+
+  def replay_event(self) -> Event:
+    """The event whose calls to run. Only a REPLAY_CALLS decision carries one.
+
+    Raises:
+      ValueError: If the decision names no event, which would mean
+        `decide_resume` returned REPLAY_CALLS without saying what to replay.
+    """
+    if self.event is None:
+      raise ValueError(f'{self.action} decision carries no event to replay')
+    return self.event
+
+
+def _branch_carries_call(
+    branch: str | None, function_calls: list[types.FunctionCall]
+) -> bool:
+  """Whether `branch` was opened by one of `function_calls`.
+
+  A branch is a dot-joined `name@run_id` path, so the run ids are parsed out and
+  compared whole: testing `id in branch` as a substring matches any id that
+  merely contains this one.
+  """
+  if not branch:
+    return False
+  run_ids = _BranchPath.from_string(branch).run_ids
+  return any(fc.id in run_ids for fc in function_calls if fc.id is not None)
+
+
+def _pause_left_calls_unanswered(
+    invocation_context: InvocationContext, events: list[Event]
+) -> bool:
+  """Whether a pause earlier in `events` is still waiting on a response.
+
+  Every event before the last is considered, not just the previous one: an LRO
+  followed by several text responses leaves the pausing call further back than
+  a two-event window can see.
+  """
+  pause_events = [
+      ev for ev in events[:-1] if invocation_context.should_pause_invocation(ev)
+  ]
+  if not pause_events:
+    return False
+  awaited = {
+      fc.id for ev in pause_events for fc in ev.get_function_calls() if fc.id
+  }
+  for ev in pause_events:
+    if ev.long_running_tool_ids:
+      awaited.update(ev.long_running_tool_ids)
+  answered = {
+      fr.id for ev in events for fr in ev.get_function_responses() if fr.id
+  }
+  # `issubset`, not `&`: this asks whether *any* awaited id is still open, so a
+  # partially answered pause keeps waiting. `decide_resume` asks the opposite
+  # question of its own ids -- whether *none* are answered -- and drops
+  # `issubset` for that reason. The two are not interchangeable.
+  return bool(awaited) and not awaited.issubset(answered)
+
+
+def _find_target_call_event(
+    events: list[Event], tools_dict: dict[str, Any]
+) -> Event | None:
+  """The most recent event before the last that calls a tool this flow owns."""
+  for ev in reversed(events[:-1]):
+    calls = ev.get_function_calls()
+    if calls and any(fc.name in tools_dict for fc in calls):
+      return ev
+  return None
+
+
+def _find_answer_event(
+    events: list[Event],
+    call_event: Event,
+    call_idx: int,
+    call_ids: set[str | None],
+    call_names: set[str | None],
+) -> Event:
+  """The event answering `call_event`, or the last event when none does.
+
+  A response counts when it carries a matching id, or a matching name with no
+  id, or is a HITL prompt raised on a branch that one of the calls opened --
+  the nested case, where the answer arrives against the sub-branch instead of
+  against the original call id.
+  """
+  # Imported here, not at module scope: google.adk.workflow imports back into
+  # the flows package.
+  # pylint: disable=g-import-not-at-top
+  from ...workflow.utils._workflow_hitl_utils import REQUEST_CREDENTIAL_FUNCTION_CALL_NAME
+  from ...workflow.utils._workflow_hitl_utils import REQUEST_INPUT_FUNCTION_CALL_NAME
+
+  # pylint: enable=g-import-not-at-top
+
+  hitl_names = {
+      REQUEST_INPUT_FUNCTION_CALL_NAME,
+      REQUEST_CREDENTIAL_FUNCTION_CALL_NAME,
+  }
+  calls = call_event.get_function_calls()
+  # `call_idx` is passed in rather than searched for again: the caller has
+  # already located `call_event`, and this runs on every resumable step.
+  start = call_idx + 1
+  for ev in reversed(events[start:]):
+    for fr in ev.get_function_responses():
+      if (
+          (fr.id is not None and fr.id in call_ids)
+          or (fr.id is None and fr.name in call_names)
+          or (fr.name in hitl_names and _branch_carries_call(ev.branch, calls))
+      ):
+        return ev
+  return events[-1]
+
+
+def _is_sub_branch_answer(answer_event: Event, call_event: Event) -> bool:
+  """Whether the answer came back from a branch the call opened."""
+  return answer_event.author == 'user' and _branch_carries_call(
+      answer_event.branch, call_event.get_function_calls()
+  )
+
+
+def _needs_call_replay(
+    call_names: set[str | None],
+    answers: list[types.FunctionResponse],
+    from_sub_branch: bool,
+) -> bool:
+  """Whether the calls named by `call_names` still have to be run.
+
+  `call_names` holds every name on the call event, not just the first: one
+  event can carry parallel calls, and an answer to the second is not evidence
+  the first never ran.
+  """
+  if not call_names:
+    return False
+  return (
+      not answers
+      or any(fr.name not in call_names for fr in answers)
+      or from_sub_branch
+  )
+
+
+def decide_resume(
+    invocation_context: InvocationContext,
+    events: list[Event],
+    tools_dict: dict[str, Any],
+) -> ResumeDecision:
+  """Decides how a resumable flow continues from `events`.
+
+  Args:
+    invocation_context: Supplies `should_pause_invocation`.
+    events: The current branch's events for this invocation, oldest first, and
+      containing at least two events.
+    tools_dict: The tools this flow can run, by name.
+
+  Returns:
+    PAUSE when a call is still unanswered, REPLAY_CALLS (naming the event whose
+    calls to run) when a call was never executed, else CONTINUE.
+  """
+  paused_by_last = invocation_context.should_pause_invocation(events[-1])
+  if not paused_by_last and _pause_left_calls_unanswered(
+      invocation_context, events
+  ):
+    return ResumeDecision(ResumeAction.PAUSE)
+
+  pause = paused_by_last
+  call_event = _find_target_call_event(events, tools_dict)
+  if call_event:
+    call_idx = next(i for i, ev in enumerate(events) if ev is call_event)
+    calls = call_event.get_function_calls()
+    call_names = {fc.name for fc in calls}
+    lro_ids = {
+        lro
+        for ev in events[call_idx:]
+        for lro in ev.long_running_tool_ids or []
+    }
+    call_ids = {fc.id for fc in calls} | lro_ids
+    answer_event = _find_answer_event(
+        events, call_event, call_idx, call_ids, call_names
+    )
+    answered_ids = {
+        fr.id
+        for ev in events[call_idx + 1 :]
+        for fr in ev.get_function_responses()
+        if fr.id is not None
+    }
+    # An answer on a sub-branch resolves the call however its ids look, so it
+    # short-circuits both unanswered tests rather than being repeated in each.
+    from_sub_branch = _is_sub_branch_answer(answer_event, call_event)
+    answers = answer_event.get_function_responses()
+    # `ids & answered` alone decides these: a set that is a subset of the
+    # answered ids necessarily intersects it, so testing `issubset` as well
+    # never changes the outcome.
+    lro_unanswered = bool(lro_ids) and not lro_ids & answered_ids
+    call_unanswered = (
+        bool(call_ids)
+        and not call_ids & answered_ids
+        and not any(fr.name in call_names for fr in answers)
+    )
+    if not from_sub_branch and (lro_unanswered or call_unanswered):
+      pause = True
+    elif _needs_call_replay(call_names, answers, from_sub_branch):
+      return ResumeDecision(ResumeAction.REPLAY_CALLS, call_event)
+
+  return ResumeDecision(ResumeAction.PAUSE if pause else ResumeAction.CONTINUE)

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -62,6 +62,8 @@
 from ._invocation_utils import require_agent as _require_agent
 from ._invocation_utils import require_run_config as _require_run_config
 from ._invocation_utils import run_config_for_new_live_session
+from ._resume_utils import decide_resume
+from ._resume_utils import ResumeAction
 from .functions import build_auth_request_event
 
 # Prefix used by toolset auth credential IDs
@@ -1271,6 +1273,28 @@ async def run_async(
           logger.warning('The last event is partial, which is not expected.')
         break
 
+  async def _replay_function_calls(
+      self,
+      invocation_context: InvocationContext,
+      model_response_event: Event,
+      llm_request: LlmRequest,
+  ) -> AsyncGenerator[Event, None]:
+    """Runs `model_response_event`'s function calls, re-issuing event ids.
+
+    A node that interrupts mid-call raises `NodeInterruptedError`, which is a
+    `BaseException` specifically so intermediate handlers do not swallow it.
+    It is left to propagate: `NodeRunner` catches it and reads the interrupt
+    ids off the context, which `ctx.run_node` populated before raising.
+    """
+    async with Aclosing(
+        self._postprocess_handle_function_calls_async(
+            invocation_context, model_response_event, llm_request
+        )
+    ) as agen:
+      async for event in agen:
+        event.id = Event.new_id()
+        yield event
+
   async def _run_one_step_async(
       self,
       invocation_context: InvocationContext,
@@ -1295,24 +1319,22 @@ async def _run_one_step_async(
         current_invocation=True, current_branch=True
     )
 
-    # Long running tool calls should have been handled before this point.
-    # If there are still long running tool calls, it means the agent is paused
-    # before, and its branch hasn't been resumed yet.
+    # For a multi-event branch, decide whether to pause (unanswered tool
+    # calls or LROs), replay unexecuted tool calls, or continue to the LLM.
     if invocation_context.is_resumable and events and len(events) > 1:
-      pause = False
-      if invocation_context.should_pause_invocation(events[-1]):
-        pause = True
-      elif invocation_context.should_pause_invocation(events[-2]):
-        # NOTE: This only checks the last 2 events. If an LRO is followed by
-        # multiple text responses, this check may not trigger correctly.
-        # This is a known limitation of the current 2-event window.
-        # Check if the function call in events[-2] is resolved by events[-1]
-        fc_ids = {fc.id for fc in events[-2].get_function_calls()}
-        fr_ids = {fr.id for fr in events[-1].get_function_responses()}
-        if fc_ids and not fc_ids.issubset(fr_ids):
-          pause = True
-
-      if pause:
+      decision = decide_resume(
+          invocation_context, events, llm_request.tools_dict
+      )
+      if decision.action is ResumeAction.PAUSE:
+        return
+      if decision.action is ResumeAction.REPLAY_CALLS:
+        async with Aclosing(
+            self._replay_function_calls(
+                invocation_context, decision.replay_event(), llm_request
+            )
+        ) as agen:
+          async for event in agen:
+            yield event
         return
 
     if (
@@ -1321,16 +1343,14 @@ async def _run_one_step_async(
         and not events[-1].partial
         and events[-1].get_function_calls()
     ):
-      model_response_event = events[-1]
       async with Aclosing(
-          self._postprocess_handle_function_calls_async(
-              invocation_context, model_response_event, llm_request
+          self._replay_function_calls(
+              invocation_context, events[-1], llm_request
           )
       ) as agen:
         async for event in agen:
-          event.id = Event.new_id()
           yield event
-        return
+      return
 
     # Calls the LLM.
     model_response_event = Event(