fix: include rows/cols in stdin messages and send Ctrl+L after resize toggle

Author: V2arK

centml/cli/shell.py

@@ -52,17 +52,22 @@ def _resolve_pod(cclient, deployment_id, pod_name=None):
                 running_pods.append(pod.name)
 
     if not running_pods:
-        raise click.ClickException(f"No running pods found for deployment {deployment_id}")
+        raise click.ClickException(
+            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 click.ClickException(f"Pod '{pod_name}' not found. Available running pods: {pods_list}")
+            raise click.ClickException(
+                f"Pod '{pod_name}' not found. Available running pods: {pods_list}"
+            )
         return pod_name
 
     if len(running_pods) > 1:
         click.echo(
-            f"Multiple running pods found, connecting to {running_pods[0]}. " f"Use --pod to specify a different pod.",
+            f"Multiple running pods found, connecting to {running_pods[0]}. "
+            f"Use --pod to specify a different pod.",
             err=True,
         )
     return running_pods[0]
@@ -110,8 +115,16 @@ def _on_stdin_ready():
                     data = await asyncio.wait_for(read_queue.get(), timeout=0.5)
                 except asyncio.TimeoutError:
                     continue
+                r, c = shutil.get_terminal_size()
                 await ws.send(
-                    json.dumps({"operation": "stdin", "data": data.decode("utf-8", errors="replace")})
+                    json.dumps(
+                        {
+                            "operation": "stdin",
+                            "data": data.decode("utf-8", errors="replace"),
+                            "rows": r,
+                            "cols": c,
+                        }
+                    )
                 )
         finally:
             loop.remove_reader(stdin_fd)
@@ -144,14 +157,20 @@ async def _interactive_session(ws_url, token):
         rows, cols = shutil.get_terminal_size()
 
         headers = {"Authorization": f"Bearer {token}"}
-        async with websockets.connect(ws_url, additional_headers=headers, close_timeout=2) as ws:
-            await ws.send(json.dumps({"operation": "resize", "rows": rows, "cols": cols}))
+        async with websockets.connect(
+            ws_url, additional_headers=headers, close_timeout=2
+        ) as ws:
+            await ws.send(
+                json.dumps({"operation": "resize", "rows": rows, "cols": cols})
+            )
 
             loop = asyncio.get_running_loop()
 
             def _send_resize():
                 r, c = shutil.get_terminal_size()
-                asyncio.ensure_future(ws.send(json.dumps({"operation": "resize", "rows": r, "cols": c})))
+                asyncio.ensure_future(
+                    ws.send(json.dumps({"operation": "resize", "rows": r, "cols": c}))
+                )
 
             async def _force_initial_redraw():
                 """Force SIGWINCH on the remote by toggling the PTY width.
@@ -164,9 +183,17 @@ async def _force_initial_redraw():
                 try:
                     await asyncio.sleep(0.5)
                     r, c = shutil.get_terminal_size()
-                    await ws.send(json.dumps({"operation": "resize", "rows": r, "cols": c + 1}))
+                    await ws.send(
+                        json.dumps({"operation": "resize", "rows": r, "cols": c + 1})
+                    )
                     await asyncio.sleep(0.05)
-                    await ws.send(json.dumps({"operation": "resize", "rows": r, "cols": c}))
+                    await ws.send(
+                        json.dumps({"operation": "resize", "rows": r, "cols": c})
+                    )
+                    await asyncio.sleep(0.1)
+                    # Ctrl+L clears screen and forces the shell to redraw
+                    # its prompt at the now-correct terminal width.
+                    await ws.send(json.dumps({"operation": "stdin", "data": "\x0c"}))
                 except Exception:
                     pass
 
@@ -208,7 +235,9 @@ async def _exec_session(ws_url, token, command):
     rows, cols = shutil.get_terminal_size(fallback=(80, 24))
     headers = {"Authorization": f"Bearer {token}"}
 
-    async with websockets.connect(ws_url, additional_headers=headers, close_timeout=2) as ws:
+    async with websockets.connect(
+        ws_url, additional_headers=headers, close_timeout=2
+    ) as ws:
         await ws.send(json.dumps({"operation": "resize", "rows": rows, "cols": cols}))
 
         # Suppress echo/bracketed-paste, emit begin marker, run command,
@@ -264,8 +293,16 @@ async def _exec_session(ws_url, token, command):
 
 @click.command(help="Open an interactive shell to a deployment pod")
 @click.argument("deployment_id", type=int)
-@click.option("--pod", default=None, help="Specific pod name (auto-selects first running pod)")
-@click.option("--shell", "shell_type", default=None, type=click.Choice(["bash", "sh", "zsh"]), help="Shell type")
+@click.option(
+    "--pod", default=None, help="Specific pod name (auto-selects first running pod)"
+)
+@click.option(
+    "--shell",
+    "shell_type",
+    default=None,
+    type=click.Choice(["bash", "sh", "zsh"]),
+    help="Shell type",
+)
 @handle_exception
 def shell(deployment_id, pod, shell_type):
     if not sys.stdin.isatty():
@@ -274,23 +311,36 @@ def shell(deployment_id, pod, shell_type):
     with get_centml_client() as cclient:
         pod_name = _resolve_pod(cclient, deployment_id, pod)
 
-    ws_url = _build_ws_url(settings.CENTML_PLATFORM_API_URL, deployment_id, pod_name, shell_type)
+    ws_url = _build_ws_url(
+        settings.CENTML_PLATFORM_API_URL, deployment_id, pod_name, shell_type
+    )
     token = auth.get_centml_token()
     exit_code = asyncio.run(_interactive_session(ws_url, token))
     sys.exit(exit_code)
 
 
-@click.command(help="Execute a command in a deployment pod", context_settings={"ignore_unknown_options": True})
+@click.command(
+    help="Execute a command in a deployment pod",
+    context_settings={"ignore_unknown_options": True},
+)
 @click.argument("deployment_id", type=int)
 @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(
+    "--shell",
+    "shell_type",
+    default=None,
+    type=click.Choice(["bash", "sh", "zsh"]),
+    help="Shell type",
+)
 @handle_exception
 def exec_cmd(deployment_id, command, pod, shell_type):
     with get_centml_client() as cclient:
         pod_name = _resolve_pod(cclient, deployment_id, pod)
 
-    ws_url = _build_ws_url(settings.CENTML_PLATFORM_API_URL, deployment_id, pod_name, shell_type)
+    ws_url = _build_ws_url(
+        settings.CENTML_PLATFORM_API_URL, deployment_id, pod_name, shell_type
+    )
     token = auth.get_centml_token()
     cmd_str = " ".join(command)
     exit_code = asyncio.run(_exec_session(ws_url, token, cmd_str))