From f3942e92dcd161f3aec2b1bb65b449f87c326ac9 Mon Sep 17 00:00:00 2001 From: SHAI-shivansh-sharma Date: Thu, 10 Sep 2026 21:45:23 +0530 Subject: [PATCH] Add state file option to adk run --- src/google/adk/cli/cli_tools_click.py | 26 ++++- .../cli/utils/test_cli_tools_click.py | 94 +++++++++++++++++++ 2 files changed, 118 insertions(+), 2 deletions(-) diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 42302f4b3d..03e37c16a0 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -993,6 +993,13 @@ def wrapper(*args, **kwargs): type=str, help="Optional. Initial state for the run as a JSON string.", ) +@click.option( + "--state_file", + type=click.Path( + exists=True, dir_okay=False, file_okay=True, resolve_path=True + ), + help="Optional. Path to a JSON file containing initial state for the run.", +) @click.option( "--timeout", type=str, @@ -1032,6 +1039,7 @@ def cli_run( replay: Optional[str], resume: Optional[str], state: Optional[str] = None, + state_file: Optional[str] = None, timeout: Optional[str] = None, in_memory: bool = False, jsonl: bool = False, @@ -1057,6 +1065,20 @@ def cli_run( agent_parent_folder = os.path.dirname(agent) agent_folder_name = os.path.basename(agent) + if state is not None and state_file is not None: + raise click.UsageError( + "Options 'state' and 'state_file' cannot be set together." + ) + + state_str = state + if state_file is not None: + try: + state_str = Path(state_file).read_text(encoding="utf-8") + except OSError as e: + raise click.ClickException( + f"Failed to read --state_file '{state_file}': {e}" + ) from e + # If query is provided, we run in single-step mode (JSONL output) if query is not None: from .cli import run_once_cli @@ -1066,7 +1088,7 @@ def cli_run( agent_parent_dir=agent_parent_folder, agent_folder_name=agent_folder_name, query=query, - state_str=state, + state_str=state_str, session_id=session_id, replay=replay, timeout=timeout, @@ -1092,7 +1114,7 @@ def cli_run( saved_session_file=resume, save_session=save_session, session_id=session_id, - state_str=state, + state_str=state_str, timeout=timeout, in_memory=in_memory, jsonl=jsonl, diff --git a/tests/unittests/cli/utils/test_cli_tools_click.py b/tests/unittests/cli/utils/test_cli_tools_click.py index 05082b1b77..30b99a7f66 100644 --- a/tests/unittests/cli/utils/test_cli_tools_click.py +++ b/tests/unittests/cli/utils/test_cli_tools_click.py @@ -647,6 +647,36 @@ def test_cli_run_interactive_with_state( assert called_kwargs.get("state_str") == '{"x": 1}' +def test_cli_run_interactive_with_state_file( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """`adk run` in interactive mode should pass state file contents.""" + # Arrange + agent_dir = tmp_path / "agent_interactive" + agent_dir.mkdir() + (agent_dir / "__init__.py").touch() + (agent_dir / "agent.py").touch() + state_file = tmp_path / "state.json" + state_file.write_text('{"x": 1}', encoding="utf-8") + + mock_run_cli = mock.AsyncMock() + monkeypatch.setattr("google.adk.cli.cli.run_cli", mock_run_cli) + + runner = CliRunner() + + # Act + result = runner.invoke( + cli_tools_click.main, + ["run", str(agent_dir), "--state_file", str(state_file)], + ) + + # Assert + assert result.exit_code == 0 + assert mock_run_cli.called + called_kwargs = mock_run_cli.call_args.kwargs + assert called_kwargs.get("state_str") == '{"x": 1}' + + def test_cli_run_options_with_query( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -685,6 +715,70 @@ def test_cli_run_options_with_query( assert called_kwargs.get("jsonl") is True +def test_cli_run_options_with_query_and_state_file( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """`adk run` with query should pass state file contents to run_once_cli.""" + # Arrange + agent_dir = tmp_path / "agent_opts" + agent_dir.mkdir() + (agent_dir / "__init__.py").touch() + state_file = tmp_path / "state.json" + state_file.write_text('{"x": 1}', encoding="utf-8") + + mock_run_once = mock.AsyncMock(return_value=0) + monkeypatch.setattr("google.adk.cli.cli.run_once_cli", mock_run_once) + + runner = CliRunner() + + # Act + result = runner.invoke( + cli_tools_click.main, + ["run", str(agent_dir), "hello", "--state_file", str(state_file)], + ) + + # Assert + assert result.exit_code == 0 + assert mock_run_once.called + called_kwargs = mock_run_once.call_args.kwargs + assert called_kwargs.get("query") == "hello" + assert called_kwargs.get("state_str") == '{"x": 1}' + + +def test_cli_run_rejects_state_and_state_file_together( + tmp_path: Path, +) -> None: + """`adk run` should reject simultaneous --state and --state_file.""" + # Arrange + agent_dir = tmp_path / "agent_opts" + agent_dir.mkdir() + (agent_dir / "__init__.py").touch() + state_file = tmp_path / "state.json" + state_file.write_text('{"x": 1}', encoding="utf-8") + + runner = CliRunner() + + # Act + result = runner.invoke( + cli_tools_click.main, + [ + "run", + str(agent_dir), + "--state", + '{"y": 2}', + "--state_file", + str(state_file), + ], + ) + + # Assert + assert result.exit_code != 0 + assert ( + "Options 'state' and 'state_file' cannot be set together." + in result.output + ) + + def test_cli_run_auto_resume_with_query( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: