feat: Use --memory_service_uri in ADK CLI run command

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

tests/unittests/cli/test_service_registry.py

@@ -165,6 +165,13 @@ def test_create_memory_service_agentengine_full(registry, mock_services):
   )
 
 
+def test_create_memory_service_memory(registry):
+  from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
+
+  memory_service = registry.create_memory_service("memory://")
+  assert isinstance(memory_service, InMemoryMemoryService)
+
+
 # General Tests
 def test_unsupported_scheme(registry, mock_services):
   session_service = registry.create_session_service("unsupported://foo")

tests/unittests/cli/utils/test_cli.py

@@ -354,9 +354,145 @@ async def test_run_cli_accepts_memory_scheme(
       save_session=False,
       session_service_uri="memory://",
       artifact_service_uri="memory://",
+      memory_service_uri="memory://",
   )
 
 
+@pytest.mark.asyncio
+async def test_run_cli_invalid_memory_uri_surfaces_value_error(
+    fake_agent, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+  """run_cli should let ValueError propagate for invalid memory service URIs."""
+  parent_dir, folder_name = fake_agent
+  input_json = {"state": {}, "queries": []}
+  input_path = tmp_path / "invalid_memory_uri.json"
+  input_path.write_text(json.dumps(input_json))
+
+  def _raise_invalid_memory_uri(
+      *,
+      base_dir: Path | str,
+      memory_service_uri: str | None = None,
+  ) -> object:
+    del base_dir, memory_service_uri
+    raise ValueError("Unsupported memory service URI: unknown://x")
+
+  monkeypatch.setattr(
+      cli, "create_memory_service_from_options", _raise_invalid_memory_uri
+  )
+
+  with pytest.raises(ValueError, match="Unsupported memory service URI"):
+    await cli.run_cli(
+        agent_parent_dir=str(parent_dir),
+        agent_folder_name=folder_name,
+        input_file=str(input_path),
+        saved_session_file=None,
+        save_session=False,
+        memory_service_uri="unknown://x",
+    )
+
+
+@pytest.mark.asyncio
+async def test_run_cli_passes_memory_service_to_input_file(
+    fake_agent, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+  """run_cli should construct and pass the configured memory service."""
+  parent_dir, folder_name = fake_agent
+  input_json = {"state": {}, "queries": []}
+  input_path = tmp_path / "memory_input.json"
+  input_path.write_text(json.dumps(input_json))
+
+  memory_service_sentinel = object()
+  captured_factory_args: dict[str, Any] = {}
+  captured_memory_service: dict[str, Any] = {}
+
+  def _memory_factory(
+      *,
+      base_dir: Path | str,
+      memory_service_uri: str | None = None,
+  ) -> object:
+    captured_factory_args["base_dir"] = base_dir
+    captured_factory_args["memory_service_uri"] = memory_service_uri
+    return memory_service_sentinel
+
+  async def _run_input_file(
+      app_name: str,
+      user_id: str,
+      agent_or_app: BaseAgent | App,
+      artifact_service: Any,
+      session_service: Any,
+      credential_service: InMemoryCredentialService,
+      input_path: str,
+      memory_service: Any = None,
+  ) -> object:
+    del app_name, user_id, agent_or_app, artifact_service
+    del session_service, credential_service, input_path
+    captured_memory_service["value"] = memory_service
+    return object()
+
+  monkeypatch.setattr(
+      cli, "create_memory_service_from_options", _memory_factory
+  )
+  monkeypatch.setattr(cli, "run_input_file", _run_input_file)
+
+  await cli.run_cli(
+      agent_parent_dir=str(parent_dir),
+      agent_folder_name=folder_name,
+      input_file=str(input_path),
+      saved_session_file=None,
+      save_session=False,
+      memory_service_uri="memory://",
+  )
+
+  assert Path(captured_factory_args["base_dir"]) == parent_dir.resolve()
+  assert captured_factory_args["memory_service_uri"] == "memory://"
+  assert captured_memory_service["value"] is memory_service_sentinel
+
+
+@pytest.mark.asyncio
+async def test_run_cli_loads_dotenv_before_memory_service_creation(
+    fake_agent, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+  """run_cli should load agent .env values before creating memory service."""
+  parent_dir, folder_name = fake_agent
+  input_json = {"state": {}, "queries": []}
+  input_path = tmp_path / "dotenv_order_input.json"
+  input_path.write_text(json.dumps(input_json))
+
+  call_order: list[str] = []
+
+  def _load_dotenv_for_agent(agent_name: str, agents_dir: str) -> None:
+    del agent_name, agents_dir
+    call_order.append("load_dotenv")
+
+  def _memory_factory(
+      *,
+      base_dir: Path | str,
+      memory_service_uri: str | None = None,
+  ) -> object:
+    del base_dir, memory_service_uri
+    call_order.append("create_memory")
+    return object()
+
+  monkeypatch.setenv("ADK_DISABLE_LOAD_DOTENV", "0")
+  monkeypatch.setattr(cli.envs, "load_dotenv_for_agent", _load_dotenv_for_agent)
+  monkeypatch.setattr(
+      cli, "create_memory_service_from_options", _memory_factory
+  )
+
+  await cli.run_cli(
+      agent_parent_dir=str(parent_dir),
+      agent_folder_name=folder_name,
+      input_file=str(input_path),
+      saved_session_file=None,
+      save_session=False,
+      memory_service_uri="memory://",
+  )
+
+  assert "create_memory" in call_order
+  assert "load_dotenv" in call_order
+  assert call_order.index("load_dotenv") < call_order.index("create_memory")
+
+
 @pytest.mark.asyncio
 async def test_run_interactively_whitespace_and_exit(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch

tests/unittests/cli/utils/test_cli_tools_click.py

@@ -23,6 +23,7 @@ from types import SimpleNamespace
 from typing import Any
 from typing import Dict
 from typing import List
+from typing import Optional
 from typing import Tuple
 from unittest import mock
 
@@ -129,7 +130,7 @@ def test_cli_create_cmd_invokes_run_cmd(
 
 # cli run
 @pytest.mark.parametrize(
-    "cli_args,expected_session_uri,expected_artifact_uri",
+    "cli_args,expected_session_uri,expected_artifact_uri,expected_memory_uri",
     [
         pytest.param(
             [
@@ -137,15 +138,19 @@ def test_cli_create_cmd_invokes_run_cmd(
                 "memory://",
                 "--artifact_service_uri",
                 "memory://",
+                "--memory_service_uri",
+                "memory://",
             ],
             "memory://",
             "memory://",
+            "memory://",
             id="memory_scheme_uris",
         ),
         pytest.param(
             [],
             None,
             None,
+            None,
             id="default_uris_none",
         ),
     ],
@@ -154,8 +159,9 @@ def test_cli_run_service_uris(
     tmp_path: Path,
     monkeypatch: pytest.MonkeyPatch,
     cli_args: list,
-    expected_session_uri: str,
-    expected_artifact_uri: str,
+    expected_session_uri: Optional[str],
+    expected_artifact_uri: Optional[str],
+    expected_memory_uri: Optional[str],
 ) -> None:
   """`adk run` should forward service URIs correctly to run_cli."""
   agent_dir = tmp_path / "agent"
@@ -186,6 +192,7 @@ def test_cli_run_service_uris(
   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.get("memory_service_uri") == expected_memory_uri
   assert coro_locals["agent_folder_name"] == "agent"
 
 

tests/unittests/cli/utils/test_service_factory.py

@@ -252,6 +252,15 @@ def test_create_memory_service_defaults_to_in_memory(tmp_path: Path):
   assert isinstance(service, InMemoryMemoryService)
 
 
+def test_create_memory_service_supports_memory_uri(tmp_path: Path):
+  service = service_factory.create_memory_service_from_options(
+      base_dir=tmp_path,
+      memory_service_uri="memory://",
+  )
+
+  assert isinstance(service, InMemoryMemoryService)
+
+
 def test_create_memory_service_raises_on_unknown_scheme(
     tmp_path: Path, monkeypatch
 ):