feat: Change service creation and add app name mapping for sessions

This change refactors how session, memory, and artifact services are created in the fast_api server, using the shared service_factory.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 839997110

Author: GWeale

src/google/adk/cli/cli.py

@@ -159,11 +159,22 @@ async def run_cli(
   load_services_module(str(agent_root))
   user_id = 'test_user'
 
+  agents_dir = str(agent_parent_path)
+  agent_loader = AgentLoader(agents_dir=agents_dir)
+  agent_or_app = agent_loader.load_agent(agent_folder_name)
+  session_app_name = (
+      agent_or_app.name if isinstance(agent_or_app, App) else agent_folder_name
+  )
+  app_name_to_dir = None
+  if isinstance(agent_or_app, App) and agent_or_app.name != agent_folder_name:
+    app_name_to_dir = {agent_or_app.name: agent_folder_name}
+
   # Create session and artifact services using factory functions
   # Sessions persist under <agents_dir>/<agent>/.adk/session.db by default.
   session_service = create_session_service_from_options(
       base_dir=agent_parent_path,
       session_service_uri=session_service_uri,
+      app_name_to_dir=app_name_to_dir,
   )
 
   artifact_service = create_artifact_service_from_options(
@@ -172,13 +183,6 @@ async def run_cli(
   )
 
   credential_service = InMemoryCredentialService()
-  agents_dir = str(agent_parent_path)
-  agent_or_app = AgentLoader(agents_dir=agents_dir).load_agent(
-      agent_folder_name
-  )
-  session_app_name = (
-      agent_or_app.name if isinstance(agent_or_app, App) else agent_folder_name
-  )
   if not is_env_enabled('ADK_DISABLE_LOAD_DOTENV'):
     envs.load_dotenv_for_agent(agent_folder_name, agents_dir)
 

src/google/adk/cli/fast_api.py

@@ -35,28 +35,24 @@
 from starlette.types import Lifespan
 from watchdog.observers import Observer
 
-from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
 from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
 from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
 from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
-from ..memory.in_memory_memory_service import InMemoryMemoryService
 from ..runners import Runner
-from ..sessions.in_memory_session_service import InMemorySessionService
 from .adk_web_server import AdkWebServer
-from .service_registry import get_service_registry
 from .service_registry import load_services_module
 from .utils import envs
 from .utils import evals
 from .utils.agent_change_handler import AgentChangeEventHandler
 from .utils.agent_loader import AgentLoader
+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
 
 logger = logging.getLogger("google_adk." + __name__)
 
 _LAZY_SERVICE_IMPORTS: dict[str, str] = {
     "AgentLoader": ".utils.agent_loader",
-    "InMemoryArtifactService": "..artifacts.in_memory_artifact_service",
-    "InMemoryMemoryService": "..memory.in_memory_memory_service",
-    "InMemorySessionService": "..sessions.in_memory_session_service",
     "LocalEvalSetResultsManager": "..evaluation.local_eval_set_results_manager",
     "LocalEvalSetsManager": "..evaluation.local_eval_sets_manager",
 }
@@ -112,48 +108,31 @@ def get_fast_api_app(
   # Load services.py from agents_dir for custom service registration.
   load_services_module(agents_dir)
 
-  service_registry = get_service_registry()
-
   # Build the Memory service
-  if memory_service_uri:
-    memory_service = service_registry.create_memory_service(
-        memory_service_uri, agents_dir=agents_dir
+  try:
+    memory_service = create_memory_service_from_options(
+        base_dir=agents_dir,
+        memory_service_uri=memory_service_uri,
     )
-    if not memory_service:
-      raise click.ClickException(
-          "Unsupported memory service URI: %s" % memory_service_uri
-      )
-  else:
-    memory_service = InMemoryMemoryService()
+  except ValueError as exc:
+    raise click.ClickException(str(exc)) from exc
 
   # Build the Session service
-  if session_service_uri:
-    session_kwargs = session_db_kwargs or {}
-    session_service = service_registry.create_session_service(
-        session_service_uri, agents_dir=agents_dir, **session_kwargs
-    )
-    if not session_service:
-      # Fallback to DatabaseSessionService if the service registry doesn't
-      # support the session service URI scheme.
-      from ..sessions.database_session_service import DatabaseSessionService
-
-      session_service = DatabaseSessionService(
-          db_url=session_service_uri, **session_kwargs
-      )
-  else:
-    session_service = InMemorySessionService()
+  session_service = create_session_service_from_options(
+      base_dir=agents_dir,
+      session_service_uri=session_service_uri,
+      session_db_kwargs=session_db_kwargs,
+  )
 
   # Build the Artifact service
-  if artifact_service_uri:
-    artifact_service = service_registry.create_artifact_service(
-        artifact_service_uri, agents_dir=agents_dir
+  try:
+    artifact_service = create_artifact_service_from_options(
+        base_dir=agents_dir,
+        artifact_service_uri=artifact_service_uri,
+        strict_uri=True,
     )
-    if not artifact_service:
-      raise click.ClickException(
-          "Unsupported artifact service URI: %s" % artifact_service_uri
-      )
-  else:
-    artifact_service = InMemoryArtifactService()
+  except ValueError as exc:
+    raise click.ClickException(str(exc)) from exc
 
   # Build  the Credential service
   credential_service = InMemoryCredentialService()

src/google/adk/cli/utils/local_storage.py

@@ -17,6 +17,7 @@
 import asyncio
 import logging
 from pathlib import Path
+from typing import Mapping
 from typing import Optional
 
 from typing_extensions import override
@@ -61,6 +62,7 @@ def create_local_session_service(
     *,
     base_dir: Path | str,
     per_agent: bool = False,
+    app_name_to_dir: Optional[Mapping[str, str]] = None,
 ) -> BaseSessionService:
   """Creates a local SQLite-backed session service.
 
@@ -69,6 +71,8 @@ def create_local_session_service(
     per_agent: If True, creates a PerAgentDatabaseSessionService that stores
       sessions in each agent's .adk folder. If False, creates a single
       SqliteSessionService at base_dir/.adk/session.db.
+    app_name_to_dir: Optional mapping from logical app name to on-disk agent
+      folder name. Only used when per_agent is True; defaults to identity.
 
   Returns:
     A BaseSessionService instance backed by SQLite.
@@ -78,7 +82,10 @@ def create_local_session_service(
         "Using per-agent session storage rooted at %s",
         base_dir,
     )
-    return PerAgentDatabaseSessionService(agents_root=base_dir)
+    return PerAgentDatabaseSessionService(
+        agents_root=base_dir,
+        app_name_to_dir=app_name_to_dir,
+    )
 
   return create_local_database_session_service(base_dir=base_dir)
 
@@ -108,23 +115,26 @@ def __init__(
       self,
       *,
       agents_root: Path | str,
+      app_name_to_dir: Optional[Mapping[str, str]] = None,
   ):
     self._agents_root = Path(agents_root).resolve()
+    self._app_name_to_dir = dict(app_name_to_dir or {})
     self._services: dict[str, BaseSessionService] = {}
     self._service_lock = asyncio.Lock()
 
   async def _get_service(self, app_name: str) -> BaseSessionService:
     async with self._service_lock:
-      service = self._services.get(app_name)
+      storage_name = self._app_name_to_dir.get(app_name, app_name)
+      service = self._services.get(storage_name)
       if service is not None:
         return service
       folder = dot_adk_folder_for_agent(
-          agents_root=self._agents_root, app_name=app_name
+          agents_root=self._agents_root, app_name=storage_name
       )
       service = create_local_database_session_service(
           base_dir=folder.agent_dir,
       )
-      self._services[app_name] = service
+      self._services[storage_name] = service
       return service
 
   @override

src/google/adk/cli/utils/service_factory.py

@@ -33,6 +33,7 @@ def create_session_service_from_options(
     base_dir: Path | str,
     session_service_uri: Optional[str] = None,
     session_db_kwargs: Optional[dict[str, Any]] = None,
+    app_name_to_dir: Optional[dict[str, str]] = None,
 ) -> BaseSessionService:
   """Creates a session service based on CLI/web options."""
   base_path = Path(base_dir)
@@ -64,7 +65,11 @@ def create_session_service_from_options(
     return DatabaseSessionService(db_url=session_service_uri, **fallback_kwargs)
 
   # Default to per-agent local SQLite storage in <agents_root>/<agent>/.adk/.
-  return create_local_session_service(base_dir=base_path, per_agent=True)
+  return create_local_session_service(
+      base_dir=base_path,
+      per_agent=True,
+      app_name_to_dir=app_name_to_dir,
+  )
 
 
 def create_memory_service_from_options(
@@ -96,6 +101,7 @@ def create_artifact_service_from_options(
     *,
     base_dir: Path | str,
     artifact_service_uri: Optional[str] = None,
+    strict_uri: bool = False,
 ) -> BaseArtifactService:
   """Creates an artifact service based on CLI/web options."""
   base_path = Path(base_dir)
@@ -108,6 +114,10 @@ def create_artifact_service_from_options(
         agents_dir=str(base_path),
     )
     if service is None:
+      if strict_uri:
+        raise ValueError(
+            f"Unsupported artifact service URI: {artifact_service_uri}"
+        )
       logger.warning(
           "Unsupported artifact service URI: %s, falling back to in-memory",
           artifact_service_uri,

tests/unittests/cli/test_fast_api.py

@@ -416,15 +416,15 @@ def test_app(
   with (
       patch("signal.signal", return_value=None),
       patch(
-          "google.adk.cli.fast_api.InMemorySessionService",
+          "google.adk.cli.fast_api.create_session_service_from_options",
           return_value=mock_session_service,
       ),
       patch(
-          "google.adk.cli.fast_api.InMemoryArtifactService",
+          "google.adk.cli.fast_api.create_artifact_service_from_options",
           return_value=mock_artifact_service,
       ),
       patch(
-          "google.adk.cli.fast_api.InMemoryMemoryService",
+          "google.adk.cli.fast_api.create_memory_service_from_options",
           return_value=mock_memory_service,
       ),
       patch(
@@ -556,15 +556,15 @@ def test_app_with_a2a(
   with (
       patch("signal.signal", return_value=None),
       patch(
-          "google.adk.cli.fast_api.InMemorySessionService",
+          "google.adk.cli.fast_api.create_session_service_from_options",
           return_value=mock_session_service,
       ),
       patch(
-          "google.adk.cli.fast_api.InMemoryArtifactService",
+          "google.adk.cli.fast_api.create_artifact_service_from_options",
           return_value=mock_artifact_service,
       ),
       patch(
-          "google.adk.cli.fast_api.InMemoryMemoryService",
+          "google.adk.cli.fast_api.create_memory_service_from_options",
           return_value=mock_memory_service,
       ),
       patch(

tests/unittests/cli/utils/test_local_storage.py

@@ -17,6 +17,7 @@
 from pathlib import Path
 
 from google.adk.cli.utils.local_storage import create_local_database_session_service
+from google.adk.cli.utils.local_storage import create_local_session_service
 from google.adk.cli.utils.local_storage import PerAgentDatabaseSessionService
 from google.adk.sessions.sqlite_session_service import SqliteSessionService
 import pytest
@@ -48,6 +49,29 @@ async def test_per_agent_session_service_creates_scoped_dot_adk(
   assert agent_b_sessions.sessions[0].app_name == "agent_b"
 
 
+@pytest.mark.asyncio
+async def test_per_agent_session_service_respects_app_name_alias(
+    tmp_path: Path,
+) -> None:
+  folder_name = "agent_folder"
+  logical_name = "custom_app"
+  (tmp_path / folder_name).mkdir()
+
+  service = create_local_session_service(
+      base_dir=tmp_path,
+      per_agent=True,
+      app_name_to_dir={logical_name: folder_name},
+  )
+
+  session = await service.create_session(
+      app_name=logical_name,
+      user_id="user",
+  )
+
+  assert session.app_name == logical_name
+  assert (tmp_path / folder_name / ".adk" / "session.db").exists()
+
+
 def test_create_local_database_session_service_returns_sqlite(
     tmp_path: Path,
 ) -> None:

tests/unittests/cli/utils/test_service_factory.py

@@ -60,6 +60,25 @@ async def test_create_session_service_defaults_to_per_agent_sqlite(
   assert (agent_dir / ".adk" / "session.db").exists()
 
 
+@pytest.mark.asyncio
+async def test_create_session_service_respects_app_name_mapping(
+    tmp_path: Path,
+) -> None:
+  agent_dir = tmp_path / "agent_folder"
+  logical_name = "custom_app"
+  agent_dir.mkdir()
+
+  service = service_factory.create_session_service_from_options(
+      base_dir=tmp_path,
+      app_name_to_dir={logical_name: "agent_folder"},
+  )
+
+  assert isinstance(service, PerAgentDatabaseSessionService)
+  session = await service.create_session(app_name=logical_name, user_id="user")
+  assert session.app_name == logical_name
+  assert (agent_dir / ".adk" / "session.db").exists()
+
+
 def test_create_session_service_fallbacks_to_database(
     tmp_path: Path, monkeypatch
 ):
@@ -101,6 +120,21 @@ def test_create_artifact_service_uses_registry(tmp_path: Path, monkeypatch):
   )
 
 
+def test_create_artifact_service_raises_on_unknown_scheme_when_strict(
+    tmp_path: Path, monkeypatch
+):
+  registry = Mock()
+  registry.create_artifact_service.return_value = None
+  monkeypatch.setattr(service_factory, "get_service_registry", lambda: registry)
+
+  with pytest.raises(ValueError):
+    service_factory.create_artifact_service_from_options(
+        base_dir=tmp_path,
+        artifact_service_uri="unknown://foo",
+        strict_uri=True,
+    )
+
+
 def test_create_memory_service_uses_registry(tmp_path: Path, monkeypatch):
   registry = Mock()
   expected = object()