From b30c2f4e139e0d4410c5f8dd61acee2056ad06ea Mon Sep 17 00:00:00 2001 From: George Weale Date: Wed, 7 Jan 2026 10:17:33 -0800 Subject: [PATCH] fix: avoid local .adk storage in Cloud Run/GKE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Default session and artifact services to in-memory when running in Cloud Run/Kubernetes (or when agents_dir isn’t writable) to prevent startup failures from attempting to create .adk under read-only/unwritable container paths (e.g. /app/agents/.adk). Local development defaults are unchanged. - ADK_FORCE_LOCAL_STORAGE=1 to always use .adk defaults - ADK_DISABLE_LOCAL_STORAGE=1 to always avoid local storage If local artifact initialization raises PermissionError, fall back to in-memory and log a warning Close #3907 Co-authored-by: George Weale PiperOrigin-RevId: 853315459 --- src/google/adk/cli/cli.py | 8 +- src/google/adk/cli/cli_deploy.py | 55 ++++-- src/google/adk/cli/cli_tools_click.py | 105 +++++------ src/google/adk/cli/fast_api.py | 3 + src/google/adk/cli/utils/service_factory.py | 168 ++++++++++++++++- tests/unittests/cli/utils/test_cli.py | 10 +- tests/unittests/cli/utils/test_cli_deploy.py | 47 ++++- .../cli/utils/test_service_factory.py | 178 ++++++++++++++++++ 8 files changed, 487 insertions(+), 87 deletions(-) diff --git a/src/google/adk/cli/cli.py b/src/google/adk/cli/cli.py index 941f1c28..7742d700 100644 --- a/src/google/adk/cli/cli.py +++ b/src/google/adk/cli/cli.py @@ -137,6 +137,7 @@ async def run_cli( session_id: Optional[str] = None, session_service_uri: Optional[str] = None, artifact_service_uri: Optional[str] = None, + use_local_storage: bool = True, ) -> None: """Runs an interactive CLI for a certain agent. @@ -153,6 +154,7 @@ async def run_cli( session_id: Optional[str], the session ID to save the session to on exit. session_service_uri: Optional[str], custom session service URI. artifact_service_uri: Optional[str], custom artifact service URI. + use_local_storage: bool, whether to use local .adk storage by default. """ agent_parent_path = Path(agent_parent_dir).resolve() agent_root = agent_parent_path / agent_folder_name @@ -169,17 +171,19 @@ async def run_cli( 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 //.adk/session.db by default. + # Create session and artifact services using factory functions. + # Sessions persist under //.adk/session.db when enabled. 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, + use_local_storage=use_local_storage, ) artifact_service = create_artifact_service_from_options( base_dir=agent_root, artifact_service_uri=artifact_service_uri, + use_local_storage=use_local_storage, ) credential_service = InMemoryCredentialService() diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index bd51e214..d36febdc 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/src/google/adk/cli/cli_deploy.py @@ -26,6 +26,7 @@ from packaging.version import parse _IS_WINDOWS = os.name == 'nt' _GCLOUD_CMD = 'gcloud.cmd' if _IS_WINDOWS else 'gcloud' +_LOCAL_STORAGE_FLAG_MIN_VERSION: Final[str] = '1.21.0' _DOCKERFILE_TEMPLATE: Final[str] = """ FROM python:3.11-slim @@ -442,26 +443,38 @@ def _get_service_option_by_adk_version( session_uri: Optional[str], artifact_uri: Optional[str], memory_uri: Optional[str], + use_local_storage: Optional[bool] = None, ) -> str: """Returns service option string based on adk_version.""" parsed_version = parse(adk_version) + options: list[str] = [] + if parsed_version >= parse('1.3.0'): - session_option = ( - f'--session_service_uri={session_uri}' if session_uri else '' - ) - artifact_option = ( - f'--artifact_service_uri={artifact_uri}' if artifact_uri else '' - ) - memory_option = f'--memory_service_uri={memory_uri}' if memory_uri else '' - return f'{session_option} {artifact_option} {memory_option}' - elif parsed_version >= parse('1.2.0'): - session_option = f'--session_db_url={session_uri}' if session_uri else '' - artifact_option = ( - f'--artifact_storage_uri={artifact_uri}' if artifact_uri else '' - ) - return f'{session_option} {artifact_option}' + if session_uri: + options.append(f'--session_service_uri={session_uri}') + if artifact_uri: + options.append(f'--artifact_service_uri={artifact_uri}') + if memory_uri: + options.append(f'--memory_service_uri={memory_uri}') else: - return f'--session_db_url={session_uri}' if session_uri else '' + if session_uri: + options.append(f'--session_db_url={session_uri}') + if parsed_version >= parse('1.2.0') and artifact_uri: + options.append(f'--artifact_storage_uri={artifact_uri}') + + if use_local_storage is not None and parsed_version >= parse( + _LOCAL_STORAGE_FLAG_MIN_VERSION + ): + # Only valid when session/artifact URIs are unset; otherwise the CLI + # rejects the combination to avoid confusing precedence. + if session_uri is None and artifact_uri is None: + options.append(( + '--use_local_storage' + if use_local_storage + else '--no_use_local_storage' + )) + + return ' '.join(options) def to_cloud_run( @@ -482,6 +495,7 @@ def to_cloud_run( session_service_uri: Optional[str] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + use_local_storage: bool = False, a2a: bool = False, extra_gcloud_args: Optional[tuple[str, ...]] = None, ): @@ -517,8 +531,12 @@ def to_cloud_run( session_service_uri: The URI of the session service. artifact_service_uri: The URI of the artifact service. memory_service_uri: The URI of the memory service. + use_local_storage: Whether to use local .adk storage in the container. """ app_name = app_name or os.path.basename(agent_folder) + if parse(adk_version) >= parse('1.3.0') and not use_local_storage: + session_service_uri = session_service_uri or 'memory://' + artifact_service_uri = artifact_service_uri or 'memory://' click.echo(f'Start generating Cloud Run source files in {temp_folder}') @@ -559,6 +577,7 @@ def to_cloud_run( session_service_uri, artifact_service_uri, memory_service_uri, + use_local_storage, ), trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '', allow_origins_option=allow_origins_option, @@ -944,6 +963,7 @@ def to_gke( session_service_uri: Optional[str] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + use_local_storage: bool = False, a2a: bool = False, ): """Deploys an agent to Google Kubernetes Engine(GKE). @@ -969,6 +989,7 @@ def to_gke( session_service_uri: The URI of the session service. artifact_service_uri: The URI of the artifact service. memory_service_uri: The URI of the memory service. + use_local_storage: Whether to use local .adk storage in the container. """ click.secho( '\n🚀 Starting ADK Agent Deployment to GKE...', fg='cyan', bold=True @@ -982,6 +1003,9 @@ def to_gke( click.echo('--------------------------------------------------\n') app_name = app_name or os.path.basename(agent_folder) + if parse(adk_version) >= parse('1.3.0') and not use_local_storage: + session_service_uri = session_service_uri or 'memory://' + artifact_service_uri = artifact_service_uri or 'memory://' click.secho('STEP 1: Preparing build environment...', bold=True) click.echo(f' - Using temporary directory: {temp_folder}') @@ -1024,6 +1048,7 @@ def to_gke( session_service_uri, artifact_service_uri, memory_service_uri, + use_local_storage, ), trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '', allow_origins_option=allow_origins_option, diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 6f01c05a..b974dad7 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -36,7 +36,6 @@ from . import cli_create from . import cli_deploy from .. import version from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE -from ..sessions.migration import migration_runner from .cli import run_cli from .fast_api import get_fast_api_app from .utils import envs @@ -368,24 +367,26 @@ def validate_exclusive(ctx, param, value): return value -def adk_services_options(): +def adk_services_options(*, default_use_local_storage: bool = True): """Decorator to add ADK services options to click commands.""" def decorator(func): @click.option( "--session_service_uri", - help=textwrap.dedent( - """\ + help=textwrap.dedent("""\ Optional. The URI of the session service. - - Leave unset to use the in-memory session service (default). + If set, ADK uses this service. + + If unset, ADK chooses a default session service (see + --use_local_storage). - Use 'agentengine://' to connect to Agent Engine sessions. can either be the full qualified resource name 'projects/abc/locations/us-central1/reasoningEngines/123' or the resource id '123'. - Use 'memory://' to run with the in-memory session service. - Use 'sqlite://' to connect to a SQLite DB. - - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported database URIs.""" - ), + - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls + for supported database URIs."""), ) @click.option( "--artifact_service_uri", @@ -393,13 +394,29 @@ def adk_services_options(): help=textwrap.dedent( """\ Optional. The URI of the artifact service. - - Leave unset to store artifacts under '.adk/artifacts' locally. + If set, ADK uses this service. + + If unset, ADK chooses a default artifact service (see + --use_local_storage). - Use 'gs://' to connect to the GCS artifact service. - Use 'memory://' to force the in-memory artifact service. - Use 'file://' to store artifacts in a custom local directory.""" ), default=None, ) + @click.option( + "--use_local_storage/--no_use_local_storage", + default=default_use_local_storage, + show_default=True, + help=( + "Optional. Whether to use local .adk storage when " + "--session_service_uri and --artifact_service_uri are unset. " + "Cannot be combined with explicit service URIs. When the agents " + "directory isn't writable (common in Cloud Run/Kubernetes), ADK " + "falls back to in-memory unless overridden by " + "ADK_FORCE_LOCAL_STORAGE=1 or ADK_DISABLE_LOCAL_STORAGE=1." + ), + ) @click.option( "--memory_service_uri", type=str, @@ -415,6 +432,17 @@ def adk_services_options(): ) @functools.wraps(func) def wrapper(*args, **kwargs): + ctx = click.get_current_context(silent=True) + if ctx is not None: + use_local_storage_source = ctx.get_parameter_source("use_local_storage") + if use_local_storage_source != ParameterSource.DEFAULT and ( + kwargs.get("session_service_uri") is not None + or kwargs.get("artifact_service_uri") is not None + ): + raise click.UsageError( + "--use_local_storage/--no_use_local_storage cannot be used with " + "--session_service_uri or --artifact_service_uri." + ) return func(*args, **kwargs) return wrapper @@ -423,7 +451,7 @@ def adk_services_options(): @main.command("run", cls=HelpfulCommand) -@adk_services_options() +@adk_services_options(default_use_local_storage=True) @click.option( "--save_session", type=bool, @@ -481,6 +509,7 @@ def cli_run( session_service_uri: Optional[str] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + use_local_storage: bool = True, ): """Runs an interactive CLI for a certain agent. @@ -513,6 +542,7 @@ def cli_run( session_id=session_id, session_service_uri=session_service_uri, artifact_service_uri=artifact_service_uri, + use_local_storage=use_local_storage, ) ) @@ -1113,7 +1143,7 @@ def fast_api_common_options(): @main.command("web") @fast_api_common_options() @web_options() -@adk_services_options() +@adk_services_options(default_use_local_storage=True) @deprecated_adk_services_options() @click.argument( "agents_dir", @@ -1136,6 +1166,7 @@ def cli_web( session_service_uri: Optional[str] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + use_local_storage: bool = True, session_db_url: Optional[str] = None, # Deprecated artifact_storage_uri: Optional[str] = None, # Deprecated a2a: bool = False, @@ -1184,6 +1215,7 @@ def cli_web( session_service_uri=session_service_uri, artifact_service_uri=artifact_service_uri, memory_service_uri=memory_service_uri, + use_local_storage=use_local_storage, eval_storage_uri=eval_storage_uri, allow_origins=allow_origins, web=True, @@ -1221,7 +1253,7 @@ def cli_web( default=os.getcwd(), ) @fast_api_common_options() -@adk_services_options() +@adk_services_options(default_use_local_storage=True) @deprecated_adk_services_options() def cli_api_server( agents_dir: str, @@ -1237,6 +1269,7 @@ def cli_api_server( session_service_uri: Optional[str] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + use_local_storage: bool = True, session_db_url: Optional[str] = None, # Deprecated artifact_storage_uri: Optional[str] = None, # Deprecated a2a: bool = False, @@ -1262,6 +1295,7 @@ def cli_api_server( session_service_uri=session_service_uri, artifact_service_uri=artifact_service_uri, memory_service_uri=memory_service_uri, + use_local_storage=use_local_storage, eval_storage_uri=eval_storage_uri, allow_origins=allow_origins, web=False, @@ -1403,7 +1437,7 @@ def cli_api_server( multiple=True, ) # TODO: Add eval_storage_uri option back when evals are supported in Cloud Run. -@adk_services_options() +@adk_services_options(default_use_local_storage=False) @deprecated_adk_services_options() @click.pass_context def cli_deploy_cloud_run( @@ -1424,6 +1458,7 @@ def cli_deploy_cloud_run( session_service_uri: Optional[str] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + use_local_storage: bool = False, session_db_url: Optional[str] = None, # Deprecated artifact_storage_uri: Optional[str] = None, # Deprecated a2a: bool = False, @@ -1501,6 +1536,7 @@ def cli_deploy_cloud_run( session_service_uri=session_service_uri, artifact_service_uri=artifact_service_uri, memory_service_uri=memory_service_uri, + use_local_storage=use_local_storage, a2a=a2a, extra_gcloud_args=tuple(gcloud_args), ) @@ -1508,47 +1544,6 @@ def cli_deploy_cloud_run( click.secho(f"Deploy failed: {e}", fg="red", err=True) -@main.group() -def migrate(): - """ADK migration commands.""" - pass - - -@migrate.command("session", cls=HelpfulCommand) -@click.option( - "--source_db_url", - required=True, - help=( - "SQLAlchemy URL of source database in database session service, e.g." - " sqlite:///source.db." - ), -) -@click.option( - "--dest_db_url", - required=True, - help=( - "SQLAlchemy URL of destination database in database session service," - " e.g. sqlite:///dest.db." - ), -) -@click.option( - "--log_level", - type=LOG_LEVELS, - default="INFO", - help="Optional. Set the logging level", -) -def cli_migrate_session( - *, source_db_url: str, dest_db_url: str, log_level: str -): - """Migrates a session database to the latest schema version.""" - logs.setup_adk_logger(getattr(logging, log_level.upper())) - try: - migration_runner.upgrade(source_db_url, dest_db_url) - click.secho("Migration check and upgrade process finished.", fg="green") - except Exception as e: - click.secho(f"Migration failed: {e}", fg="red", err=True) - - @deploy.command("agent_engine") @click.option( "--api_key", @@ -1843,7 +1838,7 @@ def cli_deploy_agent_engine( " version in the dev environment)" ), ) -@adk_services_options() +@adk_services_options(default_use_local_storage=False) @click.argument( "agent", type=click.Path( @@ -1866,6 +1861,7 @@ def cli_deploy_gke( session_service_uri: Optional[str] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + use_local_storage: bool = False, ): """Deploys an agent to GKE. @@ -1894,6 +1890,7 @@ def cli_deploy_gke( session_service_uri=session_service_uri, artifact_service_uri=artifact_service_uri, memory_service_uri=memory_service_uri, + use_local_storage=use_local_storage, ) except Exception as e: click.secho(f"Deploy failed: {e}", fg="red", err=True) diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index b5abaa50..8e87aec2 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -76,6 +76,7 @@ def get_fast_api_app( session_db_kwargs: Optional[Mapping[str, Any]] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + use_local_storage: bool = True, eval_storage_uri: Optional[str] = None, allow_origins: Optional[list[str]] = None, web: bool, @@ -122,6 +123,7 @@ def get_fast_api_app( base_dir=agents_dir, session_service_uri=session_service_uri, session_db_kwargs=session_db_kwargs, + use_local_storage=use_local_storage, ) # Build the Artifact service @@ -130,6 +132,7 @@ def get_fast_api_app( base_dir=agents_dir, artifact_service_uri=artifact_service_uri, strict_uri=True, + use_local_storage=use_local_storage, ) except ValueError as exc: raise click.ClickException(str(exc)) from exc diff --git a/src/google/adk/cli/utils/service_factory.py b/src/google/adk/cli/utils/service_factory.py index 840c5b1c..c03ac10b 100644 --- a/src/google/adk/cli/utils/service_factory.py +++ b/src/google/adk/cli/utils/service_factory.py @@ -13,7 +13,9 @@ # limitations under the License. from __future__ import annotations +import errno import logging +import os from pathlib import Path from typing import Any from typing import Optional @@ -21,12 +23,111 @@ from typing import Optional from ...artifacts.base_artifact_service import BaseArtifactService from ...memory.base_memory_service import BaseMemoryService from ...sessions.base_session_service import BaseSessionService +from ...utils.env_utils import is_env_enabled from ..service_registry import get_service_registry from .local_storage import create_local_artifact_service from .local_storage import create_local_session_service logger = logging.getLogger("google_adk." + __name__) +_DISABLE_LOCAL_STORAGE_ENV = "ADK_DISABLE_LOCAL_STORAGE" +_FORCE_LOCAL_STORAGE_ENV = "ADK_FORCE_LOCAL_STORAGE" +_LOCAL_STORAGE_ERRNOS = frozenset({ + errno.EACCES, + errno.EPERM, + errno.EROFS, +}) + +_CLOUD_RUN_SERVICE_ENV = "K_SERVICE" +_KUBERNETES_HOST_ENV = "KUBERNETES_SERVICE_HOST" + + +def _is_cloud_run() -> bool: + """Returns True when running in Cloud Run.""" + return bool(os.environ.get(_CLOUD_RUN_SERVICE_ENV)) + + +def _is_kubernetes() -> bool: + """Returns True when running in Kubernetes (including GKE).""" + return bool(os.environ.get(_KUBERNETES_HOST_ENV)) + + +def _is_dir_writable(path: Path) -> bool: + """Returns True if the directory exists and is writable/executable.""" + try: + if not path.exists() or not path.is_dir(): + return False + except OSError: + return False + return os.access(path, os.W_OK | os.X_OK) + + +def _resolve_use_local_storage( + *, + base_path: Path, + requested: bool, +) -> tuple[bool, str | None]: + """Resolves effective local storage setting with safe defaults.""" + if is_env_enabled(_DISABLE_LOCAL_STORAGE_ENV): + warning_message = ( + "Local storage is disabled by %s; using in-memory services. " + "Set --session_service_uri/--artifact_service_uri for production " + "deployments." + ) % _DISABLE_LOCAL_STORAGE_ENV + return False, warning_message + + if is_env_enabled(_FORCE_LOCAL_STORAGE_ENV): + if not _is_dir_writable(base_path): + warning_message = ( + "Local storage is forced by %s, but %s is not writable; " + "using in-memory services." + ) % (_FORCE_LOCAL_STORAGE_ENV, base_path) + return False, warning_message + return True, None + + if not requested: + return False, None + + if _is_cloud_run() or _is_kubernetes(): + warning_message = ( + "Detected Cloud Run/Kubernetes runtime; using in-memory services " + "instead of local .adk storage. Set %s=1 to force local storage." + ) % _FORCE_LOCAL_STORAGE_ENV + return False, warning_message + + if not _is_dir_writable(base_path): + warning_message = ( + "Agents directory %s is not writable; using in-memory services " + "instead of local .adk storage. Set %s=1 to force local storage." + ) % (base_path, _FORCE_LOCAL_STORAGE_ENV) + return False, warning_message + + return True, None + + +def _create_in_memory_session_service( + warning_message: str | None = None, + *warning_args: object, +) -> BaseSessionService: + """Creates an in-memory session service, optionally logging a warning.""" + if warning_message is not None: + logger.warning(warning_message, *warning_args) + from ...sessions.in_memory_session_service import InMemorySessionService + + return InMemorySessionService() + + +def _create_in_memory_artifact_service( + warning_message: str | None = None, + *warning_args: object, +) -> BaseArtifactService: + """Creates an in-memory artifact service, optionally logging a warning.""" + if warning_message is not None: + logger.warning(warning_message, *warning_args) + from ...artifacts.in_memory_artifact_service import InMemoryArtifactService + + return InMemoryArtifactService() + def create_session_service_from_options( *, @@ -34,6 +135,7 @@ def create_session_service_from_options( session_service_uri: Optional[str] = None, session_db_kwargs: Optional[dict[str, Any]] = None, app_name_to_dir: Optional[dict[str, str]] = None, + use_local_storage: bool = True, ) -> BaseSessionService: """Creates a session service based on CLI/web options.""" base_path = Path(base_dir) @@ -64,12 +166,36 @@ def create_session_service_from_options( ) return DatabaseSessionService(db_url=session_service_uri, **fallback_kwargs) - # Default to per-agent local SQLite storage in //.adk/. - return create_local_session_service( - base_dir=base_path, - per_agent=True, - app_name_to_dir=app_name_to_dir, + effective_use_local_storage, auto_warning = _resolve_use_local_storage( + base_path=base_path, + requested=use_local_storage, ) + if not effective_use_local_storage: + if auto_warning is not None: + return _create_in_memory_session_service(auto_warning) + return _create_in_memory_session_service( + "Local session storage is disabled; using in-memory session service. " + "Set --session_service_uri for production deployments." + ) + + # Default to per-agent local SQLite storage in //.adk/. + try: + return create_local_session_service( + base_dir=base_path, + per_agent=True, + app_name_to_dir=app_name_to_dir, + ) + except OSError as exc: + if exc.errno not in _LOCAL_STORAGE_ERRNOS and not isinstance( + exc, PermissionError + ): + raise + return _create_in_memory_session_service( + "Failed to initialize local session storage under %s (%r); " + "falling back to in-memory session service.", + base_path, + exc, + ) def create_memory_service_from_options( @@ -102,6 +228,7 @@ def create_artifact_service_from_options( base_dir: Path | str, artifact_service_uri: Optional[str] = None, strict_uri: bool = False, + use_local_storage: bool = True, ) -> BaseArtifactService: """Creates an artifact service based on CLI/web options.""" base_path = Path(base_dir) @@ -118,13 +245,34 @@ def create_artifact_service_from_options( raise ValueError( f"Unsupported artifact service URI: {artifact_service_uri}" ) - logger.warning( + return _create_in_memory_artifact_service( "Unsupported artifact service URI: %s, falling back to in-memory", artifact_service_uri, ) - from ...artifacts.in_memory_artifact_service import InMemoryArtifactService - - return InMemoryArtifactService() return service - return create_local_artifact_service(base_dir=base_path) + effective_use_local_storage, auto_warning = _resolve_use_local_storage( + base_path=base_path, + requested=use_local_storage, + ) + if not effective_use_local_storage: + if auto_warning is not None: + return _create_in_memory_artifact_service(auto_warning) + return _create_in_memory_artifact_service( + "Local artifact storage is disabled; using in-memory artifact service. " + "Set --artifact_service_uri for production deployments." + ) + + try: + return create_local_artifact_service(base_dir=base_path) + except OSError as exc: + if exc.errno not in _LOCAL_STORAGE_ERRNOS and not isinstance( + exc, PermissionError + ): + raise + return _create_in_memory_artifact_service( + "Failed to initialize local artifact storage under %s (%r); " + "falling back to in-memory artifact service.", + base_path, + exc, + ) diff --git a/tests/unittests/cli/utils/test_cli.py b/tests/unittests/cli/utils/test_cli.py index 73ae89a9..fc2455f6 100644 --- a/tests/unittests/cli/utils/test_cli.py +++ b/tests/unittests/cli/utils/test_cli.py @@ -293,7 +293,10 @@ async def test_run_cli_save_session( def test_create_artifact_service_defaults_to_file(tmp_path: Path) -> None: """Service factory should default to FileArtifactService when URI is unset.""" - service = create_artifact_service_from_options(base_dir=tmp_path) + service = create_artifact_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) assert isinstance(service, FileArtifactService) expected_root = Path(tmp_path) / ".adk" / "artifacts" assert service.root_dir == expected_root @@ -304,7 +307,10 @@ def test_create_artifact_service_uses_shared_root( tmp_path: Path, ) -> None: """Artifact service should use a single file artifact service.""" - service = create_artifact_service_from_options(base_dir=tmp_path) + service = create_artifact_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) assert isinstance(service, FileArtifactService) expected_root = Path(tmp_path) / ".adk" / "artifacts" assert service.root_dir == expected_root diff --git a/tests/unittests/cli/utils/test_cli_deploy.py b/tests/unittests/cli/utils/test_cli_deploy.py index 696344eb..7dd6d263 100644 --- a/tests/unittests/cli/utils/test_cli_deploy.py +++ b/tests/unittests/cli/utils/test_cli_deploy.py @@ -128,13 +128,15 @@ def test_resolve_project_from_gcloud_fails( @pytest.mark.parametrize( - "adk_version, session_uri, artifact_uri, memory_uri, expected", + "adk_version, session_uri, artifact_uri, memory_uri, use_local_storage, " + "expected", [ ( "1.3.0", "sqlite://s", "gs://a", "rag://m", + None, ( "--session_service_uri=sqlite://s --artifact_service_uri=gs://a" " --memory_service_uri=rag://m" @@ -145,6 +147,7 @@ def test_resolve_project_from_gcloud_fails( "sqlite://s", "gs://a", "rag://m", + None, "--session_db_url=sqlite://s --artifact_storage_uri=gs://a", ), ( @@ -152,6 +155,7 @@ def test_resolve_project_from_gcloud_fails( "sqlite://s", "gs://a", "rag://m", + None, "--session_db_url=sqlite://s", ), ( @@ -159,16 +163,49 @@ def test_resolve_project_from_gcloud_fails( "sqlite://s", None, None, - "--session_service_uri=sqlite://s ", + None, + "--session_service_uri=sqlite://s", ), ( "1.3.0", None, "gs://a", "rag://m", - " --artifact_service_uri=gs://a --memory_service_uri=rag://m", + None, + "--artifact_service_uri=gs://a --memory_service_uri=rag://m", + ), + ( + "1.2.0", + None, + "gs://a", + None, + None, + "--artifact_storage_uri=gs://a", + ), + ( + "1.21.0", + None, + None, + None, + False, + "--no_use_local_storage", + ), + ( + "1.21.0", + None, + None, + None, + True, + "--use_local_storage", + ), + ( + "1.21.0", + "sqlite://s", + "gs://a", + None, + False, + "--session_service_uri=sqlite://s --artifact_service_uri=gs://a", ), - ("1.2.0", None, "gs://a", None, " --artifact_storage_uri=gs://a"), ], ) def test_get_service_option_by_adk_version( @@ -176,6 +213,7 @@ def test_get_service_option_by_adk_version( session_uri: str | None, artifact_uri: str | None, memory_uri: str | None, + use_local_storage: bool | None, expected: str, ) -> None: """It should return the correct service URI flags for a given ADK version.""" @@ -184,6 +222,7 @@ def test_get_service_option_by_adk_version( session_uri=session_uri, artifact_uri=artifact_uri, memory_uri=memory_uri, + use_local_storage=use_local_storage, ) assert actual.rstrip() == expected.rstrip() diff --git a/tests/unittests/cli/utils/test_service_factory.py b/tests/unittests/cli/utils/test_service_factory.py index a8eb0fdb..87b567be 100644 --- a/tests/unittests/cli/utils/test_service_factory.py +++ b/tests/unittests/cli/utils/test_service_factory.py @@ -16,13 +16,17 @@ from __future__ import annotations +import os from pathlib import Path from unittest.mock import Mock +from google.adk.artifacts.file_artifact_service import FileArtifactService +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService from google.adk.cli.utils.local_storage import PerAgentDatabaseSessionService import google.adk.cli.utils.service_factory as service_factory from google.adk.memory.in_memory_memory_service import InMemoryMemoryService from google.adk.sessions.database_session_service import DatabaseSessionService +from google.adk.sessions.in_memory_session_service import InMemorySessionService import pytest @@ -52,6 +56,7 @@ async def test_create_session_service_defaults_to_per_agent_sqlite( agent_dir.mkdir() service = service_factory.create_session_service_from_options( base_dir=tmp_path, + use_local_storage=True, ) assert isinstance(service, PerAgentDatabaseSessionService) @@ -71,6 +76,7 @@ async def test_create_session_service_respects_app_name_mapping( service = service_factory.create_session_service_from_options( base_dir=tmp_path, app_name_to_dir={logical_name: "agent_folder"}, + use_local_storage=True, ) assert isinstance(service, PerAgentDatabaseSessionService) @@ -173,3 +179,175 @@ def test_create_memory_service_raises_on_unknown_scheme( base_dir=tmp_path, memory_service_uri="unknown://foo", ) + + +@pytest.mark.asyncio +async def test_create_session_service_defaults_to_in_memory_when_disabled( + tmp_path: Path, +) -> None: + service = service_factory.create_session_service_from_options( + base_dir=tmp_path, + use_local_storage=False, + ) + + assert isinstance(service, InMemorySessionService) + session = await service.create_session(app_name="agent_a", user_id="user") + assert session.app_name == "agent_a" + assert not (tmp_path / "agent_a" / ".adk").exists() + + +def test_create_artifact_service_defaults_to_in_memory_when_disabled( + tmp_path: Path, +) -> None: + service = service_factory.create_artifact_service_from_options( + base_dir=tmp_path, + use_local_storage=False, + ) + + assert isinstance(service, InMemoryArtifactService) + assert not (tmp_path / ".adk").exists() + + +def test_create_session_service_fallbacks_to_in_memory_on_permission_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def _raise_permission_error(*_args, **_kwargs): + raise PermissionError("nope") + + monkeypatch.setattr( + service_factory, "create_local_session_service", _raise_permission_error + ) + + service = service_factory.create_session_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) + + assert isinstance(service, InMemorySessionService) + + +@pytest.mark.skipif(os.name == "nt", reason="chmod behavior differs on Windows") +def test_create_services_default_to_in_memory_when_agents_dir_unwritable( + tmp_path: Path, +) -> None: + agents_dir = tmp_path / "agents" + agents_dir.mkdir() + try: + agents_dir.chmod(0o555) + if os.access(agents_dir, os.W_OK | os.X_OK): + pytest.skip("Test cannot make directory unwritable in this environment.") + + session_service = service_factory.create_session_service_from_options( + base_dir=agents_dir, + use_local_storage=True, + ) + assert isinstance(session_service, InMemorySessionService) + + artifact_service = service_factory.create_artifact_service_from_options( + base_dir=agents_dir, + use_local_storage=True, + ) + assert isinstance(artifact_service, InMemoryArtifactService) + finally: + agents_dir.chmod(0o755) + + +def test_adk_disable_local_storage_env_forces_in_memory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("ADK_DISABLE_LOCAL_STORAGE", "1") + + session_service = service_factory.create_session_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) + assert isinstance(session_service, InMemorySessionService) + + artifact_service = service_factory.create_artifact_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) + assert isinstance(artifact_service, InMemoryArtifactService) + + +def test_cloud_run_env_defaults_to_in_memory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("K_SERVICE", "adk-service") + + session_service = service_factory.create_session_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) + assert isinstance(session_service, InMemorySessionService) + + artifact_service = service_factory.create_artifact_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) + assert isinstance(artifact_service, InMemoryArtifactService) + + +def test_kubernetes_env_defaults_to_in_memory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "10.0.0.1") + + session_service = service_factory.create_session_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) + assert isinstance(session_service, InMemorySessionService) + + artifact_service = service_factory.create_artifact_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) + assert isinstance(artifact_service, InMemoryArtifactService) + + +@pytest.mark.asyncio +async def test_adk_force_local_storage_env_overrides_flag( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("ADK_FORCE_LOCAL_STORAGE", "1") + agent_dir = tmp_path / "agent_a" + agent_dir.mkdir() + + session_service = service_factory.create_session_service_from_options( + base_dir=tmp_path, + use_local_storage=False, + ) + assert isinstance(session_service, PerAgentDatabaseSessionService) + await session_service.create_session(app_name="agent_a", user_id="user") + assert (agent_dir / ".adk" / "session.db").exists() + + artifact_service = service_factory.create_artifact_service_from_options( + base_dir=tmp_path, + use_local_storage=False, + ) + assert isinstance(artifact_service, FileArtifactService) + + +def test_create_artifact_service_fallbacks_to_in_memory_on_permission_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def _raise_permission_error(*_args, **_kwargs): + raise PermissionError("nope") + + monkeypatch.setattr( + service_factory, "create_local_artifact_service", _raise_permission_error + ) + + service = service_factory.create_artifact_service_from_options( + base_dir=tmp_path, + use_local_storage=True, + ) + + assert isinstance(service, InMemoryArtifactService)