Skip to content

Commit 4e55d73

Browse files
raballewclaude
andauthored
feat: auto-connect jmp shell to existing leases (#330)
## Summary - When running `jmp shell` without `--selector` or `--name`, automatically resolve from active leases - If exactly one active lease exists for the current client, auto-connect to it - If multiple active leases exist and running interactively, display a numbered picker with exporter name, selector, and expiry details - Only show the current user's leases, not other clients' leases - Clean up RST markup in shell help docstring Closes #326 ## Test plan - [x] Auto-connects to single active lease without prompting - [x] Shows numbered picker with lease details when multiple leases exist (TTY) - [x] Errors with guidance when multiple leases exist (non-TTY) - [x] Filters out other clients' leases - [x] Shows "no active leases found" when user has no leases (even if others do) - [x] Existing `--lease`, `--selector`, `--name`, and `JMP_LEASE` env var still work - [x] 10 unit tests passing 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent c6cf780 commit 4e55d73

2 files changed

Lines changed: 214 additions & 15 deletions

File tree

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

Lines changed: 59 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ async def _run_shell_with_lease_async(lease, exporter_logs, config, command, can
137137
logger.debug("Exporter ready (status: %s), launching shell...", result)
138138

139139
if monitor.status_message and monitor.status_message.startswith(HOOK_WARNING_PREFIX):
140-
warning_text = monitor.status_message[len(HOOK_WARNING_PREFIX):]
140+
warning_text = monitor.status_message[len(HOOK_WARNING_PREFIX) :]
141141
click.echo(click.style(f"Warning: {warning_text}", fg="yellow", bold=True))
142142

143143
# Run the shell command
@@ -194,11 +194,9 @@ async def _run_shell_with_lease_async(lease, exporter_logs, config, command, can
194194
if monitor.status_message and monitor.status_message.startswith(
195195
HOOK_WARNING_PREFIX
196196
):
197-
warning_text = monitor.status_message[len(HOOK_WARNING_PREFIX):]
197+
warning_text = monitor.status_message[len(HOOK_WARNING_PREFIX) :]
198198
click.echo(
199-
click.style(
200-
f"Warning: {warning_text}", fg="yellow", bold=True
201-
)
199+
click.style(f"Warning: {warning_text}", fg="yellow", bold=True)
202200
)
203201
logger.info("afterLease hook completed")
204202
elif result == ExporterStatus.AFTER_LEASE_HOOK_FAILED:
@@ -294,6 +292,60 @@ async def _shell_with_signal_handling( # noqa: C901
294292
return exit_code
295293

296294

295+
def _format_lease_display(lease) -> str:
296+
parts = []
297+
if lease.exporter:
298+
parts.append(f"exporter={lease.exporter}")
299+
if lease.selector:
300+
parts.append(f"selector={lease.selector}")
301+
if lease.effective_end_time:
302+
parts.append(f"expires {lease.effective_end_time.strftime('%Y-%m-%d %H:%M')}")
303+
elif lease.effective_begin_time and lease.duration:
304+
end = lease.effective_begin_time + lease.duration
305+
parts.append(f"expires {end.strftime('%Y-%m-%d %H:%M')}")
306+
return ", ".join(parts) if parts else ""
307+
308+
309+
async def _resolve_lease_from_active_async(config) -> str:
310+
lease_list = await config.list_leases(only_active=True)
311+
client_name = config.metadata.name
312+
leases = [lease for lease in lease_list.leases if lease.client == client_name]
313+
314+
if not leases:
315+
raise click.UsageError(
316+
"no active leases found. Use --selector/-l or --name/-n to create one, "
317+
"or create a lease with 'jmp create lease'."
318+
)
319+
320+
if len(leases) == 1:
321+
return leases[0].name
322+
323+
if sys.stdin.isatty():
324+
click.echo("Multiple active leases found:\n")
325+
for i, lease in enumerate(leases, 1):
326+
info = _format_lease_display(lease)
327+
click.echo(f" {i}) {lease.name}")
328+
if info:
329+
click.echo(f" {info}")
330+
click.echo()
331+
chosen = click.prompt(
332+
"Select a lease [1-{}]".format(len(leases)),
333+
type=click.IntRange(1, len(leases)),
334+
)
335+
return leases[chosen - 1].name
336+
337+
lease_summaries = []
338+
for lease in leases:
339+
info = _format_lease_display(lease)
340+
summary = f"{lease.name} ({info})" if info else lease.name
341+
lease_summaries.append(summary)
342+
raise click.UsageError(
343+
"multiple active leases found:\n "
344+
+ "\n ".join(lease_summaries)
345+
+ "\nUse --lease to specify one, or run interactively to select."
346+
)
347+
348+
297349
@click.command("shell")
298350
@opt_config()
299351
@click.argument("command", nargs=-1)
@@ -324,16 +376,14 @@ def shell(
324376
325377
Example:
326378
327-
.. code-block:: bash
328-
329-
$ jmp shell --exporter foo -- python bar.py
379+
jmp shell --exporter foo -- python bar.py
330380
"""
331381

332382
match config:
333383
case ClientConfigV1Alpha1():
334384
has_existing_lease = bool(lease_name or os.environ.get(JMP_LEASE))
335385
if not selector and not exporter_name and not has_existing_lease:
336-
raise click.UsageError("one of --selector/-l or --name/-n is required")
386+
lease_name = anyio.run(_resolve_lease_from_active_async, config)
337387
exit_code = anyio.run(
338388
_shell_with_signal_handling,
339389
config,

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

Lines changed: 155 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,43 @@
11
import inspect
22
from contextlib import asynccontextmanager
3-
from datetime import timedelta
3+
from datetime import datetime, timedelta
44
from unittest.mock import AsyncMock, Mock, patch
55

66
import anyio
77
import click
88
import pytest
99

10-
from jumpstarter_cli.shell import _shell_with_signal_handling, shell
10+
from jumpstarter_cli.shell import _resolve_lease_from_active_async, _shell_with_signal_handling, shell
1111

12+
from jumpstarter.client.grpc import Lease, LeaseList
1213
from jumpstarter.config.client import ClientConfigV1Alpha1
1314
from jumpstarter.config.env import JMP_LEASE
1415

1516

17+
def _make_lease(name: str, client: str = "test-client") -> Lease:
18+
return Lease(
19+
namespace="default",
20+
name=name,
21+
selector="",
22+
exporter_name=None,
23+
duration=timedelta(minutes=30),
24+
effective_duration=None,
25+
begin_time=datetime.now(),
26+
client=client,
27+
exporter="test-exporter",
28+
conditions=[],
29+
effective_begin_time=None,
30+
effective_end_time=None,
31+
)
32+
33+
34+
def _make_lease_list(names: list[str]) -> LeaseList:
35+
return LeaseList(
36+
leases=[_make_lease(n) for n in names],
37+
next_page_token=None,
38+
)
39+
40+
1641
class _DummyConfig:
1742
def __init__(self):
1843
self.captured = None
@@ -48,10 +73,13 @@ def test_shell_passes_exporter_name_to_lease_async():
4873
assert config.captured[1] == "laptop-test-exporter"
4974

5075

51-
def test_shell_requires_selector_or_name():
52-
with pytest.raises(click.UsageError, match="one of --selector/-l or --name/-n is required"):
53-
inspect.unwrap(shell.callback)(
54-
config=Mock(spec=ClientConfigV1Alpha1),
76+
def test_shell_requires_selector_or_name_when_no_leases():
77+
config = Mock(spec=ClientConfigV1Alpha1)
78+
config.metadata = type("Metadata", (), {"name": "test-client"})()
79+
config.list_leases = AsyncMock(return_value=_make_lease_list([]))
80+
with pytest.raises(click.UsageError, match="no active leases found"):
81+
shell.callback.__wrapped__.__wrapped__(
82+
config=config,
5583
command=(),
5684
lease_name=None,
5785
selector=None,
@@ -81,6 +109,118 @@ def test_shell_allows_existing_lease_name_without_selector_or_name():
81109
mock_exit.assert_called_once_with(0)
82110

83111

112+
def test_shell_auto_connects_single_lease():
113+
config = Mock(spec=ClientConfigV1Alpha1)
114+
config.metadata = type("Metadata", (), {"name": "test-client"})()
115+
with (
116+
patch("jumpstarter_cli.shell.anyio.run", side_effect=["my-only-lease", 0]) as mock_run,
117+
patch("jumpstarter_cli.shell.sys.exit") as mock_exit,
118+
):
119+
shell.callback.__wrapped__.__wrapped__(
120+
config=config,
121+
command=(),
122+
lease_name=None,
123+
selector=None,
124+
exporter_name=None,
125+
duration=timedelta(minutes=1),
126+
exporter_logs=False,
127+
acquisition_timeout=None,
128+
)
129+
130+
resolve_call_args = mock_run.call_args_list[0]
131+
assert resolve_call_args[0][0] is _resolve_lease_from_active_async
132+
assert resolve_call_args[0][1] is config
133+
shell_call_args = mock_run.call_args_list[1]
134+
assert shell_call_args[0][4] == "my-only-lease"
135+
mock_exit.assert_called_once_with(0)
136+
137+
138+
def test_shell_no_leases_shows_guidance():
139+
config = Mock(spec=ClientConfigV1Alpha1)
140+
config.metadata = type("Metadata", (), {"name": "test-client"})()
141+
config.list_leases = AsyncMock(return_value=_make_lease_list([]))
142+
with pytest.raises(click.UsageError, match="no active leases found"):
143+
shell.callback.__wrapped__.__wrapped__(
144+
config=config,
145+
command=(),
146+
lease_name=None,
147+
selector=None,
148+
exporter_name=None,
149+
duration=timedelta(minutes=1),
150+
exporter_logs=False,
151+
acquisition_timeout=None,
152+
)
153+
config.list_leases.assert_called_once_with(only_active=True)
154+
155+
156+
def test_shell_multi_lease_tty_picker():
157+
config = Mock(spec=ClientConfigV1Alpha1)
158+
config.metadata = type("Metadata", (), {"name": "test-client"})()
159+
config.list_leases = AsyncMock(return_value=_make_lease_list(["lease-a", "lease-b", "lease-c"]))
160+
with (
161+
patch("jumpstarter_cli.shell.sys.stdin") as mock_stdin,
162+
patch("jumpstarter_cli.shell.click.prompt", return_value=2),
163+
):
164+
mock_stdin.isatty.return_value = True
165+
selected = anyio.run(_resolve_lease_from_active_async, config)
166+
167+
assert selected == "lease-b"
168+
config.list_leases.assert_called_once_with(only_active=True)
169+
170+
171+
def test_shell_multi_lease_no_tty_error():
172+
config = Mock(spec=ClientConfigV1Alpha1)
173+
config.metadata = type("Metadata", (), {"name": "test-client"})()
174+
config.list_leases = AsyncMock(return_value=_make_lease_list(["lease-a", "lease-b"]))
175+
with (
176+
patch("jumpstarter_cli.shell.sys.stdin") as mock_stdin,
177+
pytest.raises(click.UsageError, match="lease-a"),
178+
):
179+
mock_stdin.isatty.return_value = False
180+
shell.callback.__wrapped__.__wrapped__(
181+
config=config,
182+
command=(),
183+
lease_name=None,
184+
selector=None,
185+
exporter_name=None,
186+
duration=timedelta(minutes=1),
187+
exporter_logs=False,
188+
acquisition_timeout=None,
189+
)
190+
191+
192+
def test_shell_filters_leases_by_current_client():
193+
other_user_lease = _make_lease("other-user-lease", client="other-client")
194+
my_lease = _make_lease("my-lease", client="test-client")
195+
lease_list = LeaseList(leases=[other_user_lease, my_lease], next_page_token=None)
196+
config = Mock(spec=ClientConfigV1Alpha1)
197+
config.metadata = type("Metadata", (), {"name": "test-client"})()
198+
config.list_leases = AsyncMock(return_value=lease_list)
199+
200+
selected = anyio.run(_resolve_lease_from_active_async, config)
201+
assert selected == "my-lease"
202+
config.list_leases.assert_called_once_with(only_active=True)
203+
204+
205+
def test_shell_no_own_leases_among_others():
206+
other_lease = _make_lease("other-lease", client="other-client")
207+
lease_list = LeaseList(leases=[other_lease], next_page_token=None)
208+
config = Mock(spec=ClientConfigV1Alpha1)
209+
config.metadata = type("Metadata", (), {"name": "test-client"})()
210+
config.list_leases = AsyncMock(return_value=lease_list)
211+
with pytest.raises(click.UsageError, match="no active leases found"):
212+
shell.callback.__wrapped__.__wrapped__(
213+
config=config,
214+
command=(),
215+
lease_name=None,
216+
selector=None,
217+
exporter_name=None,
218+
duration=timedelta(minutes=1),
219+
exporter_logs=False,
220+
acquisition_timeout=None,
221+
)
222+
223+
84224
def test_shell_allows_env_lease_without_selector_or_name():
85225
with (
86226
patch("jumpstarter_cli.shell.anyio.run", return_value=0),
@@ -99,3 +239,12 @@ def test_shell_allows_env_lease_without_selector_or_name():
99239
)
100240

101241
mock_exit.assert_called_once_with(0)
242+
243+
def test_resolve_lease_handles_async_list_leases():
244+
config = Mock(spec=ClientConfigV1Alpha1)
245+
config.metadata = type("Metadata", (), {"name": "test-client"})()
246+
config.list_leases = AsyncMock(return_value=_make_lease_list(["async-lease"]))
247+
248+
selected = anyio.run(_resolve_lease_from_active_async, config)
249+
assert selected == "async-lease"
250+
config.list_leases.assert_called_once_with(only_active=True)

0 commit comments

Comments
 (0)