From 4ca904f11113c4faa3e17bb4a9662dca1f936e2e Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 23 Feb 2026 09:40:45 -0800 Subject: [PATCH] fix: Add push notification config store to agent_to_a2a This change allows users to provide a custom PushNotificationConfigStore when converting an ADK agent to an A2A Starlette application. If no custom store is provided, an InMemoryPushNotificationConfigStore is used by default, thios now lets A2A push notification configuration RPCs Close #4126 Co-authored-by: George Weale PiperOrigin-RevId: 874118109 --- src/google/adk/a2a/utils/agent_to_a2a.py | 13 +- src/google/adk/cli/fast_api.py | 7 +- .../unittests/a2a/utils/test_agent_to_a2a.py | 46 ++++++- tests/unittests/cli/test_fast_api.py | 114 +++++++++++++++--- 4 files changed, 157 insertions(+), 23 deletions(-) diff --git a/src/google/adk/a2a/utils/agent_to_a2a.py b/src/google/adk/a2a/utils/agent_to_a2a.py index 155888bc..d6a07080 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -20,7 +20,9 @@ from typing import Union from a2a.server.apps import A2AStarletteApplication from a2a.server.request_handlers import DefaultRequestHandler +from a2a.server.tasks import InMemoryPushNotificationConfigStore from a2a.server.tasks import InMemoryTaskStore +from a2a.server.tasks import PushNotificationConfigStore from a2a.types import AgentCard from starlette.applications import Starlette @@ -78,6 +80,7 @@ def to_a2a( port: int = 8000, protocol: str = "http", agent_card: Optional[Union[AgentCard, str]] = None, + push_config_store: Optional[PushNotificationConfigStore] = None, runner: Optional[Runner] = None, ) -> Starlette: """Convert an ADK agent to a A2A Starlette application. @@ -90,6 +93,9 @@ 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. + push_config_store: Optional A2A push notification config store. If not + provided, an in-memory store will be created so push-notification + config RPC methods are supported. runner: Optional pre-built Runner object. If not provided, a default runner will be created using in-memory services. @@ -127,8 +133,13 @@ def to_a2a( runner=runner or create_runner, ) + if push_config_store is None: + push_config_store = InMemoryPushNotificationConfigStore() + request_handler = DefaultRequestHandler( - agent_executor=agent_executor, task_store=task_store + agent_executor=agent_executor, + task_store=task_store, + push_config_store=push_config_store, ) # Use provided agent card or build one from the agent diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index 553629f2..8f78c15f 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -525,6 +525,7 @@ def get_fast_api_app( if a2a: from a2a.server.apps import A2AStarletteApplication from a2a.server.request_handlers import DefaultRequestHandler + from a2a.server.tasks import InMemoryPushNotificationConfigStore from a2a.server.tasks import InMemoryTaskStore from a2a.types import AgentCard from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH @@ -563,8 +564,12 @@ def get_fast_api_app( runner=create_a2a_runner_loader(app_name), ) + push_config_store = InMemoryPushNotificationConfigStore() + request_handler = DefaultRequestHandler( - agent_executor=agent_executor, task_store=a2a_task_store + agent_executor=agent_executor, + task_store=a2a_task_store, + push_config_store=push_config_store, ) with (p / "agent.json").open("r", encoding="utf-8") as f: diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index a9ff6d01..21c96d7e 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -12,12 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import ANY from unittest.mock import AsyncMock from unittest.mock import Mock from unittest.mock import patch from a2a.server.apps import A2AStarletteApplication from a2a.server.request_handlers import DefaultRequestHandler +from a2a.server.tasks import InMemoryPushNotificationConfigStore from a2a.server.tasks import InMemoryTaskStore from a2a.types import AgentCard from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor @@ -77,7 +79,9 @@ class TestToA2A: mock_task_store_class.assert_called_once() mock_agent_executor_class.assert_called_once() mock_request_handler_class.assert_called_once_with( - agent_executor=mock_agent_executor, task_store=mock_task_store + agent_executor=mock_agent_executor, + push_config_store=ANY, + task_store=mock_task_store, ) mock_card_builder_class.assert_called_once_with( agent=self.mock_agent, rpc_url="http://localhost:8000/" @@ -122,7 +126,9 @@ class TestToA2A: 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 + agent_executor=mock_agent_executor, + push_config_store=ANY, + task_store=mock_task_store, ) mock_card_builder_class.assert_called_once_with( agent=self.mock_agent, rpc_url="http://localhost:8000/" @@ -131,6 +137,42 @@ 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_passes_custom_push_config_store( + self, + mock_starlette_class, + mock_card_builder_class, + mock_task_store_class, + mock_request_handler_class, + mock_agent_executor_class, + ): + """Test to_a2a forwards a custom push config store.""" + 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_push_store = InMemoryPushNotificationConfigStore() + + result = to_a2a(self.mock_agent, push_config_store=custom_push_store) + + assert result == mock_app + mock_request_handler_class.assert_called_once_with( + agent_executor=mock_agent_executor, + push_config_store=custom_push_store, + task_store=mock_task_store, + ) + @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") diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index 913e11ae..16ee82b6 100755 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -15,7 +15,6 @@ import asyncio import json import logging -import os from pathlib import Path import signal import tempfile @@ -677,6 +676,7 @@ def test_app_with_a2a( mock_eval_sets_manager, mock_eval_set_results_manager, temp_agents_dir_with_a2a, + monkeypatch, ): """Create a TestClient for the FastAPI app with A2A enabled.""" # Mock A2A related classes @@ -728,26 +728,22 @@ def test_app_with_a2a( mock_a2a_app.return_value = mock_app_instance # Change to temp directory - original_cwd = os.getcwd() - os.chdir(temp_agents_dir_with_a2a) + monkeypatch.chdir(temp_agents_dir_with_a2a) - try: - app = get_fast_api_app( - agents_dir=".", - web=True, - session_service_uri="", - artifact_service_uri="", - memory_service_uri="", - allow_origins=["*"], - a2a=True, - host="127.0.0.1", - port=8000, - ) + app = get_fast_api_app( + agents_dir=".", + web=True, + session_service_uri="", + artifact_service_uri="", + memory_service_uri="", + allow_origins=["*"], + a2a=True, + host="127.0.0.1", + port=8000, + ) - client = TestClient(app) - yield client - finally: - os.chdir(original_cwd) + client = TestClient(app) + yield client ################################################# @@ -1406,6 +1402,86 @@ def test_a2a_agent_discovery(test_app_with_a2a): logger.info("A2A agent discovery test passed") +def test_a2a_request_handler_uses_push_config_store( + mock_session_service, + mock_artifact_service, + mock_memory_service, + mock_agent_loader, + mock_eval_sets_manager, + mock_eval_set_results_manager, + temp_agents_dir_with_a2a, + monkeypatch, +): + """Test A2A request handler gets push config store when supported.""" + with ( + patch("signal.signal", return_value=None), + patch( + "google.adk.cli.fast_api.create_session_service_from_options", + return_value=mock_session_service, + ), + patch( + "google.adk.cli.fast_api.create_artifact_service_from_options", + return_value=mock_artifact_service, + ), + patch( + "google.adk.cli.fast_api.create_memory_service_from_options", + return_value=mock_memory_service, + ), + patch( + "google.adk.cli.fast_api.AgentLoader", + return_value=mock_agent_loader, + ), + patch( + "google.adk.cli.fast_api.LocalEvalSetsManager", + return_value=mock_eval_sets_manager, + ), + patch( + "google.adk.cli.fast_api.LocalEvalSetResultsManager", + return_value=mock_eval_set_results_manager, + ), + patch("a2a.server.tasks.InMemoryTaskStore") as mock_task_store, + patch( + "a2a.server.tasks.InMemoryPushNotificationConfigStore" + ) as mock_push_config_store_class, + patch( + "google.adk.a2a.executor.a2a_agent_executor.A2aAgentExecutor" + ) as mock_executor, + patch( + "a2a.server.request_handlers.DefaultRequestHandler" + ) as mock_handler, + patch("a2a.server.apps.A2AStarletteApplication") as mock_a2a_app, + ): + mock_task_store_instance = MagicMock() + mock_task_store.return_value = mock_task_store_instance + mock_push_config_store = MagicMock() + mock_push_config_store_class.return_value = mock_push_config_store + mock_executor_instance = MagicMock() + mock_executor.return_value = mock_executor_instance + mock_handler.return_value = MagicMock() + mock_a2a_app_instance = MagicMock() + mock_a2a_app_instance.routes.return_value = [] + mock_a2a_app.return_value = mock_a2a_app_instance + + monkeypatch.chdir(temp_agents_dir_with_a2a) + _ = get_fast_api_app( + agents_dir=".", + web=True, + session_service_uri="", + artifact_service_uri="", + memory_service_uri="", + allow_origins=["*"], + a2a=True, + host="127.0.0.1", + port=8000, + ) + + mock_handler.assert_called_once_with( + agent_executor=mock_executor_instance, + push_config_store=mock_push_config_store, + task_store=mock_task_store_instance, + ) + + def test_a2a_disabled_by_default(test_app): """Test that A2A functionality is disabled by default.""" # The regular test_app fixture has a2a=False