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 <dongyuj@google.com>
PiperOrigin-RevId: 826526861
This commit is contained in:
Dongyu Jia
2025-10-31 09:54:59 -07:00
committed by Copybara-Service
parent 338c3c89c6
commit 156d235479
2 changed files with 49 additions and 1 deletions
@@ -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")