fix: unblock post-shell status probe when connection is lost

[evakhoni]

Summary

After the user exits jmp shell, a status probe is sent to detect exporter restarts before triggering the afterLease hook. If the exporter becomes unreachable without sending a TCP RST (e.g. network partition, firewall drop, suspended process), this probe blocks indefinitely — hanging the CLI after the shell exits.

The status monitor already handles this correctly: it detects stuck connections after 20 consecutive timeouts and sets connection_lost = True, logging progressive warnings to the user. However, the probe runs as an independent gRPC call and never observes the monitor's flag.

Fix

Race the probe against a lightweight watcher task that polls monitor.connection_lost every second. When the monitor marks the connection as lost, the watcher cancels the probe's task group scope, unblocking the CLI. The probe succeeds normally in the healthy case and cancels the watcher itself.

This preserves the existing monitor behaviour — users still see the progressive timeout warnings before the CLI exits cleanly.

Test

Added test_probe_cancelled_when_connection_lost to TestRunShellWithLeaseAsync:

• Verifies the probe stays blocked while connection_lost is False (5s negative phase)
• Verifies the probe is cancelled once connection_lost becomes True (positive phase)
• Guarded by fail_after(10) to fail hard instead of hanging if the fix regresses

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 8622e12f-b104-4bcd-ab8c-284518f0fa6d

📥 Commits

Reviewing files that changed from the base of the PR and between 0d2f24f and 58bf315.

📒 Files selected for processing (2)

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

🚧 Files skipped from review as they are similar to previous changes (1)

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

---

📝 Walkthrough

Walkthrough

The shell’s post-run status probe now stops when connection loss is detected and marks the monitor accordingly if no probe result is returned. Tests cover both normal probe completion and cancellation during a blocked probe.

Changes

Connection probe cancellation

Layer / File(s)	Summary
Cancellation-aware probe implementation python/packages/jumpstarter-cli/jumpstarter_cli/shell.py	Introduces _cancel_if_connection_lost() and uses it in _run_shell_with_lease_async to cancel the status probe on connection loss and treat a missing probe result as connection-lost.
Probe cancellation test coverage python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py	Adds math and the helper import, plus async tests for normal helper completion and for _run_shell_with_lease_async exiting when the probe is cancelled by connection loss.

Estimated code review effort: 2 (Simple) | ~12 minutes

Possibly related PRs

• jumpstarter-dev/jumpstarter#603: Both PRs modify _run_shell_with_lease_async's post-shell probe and the decision to continue or skip the afterLease flow based on probe outcome.

Suggested reviewers: mangelajo, raballew

Poem

A rabbit peeks as signals fade,
The probe stops gently, connection-wade.
If tunnels close before the end,
The watcher helps the task unbend,
Hop-safe tests now make the path well-laid. 🐇

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: unblocking the post-shell status probe when connection is lost.
Description check	✅ Passed	The description matches the changeset and explains the hang fix, implementation, and tests.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

python/packages/jumpstarter-cli/jumpstarter_cli/shell.py (1)

355-391: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider extracting the cancellation-aware probe into a helper.

_run_shell_with_lease_async is already flagged # noqa: C901; the new task-group + nested closure adds another layer of nesting/branching to an already large function. Pulling this block into a small helper (e.g., _probe_status_or_mark_lost(client, monitor) returning the resolved status or None) would keep the cancellation logic testable in isolation and reduce the parent function's complexity.

♻️ Sketch of extracted helper

+async def _probe_status_cancel_on_loss(client, monitor):
+    """Probe status, cancelling if the monitor detects connection loss."""
+    probe_status = None
+    async with anyio.create_task_group() as probe_tg:
+
+        async def cancel_when_connection_lost():
+            while not monitor.connection_lost:
+                await anyio.sleep(1)
+            probe_tg.cancel_scope.cancel()
+
+        probe_tg.start_soon(cancel_when_connection_lost)
+        probe_status = await client.get_status_async()
+        probe_tg.cancel_scope.cancel()
+    return probe_status

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/packages/jumpstarter-cli/jumpstarter_cli/shell.py` around lines 355 -
391, The cancellation-aware connection probe in _run_shell_with_lease_async is
adding too much nesting and branching to an already complex function. Extract
the task-group probe logic, including cancel_when_connection_lost and the
probe_status handling, into a small helper such as
_probe_status_or_mark_lost(client, monitor) that returns the resolved status or
None. Keep the lease.lease_ended and unexpected ExporterStatus handling in the
caller, and have the helper own the cancellation and connection_lost marking so
the probe logic is easier to test and the parent function stays simpler.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@python/packages/jumpstarter-cli/jumpstarter_cli/shell.py`:
- Around line 355-391: The cancellation-aware connection probe in
_run_shell_with_lease_async is adding too much nesting and branching to an
already complex function. Extract the task-group probe logic, including
cancel_when_connection_lost and the probe_status handling, into a small helper
such as _probe_status_or_mark_lost(client, monitor) that returns the resolved
status or None. Keep the lease.lease_ended and unexpected ExporterStatus
handling in the caller, and have the helper own the cancellation and
connection_lost marking so the probe logic is easier to test and the parent
function stays simpler.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 2a44bcc5-3c1b-49cf-9a69-8cee40731e87

📥 Commits

Reviewing files that changed from the base of the PR and between d9705b7 and 6c4ea12.

📒 Files selected for processing (2)

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