address comments

Author: michaelshin

centml/cli/shell.py

@@ -9,18 +9,47 @@
 from centml.sdk import auth
 from centml.sdk.api import get_centml_client
 from centml.sdk.config import settings
-from centml.sdk.shell import ShellError, build_ws_url, exec_session, interactive_session, resolve_pod
+from centml.sdk.shell import PodNotFoundError, ShellError, build_ws_url, exec_session, get_running_pods, interactive_session
 
 
-def _connect_args(deployment_id, pod, shell_type):
+def _resolve_pod(running_pods: list[str], pod_name: str) -> str:
+    """Validate that *pod_name* exists in *running_pods*."""
+    if pod_name not in running_pods:
+        pods_list = ", ".join(running_pods)
+        raise PodNotFoundError(f"Pod '{pod_name}' not found. Available running pods: {pods_list}")
+    return pod_name
+
+
+
+def _select_pod(running_pods, deployment_id):
+    click.echo(f"Multiple running pods found for deployment {deployment_id}:")
+    for i, name in enumerate(running_pods, 1):
+        click.echo(f"  [{i}] {name}")
+
+    choice = click.prompt(
+        "Select a pod",
+        type=click.IntRange(1, len(running_pods)),
+        prompt_suffix=f" [1-{len(running_pods)}]: ",
+    )
+    return running_pods[choice - 1]
+
+
+def _connect_args(deployment_id, pod, shell_type, first_pod=False):
     """Resolve pod, build WebSocket URL, and obtain auth token."""
     with get_centml_client() as cclient:
-        try:
-            pod_name, warning = resolve_pod(cclient, deployment_id, pod)
-        except ShellError as exc:
-            raise click.ClickException(str(exc)) from exc
-    if warning is not None:
-        click.echo(f"{warning} Use --pod to specify a different pod.", err=True)
+        running_pods = get_running_pods(cclient, deployment_id)
+        if not running_pods:
+            raise click.ClickException(f"No running pods found for deployment {deployment_id}")
+
+        if pod is not None:
+            try:
+                pod_name = _resolve_pod(running_pods, pod)
+            except ShellError as exc:
+                raise click.ClickException(str(exc)) from exc
+        elif len(running_pods) == 1 or first_pod:
+            pod_name = running_pods[0]
+        else:
+            pod_name = _select_pod(running_pods, deployment_id)
 
     ws_url = build_ws_url(settings.CENTML_PLATFORM_API_URL, deployment_id, pod_name, shell_type)
     token = auth.get_centml_token()
@@ -29,14 +58,15 @@ def _connect_args(deployment_id, pod, shell_type):
 
 @click.command(help="Open an interactive shell to a deployment pod")
 @click.argument("deployment_id", type=int)
-@click.option("--pod", default=None, help="Specify a pod name (Default: auto-selects first running pod)")
+@click.option("--pod", default=None, help="Specify a pod name")
 @click.option("--shell", "shell_type", default=None, type=click.Choice(["bash", "sh", "zsh"]), help="Shell type")
+@click.option("--first-pod", is_flag=True, default=False, help="Auto-select the first running pod (skip interactive selection)")
 @handle_exception
-def shell(deployment_id, pod, shell_type):
+def shell(deployment_id, pod, shell_type, first_pod):
     if not sys.stdin.isatty():
         raise click.ClickException("Interactive shell requires a terminal (TTY)")
 
-    ws_url, token = _connect_args(deployment_id, pod, shell_type)
+    ws_url, token = _connect_args(deployment_id, pod, shell_type, first_pod)
     exit_code = asyncio.run(interactive_session(ws_url, token))
     sys.exit(exit_code)
 
@@ -46,8 +76,9 @@ def shell(deployment_id, pod, shell_type):
 @click.argument("command", nargs=-1, required=True, type=click.UNPROCESSED)
 @click.option("--pod", default=None, help="Specific pod name")
 @click.option("--shell", "shell_type", default=None, type=click.Choice(["bash", "sh", "zsh"]), help="Shell type")
+@click.option("--first-pod", is_flag=True, default=False, help="Auto-select the first running pod (skip interactive selection)")
 @handle_exception
-def exec_cmd(deployment_id, command, pod, shell_type):
-    ws_url, token = _connect_args(deployment_id, pod, shell_type)
+def exec_cmd(deployment_id, command, pod, shell_type, first_pod):
+    ws_url, token = _connect_args(deployment_id, pod, shell_type, first_pod)
     exit_code = asyncio.run(exec_session(ws_url, token, " ".join(command)))
     sys.exit(exit_code)

centml/sdk/shell/__init__.py

@@ -1,12 +1,12 @@
 from centml.sdk.shell.exceptions import NoPodAvailableError, PodNotFoundError, ShellError
-from centml.sdk.shell.session import build_ws_url, exec_session, interactive_session, resolve_pod
+from centml.sdk.shell.session import build_ws_url, exec_session, get_running_pods, interactive_session
 
 __all__ = [
     "ShellError",
     "NoPodAvailableError",
     "PodNotFoundError",
     "build_ws_url",
-    "resolve_pod",
+    "get_running_pods",
     "interactive_session",
     "exec_session",
 ]

centml/sdk/shell/session.py

@@ -6,13 +6,11 @@
 import termios
 import tty
 import urllib.parse
-from typing import Optional, Tuple
 
 import pyte
 import websockets
 
 from centml.sdk import PodStatus
-from centml.sdk.shell.exceptions import NoPodAvailableError, PodNotFoundError
 
 # exec_session wraps commands between BEGIN/END markers so it can separate
 # real command output from shell noise (prompt, echoed input, login banners).
@@ -35,27 +33,14 @@ def build_ws_url(api_url, deployment_id, pod_name, shell_type=None):
     return url
 
 
-def resolve_pod(cclient, deployment_id, pod_name=None) -> Tuple[str, Optional[str]]:
+def get_running_pods(cclient, deployment_id) -> list[str]:
     status = cclient.get_status_v3(deployment_id)
     running_pods = []
     for revision in status.revision_pod_details_list or []:
         for pod in revision.pod_details_list or []:
             if pod.status == PodStatus.RUNNING and pod.name:
                 running_pods.append(pod.name)
-
-    if not running_pods:
-        raise NoPodAvailableError(f"No running pods found for deployment {deployment_id}")
-
-    if pod_name is not None:
-        if pod_name not in running_pods:
-            pods_list = ", ".join(running_pods)
-            raise PodNotFoundError(f"Pod '{pod_name}' not found. Available running pods: {pods_list}")
-        return pod_name, None
-
-    warning = None
-    if len(running_pods) > 1:
-        warning = f"Multiple running pods found, connecting to {running_pods[0]}."
-    return running_pods[0], warning
+    return running_pods
 
 
 async def forward_io(ws, term_size, shutdown):