refactor: use pyte for exec ANSI stripping and add ConnectionClosed handling

Replace regex-based _strip_ansi with pyte single-row screen for marker
detection. pyte interprets all VT100/VT220 sequences including OSC and
truecolor escapes that the regex could miss.

Author: V2arK

centml/cli/shell.py

@@ -2,7 +2,6 @@
 
 import asyncio
 import json
-import re
 import shutil
 import signal
 import sys
@@ -337,8 +336,6 @@ def _send_resize():
         sys.stdout.buffer.flush()
 
 
-_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;?]*[a-zA-Z]|\x1b\].*?\x07|\x1b\[.*?\x1b\\")
-
 _BEGIN_MARKER = "__CENTML_BEGIN_5f3a__"
 _END_MARKER = "__CENTML_END_5f3a__"
 
@@ -348,9 +345,17 @@ def _send_resize():
 _PRINTF_END = r"\137\137CENTML_END_5f3a\137\137"
 
 
-def _strip_ansi(text):
-    """Remove ANSI escape sequences from text."""
-    return _ANSI_ESCAPE_RE.sub("", text)
+def _pyte_extract_text(line_stream, line_screen, text):
+    """Feed text through a single-row pyte screen and return visible characters.
+
+    More robust than regex ANSI stripping: pyte interprets all VT100/VT220
+    sequences including OSC, cursor repositioning, and truecolor escapes.
+    """
+    line_screen.reset()
+    line_stream.feed(text)
+    return "".join(
+        line_screen.buffer[0][col].data for col in range(line_screen.columns)
+    ).rstrip()
 
 
 async def _exec_session(ws_url, token, command):
@@ -360,6 +365,9 @@ async def _exec_session(ws_url, token, command):
     Suppresses shell echo and uses markers to capture only command output.
     """
     cols, rows = shutil.get_terminal_size(fallback=(80, 24))
+    # Single-row screen for interpreting escape sequences in marker detection.
+    line_screen = pyte.Screen(cols, 1)
+    line_stream = pyte.Stream(line_screen)
     headers = {"Authorization": f"Bearer {token}"}
 
     async with websockets.connect(
@@ -386,35 +394,40 @@ async def _exec_session(ws_url, token, command):
         buffer = ""
         is_capturing = False
         is_done = False
-        async for raw_msg in ws:
-            msg = json.loads(raw_msg)
-            if msg.get("data"):
-                buffer += msg["data"]
-                while "\n" in buffer:
-                    line, buffer = buffer.split("\n", 1)
-                    clean = _strip_ansi(line).rstrip("\r")
-                    if _BEGIN_MARKER in clean:
-                        is_capturing = True
-                        continue
-                    if _END_MARKER in clean:
-                        parts = clean.split(_END_MARKER + ":")
-                        if len(parts) > 1:
-                            try:
-                                exit_code = int(parts[1].strip())
-                            except ValueError:
-                                pass
-                        is_done = True
-                        break
-                    if is_capturing:
-                        sys.stdout.write(line + "\n")
-                        sys.stdout.flush()
-            elif msg.get("error"):
-                sys.stderr.write(f"Error: {msg['error']}\n")
-                return 1
-            if is_done or "Code" in msg:
-                if "Code" in msg:
-                    exit_code = msg["Code"]
-                break
+        try:
+            async for raw_msg in ws:
+                msg = json.loads(raw_msg)
+                if msg.get("data"):
+                    buffer += msg["data"]
+                    while "\n" in buffer:
+                        line, buffer = buffer.split("\n", 1)
+                        clean = _pyte_extract_text(
+                            line_stream, line_screen, line.rstrip("\r")
+                        )
+                        if _BEGIN_MARKER in clean:
+                            is_capturing = True
+                            continue
+                        if _END_MARKER in clean:
+                            parts = clean.split(_END_MARKER + ":")
+                            if len(parts) > 1:
+                                try:
+                                    exit_code = int(parts[1].strip())
+                                except ValueError:
+                                    pass
+                            is_done = True
+                            break
+                        if is_capturing:
+                            sys.stdout.write(line + "\n")
+                            sys.stdout.flush()
+                elif msg.get("error"):
+                    sys.stderr.write(f"Error: {msg['error']}\n")
+                    return 1
+                if is_done or "Code" in msg:
+                    if "Code" in msg:
+                        exit_code = msg["Code"]
+                    break
+        except websockets.ConnectionClosed:
+            pass
         return exit_code
 
 

tests/test_shell.py

@@ -52,18 +52,6 @@ def _make_status_response(revisions):
 # ===========================================================================
 
 
-class TestStripAnsi:
-    def test_strips_csi_sequences(self):
-        from centml.cli.shell import _strip_ansi
-
-        assert _strip_ansi("\x1b[?2004htext\x1b[0m") == "text"
-
-    def test_preserves_plain_text(self):
-        from centml.cli.shell import _strip_ansi
-
-        assert _strip_ansi("hello world") == "hello world"
-
-
 class TestBuildWsUrl:
     def test_https_to_wss(self):
         from centml.cli.shell import _build_ws_url
@@ -339,6 +327,67 @@ def test_filters_noise_before_marker(self):
         assert "real output" in output
         assert "prompt$" not in output
 
+    def test_connection_closed_returns_zero(self):
+        """Graceful exit when server closes connection without Code message."""
+        from centml.cli.shell import _exec_session
+
+        import websockets as _ws_lib
+
+        ws = AsyncMock()
+
+        async def _raise_closed():
+            yield json.dumps({"data": "partial\n"})
+            raise _ws_lib.ConnectionClosed(None, None)
+
+        ws.__aiter__ = MagicMock(return_value=_raise_closed())
+
+        with patch("centml.cli.shell.websockets") as mock_ws_mod:
+            mock_ws_mod.connect = MagicMock(
+                return_value=AsyncMock(
+                    __aenter__=AsyncMock(return_value=ws),
+                    __aexit__=AsyncMock(return_value=False),
+                )
+            )
+            mock_ws_mod.ConnectionClosed = _ws_lib.ConnectionClosed
+
+            exit_code = asyncio.run(
+                _exec_session("wss://test/ws", "fake-token", "exit")
+            )
+
+        assert exit_code == 0
+
+    def test_handles_ansi_around_markers(self):
+        """Markers wrapped in ANSI codes are still detected via pyte."""
+        from centml.cli.shell import _exec_session, _BEGIN_MARKER, _END_MARKER
+
+        ws = AsyncMock()
+        # Markers surrounded by ANSI color codes.
+        data = f"\x1b[32m{_BEGIN_MARKER}\x1b[0m\noutput\n\x1b[32m{_END_MARKER}:0\x1b[0m\n"
+        messages = [json.dumps({"data": data}), json.dumps({"Code": 0})]
+        ws.__aiter__ = MagicMock(return_value=_async_iter_from_list(messages))
+
+        captured = []
+        with patch("centml.cli.shell.websockets") as mock_ws_mod, patch(
+            "centml.cli.shell.sys"
+        ) as mock_sys:
+            mock_ws_mod.connect = MagicMock(
+                return_value=AsyncMock(
+                    __aenter__=AsyncMock(return_value=ws),
+                    __aexit__=AsyncMock(return_value=False),
+                )
+            )
+            mock_sys.stdout.write = lambda s: captured.append(s)
+            mock_sys.stdout.flush = MagicMock()
+            mock_sys.stderr.write = MagicMock()
+
+            exit_code = asyncio.run(
+                _exec_session("wss://test/ws", "fake-token", "echo test")
+            )
+
+        assert exit_code == 0
+        output = "".join(captured)
+        assert "output" in output
+
 
 # ===========================================================================
 # _interactive_session -- terminal restore