Skip to content

Commit 468bb22

Browse files
STHITAPRAJNASclaude
authored andcommitted
fixup: code style and convention alignment
- hoist a2a_push_config_store init outside per-agent loop, matching the task_store pattern (both now use if-None guard at the top) - remove redundant local push_config_store variable; use a2a_push_config_store directly in DefaultRequestHandler call - trim docstring: drop inline Example block, shorten to match the one-liner style used by the other params in the same function - rename test to test_a2a_custom_task_store_is_used (shorter, matches existing naming style) - fix test mock setup to use an explicit mock_a2a_app_instance like the test_app_with_a2a fixture does - simplify test assertion: drop intermediate call_kwargs variable
1 parent 004318a commit 468bb22

2 files changed

Lines changed: 18 additions & 33 deletions

File tree

src/google/adk/cli/fast_api.py

Lines changed: 8 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -131,22 +131,11 @@ def get_fast_api_app(
131131
web: Whether to enable the web UI and serve its assets.
132132
a2a: Whether to enable Agent-to-Agent (A2A) protocol support.
133133
a2a_task_store: Optional A2A TaskStore instance. Defaults to
134-
InMemoryTaskStore when a2a=True. Pass a DatabaseTaskStore (from the
135-
a2a-sdk) for persistence across server restarts and horizontal replicas.
136-
Example::
137-
138-
from a2a.server.tasks import DatabaseTaskStore
139-
from sqlalchemy.ext.asyncio import create_async_engine
140-
141-
engine = create_async_engine("postgresql+asyncpg://user:pw@host/db")
142-
app = get_fast_api_app(
143-
agents_dir="agents/", web=True, a2a=True,
144-
a2a_task_store=DatabaseTaskStore(engine),
145-
)
146-
147-
a2a_push_config_store: Optional A2A PushNotificationConfigStore instance.
148-
Defaults to InMemoryPushNotificationConfigStore when a2a=True. Pass a
149-
DatabasePushNotificationConfigStore for persistence across restarts.
134+
InMemoryTaskStore. Pass a custom store (e.g. DatabaseTaskStore) to
135+
persist task state across restarts or share it across replicas.
136+
a2a_push_config_store: Optional A2A PushNotificationConfigStore.
137+
Defaults to InMemoryPushNotificationConfigStore. Pass a custom store
138+
for persistence across restarts.
150139
host: Host address for the server (defaults to 127.0.0.1).
151140
port: Port number for the server (defaults to 8000).
152141
url_prefix: Optional prefix for all URL routes.
@@ -619,6 +608,8 @@ async def get_agent_builder(
619608
if base_path.exists() and base_path.is_dir():
620609
if a2a_task_store is None:
621610
a2a_task_store = InMemoryTaskStore()
611+
if a2a_push_config_store is None:
612+
a2a_push_config_store = InMemoryPushNotificationConfigStore()
622613

623614
def create_a2a_runner_loader(captured_app_name: str):
624615
"""Factory function to create A2A runner with proper closure."""
@@ -646,16 +637,10 @@ async def _get_a2a_runner_async() -> Runner:
646637
runner=create_a2a_runner_loader(app_name),
647638
)
648639

649-
push_config_store = (
650-
a2a_push_config_store
651-
if a2a_push_config_store is not None
652-
else InMemoryPushNotificationConfigStore()
653-
)
654-
655640
request_handler = DefaultRequestHandler(
656641
agent_executor=agent_executor,
657642
task_store=a2a_task_store,
658-
push_config_store=push_config_store,
643+
push_config_store=a2a_push_config_store,
659644
)
660645

661646
with (p / "agent.json").open("r", encoding="utf-8") as f:

tests/unittests/cli/test_fast_api.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1876,7 +1876,7 @@ def test_a2a_uses_in_memory_task_store_by_default(
18761876
temp_agents_dir_with_a2a,
18771877
monkeypatch,
18781878
):
1879-
"""Test that InMemoryTaskStore is created when no task_store is provided."""
1879+
"""Test that InMemoryTaskStore is used when no task store is provided."""
18801880
with (
18811881
patch("signal.signal", return_value=None),
18821882
patch(
@@ -1909,7 +1909,9 @@ def test_a2a_uses_in_memory_task_store_by_default(
19091909
patch("a2a.server.request_handlers.DefaultRequestHandler"),
19101910
patch("a2a.server.apps.A2AStarletteApplication") as mock_a2a_app,
19111911
):
1912-
mock_a2a_app.return_value.routes.return_value = []
1912+
mock_a2a_app_instance = MagicMock()
1913+
mock_a2a_app_instance.routes.return_value = []
1914+
mock_a2a_app.return_value = mock_a2a_app_instance
19131915
monkeypatch.chdir(temp_agents_dir_with_a2a)
19141916

19151917
_ = get_fast_api_app(
@@ -1927,7 +1929,7 @@ def test_a2a_uses_in_memory_task_store_by_default(
19271929
mock_task_store_class.assert_called_once()
19281930

19291931

1930-
def test_a2a_custom_task_store_bypasses_in_memory_default(
1932+
def test_a2a_custom_task_store_is_used(
19311933
mock_session_service,
19321934
mock_artifact_service,
19331935
mock_memory_service,
@@ -1937,7 +1939,7 @@ def test_a2a_custom_task_store_bypasses_in_memory_default(
19371939
temp_agents_dir_with_a2a,
19381940
monkeypatch,
19391941
):
1940-
"""Test that a custom task_store is forwarded and InMemoryTaskStore is not created."""
1942+
"""Test that a custom task store is forwarded to DefaultRequestHandler."""
19411943
custom_task_store = MagicMock()
19421944

19431945
with (
@@ -1974,7 +1976,9 @@ def test_a2a_custom_task_store_bypasses_in_memory_default(
19741976
) as mock_handler,
19751977
patch("a2a.server.apps.A2AStarletteApplication") as mock_a2a_app,
19761978
):
1977-
mock_a2a_app.return_value.routes.return_value = []
1979+
mock_a2a_app_instance = MagicMock()
1980+
mock_a2a_app_instance.routes.return_value = []
1981+
mock_a2a_app.return_value = mock_a2a_app_instance
19781982
monkeypatch.chdir(temp_agents_dir_with_a2a)
19791983

19801984
_ = get_fast_api_app(
@@ -1990,12 +1994,8 @@ def test_a2a_custom_task_store_bypasses_in_memory_default(
19901994
port=8000,
19911995
)
19921996

1993-
# InMemoryTaskStore must NOT be instantiated when a custom store is supplied
19941997
mock_task_store_class.assert_not_called()
1995-
1996-
# The custom store must be passed through to DefaultRequestHandler
1997-
call_kwargs = mock_handler.call_args.kwargs
1998-
assert call_kwargs["task_store"] is custom_task_store
1998+
assert mock_handler.call_args.kwargs["task_store"] is custom_task_store
19991999

20002000

20012001
def test_patch_memory(test_app, create_test_session, mock_memory_service):

0 commit comments

Comments
 (0)