From 9d776cace5a3224a306628b8e8d10bd13ae2b330 Mon Sep 17 00:00:00 2001 From: Junyi Zheng <279671317+Junyi-Zheng@users.noreply.github.com> Date: Fri, 18 Sep 2026 22:48:15 -0400 Subject: [PATCH 1/2] fix: route authentication responses to the requesting sub-agent --- src/google/adk/agents/_agent_router.py | 83 ++++++++- src/google/adk/runners.py | 12 +- tests/unittests/agents/test_agent_router.py | 163 +++++++++++++++++ .../runners/test_resume_invocation.py | 168 ++++++++++++++++++ 4 files changed, 423 insertions(+), 3 deletions(-) diff --git a/src/google/adk/agents/_agent_router.py b/src/google/adk/agents/_agent_router.py index 4cbc7faff2..7736db9092 100644 --- a/src/google/adk/agents/_agent_router.py +++ b/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, diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 5c6c9cec0a..4ee035a7be 100644 --- a/src/google/adk/runners.py +++ b/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( diff --git a/tests/unittests/agents/test_agent_router.py b/tests/unittests/agents/test_agent_router.py index d523353828..dfe1b06e5a 100644 --- a/tests/unittests/agents/test_agent_router.py +++ b/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) diff --git a/tests/unittests/runners/test_resume_invocation.py b/tests/unittests/runners/test_resume_invocation.py index 6b516ff06d..ea930fa1d5 100644 --- a/tests/unittests/runners/test_resume_invocation.py +++ b/tests/unittests/runners/test_resume_invocation.py @@ -426,3 +426,171 @@ async def _run_async_impl( assert not await runner.run_async( invocation_id=invocation_events[0].invocation_id ) + + +@pytest.mark.parametrize("resumable", [False, True]) +@pytest.mark.parametrize("auth_stage", ["tool", "toolset"]) +@pytest.mark.parametrize("nested", [False, True]) +async def test_auth_response_resumes_restricted_sub_agent( + resumable, auth_stage, nested +): + """Authentication resumes its owner even when transfer to its parent is disabled. + + Setup: transfer to a restricted child that requests an OIDC credential. + Act: return the credential using the emitted authentication call ID. + Assert: the protected tool succeeds once, without another auth request; + a subsequent ordinary user message still returns to the root agent. + """ + from google.adk.auth.auth_credential import AuthCredential + from google.adk.auth.auth_credential import OAuth2Auth + from google.adk.auth.auth_schemes import OpenIdConnectWithConfig + from google.adk.auth.auth_tool import AuthConfig + from google.adk.runners import Runner + from google.adk.sessions.in_memory_session_service import InMemorySessionService + from google.adk.tools.base_toolset import BaseToolset + from google.adk.tools.function_tool import FunctionTool + from google.adk.tools.tool_context import ToolContext + + auth_config = AuthConfig( + auth_scheme=OpenIdConnectWithConfig( + authorization_endpoint="https://issuer.example/authorize", + token_endpoint="https://issuer.example/token", + scopes=["openid"], + ), + raw_auth_credential=AuthCredential( + auth_type="oauth2", + oauth2=OAuth2Auth(client_id="client", client_secret="secret"), + ), + credential_key="test-credential", + ) + successful_calls = [] + + def protected_tool(tool_context: ToolContext) -> dict: + if auth_stage == "tool": + credential = tool_context.get_auth_response(auth_config) + else: + credential = tool_context.get_invocation_context().credential_by_key.get( + auth_config.credential_key + ) + if not credential or not credential.oauth2.access_token: + tool_context.request_credential(auth_config) + return {"status": "authorization_required"} + successful_calls.append(credential.oauth2.access_token) + return {"status": "ok"} + + class AuthToolset(BaseToolset): + + def get_auth_config(self): + return auth_config + + async def get_tools(self, readonly_context=None): + return [FunctionTool(protected_tool)] + + async def close(self): + pass + + child = LlmAgent( + name="worker", + disallow_transfer_to_parent=True, + disallow_transfer_to_peers=True, + model=testing_utils.MockModel.create([ + Part.from_function_call(name="protected_tool", args={}), + "child completed", + ]), + tools=[protected_tool] if auth_stage == "tool" else [AuthToolset()], + ) + delegate = child + if nested: + delegate = LlmAgent( + name="middle", + disallow_transfer_to_parent=True, + model=testing_utils.MockModel.create([transfer_call_part(child.name)]), + sub_agents=[child], + ) + root = LlmAgent( + name="root", + model=testing_utils.MockModel.create([ + transfer_call_part(delegate.name), + "root handles next message", + ]), + sub_agents=[delegate], + ) + app = App( + name="test_app", + root_agent=root, + resumability_config=ResumabilityConfig(is_resumable=resumable), + ) + runner = Runner(app=app, session_service=InMemorySessionService()) + session = await runner.session_service.create_session( + app_name=app.name, user_id="user" + ) + + async def collect(message): + return [ + event + async for event in runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=message, + ) + ] + + try: + initial = await collect(testing_utils.UserContent("run protected tool")) + auth_events = [ + (event, call) + for event in initial + for call in event.get_function_calls() + if call.name == "adk_request_credential" + ] + assert len(auth_events) == 1 + request_event, call = auth_events[0] + assert call.id in request_event.long_running_tool_ids + assert not successful_calls + config = AuthConfig.model_validate(call.args["authConfig"]) + config.exchanged_auth_credential = AuthCredential( + auth_type="oauth2", oauth2=OAuth2Auth(access_token="test-token") + ) + resumed = await collect( + testing_utils.UserContent( + Part( + function_response=FunctionResponse( + id=call.id, + name=call.name, + response=config.model_dump(mode="json", by_alias=True), + ) + ) + ) + ) + + assert successful_calls == ["test-token"] + assert not any( + call.name == "adk_request_credential" + for event in resumed + for call in event.get_function_calls() + ) + assert any( + event.author == child.name + and event.content + and any(part.text == "child completed" for part in event.content.parts) + for event in resumed + ) + # The resumed child keeps the original branch, including nested transfers. + assert all( + event.branch == request_event.branch + for event in resumed + if event.author == child.name and event.content + ) + following = await collect(testing_utils.UserContent("a new request")) + assert any( + event.author == root.name + and event.content + and any( + part.text == "root handles next message" + for part in event.content.parts + ) + for event in following + ) + assert successful_calls == ["test-token"] + finally: + await runner.close() From 2028cbd6ff9ad713d02f0985f6a4a1210840a5a5 Mon Sep 17 00:00:00 2001 From: Junyi Zheng <279671317+Junyi-Zheng@users.noreply.github.com> Date: Fri, 18 Sep 2026 23:11:47 -0400 Subject: [PATCH 2/2] test: cover authentication routing with SQLite sessions --- tests/unittests/runners/test_resume_invocation.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/unittests/runners/test_resume_invocation.py b/tests/unittests/runners/test_resume_invocation.py index ea930fa1d5..517c40cbf9 100644 --- a/tests/unittests/runners/test_resume_invocation.py +++ b/tests/unittests/runners/test_resume_invocation.py @@ -431,8 +431,9 @@ async def _run_async_impl( @pytest.mark.parametrize("resumable", [False, True]) @pytest.mark.parametrize("auth_stage", ["tool", "toolset"]) @pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("session_storage", ["memory", "sqlite"]) async def test_auth_response_resumes_restricted_sub_agent( - resumable, auth_stage, nested + resumable, auth_stage, nested, session_storage, tmp_path ): """Authentication resumes its owner even when transfer to its parent is disabled. @@ -446,6 +447,7 @@ async def test_auth_response_resumes_restricted_sub_agent( from google.adk.auth.auth_schemes import OpenIdConnectWithConfig from google.adk.auth.auth_tool import AuthConfig from google.adk.runners import Runner + from google.adk.sessions.database_session_service import DatabaseSessionService from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.tools.base_toolset import BaseToolset from google.adk.tools.function_tool import FunctionTool @@ -520,7 +522,14 @@ async def close(self): root_agent=root, resumability_config=ResumabilityConfig(is_resumable=resumable), ) - runner = Runner(app=app, session_service=InMemorySessionService()) + session_service = ( + InMemorySessionService() + if session_storage == "memory" + else DatabaseSessionService( + db_url=f"sqlite+aiosqlite:///{tmp_path / 'sessions.sqlite'}" + ) + ) + runner = Runner(app=app, session_service=session_service) session = await runner.session_service.create_session( app_name=app.name, user_id="user" ) @@ -594,3 +603,5 @@ async def collect(message): assert successful_calls == ["test-token"] finally: await runner.close() + if session_storage == "sqlite": + await session_service.close()