diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 379c314a5a..dbe3407998 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -1714,6 +1714,20 @@ def decorator(func): ), default=None, ) + @click.option( + "--app_name", + type=str, + help=( + "Optional. Overrides the application name used to store and look up" + " sessions and artifacts. By default the app_name is the agent's" + " folder name, so multiple servers loading an identically-named" + " agent folder against a shared backend share the same sessions." + " Set a distinct value per server to keep their persisted data" + " isolated without renaming the folder. Also settable via the" + " ADK_APP_NAME environment variable." + ), + default=None, + ) @functools.wraps(func) @click.pass_context def wrapper(ctx, *args, **kwargs): @@ -1788,6 +1802,7 @@ def cli_web( logo_text: str | None = None, logo_image_url: str | None = None, trigger_sources: list[str] | None = None, + app_name: str | None = None, ): """Starts a FastAPI server with Web UI for agents. @@ -1828,6 +1843,7 @@ async def _lifespan(app: FastAPI): app = get_fast_api_app( agents_dir=agents_dir, + app_name=app_name, session_service_uri=session_service_uri, artifact_service_uri=artifact_service_uri, memory_service_uri=memory_service_uri, @@ -1929,6 +1945,7 @@ def cli_api_server( with_ui: bool = False, gemini_enterprise_app_name: str | None = None, express_mode: bool = False, + app_name: str | None = None, ): """Starts a FastAPI server for agents. @@ -1954,6 +1971,7 @@ def cli_api_server( config = uvicorn.Config( get_fast_api_app( agents_dir=agents_dir, + app_name=app_name, session_service_uri=session_service_uri, artifact_service_uri=artifact_service_uri, memory_service_uri=memory_service_uri, diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index f096717dc2..99f94d4c86 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -404,6 +404,7 @@ def get_fast_api_app( *, agents_dir: str, agent_loader: BaseAgentLoader | None = None, + app_name: str | None = None, session_service_uri: str | None = None, session_db_kwargs: Mapping[str, Any] | None = None, artifact_service_uri: str | None = None, @@ -443,6 +444,13 @@ def get_fast_api_app( services.py/yaml), and as a base for local storage. agent_loader: An optional custom loader for retrieving agent instances. If not provided, a default AgentLoader targeting agents_dir is used. + app_name: Optional override for the application name used to store and look + up sessions and artifacts. By default the app_name is the agent's folder + name, which means two servers loading an identically-named agent folder + against a shared backend share the same session/artifact namespace. Set a + distinct value per server to keep their persisted data isolated without + renaming the folder. Falls back to the ``ADK_APP_NAME`` environment + variable when not provided. session_service_uri: A URI defining the backend for session persistence. Supports schemes like 'memory://', 'sqlite://', 'postgresql://', 'mysql://', or 'agentengine://'. Defaults to per-agent local SQLite @@ -562,6 +570,20 @@ def get_fast_api_app( except ValueError as exc: raise click.ClickException(str(exc)) from exc + # Pin the app_name used for persisted storage when an override is supplied + # (via the app_name argument or the ADK_APP_NAME environment variable). This + # keeps sessions/artifacts isolated across servers that load an + # identically-named agent folder against a shared backend. + app_name = app_name or os.environ.get("ADK_APP_NAME") + if app_name: + from .utils.app_name_override import maybe_override_app_name + + session_service, artifact_service = maybe_override_app_name( + app_name, + session_service=session_service, + artifact_service=artifact_service, + ) + # Build the Credential service credential_service = InMemoryCredentialService() diff --git a/src/google/adk/cli/utils/app_name_override.py b/src/google/adk/cli/utils/app_name_override.py new file mode 100644 index 0000000000..1d6dae1cf7 --- /dev/null +++ b/src/google/adk/cli/utils/app_name_override.py @@ -0,0 +1,147 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Service wrappers that pin the ``app_name`` used for persisted storage. + +When agents are served via ``adk web`` / ``adk api_server`` by pointing at a +directory, the ``app_name`` used to key sessions and artifacts is derived from +the agent's *folder name*. Two servers that load an identically-named agent +folder against a shared backend therefore write to the same ``app_name`` and +share each other's sessions. + +These thin decorators let a server pin a distinct ``app_name`` for its +persisted data without renaming the folder (so the loader, the REST routes and +the web UI keep using the folder name, while storage is namespaced by the +override). They are wired up from :func:`get_fast_api_app` when an override is +supplied via the ``app_name`` argument, the ``--app_name`` CLI flag, or the +``ADK_APP_NAME`` environment variable. + +The wrappers only rewrite the ``app_name`` keyword argument; every other +argument and any implementation-specific attribute is delegated unchanged. +""" + +from __future__ import annotations + +from typing import Any + +from ...artifacts.base_artifact_service import BaseArtifactService +from ...events.event import Event +from ...sessions.base_session_service import BaseSessionService +from ...sessions.session import Session + + +class _AppNameOverrideBase: + """Holds the wrapped service and the override, delegating unknown attrs.""" + + def __init__(self, delegate: Any, app_name: str) -> None: + self._delegate = delegate + self._app_name = app_name + + def __getattr__(self, name: str) -> Any: + # __getattr__ only runs for attributes not found normally. Guard the two + # private attributes so a lookup before __init__ completes (e.g. during + # copy/pickle) raises AttributeError instead of recursing forever. + if name in ("_delegate", "_app_name"): + raise AttributeError(name) + return getattr(self._delegate, name) + + +class AppNameOverrideSessionService(_AppNameOverrideBase, BaseSessionService): + """A :class:`BaseSessionService` that forces a fixed ``app_name``.""" + + async def create_session(self, **kwargs: Any) -> Session: + kwargs["app_name"] = self._app_name + return await self._delegate.create_session(**kwargs) + + async def get_session(self, **kwargs: Any): + kwargs["app_name"] = self._app_name + return await self._delegate.get_session(**kwargs) + + async def list_sessions(self, **kwargs: Any): + kwargs["app_name"] = self._app_name + return await self._delegate.list_sessions(**kwargs) + + async def delete_session(self, **kwargs: Any) -> None: + kwargs["app_name"] = self._app_name + return await self._delegate.delete_session(**kwargs) + + async def get_user_state(self, **kwargs: Any) -> dict[str, Any]: + kwargs["app_name"] = self._app_name + return await self._delegate.get_user_state(**kwargs) + + async def append_event(self, session: Session, event: Event) -> Event: + # ``session.app_name`` already carries the override, because the session + # was created or fetched through this wrapper, so no rewrite is needed. + return await self._delegate.append_event(session=session, event=event) + + async def flush(self) -> None: + return await self._delegate.flush() + + +class AppNameOverrideArtifactService(_AppNameOverrideBase, BaseArtifactService): + """A :class:`BaseArtifactService` that forces a fixed ``app_name``.""" + + async def save_artifact(self, **kwargs: Any) -> int: + kwargs["app_name"] = self._app_name + return await self._delegate.save_artifact(**kwargs) + + async def load_artifact(self, **kwargs: Any): + kwargs["app_name"] = self._app_name + return await self._delegate.load_artifact(**kwargs) + + async def list_artifact_keys(self, **kwargs: Any) -> list[str]: + kwargs["app_name"] = self._app_name + return await self._delegate.list_artifact_keys(**kwargs) + + async def delete_artifact(self, **kwargs: Any) -> None: + kwargs["app_name"] = self._app_name + return await self._delegate.delete_artifact(**kwargs) + + async def list_versions(self, **kwargs: Any) -> list[int]: + kwargs["app_name"] = self._app_name + return await self._delegate.list_versions(**kwargs) + + async def list_artifact_versions(self, **kwargs: Any): + kwargs["app_name"] = self._app_name + return await self._delegate.list_artifact_versions(**kwargs) + + async def get_artifact_version(self, **kwargs: Any): + kwargs["app_name"] = self._app_name + return await self._delegate.get_artifact_version(**kwargs) + + +def maybe_override_app_name( + app_name: str | None, + *, + session_service: BaseSessionService, + artifact_service: BaseArtifactService, +) -> tuple[BaseSessionService, BaseArtifactService]: + """Wraps the services to pin ``app_name`` when an override is provided. + + Args: + app_name: The override application name. When falsy the services are + returned unchanged. + session_service: The session service to wrap. + artifact_service: The artifact service to wrap. + + Returns: + A ``(session_service, artifact_service)`` tuple, wrapped when an override + is supplied and unchanged otherwise. + """ + if not app_name: + return session_service, artifact_service + return ( + AppNameOverrideSessionService(session_service, app_name), + AppNameOverrideArtifactService(artifact_service, app_name), + ) diff --git a/tests/unittests/cli/utils/test_app_name_override.py b/tests/unittests/cli/utils/test_app_name_override.py new file mode 100644 index 0000000000..32a6d4ad7c --- /dev/null +++ b/tests/unittests/cli/utils/test_app_name_override.py @@ -0,0 +1,242 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.cli.utils.app_name_override import AppNameOverrideArtifactService +from google.adk.cli.utils.app_name_override import AppNameOverrideSessionService +from google.adk.cli.utils.app_name_override import maybe_override_app_name +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.genai import types +import pytest + +# The "folder name" a caller/route uses; the override is what should be stored. +FOLDER = "assistant" +OVERRIDE = "tenant_a" +USER = "u1" + + +@pytest.mark.asyncio +async def test_session_override_stores_under_override_name(): + delegate = InMemorySessionService() + service = AppNameOverrideSessionService(delegate, OVERRIDE) + + created = await service.create_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + # The returned session is keyed by the override, not the folder name. + assert created.app_name == OVERRIDE + + # It is reachable through the override on the underlying service ... + assert ( + await delegate.get_session( + app_name=OVERRIDE, user_id=USER, session_id="s1" + ) + is not None + ) + # ... and NOT under the original folder name. + assert ( + await delegate.get_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_session_override_get_list_delete_roundtrip(): + service = AppNameOverrideSessionService(InMemorySessionService(), OVERRIDE) + + await service.create_session(app_name=FOLDER, user_id=USER, session_id="s1") + + # get/list/delete all accept the folder name but operate on the override. + assert ( + await service.get_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + is not None + ) + listed = await service.list_sessions(app_name=FOLDER, user_id=USER) + assert [s.id for s in listed.sessions] == ["s1"] + + await service.delete_session(app_name=FOLDER, user_id=USER, session_id="s1") + assert ( + await service.get_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_two_overrides_do_not_share_a_backend(): + """Two servers, same folder + backend, distinct overrides -> isolated.""" + delegate = InMemorySessionService() + service_a = AppNameOverrideSessionService(delegate, "tenant_a") + service_b = AppNameOverrideSessionService(delegate, "tenant_b") + + await service_a.create_session(app_name=FOLDER, user_id=USER, session_id="s1") + + # B uses the same folder name but a different override -> cannot see A's. + assert ( + await service_b.get_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + is None + ) + # A second server with the SAME override does share, on purpose. + service_a2 = AppNameOverrideSessionService(delegate, "tenant_a") + assert ( + await service_a2.get_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + is not None + ) + + +@pytest.mark.asyncio +async def test_artifact_override_stores_under_override_name(): + delegate = InMemoryArtifactService() + service = AppNameOverrideArtifactService(delegate, OVERRIDE) + part = types.Part.from_text(text="hello") + + await service.save_artifact( + app_name=FOLDER, + user_id=USER, + session_id="s1", + filename="a.txt", + artifact=part, + ) + + # Stored under the override, not the folder name. + assert ( + await delegate.load_artifact( + app_name=OVERRIDE, user_id=USER, session_id="s1", filename="a.txt" + ) + is not None + ) + assert ( + await delegate.load_artifact( + app_name=FOLDER, user_id=USER, session_id="s1", filename="a.txt" + ) + is None + ) + + +def test_maybe_override_is_noop_without_app_name(): + session_service = InMemorySessionService() + artifact_service = InMemoryArtifactService() + + s, a = maybe_override_app_name( + None, session_service=session_service, artifact_service=artifact_service + ) + assert s is session_service + assert a is artifact_service + + +def test_maybe_override_wraps_when_app_name_given(): + s, a = maybe_override_app_name( + OVERRIDE, + session_service=InMemorySessionService(), + artifact_service=InMemoryArtifactService(), + ) + assert isinstance(s, AppNameOverrideSessionService) + assert isinstance(a, AppNameOverrideArtifactService) + + +def test_unknown_attributes_delegate_through(): + delegate = InMemorySessionService() + service = AppNameOverrideSessionService(delegate, OVERRIDE) + # sessions dict is an impl detail of InMemorySessionService; it must be + # reachable via delegation rather than raising AttributeError. + assert service.sessions is delegate.sessions + + +def _make_session_backend(name: str, tmp_path): + """Builds a concrete session backend, skipping when drivers are absent. + + The override wrapper operates at the ``BaseSessionService`` abstraction, so it + must behave identically regardless of the underlying storage. ``postgresql`` + and ``mysql`` use the very same ``DatabaseSessionService`` class as the + ``sqlite+aiosqlite`` case below (only the connection URL differs), so covering + the SQLAlchemy path with sqlite exercises the same code that runs on Postgres. + """ + if name == "in_memory": + return InMemorySessionService() + if name == "database_sqlite": + pytest.importorskip("sqlalchemy") + pytest.importorskip("aiosqlite") + from google.adk.sessions.database_session_service import ( + DatabaseSessionService, + ) + + db_path = (tmp_path / "database.sqlite").as_posix() + return DatabaseSessionService(db_url=f"sqlite+aiosqlite:///{db_path}") + if name == "sqlite_service": + pytest.importorskip("aiosqlite") + from google.adk.sessions.sqlite_session_service import SqliteSessionService + + return SqliteSessionService(db_path=(tmp_path / "sqlite.db").as_posix()) + raise ValueError(name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "backend_name", + ["in_memory", "database_sqlite", "sqlite_service"], +) +async def test_override_is_backend_agnostic(backend_name, tmp_path): + """The override must isolate storage identically on every backend. + + This is the guard that keeps the feature non-breaking across databases: the + wrapper only rewrites the ``app_name`` string, so create/get/list/delete must + round-trip under the override (and never under the folder name) whether the + backend is in-memory, a SQLAlchemy database, or the sqlite service. + """ + delegate = _make_session_backend(backend_name, tmp_path) + service = AppNameOverrideSessionService(delegate, OVERRIDE) + + created = await service.create_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + assert created.app_name == OVERRIDE + + fetched = await service.get_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + assert fetched is not None and fetched.app_name == OVERRIDE + + listed = await service.list_sessions(app_name=FOLDER, user_id=USER) + assert [s.id for s in listed.sessions] == ["s1"] + + # Persisted under the override, never under the caller-supplied folder name. + assert ( + await delegate.get_session( + app_name=OVERRIDE, user_id=USER, session_id="s1" + ) + is not None + ) + assert ( + await delegate.get_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + is None + ) + + await service.delete_session(app_name=FOLDER, user_id=USER, session_id="s1") + assert ( + await service.get_session( + app_name=FOLDER, user_id=USER, session_id="s1" + ) + is None + )