From 76f5f813fdccaa7fb9291f6dbf246d5ebe31cc57 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 13:51:55 +0000 Subject: [PATCH 01/13] feat: Add support for custom task store in to_a2a function This update introduces an optional `task_store` parameter to the `to_a2a` function, allowing users to provide a custom task store instance. If not specified, the function defaults to using `InMemoryTaskStore`. Additionally, unit tests have been added to verify the functionality with a custom task store. --- src/google/adk/a2a/utils/agent_to_a2a.py | 13 +++++- .../unittests/a2a/utils/test_agent_to_a2a.py | 45 +++++++++++++++++++ 2 files changed, 57 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 72a2292fb3..a817a93095 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -30,6 +30,7 @@ else: raise e +from typing import Any from typing import Optional from typing import Union @@ -91,6 +92,7 @@ def to_a2a( protocol: str = "http", agent_card: Optional[Union[AgentCard, str]] = None, runner: Optional[Runner] = None, + task_store: Optional[Any] = None, ) -> Starlette: """Convert an ADK agent to a A2A Starlette application. @@ -104,6 +106,9 @@ def to_a2a( agent. runner: Optional pre-built Runner object. If not provided, a default runner will be created using in-memory services. + task_store: Optional task store instance. If not provided, an + InMemoryTaskStore will be created. Must be compatible with + DefaultRequestHandler's task_store parameter. Returns: A Starlette application that can be run with uvicorn @@ -115,6 +120,11 @@ def to_a2a( # Or with custom agent card: app = to_a2a(agent, agent_card=my_custom_agent_card) + + # Or with custom task store: + from a2a.server.tasks import MyCustomTaskStore + custom_store = MyCustomTaskStore() + app = to_a2a(agent, task_store=custom_store) """ # Set up ADK logging to ensure logs are visible when using uvicorn directly setup_adk_logger(logging.INFO) @@ -132,7 +142,8 @@ async def create_runner() -> Runner: ) # Create A2A components - task_store = InMemoryTaskStore() + if task_store is None: + task_store = InMemoryTaskStore() agent_executor = A2aAgentExecutor( runner=runner or create_runner, diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index ee80b0233b..af800c4754 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -148,6 +148,51 @@ def test_to_a2a_with_custom_runner( "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_task_store( + 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 task store.""" + # Arrange + mock_app = Mock(spec=Starlette) + mock_starlette_class.return_value = mock_app + custom_task_store = Mock(spec=InMemoryTaskStore) + 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 + + # Act + result = to_a2a(self.mock_agent, task_store=custom_task_store) + + # Assert + assert result == mock_app + mock_starlette_class.assert_called_once() + # Verify InMemoryTaskStore was NOT created since we provided a custom one + mock_task_store_class.assert_not_called() + mock_agent_executor_class.assert_called_once() + # Verify the custom task store was used + mock_request_handler_class.assert_called_once_with( + agent_executor=mock_agent_executor, task_store=custom_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") From ff62647264ef9e0f1b721d9a2419d3d4879c6e74 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 13:56:33 +0000 Subject: [PATCH 02/13] Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/a2a/utils/agent_to_a2a.py | 2 +- 1 file changed, 1 insertion(+), 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 a817a93095..9344c5ba37 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -92,7 +92,7 @@ def to_a2a( protocol: str = "http", agent_card: Optional[Union[AgentCard, str]] = None, runner: Optional[Runner] = None, - task_store: Optional[Any] = None, + task_store: Optional["InMemoryTaskStore"] = None, ) -> Starlette: """Convert an ADK agent to a A2A Starlette application. From 8c065dbf879bda04b113e77733a4c4fda12a4f48 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 13:58:38 +0000 Subject: [PATCH 03/13] Remove unused import 'Any' from agent_to_a2a.py Removed unused import of 'Any' from typing. --- src/google/adk/a2a/utils/agent_to_a2a.py | 1 - 1 file changed, 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 9344c5ba37..2a716b6ccd 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -30,7 +30,6 @@ else: raise e -from typing import Any from typing import Optional from typing import Union From 933ad25a98d99fa1cb8c8cccca4448a66c1884f8 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:10:43 +0000 Subject: [PATCH 04/13] Refactor task_store type to use TaskStore --- src/google/adk/a2a/utils/agent_to_a2a.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/google/adk/a2a/utils/agent_to_a2a.py b/src/google/adk/a2a/utils/agent_to_a2a.py index 2a716b6ccd..9e5617dbe5 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -20,7 +20,7 @@ try: from a2a.server.apps import A2AStarletteApplication from a2a.server.request_handlers import DefaultRequestHandler - from a2a.server.tasks import InMemoryTaskStore + from a2a.server.tasks import TaskStore, InMemoryTaskStore from a2a.types import AgentCard except ImportError as e: if sys.version_info < (3, 10): @@ -91,7 +91,7 @@ def to_a2a( protocol: str = "http", agent_card: Optional[Union[AgentCard, str]] = None, runner: Optional[Runner] = None, - task_store: Optional["InMemoryTaskStore"] = None, + task_store: Optional["TaskStore"] = None, ) -> Starlette: """Convert an ADK agent to a A2A Starlette application. From 2953b2803dea1cdf16f493068d941d877a669540 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:12:21 +0000 Subject: [PATCH 05/13] Correct type hint for task_store in agent_to_a2a Fix type hint for task_store parameter. --- src/google/adk/a2a/utils/agent_to_a2a.py | 2 +- 1 file changed, 1 insertion(+), 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 9e5617dbe5..abcc6cf261 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -91,7 +91,7 @@ def to_a2a( protocol: str = "http", agent_card: Optional[Union[AgentCard, str]] = None, runner: Optional[Runner] = None, - task_store: Optional["TaskStore"] = None, + task_store: Optional[TaskStore] = None, ) -> Starlette: """Convert an ADK agent to a A2A Starlette application. From 019001b6e1915bcf032582c143eb76a7b4a66810 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:18:28 +0000 Subject: [PATCH 06/13] Update tests/unittests/a2a/utils/test_agent_to_a2a.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/unittests/a2a/utils/test_agent_to_a2a.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index af800c4754..db165bad7a 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -165,7 +165,7 @@ def test_to_a2a_with_custom_task_store( # Arrange mock_app = Mock(spec=Starlette) mock_starlette_class.return_value = mock_app - custom_task_store = Mock(spec=InMemoryTaskStore) + custom_task_store = Mock(spec=TaskStore) mock_agent_executor = Mock(spec=A2aAgentExecutor) mock_agent_executor_class.return_value = mock_agent_executor mock_request_handler = Mock(spec=DefaultRequestHandler) From b02954d2f3f277ca9c866960a0ee7eccb2f7f611 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:22:41 +0000 Subject: [PATCH 07/13] Update test_agent_to_a2a to use TaskStore for custom task store mock Refactor the test to replace the mock of InMemoryTaskStore with TaskStore, ensuring consistency with the updated type hint for task_store. This change enhances the clarity and accuracy of the unit tests. --- tests/unittests/a2a/utils/test_agent_to_a2a.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index af800c4754..87a778adaa 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -28,7 +28,7 @@ try: from a2a.server.apps import A2AStarletteApplication from a2a.server.request_handlers import DefaultRequestHandler - from a2a.server.tasks import InMemoryTaskStore + from a2a.server.tasks import TaskStore, InMemoryTaskStore from a2a.types import AgentCard from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor from google.adk.a2a.utils.agent_card_builder import AgentCardBuilder @@ -165,7 +165,7 @@ def test_to_a2a_with_custom_task_store( # Arrange mock_app = Mock(spec=Starlette) mock_starlette_class.return_value = mock_app - custom_task_store = Mock(spec=InMemoryTaskStore) + custom_task_store = Mock(spec=TaskStore) mock_agent_executor = Mock(spec=A2aAgentExecutor) mock_agent_executor_class.return_value = mock_agent_executor mock_request_handler = Mock(spec=DefaultRequestHandler) From 870f1a732f7096a85ea54d2245df604603ddd356 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:25:36 +0000 Subject: [PATCH 08/13] Update tests/unittests/a2a/utils/test_agent_to_a2a.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/unittests/a2a/utils/test_agent_to_a2a.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index 87a778adaa..d841243e5b 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -155,11 +155,11 @@ def test_to_a2a_with_custom_runner( @patch("google.adk.a2a.utils.agent_to_a2a.Starlette") def test_to_a2a_with_custom_task_store( self, - mock_starlette_class, - mock_card_builder_class, - mock_task_store_class, - mock_request_handler_class, mock_agent_executor_class, + mock_request_handler_class, + mock_task_store_class, + mock_card_builder_class, + mock_starlette_class, ): """Test to_a2a with a custom task store.""" # Arrange From f84d2e500ce2ec6e1a27431f56dde9f27f8f81d1 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:26:09 +0000 Subject: [PATCH 09/13] Update src/google/adk/a2a/utils/agent_to_a2a.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/a2a/utils/agent_to_a2a.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/google/adk/a2a/utils/agent_to_a2a.py b/src/google/adk/a2a/utils/agent_to_a2a.py index abcc6cf261..a5639ee09f 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -121,9 +121,9 @@ def to_a2a( app = to_a2a(agent, agent_card=my_custom_agent_card) # Or with custom task store: - from a2a.server.tasks import MyCustomTaskStore - custom_store = MyCustomTaskStore() - app = to_a2a(agent, task_store=custom_store) + from a2a.server.tasks import TaskStore + class MyCustomTaskStore(TaskStore): pass # A user-defined TaskStore + app = to_a2a(agent, task_store=MyCustomTaskStore()) """ # Set up ADK logging to ensure logs are visible when using uvicorn directly setup_adk_logger(logging.INFO) From 6f265285bb0e249709e43923230cce6437d0dfa7 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:29:59 +0000 Subject: [PATCH 10/13] Update tests/unittests/a2a/utils/test_agent_to_a2a.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/unittests/a2a/utils/test_agent_to_a2a.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index d841243e5b..87a778adaa 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -155,11 +155,11 @@ def test_to_a2a_with_custom_runner( @patch("google.adk.a2a.utils.agent_to_a2a.Starlette") def test_to_a2a_with_custom_task_store( self, - mock_agent_executor_class, - mock_request_handler_class, - mock_task_store_class, - mock_card_builder_class, 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 task store.""" # Arrange From bfdac2d4ae00cf3c01f3115919f428729285a4f7 Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:30:20 +0000 Subject: [PATCH 11/13] Update src/google/adk/a2a/utils/agent_to_a2a.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/a2a/utils/agent_to_a2a.py | 2 +- 1 file changed, 1 insertion(+), 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 a5639ee09f..b01e39f23c 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -122,7 +122,7 @@ def to_a2a( # Or with custom task store: from a2a.server.tasks import TaskStore - class MyCustomTaskStore(TaskStore): pass # A user-defined TaskStore + class MyCustomTaskStore(TaskStore): ... # A user-defined TaskStore; abstract methods must be implemented app = to_a2a(agent, task_store=MyCustomTaskStore()) """ # Set up ADK logging to ensure logs are visible when using uvicorn directly From 9454d4e248b5638b3a702f3a6a609a234822acbb Mon Sep 17 00:00:00 2001 From: secprog Date: Thu, 20 Nov 2025 14:35:24 +0000 Subject: [PATCH 12/13] Remove unused mock imports in test_agent_to_a2a.py to streamline the test setup. --- tests/unittests/a2a/utils/test_agent_to_a2a.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index 87a778adaa..7f0e8e321c 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -168,10 +168,6 @@ def test_to_a2a_with_custom_task_store( custom_task_store = Mock(spec=TaskStore) 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 # Act result = to_a2a(self.mock_agent, task_store=custom_task_store) From 429a928b8c98f3a029b1679b8ee0ae3e0b249414 Mon Sep 17 00:00:00 2001 From: secprog Date: Fri, 21 Nov 2025 15:54:31 +0000 Subject: [PATCH 13/13] refactor: Clean up import statements in experiment and run_experiment modules Removed unnecessary blank lines and organized imports for better readability in `experiment.py` and `run_experiment.py`. Also, updated import statements in `agent_to_a2a.py` and its corresponding test file to maintain consistency and clarity. --- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - src/google/adk/a2a/utils/agent_to_a2a.py | 3 ++- tests/unittests/a2a/utils/test_agent_to_a2a.py | 3 ++- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index 2f5d03a772..f68b349d9c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index cfd850b3a3..1bc4ee58c8 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/src/google/adk/a2a/utils/agent_to_a2a.py b/src/google/adk/a2a/utils/agent_to_a2a.py index b01e39f23c..6e664ca7f8 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -20,7 +20,8 @@ try: from a2a.server.apps import A2AStarletteApplication from a2a.server.request_handlers import DefaultRequestHandler - from a2a.server.tasks import TaskStore, InMemoryTaskStore + from a2a.server.tasks import InMemoryTaskStore + from a2a.server.tasks import TaskStore from a2a.types import AgentCard except ImportError as e: if sys.version_info < (3, 10): diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index 7f0e8e321c..69040701b3 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -28,7 +28,8 @@ try: from a2a.server.apps import A2AStarletteApplication from a2a.server.request_handlers import DefaultRequestHandler - from a2a.server.tasks import TaskStore, InMemoryTaskStore + from a2a.server.tasks import InMemoryTaskStore + from a2a.server.tasks import TaskStore from a2a.types import AgentCard from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor from google.adk.a2a.utils.agent_card_builder import AgentCardBuilder