From efbffd703d47f2d0bf7d841483cffc48b9cc77f0 Mon Sep 17 00:00:00 2001 From: Sam Schillace Date: Sat, 25 Jul 2026 12:52:18 -0700 Subject: [PATCH] feat(run): support pre-allocated session IDs Add a --session-id option to the run command so a caller-supplied id is used for a new session instead of a freshly minted UUID, and add an 'amplifier session new' subcommand that prints a fresh UUID for pre-allocation. --session-id is rejected together with --resume (which already selects the session to continue). ID=$(amplifier session new) amplifier run --session-id "$ID" "your prompt" --- amplifier_app_cli/commands/run.py | 18 +++- amplifier_app_cli/commands/session.py | 14 +++ tests/test_run_session_id_flag.py | 149 ++++++++++++++++++++++++++ 3 files changed, 179 insertions(+), 2 deletions(-) create mode 100644 tests/test_run_session_id_flag.py diff --git a/amplifier_app_cli/commands/run.py b/amplifier_app_cli/commands/run.py index d2dd85e..228e888 100644 --- a/amplifier_app_cli/commands/run.py +++ b/amplifier_app_cli/commands/run.py @@ -55,6 +55,11 @@ def register_run_command( help="Execution mode", ) @click.option("--resume", help="Resume specific session with new prompt") + @click.option( + "--session-id", + default=None, + help="Use a pre-allocated session ID for a new session (see 'amplifier session new')", + ) @click.option("--verbose", "-v", is_flag=True, help="Verbose output") @click.option( "--output-format", @@ -70,12 +75,21 @@ def run( max_tokens: int | None, mode: str, resume: str | None, + session_id: str | None, verbose: bool, output_format: str, ): """Execute a prompt or start an interactive session.""" from ..session_store import SessionStore + # --session-id applies only to new sessions; --resume already carries an id + if session_id and resume: + console.print( + "[red]Error:[/red] --session-id cannot be combined with --resume\n" + "--resume already selects the session to continue" + ) + sys.exit(1) + # Handle --resume flag if resume: store = SessionStore() @@ -362,7 +376,7 @@ def run( ) else: # New session - banner displayed by interactive_chat - session_id = str(uuid.uuid4()) + session_id = session_id or str(uuid.uuid4()) asyncio.run( interactive_chat( config_data, @@ -407,7 +421,7 @@ def run( ) else: # Create new session - session_id = str(uuid.uuid4()) + session_id = session_id or str(uuid.uuid4()) if output_format == "text": config_summary = get_effective_config_summary( config_data, config_source_name diff --git a/amplifier_app_cli/commands/session.py b/amplifier_app_cli/commands/session.py index 7c82579..a701a15 100644 --- a/amplifier_app_cli/commands/session.py +++ b/amplifier_app_cli/commands/session.py @@ -530,6 +530,20 @@ def session(ctx: click.Context): click.echo("\n" + ctx.get_help()) ctx.exit() + @session.command(name="new") + def sessions_new(): + """Pre-allocate a new session ID and print it. + + The printed UUID can be passed to a later launch to give the session a + known id in advance, e.g.: + + ID=$(amplifier session new) + amplifier run --session-id "$ID" "your prompt" + """ + import uuid + + click.echo(str(uuid.uuid4())) + @session.command(name="list") @click.option("--limit", "-n", default=20, help="Number of sessions to show") @click.option( diff --git a/tests/test_run_session_id_flag.py b/tests/test_run_session_id_flag.py new file mode 100644 index 0000000..5e82545 --- /dev/null +++ b/tests/test_run_session_id_flag.py @@ -0,0 +1,149 @@ +"""Tests for pre-allocated session IDs. + +Covers the ``--session-id`` flag on the ``run`` command (use a caller-supplied +id for a new session instead of minting a UUID) and the ``session new`` +subcommand (print a fresh UUID for pre-allocation). + +The run tests build a minimal Click group by registering the run command with +mocked collaborators (so no chat/session is actually launched), then invoke it +and inspect the session id handed to the launch coroutine. +""" + +from __future__ import annotations + +import uuid +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import click +from click.testing import CliRunner + +from amplifier_app_cli.commands.run import register_run_command +from amplifier_app_cli.commands.session import register_session_commands + +_MODULE = "amplifier_app_cli.commands.run" + + +def _make_prepared_bundle() -> SimpleNamespace: + return SimpleNamespace(mount_plan={"hooks": []}) + + +def _build_cli() -> tuple[click.Group, AsyncMock]: + """Register the run command with mocked collaborators. + + Returns the CLI group and the execute_single AsyncMock so tests can read + the ``session_id`` kwarg the launch path was called with. + """ + execute_single = AsyncMock() + cli = click.Group() + register_run_command( + cli, + interactive_chat=AsyncMock(), + execute_single=execute_single, + get_module_search_paths=lambda: [], + check_first_run=lambda: False, + prompt_first_run_init=lambda console: True, + ) + return cli, execute_single + + +def _invoke(cli: click.Group, prepared_bundle: SimpleNamespace, *args: str): + with ( + patch(f"{_MODULE}.resolve_config", return_value=({}, prepared_bundle)), + patch(f"{_MODULE}.create_config_manager", return_value=MagicMock()), + patch(f"{_MODULE}.AppSettings", return_value=MagicMock()), + patch( + "amplifier_app_cli.utils.startup_checker.check_and_notify", + new=AsyncMock(), + ), + ): + runner = CliRunner() + return runner.invoke(cli, ["run", *args]) + + +def test_session_id_flag_used_for_new_session(): + """``run --session-id X`` must launch the new session with that exact id.""" + prepared_bundle = _make_prepared_bundle() + cli, execute_single = _build_cli() + + result = _invoke( + cli, + prepared_bundle, + "hello", + "--session-id", + "my-preallocated-id", + "--mode", + "single", + "--output-format", + "json", + ) + + assert result.exit_code == 0, result.output + execute_single.assert_awaited_once() + await_args = execute_single.await_args + assert await_args is not None + assert await_args.kwargs["session_id"] == "my-preallocated-id" + + +def test_no_session_id_flag_mints_uuid(): + """Omitting ``--session-id`` must fall back to a generated UUID.""" + prepared_bundle = _make_prepared_bundle() + cli, execute_single = _build_cli() + + result = _invoke( + cli, + prepared_bundle, + "hello", + "--mode", + "single", + "--output-format", + "json", + ) + + assert result.exit_code == 0, result.output + execute_single.assert_awaited_once() + await_args = execute_single.await_args + assert await_args is not None + generated = await_args.kwargs["session_id"] + # A valid UUID4 string was minted (raises if not parseable). + assert str(uuid.UUID(generated)) == generated + + +def test_session_id_conflicts_with_resume(): + """``--session-id`` combined with ``--resume`` is rejected before launch.""" + prepared_bundle = _make_prepared_bundle() + cli, execute_single = _build_cli() + + result = _invoke( + cli, + prepared_bundle, + "hello", + "--session-id", + "abc", + "--resume", + "some-other-session", + "--mode", + "single", + ) + + assert result.exit_code != 0 + assert "cannot be combined with --resume" in result.output + execute_single.assert_not_awaited() + + +def test_session_new_prints_valid_uuid(): + """``session new`` prints a single parseable UUID and nothing else.""" + cli = click.Group() + register_session_commands( + cli, + interactive_chat=AsyncMock(), + execute_single=AsyncMock(), + get_module_search_paths=lambda: [], + ) + + runner = CliRunner() + result = runner.invoke(cli, ["session", "new"]) + + assert result.exit_code == 0, result.output + printed = result.output.strip() + assert str(uuid.UUID(printed)) == printed