feat: auto-connect jmp shell to existing leases

python/packages/jumpstarter-cli/jumpstarter_cli/shell.py

@@ -137,7 +137,7 @@ async def _run_shell_with_lease_async(lease, exporter_logs, config, command, can
                             logger.debug("Exporter ready (status: %s), launching shell...", result)
 
                             if monitor.status_message and monitor.status_message.startswith(HOOK_WARNING_PREFIX):
-                                warning_text = monitor.status_message[len(HOOK_WARNING_PREFIX):]
+                                warning_text = monitor.status_message[len(HOOK_WARNING_PREFIX) :]
                                 click.echo(click.style(f"Warning: {warning_text}", fg="yellow", bold=True))
 
                             # Run the shell command
@@ -194,11 +194,9 @@ async def _run_shell_with_lease_async(lease, exporter_logs, config, command, can
                                                 if monitor.status_message and monitor.status_message.startswith(
                                                     HOOK_WARNING_PREFIX
                                                 ):
-                                                    warning_text = monitor.status_message[len(HOOK_WARNING_PREFIX):]
+                                                    warning_text = monitor.status_message[len(HOOK_WARNING_PREFIX) :]
                                                     click.echo(
-                                                        click.style(
-                                                            f"Warning: {warning_text}", fg="yellow", bold=True
-                                                        )
+                                                        click.style(f"Warning: {warning_text}", fg="yellow", bold=True)
                                                     )
                                                 logger.info("afterLease hook completed")
                                             elif result == ExporterStatus.AFTER_LEASE_HOOK_FAILED:
@@ -294,6 +292,60 @@ async def _shell_with_signal_handling(  # noqa: C901
     return exit_code
 
 
+def _format_lease_display(lease) -> str:
+    parts = []
+    if lease.exporter:
+        parts.append(f"exporter={lease.exporter}")
+    if lease.selector:
+        parts.append(f"selector={lease.selector}")
+    if lease.effective_end_time:
+        parts.append(f"expires {lease.effective_end_time.strftime('%Y-%m-%d %H:%M')}")
+    elif lease.effective_begin_time and lease.duration:
+        end = lease.effective_begin_time + lease.duration
+        parts.append(f"expires {end.strftime('%Y-%m-%d %H:%M')}")
+    return ", ".join(parts) if parts else ""
+
+
+async def _resolve_lease_from_active_async(config) -> str:
+    lease_list = await config.list_leases(only_active=True)
+    client_name = config.metadata.name
+    leases = [lease for lease in lease_list.leases if lease.client == client_name]
+
+    if not leases:
+        raise click.UsageError(
+            "no active leases found. Use --selector/-l or --name/-n to create one, "
+            "or create a lease with 'jmp create lease'."
+        )
+
+    if len(leases) == 1:
+        return leases[0].name
+
+    if sys.stdin.isatty():
+        click.echo("Multiple active leases found:\n")
+        for i, lease in enumerate(leases, 1):
+            info = _format_lease_display(lease)
+            click.echo(f"  {i}) {lease.name}")
+            if info:
+                click.echo(f"     {info}")
+        click.echo()
+        chosen = click.prompt(
+            "Select a lease [1-{}]".format(len(leases)),
+            type=click.IntRange(1, len(leases)),
+        )
+        return leases[chosen - 1].name
+
+    lease_summaries = []
+    for lease in leases:
+        info = _format_lease_display(lease)
+        summary = f"{lease.name} ({info})" if info else lease.name
+        lease_summaries.append(summary)
+    raise click.UsageError(
+        "multiple active leases found:\n  "
+        + "\n  ".join(lease_summaries)
+        + "\nUse --lease to specify one, or run interactively to select."
+    )
+
+
 @click.command("shell")
 @opt_config()
 @click.argument("command", nargs=-1)
@@ -324,16 +376,14 @@ def shell(
 
     Example:
 
-    .. code-block:: bash
-
-        $ jmp shell --exporter foo -- python bar.py
+        jmp shell --exporter foo -- python bar.py
     """
 
     match config:
         case ClientConfigV1Alpha1():
             has_existing_lease = bool(lease_name or os.environ.get(JMP_LEASE))
             if not selector and not exporter_name and not has_existing_lease:
-                raise click.UsageError("one of --selector/-l or --name/-n is required")
+                lease_name = anyio.run(_resolve_lease_from_active_async, config)
             exit_code = anyio.run(
                 _shell_with_signal_handling,
                 config,

python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py

@@ -1,18 +1,43 @@
 import inspect
 from contextlib import asynccontextmanager
-from datetime import timedelta
+from datetime import datetime, timedelta
 from unittest.mock import AsyncMock, Mock, patch
 
 import anyio
 import click
 import pytest
 
-from jumpstarter_cli.shell import _shell_with_signal_handling, shell
+from jumpstarter_cli.shell import _resolve_lease_from_active_async, _shell_with_signal_handling, shell
 
+from jumpstarter.client.grpc import Lease, LeaseList
 from jumpstarter.config.client import ClientConfigV1Alpha1
 from jumpstarter.config.env import JMP_LEASE
 
 
+def _make_lease(name: str, client: str = "test-client") -> Lease:
+    return Lease(
+        namespace="default",
+        name=name,
+        selector="",
+        exporter_name=None,
+        duration=timedelta(minutes=30),
+        effective_duration=None,
+        begin_time=datetime.now(),
+        client=client,
+        exporter="test-exporter",
+        conditions=[],
+        effective_begin_time=None,
+        effective_end_time=None,
+    )
+
+
+def _make_lease_list(names: list[str]) -> LeaseList:
+    return LeaseList(
+        leases=[_make_lease(n) for n in names],
+        next_page_token=None,
+    )
+
+
 class _DummyConfig:
     def __init__(self):
         self.captured = None
@@ -48,10 +73,13 @@ def test_shell_passes_exporter_name_to_lease_async():
     assert config.captured[1] == "laptop-test-exporter"
 
 
-def test_shell_requires_selector_or_name():
-    with pytest.raises(click.UsageError, match="one of --selector/-l or --name/-n is required"):
-        inspect.unwrap(shell.callback)(
-            config=Mock(spec=ClientConfigV1Alpha1),
+def test_shell_requires_selector_or_name_when_no_leases():
+    config = Mock(spec=ClientConfigV1Alpha1)
+    config.metadata = type("Metadata", (), {"name": "test-client"})()
+    config.list_leases = AsyncMock(return_value=_make_lease_list([]))
+    with pytest.raises(click.UsageError, match="no active leases found"):
+        shell.callback.__wrapped__.__wrapped__(
+            config=config,
             command=(),
             lease_name=None,
             selector=None,
@@ -81,6 +109,118 @@ def test_shell_allows_existing_lease_name_without_selector_or_name():
     mock_exit.assert_called_once_with(0)
 
 
+def test_shell_auto_connects_single_lease():
+    config = Mock(spec=ClientConfigV1Alpha1)
+    config.metadata = type("Metadata", (), {"name": "test-client"})()
+    with (
+        patch("jumpstarter_cli.shell.anyio.run", side_effect=["my-only-lease", 0]) as mock_run,
+        patch("jumpstarter_cli.shell.sys.exit") as mock_exit,
+    ):
+        shell.callback.__wrapped__.__wrapped__(
+            config=config,
+            command=(),
+            lease_name=None,
+            selector=None,
+            exporter_name=None,
+            duration=timedelta(minutes=1),
+            exporter_logs=False,
+            acquisition_timeout=None,
+        )
+
+    resolve_call_args = mock_run.call_args_list[0]
+    assert resolve_call_args[0][0] is _resolve_lease_from_active_async
+    assert resolve_call_args[0][1] is config
+    shell_call_args = mock_run.call_args_list[1]
+    assert shell_call_args[0][4] == "my-only-lease"
+    mock_exit.assert_called_once_with(0)
+
+
+def test_shell_no_leases_shows_guidance():
+    config = Mock(spec=ClientConfigV1Alpha1)
+    config.metadata = type("Metadata", (), {"name": "test-client"})()
+    config.list_leases = AsyncMock(return_value=_make_lease_list([]))
+    with pytest.raises(click.UsageError, match="no active leases found"):
+        shell.callback.__wrapped__.__wrapped__(
+            config=config,
+            command=(),
+            lease_name=None,
+            selector=None,
+            exporter_name=None,
+            duration=timedelta(minutes=1),
+            exporter_logs=False,
+            acquisition_timeout=None,
+        )
+    config.list_leases.assert_called_once_with(only_active=True)
+
+
+def test_shell_multi_lease_tty_picker():
+    config = Mock(spec=ClientConfigV1Alpha1)
+    config.metadata = type("Metadata", (), {"name": "test-client"})()
+    config.list_leases = AsyncMock(return_value=_make_lease_list(["lease-a", "lease-b", "lease-c"]))
+    with (
+        patch("jumpstarter_cli.shell.sys.stdin") as mock_stdin,
+        patch("jumpstarter_cli.shell.click.prompt", return_value=2),
+    ):
+        mock_stdin.isatty.return_value = True
+        selected = anyio.run(_resolve_lease_from_active_async, config)
+
+    assert selected == "lease-b"
+    config.list_leases.assert_called_once_with(only_active=True)
+
+
+def test_shell_multi_lease_no_tty_error():
+    config = Mock(spec=ClientConfigV1Alpha1)
+    config.metadata = type("Metadata", (), {"name": "test-client"})()
+    config.list_leases = AsyncMock(return_value=_make_lease_list(["lease-a", "lease-b"]))
+    with (
+        patch("jumpstarter_cli.shell.sys.stdin") as mock_stdin,
+        pytest.raises(click.UsageError, match="lease-a"),
+    ):
+        mock_stdin.isatty.return_value = False
+        shell.callback.__wrapped__.__wrapped__(
+            config=config,
+            command=(),
+            lease_name=None,
+            selector=None,
+            exporter_name=None,
+            duration=timedelta(minutes=1),
+            exporter_logs=False,
+            acquisition_timeout=None,
+        )
+
+
+def test_shell_filters_leases_by_current_client():
+    other_user_lease = _make_lease("other-user-lease", client="other-client")
+    my_lease = _make_lease("my-lease", client="test-client")
+    lease_list = LeaseList(leases=[other_user_lease, my_lease], next_page_token=None)
+    config = Mock(spec=ClientConfigV1Alpha1)
+    config.metadata = type("Metadata", (), {"name": "test-client"})()
+    config.list_leases = AsyncMock(return_value=lease_list)
+
+    selected = anyio.run(_resolve_lease_from_active_async, config)
+    assert selected == "my-lease"
+    config.list_leases.assert_called_once_with(only_active=True)
+
+
+def test_shell_no_own_leases_among_others():
+    other_lease = _make_lease("other-lease", client="other-client")
+    lease_list = LeaseList(leases=[other_lease], next_page_token=None)
+    config = Mock(spec=ClientConfigV1Alpha1)
+    config.metadata = type("Metadata", (), {"name": "test-client"})()
+    config.list_leases = AsyncMock(return_value=lease_list)
+    with pytest.raises(click.UsageError, match="no active leases found"):
+        shell.callback.__wrapped__.__wrapped__(
+            config=config,
+            command=(),
+            lease_name=None,
+            selector=None,
+            exporter_name=None,
+            duration=timedelta(minutes=1),
+            exporter_logs=False,
+            acquisition_timeout=None,
+        )
+
+
 def test_shell_allows_env_lease_without_selector_or_name():
     with (
         patch("jumpstarter_cli.shell.anyio.run", return_value=0),
@@ -99,3 +239,12 @@ def test_shell_allows_env_lease_without_selector_or_name():
         )
 
     mock_exit.assert_called_once_with(0)
+
+def test_resolve_lease_handles_async_list_leases():
+    config = Mock(spec=ClientConfigV1Alpha1)
+    config.metadata = type("Metadata", (), {"name": "test-client"})()
+    config.list_leases = AsyncMock(return_value=_make_lease_list(["async-lease"]))
+
+    selected = anyio.run(_resolve_lease_from_active_async, config)
+    assert selected == "async-lease"
+    config.list_leases.assert_called_once_with(only_active=True)