feat: expose service URI flags

Adds the shared adk_services_options decorator to adk run and other commands so developers can pass session/artifact URIs from the CLI

Has new warning for the unsupported memory service on adk run, and removes the legacy --session_db_url/--artifact_storage_uri flags with tests

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

Author: GWeale

src/google/adk/cli/cli_tools_click.py

@@ -24,6 +24,7 @@
 import os
 from pathlib import Path
 import tempfile
+import textwrap
 from typing import Optional
 
 import click
@@ -354,7 +355,62 @@ def validate_exclusive(ctx, param, value):
   return value
 
 
+def adk_services_options():
+  """Decorator to add ADK services options to click commands."""
+
+  def decorator(func):
+    @click.option(
+        "--session_service_uri",
+        help=textwrap.dedent(
+            """\
+            Optional. The URI of the session service.
+            - Leave unset to use the in-memory session service (default).
+            - 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."""
+        ),
+    )
+    @click.option(
+        "--artifact_service_uri",
+        type=str,
+        help=textwrap.dedent(
+            """\
+            Optional. The URI of the artifact service.
+            - Leave unset to store artifacts under '.adk/artifacts' locally.
+            - 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(
+        "--memory_service_uri",
+        type=str,
+        help=textwrap.dedent("""\
+            Optional. The URI of the memory service.
+            - Use 'rag://<rag_corpus_id>' to connect to Vertex AI Rag Memory Service.
+            - 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 force the in-memory memory service."""),
+        default=None,
+    )
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+      return func(*args, **kwargs)
+
+    return wrapper
+
+  return decorator
+
+
 @main.command("run", cls=HelpfulCommand)
+@adk_services_options()
 @click.option(
     "--save_session",
     type=bool,
@@ -409,6 +465,9 @@ def cli_run(
     session_id: Optional[str],
     replay: Optional[str],
     resume: Optional[str],
+    session_service_uri: Optional[str] = None,
+    artifact_service_uri: Optional[str] = None,
+    memory_service_uri: Optional[str] = None,
 ):
   """Runs an interactive CLI for a certain agent.
 
@@ -420,6 +479,14 @@ def cli_run(
   """
   logs.log_to_tmp_folder()
 
+  # Validation warning for memory_service_uri (not supported for adk run)
+  if memory_service_uri:
+    click.secho(
+        "WARNING: --memory_service_uri is not supported for adk run.",
+        fg="yellow",
+        err=True,
+    )
+
   agent_parent_folder = os.path.dirname(agent)
   agent_folder_name = os.path.basename(agent)
 
@@ -431,6 +498,8 @@ def cli_run(
           saved_session_file=resume,
           save_session=save_session,
           session_id=session_id,
+          session_service_uri=session_service_uri,
+          artifact_service_uri=artifact_service_uri,
       )
   )
 
@@ -865,63 +934,14 @@ def wrapper(*args, **kwargs):
   return decorator
 
 
-def adk_services_options():
-  """Decorator to add ADK services options to click commands."""
-
-  def decorator(func):
-    @click.option(
-        "--session_service_uri",
-        help=(
-            """Optional. The URI of the session service.
-          - 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 'sqlite://<path_to_sqlite_file>' to connect to an aio-sqlite
-            based session service, which is good for local development.
-          - Use 'postgresql://<user>:<password>@<host>:<port>/<database_name>'
-            to connect to a PostgreSQL DB.
-          - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls
-            for more details on other database URIs supported by SQLAlchemy."""
-        ),
-    )
-    @click.option(
-        "--artifact_service_uri",
-        type=str,
-        help=(
-            "Optional. The URI of the artifact service,"
-            " supported URIs: gs://<bucket name> for GCS artifact service."
-        ),
-        default=None,
-    )
-    @click.option(
-        "--memory_service_uri",
-        type=str,
-        help=("""Optional. The URI of the memory service.
-            - Use 'rag://<rag_corpus_id>' to connect to Vertex AI Rag Memory Service.
-            - 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'."""),
-        default=None,
-    )
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-      return func(*args, **kwargs)
-
-    return wrapper
-
-  return decorator
-
-
 def deprecated_adk_services_options():
   """Deprecated ADK services options."""
 
   def warn(alternative_param, ctx, param, value):
     if value:
       click.echo(
           click.style(
-              f"WARNING: Deprecated option {param.name} is used. Please use"
+              f"WARNING: Deprecated option --{param.name} is used. Please use"
               f" {alternative_param} instead.",
               fg="yellow",
           ),
@@ -1116,6 +1136,8 @@ def cli_web(
 
     adk web --session_service_uri=[uri] --port=[port] path/to/agents_dir
   """
+  session_service_uri = session_service_uri or session_db_url
+  artifact_service_uri = artifact_service_uri or artifact_storage_uri
   logs.setup_adk_logger(getattr(logging, log_level.upper()))
 
   @asynccontextmanager
@@ -1140,8 +1162,6 @@ async def _lifespan(app: FastAPI):
         fg="green",
     )
 
-  session_service_uri = session_service_uri or session_db_url
-  artifact_service_uri = artifact_service_uri or artifact_storage_uri
   app = get_fast_api_app(
       agents_dir=agents_dir,
       session_service_uri=session_service_uri,
@@ -1215,10 +1235,10 @@ def cli_api_server(
 
     adk api_server --session_service_uri=[uri] --port=[port] path/to/agents_dir
   """
-  logs.setup_adk_logger(getattr(logging, log_level.upper()))
-
   session_service_uri = session_service_uri or session_db_url
   artifact_service_uri = artifact_service_uri or artifact_storage_uri
+  logs.setup_adk_logger(getattr(logging, log_level.upper()))
+
   config = uvicorn.Config(
       get_fast_api_app(
           agents_dir=agents_dir,

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

@@ -18,8 +18,10 @@
 
 from ...agents.base_agent import BaseAgent
 from ...agents.llm_agent import LlmAgent
+from .dot_adk_folder import DotAdkFolder
 from .state import create_empty_state
 
 __all__ = [
     'create_empty_state',
+    'DotAdkFolder',
 ]

tests/unittests/cli/utils/test_cli_tools_click.py

@@ -76,8 +76,11 @@ def __call__(self, *args: Any, **kwargs: Any) -> None:  # noqa: D401
 
 # Fixtures
 @pytest.fixture(autouse=True)
-def _mute_click(monkeypatch: pytest.MonkeyPatch) -> None:
+def _mute_click(request, monkeypatch: pytest.MonkeyPatch) -> None:
   """Suppress click output during tests."""
+  # Allow tests to opt-out of muting by using the 'unmute_click' marker
+  if "unmute_click" in request.keywords:
+    return
   monkeypatch.setattr(click, "echo", lambda *a, **k: None)
   # Keep secho for error messages
   # monkeypatch.setattr(click, "secho", lambda *a, **k: None)
@@ -121,32 +124,70 @@ def test_cli_create_cmd_invokes_run_cmd(
       cli_tools_click.main,
       ["create", "--model", "gemini", "--api_key", "key123", str(app_dir)],
   )
-  assert result.exit_code == 0
+  assert result.exit_code == 0, (result.output, repr(result.exception))
   assert rec.calls, "cli_create.run_cmd must be called"
 
 
 # cli run
-@pytest.mark.asyncio
-async def test_cli_run_invokes_run_cli(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+@pytest.mark.parametrize(
+    "cli_args,expected_session_uri,expected_artifact_uri",
+    [
+        pytest.param(
+            [
+                "--session_service_uri",
+                "memory://",
+                "--artifact_service_uri",
+                "memory://",
+            ],
+            "memory://",
+            "memory://",
+            id="memory_scheme_uris",
+        ),
+        pytest.param(
+            [],
+            None,
+            None,
+            id="default_uris_none",
+        ),
+    ],
+)
+def test_cli_run_service_uris(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    cli_args: list,
+    expected_session_uri: str,
+    expected_artifact_uri: str,
 ) -> None:
-  """`adk run` should call run_cli via asyncio.run with correct parameters."""
-  rec = _Recorder()
-  monkeypatch.setattr(cli_tools_click, "run_cli", lambda **kwargs: rec(kwargs))
-  monkeypatch.setattr(
-      cli_tools_click.asyncio, "run", lambda coro: coro
-  )  # pass-through
-
-  # create dummy agent directory
+  """`adk run` should forward service URIs correctly to run_cli."""
   agent_dir = tmp_path / "agent"
   agent_dir.mkdir()
   (agent_dir / "__init__.py").touch()
   (agent_dir / "agent.py").touch()
 
+  # Capture the coroutine's locals before closing it
+  captured_locals = []
+
+  def capture_asyncio_run(coro):
+    # Extract the locals before closing the coroutine
+    if coro.cr_frame is not None:
+      captured_locals.append(dict(coro.cr_frame.f_locals))
+    coro.close()  # Properly close the coroutine to avoid warnings
+
+  monkeypatch.setattr(cli_tools_click.asyncio, "run", capture_asyncio_run)
+
   runner = CliRunner()
-  result = runner.invoke(cli_tools_click.main, ["run", str(agent_dir)])
-  assert result.exit_code == 0
-  assert rec.calls and rec.calls[0][0][0]["agent_folder_name"] == "agent"
+  result = runner.invoke(
+      cli_tools_click.main,
+      ["run", *cli_args, str(agent_dir)],
+  )
+  assert result.exit_code == 0, (result.output, repr(result.exception))
+  assert len(captured_locals) == 1, "Expected asyncio.run to be called once"
+
+  # Verify the kwargs passed to run_cli
+  coro_locals = captured_locals[0]
+  assert coro_locals.get("session_service_uri") == expected_session_uri
+  assert coro_locals.get("artifact_service_uri") == expected_artifact_uri
+  assert coro_locals["agent_folder_name"] == "agent"
 
 
 # cli deploy cloud_run
@@ -520,10 +561,13 @@ def test_cli_web_passes_service_uris(
   assert called_kwargs.get("memory_service_uri") == "rag://mycorpus"
 
 
-def test_cli_web_passes_deprecated_uris(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, _patch_uvicorn: _Recorder
+@pytest.mark.unmute_click
+def test_cli_web_warns_and_maps_deprecated_uris(
+    tmp_path: Path,
+    _patch_uvicorn: _Recorder,
+    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-  """`adk web` should use deprecated URIs if new ones are not provided."""
+  """`adk web` should accept deprecated URI flags with warnings."""
   agents_dir = tmp_path / "agents"
   agents_dir.mkdir()
 
@@ -542,11 +586,14 @@ def test_cli_web_passes_deprecated_uris(
           "gs://deprecated",
       ],
   )
+
   assert result.exit_code == 0
-  assert mock_get_app.calls
   called_kwargs = mock_get_app.calls[0][1]
   assert called_kwargs.get("session_service_uri") == "sqlite:///deprecated.db"
   assert called_kwargs.get("artifact_service_uri") == "gs://deprecated"
+  # Check output for deprecation warnings (CliRunner captures both stdout and stderr)
+  assert "--session_db_url" in result.output
+  assert "--artifact_storage_uri" in result.output
 
 
 def test_cli_eval_with_eval_set_file_path(