From 99fc17b3365e52e869f14cd7cd39e1e39fda00c0 Mon Sep 17 00:00:00 2001 From: George Weale Date: Tue, 11 Nov 2025 21:38:17 -0800 Subject: [PATCH] feat: add adk folder manager and per agent local storage helpers Creates AdkFolderManager for creating/resetting the .adk layout, helper builders that return SQLite- and filesystembacked services for each agent Co-authored-by: George Weale PiperOrigin-RevId: 831206377 --- src/google/adk/apps/app.py | 17 ++ src/google/adk/cli/utils/dot_adk_folder.py | 74 ++++++++ src/google/adk/cli/utils/local_storage.py | 171 ++++++++++++++++++ tests/unittests/apps/test_apps.py | 26 +++ .../cli/utils/test_dot_adk_folder.py | 47 +++++ .../unittests/cli/utils/test_local_storage.py | 56 ++++++ 6 files changed, 391 insertions(+) create mode 100644 src/google/adk/cli/utils/dot_adk_folder.py create mode 100644 src/google/adk/cli/utils/local_storage.py create mode 100644 tests/unittests/cli/utils/test_dot_adk_folder.py create mode 100644 tests/unittests/cli/utils/test_local_storage.py diff --git a/src/google/adk/apps/app.py b/src/google/adk/apps/app.py index 63725889..5382eb5a 100644 --- a/src/google/adk/apps/app.py +++ b/src/google/adk/apps/app.py @@ -18,6 +18,7 @@ from typing import Optional from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field +from pydantic import model_validator from ..agents.base_agent import BaseAgent from ..agents.context_cache_config import ContextCacheConfig @@ -26,6 +27,17 @@ from ..plugins.base_plugin import BasePlugin from ..utils.feature_decorator import experimental +def validate_app_name(name: str) -> None: + """Ensures the provided application name is safe and intuitive.""" + if not name.isidentifier(): + raise ValueError( + f"Invalid app name '{name}': must be a valid identifier consisting of" + " letters, digits, and underscores." + ) + if name == "user": + raise ValueError("App name cannot be 'user'; reserved for end-user input.") + + @experimental class ResumabilityConfig(BaseModel): """The config of the resumability for an application. @@ -105,3 +117,8 @@ class App(BaseModel): The config of the resumability for the application. If configured, will be applied to all agents in the app. """ + + @model_validator(mode="after") + def _validate_name(self) -> App: + validate_app_name(self.name) + return self diff --git a/src/google/adk/cli/utils/dot_adk_folder.py b/src/google/adk/cli/utils/dot_adk_folder.py new file mode 100644 index 00000000..39bbb4f3 --- /dev/null +++ b/src/google/adk/cli/utils/dot_adk_folder.py @@ -0,0 +1,74 @@ +# Copyright 2025 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. +"""Helpers for managing an agent's `.adk` folder.""" + +from __future__ import annotations + +from functools import cached_property +from pathlib import Path + + +def _resolve_agent_dir(*, agents_root: Path | str, app_name: str) -> Path: + """Resolves the agent directory with safety checks.""" + agents_root_path = Path(agents_root).resolve() + agent_dir = (agents_root_path / app_name).resolve() + if not str(agent_dir).startswith(str(agents_root_path)): + raise ValueError( + f"Invalid app_name '{app_name}': resolves outside base directory" + ) + + return agent_dir + + +class DotAdkFolder: + """Manages the lifecycle of the `.adk` folder for a single agent.""" + + def __init__(self, agent_dir: Path | str): + self._agent_dir = Path(agent_dir).resolve() + + @property + def agent_dir(self) -> Path: + return self._agent_dir + + @cached_property + def dot_adk_dir(self) -> Path: + return self._agent_dir / ".adk" + + @cached_property + def artifacts_dir(self) -> Path: + return self.dot_adk_dir / "artifacts" + + @cached_property + def session_db_path(self) -> Path: + return self.dot_adk_dir / "session.db" + + +def dot_adk_folder_for_agent( + *, agents_root: Path | str, app_name: str +) -> DotAdkFolder: + """Creates a manager for an agent rooted under `agents_root`. + + Args: + agents_root: Directory that contains all agents. + app_name: Name of the agent directory. + + Returns: + A `DotAdkFolder` scoped to the given agent. + + Raises: + ValueError: If `app_name` traverses outside of `agents_root`. + """ + return DotAdkFolder( + _resolve_agent_dir(agents_root=agents_root, app_name=app_name) + ) diff --git a/src/google/adk/cli/utils/local_storage.py b/src/google/adk/cli/utils/local_storage.py new file mode 100644 index 00000000..9e6b3f3d --- /dev/null +++ b/src/google/adk/cli/utils/local_storage.py @@ -0,0 +1,171 @@ +# Copyright 2025 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. +"""Utilities for local .adk folder persistence.""" +from __future__ import annotations + +import asyncio +import logging +from pathlib import Path +from typing import Optional + +from typing_extensions import override + +from ...artifacts.base_artifact_service import BaseArtifactService +from ...artifacts.file_artifact_service import FileArtifactService +from ...events.event import Event +from ...sessions.base_session_service import BaseSessionService +from ...sessions.base_session_service import GetSessionConfig +from ...sessions.base_session_service import ListSessionsResponse +from ...sessions.session import Session +from .dot_adk_folder import dot_adk_folder_for_agent +from .dot_adk_folder import DotAdkFolder + +logger = logging.getLogger("google_adk." + __name__) + + +def create_local_database_session_service( + *, + base_dir: Path | str, +) -> BaseSessionService: + """Creates a SQLite-backed session service at .adk/session.db. + + Args: + base_dir: The base directory for the agent (parent of .adk folder). + + Returns: + A SqliteSessionService instance. + """ + from ...sessions.sqlite_session_service import SqliteSessionService + + manager = DotAdkFolder(base_dir) + manager.dot_adk_dir.mkdir(parents=True, exist_ok=True) + + session_db_path = manager.session_db_path + + logger.info("Creating local session service at %s", session_db_path) + return SqliteSessionService(db_path=str(session_db_path)) + + +def create_local_artifact_service( + *, base_dir: Path | str, per_agent: bool = False +) -> BaseArtifactService: + """Creates a file-backed artifact service rooted in `.adk/artifacts`. + + Args: + base_dir: Directory whose `.adk` folder will store artifacts. + per_agent: Indicates whether the service is being used in multi-agent mode. + + Returns: + A `FileArtifactService` scoped to the derived root directory. + """ + manager = DotAdkFolder(base_dir) + artifact_root = manager.artifacts_dir + artifact_root.mkdir(parents=True, exist_ok=True) + if per_agent: + logger.info( + "Using shared file artifact service rooted at %s for multi-agent mode", + artifact_root, + ) + else: + logger.info("Using file artifact service at %s", artifact_root) + return FileArtifactService(root_dir=artifact_root) + + +class PerAgentDatabaseSessionService(BaseSessionService): + """Routes session storage to per-agent `.adk/session.db` files.""" + + def __init__( + self, + *, + agents_root: Path | str, + ): + self._agents_root = Path(agents_root).resolve() + 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) + if service is not None: + return service + folder = dot_adk_folder_for_agent( + agents_root=self._agents_root, app_name=app_name + ) + service = create_local_database_session_service( + base_dir=folder.agent_dir, + ) + self._services[app_name] = service + return service + + @override + async def create_session( + self, + *, + app_name: str, + user_id: str, + state: Optional[dict[str, object]] = None, + session_id: Optional[str] = None, + ) -> Session: + service = await self._get_service(app_name) + return await service.create_session( + app_name=app_name, + user_id=user_id, + state=state, + session_id=session_id, + ) + + @override + async def get_session( + self, + *, + app_name: str, + user_id: str, + session_id: str, + config: Optional[GetSessionConfig] = None, + ) -> Optional[Session]: + service = await self._get_service(app_name) + return await service.get_session( + app_name=app_name, + user_id=user_id, + session_id=session_id, + config=config, + ) + + @override + async def list_sessions( + self, + *, + app_name: str, + user_id: Optional[str] = None, + ) -> ListSessionsResponse: + service = await self._get_service(app_name) + return await service.list_sessions(app_name=app_name, user_id=user_id) + + @override + async def delete_session( + self, + *, + app_name: str, + user_id: str, + session_id: str, + ) -> None: + service = await self._get_service(app_name) + await service.delete_session( + app_name=app_name, user_id=user_id, session_id=session_id + ) + + @override + async def append_event(self, session: Session, event: Event) -> Event: + service = await self._get_service(session.app_name) + return await service.append_event(session, event) diff --git a/tests/unittests/apps/test_apps.py b/tests/unittests/apps/test_apps.py index d5b23d5a..bfbc368b 100644 --- a/tests/unittests/apps/test_apps.py +++ b/tests/unittests/apps/test_apps.py @@ -19,6 +19,7 @@ from google.adk.agents.context_cache_config import ContextCacheConfig from google.adk.apps.app import App from google.adk.apps.app import ResumabilityConfig from google.adk.plugins.base_plugin import BasePlugin +import pytest class TestApp: @@ -168,3 +169,28 @@ class TestApp: resumability_config=None, ) assert app.resumability_config is None + + def test_app_rejects_invalid_name(self): + """Test that invalid application names are rejected.""" + mock_agent = Mock(spec=BaseAgent) + + with pytest.raises(ValueError): + App(name="../escape_attempt", root_agent=mock_agent) + + with pytest.raises(ValueError): + App(name="nested/path", root_agent=mock_agent) + + with pytest.raises(ValueError): + App(name="windows\\path", root_agent=mock_agent) + + def test_app_name_must_be_identifier(self): + mock_agent = Mock(spec=BaseAgent) + + with pytest.raises(ValueError): + App(name="invalid-name", root_agent=mock_agent) + + def test_app_name_cannot_be_user(self): + mock_agent = Mock(spec=BaseAgent) + + with pytest.raises(ValueError): + App(name="user", root_agent=mock_agent) diff --git a/tests/unittests/cli/utils/test_dot_adk_folder.py b/tests/unittests/cli/utils/test_dot_adk_folder.py new file mode 100644 index 00000000..249f9ac1 --- /dev/null +++ b/tests/unittests/cli/utils/test_dot_adk_folder.py @@ -0,0 +1,47 @@ +# Copyright 2025 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 __future__ import annotations + +from pathlib import Path + +from google.adk.cli.utils.dot_adk_folder import dot_adk_folder_for_agent +from google.adk.cli.utils.dot_adk_folder import DotAdkFolder +import pytest + + +def test_paths_are_relative_to_agent_dir(tmp_path: Path): + agent_dir = tmp_path / "agent_a" + folder = DotAdkFolder(agent_dir) + + assert folder.dot_adk_dir == agent_dir.resolve() / ".adk" + assert folder.artifacts_dir == folder.dot_adk_dir / "artifacts" + assert folder.session_db_path == folder.dot_adk_dir / "session.db" + + +def test_for_agent_validates_app_name(tmp_path: Path): + agents_root = tmp_path / "agents" + agents_root.mkdir() + + with pytest.raises(ValueError): + dot_adk_folder_for_agent( + agents_root=agents_root, app_name="../escape_attempt" + ) + + folder = dot_adk_folder_for_agent( + agents_root=agents_root, app_name="valid_agent" + ) + + expected_dir = (agents_root / "valid_agent").resolve() + assert folder.agent_dir == expected_dir diff --git a/tests/unittests/cli/utils/test_local_storage.py b/tests/unittests/cli/utils/test_local_storage.py new file mode 100644 index 00000000..324ce5e0 --- /dev/null +++ b/tests/unittests/cli/utils/test_local_storage.py @@ -0,0 +1,56 @@ +# Copyright 2025 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 __future__ import annotations + +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 PerAgentDatabaseSessionService +from google.adk.sessions.sqlite_session_service import SqliteSessionService +import pytest + + +@pytest.mark.asyncio +async def test_per_agent_session_service_creates_scoped_dot_adk( + tmp_path: Path, +) -> None: + agent_a = tmp_path / "agent_a" + agent_b = tmp_path / "agent_b" + agent_a.mkdir() + agent_b.mkdir() + + service = PerAgentDatabaseSessionService(agents_root=tmp_path) + + await service.create_session(app_name="agent_a", user_id="user_a") + await service.create_session(app_name="agent_b", user_id="user_b") + + assert (agent_a / ".adk" / "session.db").exists() + assert (agent_b / ".adk" / "session.db").exists() + + agent_a_sessions = await service.list_sessions(app_name="agent_a") + agent_b_sessions = await service.list_sessions(app_name="agent_b") + + assert len(agent_a_sessions.sessions) == 1 + assert agent_a_sessions.sessions[0].app_name == "agent_a" + assert len(agent_b_sessions.sessions) == 1 + assert agent_b_sessions.sessions[0].app_name == "agent_b" + + +def test_create_local_database_session_service_returns_sqlite( + tmp_path: Path, +) -> None: + service = create_local_database_session_service(base_dir=tmp_path) + + assert isinstance(service, SqliteSessionService)