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 <gweale@google.com>
PiperOrigin-RevId: 874118109

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

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:

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")

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