fix: route authentication responses to requesting sub-agents

src/google/adk/agents/_agent_router.py

@@ -28,8 +28,11 @@
 from ..flows.llm_flows.extensions._agent_transfer import _get_transfer_targets
 from ..flows.llm_flows.functions import _collect_function_call_ids
 from ..flows.llm_flows.functions import find_matching_function_call
+from ..flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME
 
 if TYPE_CHECKING:
+  from google.genai import types
+
   from ..agents.base_agent import BaseAgent
   from ..agents.invocation_context import InvocationContext
   from ..apps.app import ResumabilityConfig
@@ -82,11 +85,14 @@ def find_agent_to_run(
     session: Session,
     root_agent: BaseAgent,
     resumability_config: Optional[ResumabilityConfig] = None,
+    *,
+    new_message: Optional[types.Content] = None,
 ) -> BaseAgent:
   """Finds the agent to run to continue the session.
 
   A qualified agent must be either of:
 
+  - The owner of an outstanding authentication request answered by new_message.
   - The agent that returned a function call and the last user message is a
     function response to this function call.
   - The root agent.
@@ -99,6 +105,7 @@ def find_agent_to_run(
       session: The session to find the agent for.
       root_agent: The root agent of the runner.
       resumability_config: Optional resumability configuration.
+      new_message: Incoming message not yet appended to the session.
 
   Returns:
     The agent to run. (the active agent that should reply to the latest user
@@ -111,11 +118,19 @@ def find_agent_to_run(
   if isinstance(root_agent, Workflow):
     return root_agent
 
+  filtered_events = _apply_rewinds(session.events)
+  # The node runtime selects an agent before appending the incoming message.
+  # Resolve credential replies first, independently of resumability settings.
+  if new_message and (
+      auth_agent := _find_auth_response_agent(
+          filtered_events, root_agent, new_message
+      )
+  ):
+    return auth_agent
   # If the last event is a function response, should send this response to
   # 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 = _apply_rewinds(session.events)
   event = find_matching_function_call(filtered_events)
   is_resumable = resumability_config and resumability_config.is_resumable
   # Only route based on a past function response if resumability is enabled.
@@ -160,6 +175,72 @@ def _event_filter(event: Event) -> bool:
   return root_agent
 
 
+def _find_auth_response_agent(
+    events: list[Event], root_agent: BaseAgent, message: types.Content
+) -> Optional[BaseAgent]:
+  """Resolves an incoming auth response to its outstanding request's owner.
+
+  Transfer restrictions apply to new conversation turns, not to completing an
+  authentication request. Only use requests recorded in this session; node paths
+  disambiguate agents with the same name in different branches. A batch must
+  belong to one agent and invocation to select a single entry point.
+  """
+  responses = [
+      p.function_response for p in message.parts or [] if p.function_response
+  ]
+  if not responses or any(
+      not response.id or response.name != REQUEST_EUC_FUNCTION_CALL_NAME
+      for response in responses
+  ):
+    return None
+  pending_ids = {response.id for response in responses}
+  answered_ids = {
+      response.id
+      for event in events
+      for response in event.get_function_responses()
+  }
+  if pending_ids & answered_ids:
+    return None
+
+  owner = None
+  invocation_id = None
+  for event in reversed(events):
+    matching_ids = {
+        call.id
+        for call in event.get_function_calls()
+        if call.id in pending_ids
+        and call.name == REQUEST_EUC_FUNCTION_CALL_NAME
+    }
+    if not matching_ids:
+      continue
+    event_path = (
+        _NodePathBuilder.from_string(event.node_info.path).static_path
+        if event.node_info.path
+        else None
+    )
+    candidates = []
+    pending_agents = [(root_agent, root_agent.name)]
+    while pending_agents:
+      agent, path = pending_agents.pop()
+      if agent.name == event.author and (not event_path or path == event_path):
+        candidates.append(agent)
+      pending_agents.extend(
+          (child, f"{path}/{child.name}") for child in agent.sub_agents
+      )
+    if len(candidates) != 1:
+      return None
+    agent = candidates[0]
+    if owner is not None and (
+        owner is not agent or invocation_id != event.invocation_id
+    ):
+      return None
+    owner, invocation_id = agent, event.invocation_id
+    pending_ids.difference_update(matching_ids)
+    if not pending_ids:
+      return owner
+  return None
+
+
 def restore_branch_from_history(
     invocation_context: InvocationContext,
     node: BaseNode,

src/google/adk/runners.py

@@ -1114,7 +1114,9 @@ async def run_async(
           if has_task_subagent:
             agent_to_run = self.agent
           else:
-            agent_to_run = self._find_agent_to_run(session, self.agent)
+            agent_to_run = self._find_agent_to_run(
+                session, self.agent, new_message=new_message
+            )
         else:
           agent_to_run = self.agent
 
@@ -1702,7 +1704,11 @@ async def _pump_queued_events() -> None:
       await self._cleanup_root_task(queue_task, self.agent.name)
 
   def _find_agent_to_run(
-      self, session: Session, root_agent: BaseAgent
+      self,
+      session: Session,
+      root_agent: BaseAgent,
+      *,
+      new_message: Optional[types.Content] = None,
   ) -> BaseAgent:
     """Finds the agent to run to continue the session.
 
@@ -1719,6 +1725,7 @@ def _find_agent_to_run(
     Args:
         session: The session to find the agent for.
         root_agent: The root agent of the runner.
+        new_message: Incoming message not yet appended to the session.
 
     Returns:
       The agent to run. (the active agent that should reply to the latest user
@@ -1730,6 +1737,7 @@ def _find_agent_to_run(
         session=session,
         root_agent=root_agent,
         resumability_config=self.resumability_config,
+        new_message=new_message,
     )
 
   async def run_debug(

tests/unittests/agents/test_agent_router.py

@@ -25,9 +25,11 @@
 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
+import pytest
 
 
 class _MockLlmAgent(LlmAgent):
@@ -495,3 +497,164 @@ 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 _auth_request(
+    author="non_transferable", call_id="auth-1", invocation_id="inv1"
+):
+  return Event(
+      author=author,
+      invocation_id=invocation_id,
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part(
+                  function_call=types.FunctionCall(
+                      id=call_id, name="adk_request_credential", args={}
+                  )
+              )
+          ],
+      ),
+      long_running_tool_ids={call_id},
+  )
+
+
+def _auth_response(call_id="auth-1", name="adk_request_credential"):
+  return types.Part(
+      function_response=types.FunctionResponse(
+          id=call_id, name=name, response={}
+      )
+  )
+
+
+@pytest.mark.parametrize("resumable", [False, True])
+@pytest.mark.parametrize("node_path", [None, "root_agent@1/non_transferable@2"])
+def test_incoming_auth_response_resumes_owner(resumable, node_path):
+  """An outstanding auth response reaches its restricted owner in either mode."""
+  root, _, _, child = _make_agent_tree()
+  request = _auth_request()
+  request.node_info.path = node_path
+  session = Session(id="s", app_name="app", user_id="u", events=[request])
+  message = types.Content(role="user", parts=[_auth_response()])
+
+  assert (
+      _agent_router.find_agent_to_run(
+          session,
+          root,
+          ResumabilityConfig(is_resumable=resumable),
+          new_message=message,
+      )
+      is child
+  )
+  # Routing must not append the incoming message before runner callbacks run.
+  assert session.events == [request]
+
+
+@pytest.mark.parametrize("resumable", [False, True])
+@pytest.mark.parametrize(
+    "case",
+    [
+        "text",
+        "missing_id",
+        "unknown_id",
+        "wrong_response_name",
+        "wrong_call_name",
+        "unknown_author",
+        "foreign_path",
+        "answered",
+        "rewound",
+    ],
+)
+def test_invalid_auth_response_does_not_select_restricted_agent(
+    case, resumable
+):
+  """Only a matching outstanding auth request can override transfer routing."""
+  root, _, _, _ = _make_agent_tree()
+  request = _auth_request()
+  message = types.Content(role="user", parts=[_auth_response()])
+  events = [request]
+  if case == "text":
+    message.parts = [types.Part(text="new question")]
+  elif case == "missing_id":
+    message.parts = [_auth_response(None)]
+  elif case == "unknown_id":
+    message.parts = [_auth_response("unknown")]
+  elif case == "wrong_response_name":
+    message.parts = [_auth_response(name="another_tool")]
+  elif case == "wrong_call_name":
+    request.content.parts[0].function_call.name = "another_tool"
+  elif case == "unknown_author":
+    request.author = "missing_agent"
+  elif case == "foreign_path":
+    request.node_info.path = "other_root@1/non_transferable@1"
+  elif case == "answered":
+    events.extend([
+        Event(author="user", content=message),
+        Event(
+            author=root.name,
+            content=types.Content(parts=[types.Part(text="done")]),
+        ),
+    ])
+  elif case == "rewound":
+    events.append(
+        Event(
+            author="user",
+            invocation_id="inv2",
+            actions=EventActions(rewind_before_invocation_id="inv1"),
+        )
+    )
+  session = Session(id="s", app_name="app", user_id="u", events=events)
+
+  assert (
+      _agent_router.find_agent_to_run(
+          session,
+          root,
+          ResumabilityConfig(is_resumable=resumable),
+          new_message=message,
+      )
+      is root
+  )
+
+
+@pytest.mark.parametrize("path_present", [False, True])
+def test_auth_response_disambiguates_same_named_agents(path_present):
+  """A persisted node path selects the correct branch when names are repeated."""
+  left_child = _MockLlmAgent("worker", disallow_transfer_to_parent=True)
+  right_child = _MockLlmAgent("worker", disallow_transfer_to_parent=True)
+  left = LlmAgent(name="left", sub_agents=[left_child])
+  right = LlmAgent(name="right", sub_agents=[right_child])
+  root = LlmAgent(name="root", sub_agents=[left, right])
+  request = _auth_request(author="worker")
+  if path_present:
+    request.node_info.path = "root@1/right@2/worker@3"
+  session = Session(id="s", app_name="app", user_id="u", events=[request])
+  message = types.Content(role="user", parts=[_auth_response()])
+
+  assert _agent_router.find_agent_to_run(
+      session, root, new_message=message
+  ) is (right_child if path_present else root)
+
+
+@pytest.mark.parametrize(
+    "second_request",
+    ["same_owner", "other_owner", "other_invocation", "unknown"],
+)
+def test_auth_response_batch_requires_one_owner_and_invocation(second_request):
+  """A response batch is routed only when all requests have one resume target."""
+  root, _, other, child = _make_agent_tree()
+  other.disallow_transfer_to_parent = True
+  first = _auth_request()
+  second = _auth_request(call_id="auth-2")
+  if second_request == "other_owner":
+    second.author = other.name
+  elif second_request == "other_invocation":
+    second.invocation_id = "inv2"
+  events = [first, second] if second_request != "unknown" else [first]
+  session = Session(id="s", app_name="app", user_id="u", events=events)
+  message = types.Content(
+      role="user", parts=[_auth_response(), _auth_response("auth-2")]
+  )
+
+  assert _agent_router.find_agent_to_run(
+      session, root, new_message=message
+  ) is (child if second_request == "same_owner" else root)