feat(a2a): add support for persistent task stores

Merge #5597

**Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.**

### Link to Issue or Description of Change

**1. Link to an existing issue (if applicable):**

- Related: #4971

**2. Or, if no issue exists, describe the change:**

_If applicable, please follow the issue templates to provide as much detail as
possible._

**Problem:**
A2A support lacked pluggable or persistent task store backends, forcing a strict default to `InMemoryTaskStore`.

**Solution:**
Extended `ServiceRegistry` and `ServiceFactory` to support URI-driven configuration for A2A task stores. Added built-in support for `memory://`, `postgresql://`, `mysql://`, and `sqlite://` schemes. Plumbed the options down into `to_a2a()` and `get_fast_api_app()`, ensuring connection strings are securely redacted from application logs.

### Testing Plan

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

Summary: Ran `pytest` against `test_agent_to_a2a.py`, `test_fast_api.py`, `test_service_registry.py`, and `test_service_factory.py`. All 131 unit tests passed successfully.

**Manual End-to-End (E2E) Tests:**

Compiled the package distribution wheel via `uv build` and successfully validated its installation inside a clean isolated sandbox environment via `uv pip install dist/google_adk-1.32.0-py3-none-any.whl[a2a]`. Verified that the application instantiates default in-memory stores and custom persistent URI stores

### Checklist

- [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document.
- [x] I have performed a self-review of my own code.
- [x] I have commented my code, particularly in hard-to-understand areas.
- [x] I have added tests that prove my fix is effective or that my feature works.
- [x] New and existing unit tests pass locally with my changes.
- [x] I have manually tested my changes end-to-end.
- [x] Any dependent changes have been merged and published in downstream modules.

COPYBARA_INTEGRATE_REVIEW=#5597 from allen-stephen:feat/a2a-task-store-support f53ba0f
PiperOrigin-RevId: 914496800

Author: allen-stephen

src/google/adk/a2a/utils/agent_to_a2a.py

@@ -18,14 +18,13 @@
 import logging
 from typing import AsyncIterator
 from typing import Callable
-from typing import Optional
-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.server.tasks import TaskStore
 from a2a.types import AgentCard
 from starlette.applications import Starlette
 
@@ -41,8 +40,8 @@
 
 
 def _load_agent_card(
-    agent_card: Optional[Union[AgentCard, str]],
-) -> Optional[AgentCard]:
+    agent_card: AgentCard | str | None,
+) -> AgentCard | None:
   """Load agent card from various sources.
 
   Args:
@@ -82,10 +81,11 @@ def to_a2a(
     host: str = "localhost",
     port: int = 8000,
     protocol: str = "http",
-    agent_card: Optional[Union[AgentCard, str]] = None,
-    push_config_store: Optional[PushNotificationConfigStore] = None,
-    runner: Optional[Runner] = None,
-    lifespan: Optional[Callable[[Starlette], AsyncIterator[None]]] = None,
+    agent_card: AgentCard | str | None = None,
+    push_config_store: PushNotificationConfigStore | None = None,
+    task_store: TaskStore | None = None,
+    runner: Runner | None = None,
+    lifespan: Callable[[Starlette], AsyncIterator[None]] | None = None,
 ) -> Starlette:
   """Convert an ADK agent to a A2A Starlette application.
 
@@ -100,6 +100,8 @@ def to_a2a(
       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.
+      task_store: Optional A2A task store for persisting task state. If not
+        provided, an in-memory store will be created.
       runner: Optional pre-built Runner object. If not provided, a default
               runner will be created using in-memory services.
       lifespan: Optional async context manager for Starlette lifespan
@@ -127,6 +129,20 @@ async def lifespan(app):
           await app.state.db.close()
 
       app = to_a2a(agent, lifespan=lifespan)
+
+      # Or with a persistent task store (the caller owns engine disposal):
+      from a2a.server.tasks import DatabaseTaskStore
+      from sqlalchemy.ext.asyncio import create_async_engine
+
+      engine = create_async_engine("postgresql+asyncpg://...")
+      task_store = DatabaseTaskStore(engine=engine)
+
+      @asynccontextmanager
+      async def lifespan(app):
+          yield
+          await engine.dispose()
+
+      app = to_a2a(agent, task_store=task_store, lifespan=lifespan)
   """
   # Set up ADK logging to ensure logs are visible when using uvicorn directly
   adk_logger = logging.getLogger("google_adk")
@@ -145,7 +161,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,

src/google/adk/cli/fast_api.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from contextlib import asynccontextmanager
 import importlib
 import json
 import logging
@@ -24,7 +25,6 @@
 from typing import Any
 from typing import Literal
 from typing import Mapping
-from typing import Optional
 
 import click
 from fastapi import FastAPI
@@ -48,6 +48,7 @@
 from .utils.agent_change_handler import AgentChangeEventHandler
 from .utils.agent_loader import AgentLoader
 from .utils.base_agent_loader import BaseAgentLoader
+from .utils.service_factory import _create_task_store_from_options
 from .utils.service_factory import create_artifact_service_from_options
 from .utils.service_factory import create_memory_service_from_options
 from .utils.service_factory import create_session_service_from_options
@@ -75,28 +76,29 @@ def __getattr__(name: str):
 def get_fast_api_app(
     *,
     agents_dir: str,
-    agent_loader: Optional[BaseAgentLoader] = None,
-    session_service_uri: Optional[str] = None,
-    session_db_kwargs: Optional[Mapping[str, Any]] = None,
-    artifact_service_uri: Optional[str] = None,
-    memory_service_uri: Optional[str] = None,
+    agent_loader: BaseAgentLoader | None = None,
+    session_service_uri: str | None = None,
+    session_db_kwargs: Mapping[str, Any] | None = None,
+    artifact_service_uri: str | None = None,
+    memory_service_uri: str | None = None,
     use_local_storage: bool = True,
-    eval_storage_uri: Optional[str] = None,
-    allow_origins: Optional[list[str]] = None,
+    eval_storage_uri: str | None = None,
+    allow_origins: list[str] | None = None,
     web: bool,
     a2a: bool = False,
+    task_store_uri: str | None = None,
     host: str = "127.0.0.1",
     port: int = 8000,
-    url_prefix: Optional[str] = None,
+    url_prefix: str | None = None,
     trace_to_cloud: bool = False,
     otel_to_cloud: bool = False,
     reload_agents: bool = False,
-    lifespan: Optional[Lifespan[FastAPI]] = None,
-    extra_plugins: Optional[list[str]] = None,
-    logo_text: Optional[str] = None,
-    logo_image_url: Optional[str] = None,
+    lifespan: Lifespan[FastAPI] | None = None,
+    extra_plugins: list[str] | None = None,
+    logo_text: str | None = None,
+    logo_image_url: str | None = None,
     auto_create_session: bool = False,
-    trigger_sources: Optional[list[Literal["pubsub", "eventarc"]]] = None,
+    trigger_sources: list[Literal["pubsub", "eventarc"]] | None = None,
 ) -> FastAPI:
   """Constructs and returns a FastAPI application for serving ADK agents.
 
@@ -128,6 +130,8 @@ def get_fast_api_app(
     allow_origins: List of allowed origins for CORS.
     web: Whether to enable the web UI and serve its assets.
     a2a: Whether to enable Agent-to-Agent (A2A) protocol support.
+    task_store_uri: URI for the A2A task store. Uses in-memory task store if
+      None. Only used when ``a2a=True``.
     host: Host address for the server (defaults to 127.0.0.1).
     port: Port number for the server (defaults to 8000).
     url_prefix: Optional prefix for all URL routes.
@@ -272,6 +276,33 @@ def tear_down_observer(observer: Observer, _: AdkWebServer):
         web_assets_dir=ANGULAR_DIST_PATH,
     )
 
+  # Create the task store early so its engine can be disposed via the
+  # lifespan, preventing connection pool leaks on shutdown.
+  a2a_task_store = None
+  if a2a:
+    base_path = Path.cwd() / agents_dir
+    if base_path.exists() and base_path.is_dir():
+      a2a_task_store = _create_task_store_from_options(
+          task_store_uri=task_store_uri,
+      )
+
+  if a2a_task_store is not None and hasattr(a2a_task_store, "engine"):
+    outer_lifespan = lifespan
+
+    @asynccontextmanager
+    async def _a2a_lifespan(app_instance: FastAPI):
+      try:
+        if outer_lifespan:
+          async with outer_lifespan(app_instance) as ctx:
+            yield ctx
+        else:
+          yield
+      finally:
+        logger.info("Disposing A2A task store engine")
+        await a2a_task_store.engine.dispose()
+
+    lifespan = _a2a_lifespan
+
   app = adk_web_server.get_fast_api_app(
       lifespan=lifespan,
       allow_origins=allow_origins,
@@ -339,7 +370,7 @@ def _walk(node: Any) -> None:
       for doc in docs:
         _walk(doc)
 
-    def _parse_upload_filename(filename: Optional[str]) -> tuple[str, str]:
+    def _parse_upload_filename(filename: str | None) -> tuple[str, str]:
       if not filename:
         raise ValueError("Upload filename is missing.")
       filename = _normalize_relative_path(filename)
@@ -473,7 +504,7 @@ def ensure_tmp_exists(app_name: str) -> bool:
 
     @app.post("/builder/save", response_model_exclude_none=True)
     async def builder_build(
-        files: list[UploadFile], tmp: Optional[bool] = False
+        files: list[UploadFile], tmp: bool | None = False
     ) -> bool:
       try:
         # Phase 1: parse filenames and read content into memory.
@@ -544,8 +575,8 @@ async def builder_cancel(app_name: str) -> bool:
     )
     async def get_agent_builder(
         app_name: str,
-        file_path: Optional[str] = None,
-        tmp: Optional[bool] = False,
+        file_path: str | None = None,
+        tmp: bool | None = False,
     ):
       try:
         app_root = _get_app_root(app_name)
@@ -584,11 +615,10 @@ async def get_agent_builder(
           headers={"Cache-Control": "no-store"},
       )
 
-  if a2a:
+  if a2a and a2a_task_store is not None:
     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
 
@@ -598,7 +628,6 @@ async def get_agent_builder(
     base_path = Path.cwd() / agents_dir
     # the root agents directory should be an existing folder
     if base_path.exists() and base_path.is_dir():
-      a2a_task_store = InMemoryTaskStore()
 
       def create_a2a_runner_loader(captured_app_name: str):
         """Factory function to create A2A runner with proper closure."""

src/google/adk/cli/service_registry.py

@@ -69,7 +69,6 @@ def my_session_factory(uri: str, **kwargs):
 from pathlib import Path
 import sys
 from typing import Any
-from typing import Optional
 from typing import Protocol
 from urllib.parse import unquote
 from urllib.parse import urlparse
@@ -98,6 +97,7 @@ def __init__(self):
     self._session_factories: dict[str, ServiceFactory] = {}
     self._artifact_factories: dict[str, ServiceFactory] = {}
     self._memory_factories: dict[str, ServiceFactory] = {}
+    self._task_store_factories: dict[str, ServiceFactory] = {}
 
   def register_session_service(
       self, scheme: str, factory: ServiceFactory
@@ -123,6 +123,12 @@ def register_memory_service(
     """Register a factory for a custom memory service URI scheme."""
     self._memory_factories[scheme] = factory
 
+  def _register_task_store_service(
+      self, scheme: str, factory: ServiceFactory
+  ) -> None:
+    """Register a factory for a custom A2A task store URI scheme."""
+    self._task_store_factories[scheme] = factory
+
   def create_session_service(
       self, uri: str, **kwargs
   ) -> BaseSessionService | None:
@@ -150,6 +156,17 @@ def create_memory_service(
       return self._memory_factories[scheme](uri, **kwargs)
     return None
 
+  def _create_task_store_service(self, uri: str, **kwargs: Any) -> Any:
+    """Create A2A task store from URI using registered factories."""
+    scheme = urlparse(uri).scheme
+    if scheme and scheme in self._task_store_factories:
+      return self._task_store_factories[scheme](uri, **kwargs)
+    supported = sorted(self._task_store_factories.keys())
+    raise ValueError(
+        f"Unsupported A2A task store URI scheme: '{scheme}'."
+        f" Supported schemes: {supported}"
+    )
+
 
 def get_service_registry() -> ServiceRegistry:
   """Gets the singleton ServiceRegistry instance, initializing it if needed."""
@@ -333,9 +350,42 @@ def agentengine_memory_factory(uri: str, **kwargs):
   registry.register_memory_service("rag", rag_memory_factory)
   registry.register_memory_service("agentengine", agentengine_memory_factory)
 
+  # -- A2A Task Store Services --
+  def memory_task_store_factory(uri: str, **kwargs: Any) -> Any:
+    try:
+      from a2a.server.tasks import InMemoryTaskStore
+    except ImportError as e:
+      raise ImportError(
+          "A2A task store support requires the 'a2a' package."
+          " Install it with: pip install google-adk[a2a]"
+      ) from e
+
+    return InMemoryTaskStore()
+
+  def database_task_store_factory(uri: str, **kwargs: Any) -> Any:
+    try:
+      from a2a.server.tasks import DatabaseTaskStore
+    except ImportError as e:
+      raise ImportError(
+          "A2A task store support requires the 'a2a' package."
+          " Install it with: pip install google-adk[a2a]"
+      ) from e
+    from sqlalchemy.ext.asyncio import create_async_engine
+
+    engine = create_async_engine(uri)
+    return DatabaseTaskStore(engine=engine)
+
+  registry._register_task_store_service("memory", memory_task_store_factory)
+  for scheme in [
+      "postgresql+asyncpg",
+      "mysql+aiomysql",
+      "sqlite+aiosqlite",
+  ]:
+    registry._register_task_store_service(scheme, database_task_store_factory)
+
 
 def _load_gcp_config(
-    agents_dir: Optional[str], service_name: str
+    agents_dir: str | None, service_name: str
 ) -> tuple[str, str]:
   """Loads GCP project and location from environment."""
   if not agents_dir:
@@ -355,7 +405,7 @@ def _load_gcp_config(
 
 
 def _parse_agent_engine_kwargs(
-    uri_part: str, agents_dir: Optional[str]
+    uri_part: str, agents_dir: str | None
 ) -> dict[str, Any]:
   """Helper to parse agent engine resource name."""
   if not uri_part:
@@ -437,5 +487,7 @@ def _register_services_from_yaml_config(
       registry.register_artifact_service(scheme, factory)
     elif service_type == "memory":
       registry.register_memory_service(scheme, factory)
+    elif service_type == "task_store":
+      registry._register_task_store_service(scheme, factory)
     else:
       logger.warning("Unknown service type in YAML: %s", service_type)