fix: avoid local .adk storage in Cloud Run/GKE

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 <gweale@google.com>
PiperOrigin-RevId: 853315459
This commit is contained in:
George Weale
2026-01-07 10:18:11 -08:00
committed by Copybara-Service
parent d5bd8d939e
commit b30c2f4e13
8 changed files with 487 additions and 87 deletions
+6 -2
View File
@@ -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 <agents_dir>/<agent>/.adk/session.db by default.
# Create session and artifact services using factory functions.
# Sessions persist under <agents_dir>/<agent>/.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()
+40 -15
View File
@@ -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,
+51 -54
View File
@@ -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://<agent_engine>' to connect to Agent Engine
sessions. <agent_engine> 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://<path_to_sqlite_file>' 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://<bucket_name>' to connect to the GCS artifact service.
- Use 'memory://' to force the in-memory artifact service.
- Use 'file://<path>' 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)
+3
View File
@@ -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
+158 -10
View File
@@ -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 <agents_root>/<agent>/.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 <agents_root>/<agent>/.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,
)
+8 -2
View File
@@ -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
+43 -4
View File
@@ -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()
@@ -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)