add interactive picker to prime sandbox ssh (#689)

Author: kevinjosethomas

packages/prime/src/prime_cli/commands/sandbox.py

@@ -41,14 +41,15 @@
     obfuscate_env_vars,
     obfuscate_secrets,
     output_data_as_json,
+    require_selection,
     sort_by_created,
     status_color,
     validate_output_format,
 )
 from ..utils.display import SANDBOX_STATUS_COLORS
 from ..utils.time_utils import now_utc, to_utc
 
-app = PlainTyper(help="Manage code sandboxes", no_args_is_help=True)
+app = PlainTyper(help="Manage sandboxes", no_args_is_help=True)
 console = get_console()
 
 
@@ -195,6 +196,43 @@ def _guard_vm_unsupported(sandbox: Sandbox, feature_name: str) -> None:
         raise typer.Exit(1)
 
 
+def _ssh_sandbox_display(item: Dict[str, Any]) -> str:
+    """Format a sandbox row for the interactive SSH picker."""
+    return f"{item['id']}  {item['name']}  [dim]({item['image']})[/dim]"
+
+
+def _select_sandbox_for_ssh(sandbox_client: SandboxClient) -> str:
+    """Let the user pick a running sandbox to SSH into when no ID is given."""
+    # Page through every running sandbox before filtering, so SSH-able containers
+    # on later pages aren't dropped when a user has many running sandboxes.
+    sandboxes: List[Sandbox] = []
+    with console.status("[bold blue]Loading sandboxes...", spinner="dots"):
+        page = 1
+        while True:
+            sandbox_list = sandbox_client.list(status="RUNNING", per_page=100, page=page)
+            sandboxes.extend(sandbox_list.sandboxes)
+            if not sandbox_list.has_next:
+                break
+            page += 1
+
+    # SSH is only supported for non-VM sandboxes (see _guard_vm_unsupported).
+    items = [
+        {"id": sb.id, "name": sb.name, "image": sb.docker_image}
+        for sb in sort_by_created(sandboxes)
+        if not sb.vm
+    ]
+
+    selected = require_selection(
+        items,
+        "SSH into",
+        "No running sandboxes available to SSH into.",
+        item_type="sandbox",
+        display_fn=_ssh_sandbox_display,
+        page_size=50,
+    )
+    return str(selected.get("id"))
+
+
 @app.command("list", epilog=LIST_SANDBOXES_JSON_HELP)
 def list_sandboxes_cmd(
     team_id: Optional[str] = typer.Option(
@@ -236,7 +274,7 @@ def list_sandboxes_cmd(
         )
 
         table = build_table(
-            f"Code Sandboxes (Total: {sandbox_list.total})",
+            f"Sandboxes (Total: {sandbox_list.total})",
             [
                 ("ID", "cyan"),
                 ("Name", "blue"),
@@ -1443,9 +1481,11 @@ def list_ports(
         raise typer.Exit(1)
 
 
-@app.command("ssh", no_args_is_help=True)
+@app.command("ssh")
 def ssh_connect(
-    sandbox_id: str = typer.Argument(..., help="Sandbox ID to SSH into"),
+    sandbox_id: Optional[str] = typer.Argument(
+        None, help="Sandbox ID to SSH into (interactive selection if not provided)"
+    ),
     ssh_args: Optional[List[str]] = typer.Argument(
         None, help="Additional SSH arguments (e.g., -- -v for verbose)"
     ),
@@ -1459,9 +1499,11 @@ def ssh_connect(
     """Connect to a sandbox via SSH.
 
     This command creates a SSH session with an ephemeral key and cleans up on disconnect.
+    Run without a sandbox ID to pick from your running sandboxes interactively.
 
     \b
     Examples:
+        prime sandbox ssh
         prime sandbox ssh sb_abc123
         prime sandbox ssh sb_abc123 --shell bash
         prime sandbox ssh sb_abc123 -- -L 3000:localhost:3000
@@ -1473,7 +1515,7 @@ def ssh_connect(
 
     def cleanup() -> None:
         """Clean up the SSH session and temporary keys."""
-        if session_id and sandbox_client:
+        if session_id and sandbox_client and sandbox_id:
             try:
                 console.print("\n[bold blue]Cleaning up SSH session...[/bold blue]")
                 sandbox_client.close_ssh_session(sandbox_id, session_id)
@@ -1495,6 +1537,10 @@ def cleanup() -> None:
         base_client = APIClient()
         sandbox_client = SandboxClient(base_client)
 
+        # Pick a sandbox interactively when no ID was provided
+        if sandbox_id is None:
+            sandbox_id = _select_sandbox_for_ssh(sandbox_client)
+
         # Check if sandbox is running
         with console.status("[bold blue]Checking sandbox status...", spinner="dots"):
             sandbox = sandbox_client.get(sandbox_id)

packages/prime/src/prime_cli/utils/prompt.py

@@ -2,6 +2,7 @@
 from typing import Any, Callable, Dict, List, Optional
 
 import typer
+from rich.markup import escape
 
 from .plain import get_console
 
@@ -46,12 +47,15 @@ def require_selection(
     empty_message: str,
     item_type: str = "item",
     display_fn: Optional[Callable[[Dict[str, Any]], str]] = None,
+    page_size: Optional[int] = None,
 ) -> Dict[str, Any]:
     if not items:
         console.print(f"[yellow]{empty_message}[/yellow]")
         raise typer.Exit()
 
-    selected = select_item_interactive(items, action, item_type=item_type, display_fn=display_fn)
+    selected = select_item_interactive(
+        items, action, item_type=item_type, display_fn=display_fn, page_size=page_size
+    )
     if not selected:
         console.print("\n[dim]Cancelled.[/dim]")
         raise typer.Exit()
@@ -64,6 +68,7 @@ def select_item_interactive(
     action: str = "select",
     item_type: str = "item",
     display_fn: Optional[Callable[[Dict[str, Any]], str]] = None,
+    page_size: Optional[int] = None,
 ) -> Optional[Dict[str, Any]]:
     """Display items and let user select one interactively.
 
@@ -73,6 +78,8 @@ def select_item_interactive(
         item_type: Type of item being selected (e.g., "secret", "variable")
         display_fn: Function to format each item for display.
                    Defaults to showing 'name' and 'description' fields.
+        page_size: When set and there are more items than this, show them
+                   page_size at a time with [n]ext / [p]rev navigation.
 
     Returns:
         Selected item or None if cancelled
@@ -82,6 +89,9 @@ def select_item_interactive(
 
     formatter = display_fn or _default_display_fn
 
+    if page_size is not None and len(items) > page_size:
+        return _select_item_paged(items, action, item_type, formatter, page_size)
+
     console.print(f"\n[bold]Select a {item_type} to {action}:[/bold]\n")
     for i, item in enumerate(items, 1):
         console.print(f"  {i}. {formatter(item)}")
@@ -102,6 +112,72 @@ def select_item_interactive(
             return None
 
 
+def _select_item_paged(
+    items: List[Dict[str, Any]],
+    action: str,
+    item_type: str,
+    formatter: Callable[[Dict[str, Any]], str],
+    page_size: int,
+) -> Optional[Dict[str, Any]]:
+    """Paged variant of the interactive selector with next/prev navigation.
+
+    Selection numbers are global (1..len(items)) — the number printed next to an
+    item is what you type, on any page.
+    """
+    total = len(items)
+    total_pages = (total + page_size - 1) // page_size
+    page = 0
+
+    while True:
+        start = page * page_size
+        end = min(start + page_size, total)
+
+        console.print(
+            f"\n[bold]Select a {item_type} to {action} "
+            f"(page {page + 1}/{total_pages}):[/bold]\n"
+        )
+        for i in range(start, end):
+            console.print(f"  {i + 1}. {formatter(items[i])}")
+
+        nav = []
+        if page + 1 < total_pages:
+            nav.append("[n] next page")
+        if page > 0:
+            nav.append("[p] prev page")
+        console.print(f"\n[dim]{escape('  '.join(nav))}[/dim]\n")
+
+        try:
+            choice = typer.prompt("Select (empty to cancel)", default="").strip()
+        except KeyboardInterrupt:
+            return None
+
+        if not choice:
+            return None
+
+        lowered = choice.lower()
+        if lowered == "n":
+            if page + 1 < total_pages:
+                page += 1
+            else:
+                console.print("[red]Already on the last page[/red]")
+            continue
+        if lowered == "p":
+            if page > 0:
+                page -= 1
+            else:
+                console.print("[red]Already on the first page[/red]")
+            continue
+
+        try:
+            selection = int(choice)
+        except ValueError:
+            console.print("[red]Please enter a number, or 'n'/'p' to navigate[/red]")
+            continue
+        if 1 <= selection <= total:
+            return items[selection - 1]
+        console.print(f"[red]Please enter a number between 1 and {total}[/red]")
+
+
 def prompt_for_value(
     prompt_text: str,
     required: bool = True,