From 156d23547915e8f7f5c6ba55e0362f4b133c3968 Mon Sep 17 00:00:00 2001 From: Dongyu Jia Date: Fri, 31 Oct 2025 09:54:27 -0700 Subject: [PATCH] feat: Allow injecting a custom Runner into `agent_to_a2a` This change adds an optional `runner` parameter to the `to_a2a` function, enabling users to provide a pre-configured `Runner` instance instead of always using the default in-memory services. A new test case has been added to verify this functionality. closes #3104 Co-authored-by: Dongyu Jia PiperOrigin-RevId: 826526861 --- src/google/adk/a2a/utils/agent_to_a2a.py | 5 ++- .../unittests/a2a/utils/test_agent_to_a2a.py | 45 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/src/google/adk/a2a/utils/agent_to_a2a.py b/src/google/adk/a2a/utils/agent_to_a2a.py index e550dc76..72a2292f 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -90,6 +90,7 @@ def to_a2a( port: int = 8000, protocol: str = "http", agent_card: Optional[Union[AgentCard, str]] = None, + runner: Optional[Runner] = None, ) -> Starlette: """Convert an ADK agent to a A2A Starlette application. @@ -101,6 +102,8 @@ def to_a2a( agent_card: Optional pre-built AgentCard object or path to agent card JSON. If not provided, will be built automatically from the agent. + runner: Optional pre-built Runner object. If not provided, a default + runner will be created using in-memory services. Returns: A Starlette application that can be run with uvicorn @@ -132,7 +135,7 @@ def to_a2a( task_store = InMemoryTaskStore() agent_executor = A2aAgentExecutor( - runner=create_runner, + runner=runner or create_runner, ) request_handler = DefaultRequestHandler( diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index 1106bf96..ee80b023 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -103,6 +103,51 @@ class TestToA2A: "startup", mock_app.add_event_handler.call_args[0][1] ) + @patch("google.adk.a2a.utils.agent_to_a2a.A2aAgentExecutor") + @patch("google.adk.a2a.utils.agent_to_a2a.DefaultRequestHandler") + @patch("google.adk.a2a.utils.agent_to_a2a.InMemoryTaskStore") + @patch("google.adk.a2a.utils.agent_to_a2a.AgentCardBuilder") + @patch("google.adk.a2a.utils.agent_to_a2a.Starlette") + def test_to_a2a_with_custom_runner( + self, + mock_starlette_class, + mock_card_builder_class, + mock_task_store_class, + mock_request_handler_class, + mock_agent_executor_class, + ): + """Test to_a2a with a custom runner.""" + # Arrange + mock_app = Mock(spec=Starlette) + mock_starlette_class.return_value = mock_app + mock_task_store = Mock(spec=InMemoryTaskStore) + mock_task_store_class.return_value = mock_task_store + mock_agent_executor = Mock(spec=A2aAgentExecutor) + mock_agent_executor_class.return_value = mock_agent_executor + mock_request_handler = Mock(spec=DefaultRequestHandler) + mock_request_handler_class.return_value = mock_request_handler + mock_card_builder = Mock(spec=AgentCardBuilder) + mock_card_builder_class.return_value = mock_card_builder + custom_runner = Mock(spec=Runner) + + # Act + result = to_a2a(self.mock_agent, runner=custom_runner) + + # Assert + assert result == mock_app + mock_starlette_class.assert_called_once() + mock_task_store_class.assert_called_once() + mock_agent_executor_class.assert_called_once_with(runner=custom_runner) + mock_request_handler_class.assert_called_once_with( + agent_executor=mock_agent_executor, task_store=mock_task_store + ) + mock_card_builder_class.assert_called_once_with( + agent=self.mock_agent, rpc_url="http://localhost:8000/" + ) + mock_app.add_event_handler.assert_called_once_with( + "startup", mock_app.add_event_handler.call_args[0][1] + ) + @patch("google.adk.a2a.utils.agent_to_a2a.A2aAgentExecutor") @patch("google.adk.a2a.utils.agent_to_a2a.DefaultRequestHandler") @patch("google.adk.a2a.utils.agent_to_a2a.InMemoryTaskStore")