ControlModeEngine(fix[robustness]): Harden control mode I/O

Author: tony

src/libtmux/_internal/engines/control_mode.py

@@ -2,11 +2,14 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
+import queue
 import shlex
 import shutil
 import subprocess
 import threading
+import time
 import typing as t
 
 from libtmux import exc
@@ -24,6 +27,10 @@ def __init__(self, command_timeout: float | None = 10.0) -> None:
         self._lock = threading.Lock()
         self._server_args: tuple[str | int, ...] | None = None
         self.command_timeout = command_timeout
+        self._queue: queue.Queue[str | None] = queue.Queue(maxsize=1024)
+        self._reader_thread: threading.Thread | None = None
+        self._stderr_thread: threading.Thread | None = None
+        self.tmux_bin: str | None = None
 
     def close(self) -> None:
         """Terminate the tmux control mode process."""
@@ -36,16 +43,53 @@ def close(self) -> None:
                 self.process.wait()
             self.process = None
             self._server_args = None
+            # Unblock any waiters
+            with contextlib.suppress(queue.Full):
+                self._queue.put_nowait(None)
 
     def __del__(self) -> None:
         """Cleanup the process on destruction."""
         self.close()
 
+    def _reader(self, process: subprocess.Popen[str]) -> None:
+        """Background thread to read stdout and push to queue."""
+        assert process.stdout is not None
+        while True:
+            try:
+                line = process.stdout.readline()
+            except (ValueError, OSError):
+                break
+
+            if not line:
+                # EOF
+                with contextlib.suppress(queue.Full):
+                    self._queue.put_nowait(None)
+                break
+
+            if line.startswith("%") and not line.startswith(
+                ("%begin", "%end", "%error"),
+            ):
+                logger.debug("Control Mode Notification: %s", line.rstrip("\n"))
+                continue
+
+            try:
+                self._queue.put(line, timeout=1)
+            except queue.Full:
+                logger.warning("Control Mode queue full; dropping line")
+
+    def _drain_stderr(self, process: subprocess.Popen[str]) -> None:
+        """Continuously drain stderr to prevent child blocking."""
+        if process.stderr is None:
+            return
+        for err_line in process.stderr:
+            logger.debug("Control Mode stderr: %s", err_line.rstrip("\n"))
+
     def _start_process(self, server_args: t.Sequence[str | int] | None) -> None:
         """Start the tmux control mode process."""
         tmux_bin = shutil.which("tmux")
         if not tmux_bin:
             raise exc.TmuxCommandNotFound
+        self.tmux_bin = tmux_bin
 
         normalized_args: tuple[str | int, ...] = tuple(server_args or ())
 
@@ -67,31 +111,95 @@ def _start_process(self, server_args: t.Sequence[str | int] | None) -> None:
         )
         self._server_args = normalized_args
 
+        # reset queues from prior runs
+        while not self._queue.empty():
+            self._queue.get_nowait()
+
+        # Start IO threads
+        self._reader_thread = threading.Thread(
+            target=self._reader,
+            args=(self.process,),
+            daemon=True,
+        )
+        self._reader_thread.start()
+
+        self._stderr_thread = threading.Thread(
+            target=self._drain_stderr,
+            args=(self.process,),
+            daemon=True,
+        )
+        self._stderr_thread.start()
+
         # Consume startup command output with the same timeout used for commands.
-        assert self.process.stdout is not None
+        try:
+            self._read_response(cmd="new-session", timeout=self.command_timeout)
+        except exc.ControlModeTimeout:
+            logger.exception("Control Mode bootstrap command timed out")
+            self.close()
+            raise
+
+    def _read_response(self, cmd: str, timeout: float | None) -> tmux_cmd:
+        """Read response from the queue."""
+        stdout_lines: list[str] = []
+        stderr_lines: list[str] = []
+        returncode = 0
+
+        start_time = time.monotonic()
+
+        while True:
+            if timeout is None:
+                remaining: float | None = None
+            else:
+                elapsed = time.monotonic() - start_time
+                remaining = timeout - elapsed
+                if remaining <= 0:
+                    msg = "tmux control mode command timed out"
+                    raise exc.ControlModeTimeout(msg)
 
-        def bootstrap_reader() -> None:
-            assert self.process is not None
-            assert self.process.stdout is not None
-            while True:
-                line = self.process.stdout.readline()
-                if not line:
-                    logger.warning("Control Mode process exited immediately")
-                    break
-                if line.startswith("%end") or line.startswith("%error"):
-                    break
-
-        if self.command_timeout is None:
-            bootstrap_reader()
-        else:
-            bootstrap_thread = threading.Thread(target=bootstrap_reader, daemon=True)
-            bootstrap_thread.start()
-            bootstrap_thread.join(timeout=self.command_timeout)
-            if bootstrap_thread.is_alive():
-                logger.error("Control Mode bootstrap command timed out")
-                self.close()
+            try:
+                line = self._queue.get(timeout=remaining)
+            except queue.Empty:
                 msg = "tmux control mode command timed out"
-                raise exc.ControlModeTimeout(msg)
+                raise exc.ControlModeTimeout(msg) from None
+
+            if line is None:
+                logger.error("Unexpected EOF from Control Mode process")
+                returncode = 1
+                break
+
+            line = line.rstrip("\n")
+
+            if line.startswith("%begin"):
+                parts = line.split()
+                if len(parts) > 3:
+                    flags = int(parts[3])
+                    if flags & 1:
+                        returncode = 1
+                continue
+            if line.startswith("%end"):
+                parts = line.split()
+                if len(parts) > 3:
+                    flags = int(parts[3])
+                    if flags & 1:
+                        returncode = 1
+                break
+            if line.startswith("%error"):
+                returncode = 1
+                stderr_lines = stdout_lines
+                stdout_lines = []
+                break
+
+            stdout_lines.append(line)
+
+        if cmd == "has-session" and stderr_lines and not stdout_lines:
+            stdout_lines = [stderr_lines[0]]
+
+        return tmux_cmd(
+            cmd=[],
+            stdout=stdout_lines,
+            stderr=stderr_lines,
+            returncode=returncode,
+        )
 
     def run(
         self,
@@ -120,11 +228,7 @@ def run(
 
             assert self.process is not None
             assert self.process.stdin is not None
-            assert self.process.stdout is not None
-
             # Construct the command line
-            # We use shlex.join for correct shell-like quoting, required by tmux control
-            # mode.
             full_args = [cmd]
             if cmd_args:
                 full_args.extend(str(a) for a in cmd_args)
@@ -146,89 +250,19 @@ def run(
                 self.process.stdin.write(command_line + "\n")
                 self.process.stdin.flush()
 
-            # Read response
-            stdout_lines: list[str] = []
-            stderr_lines: list[str] = []
-            returncode = 0
+            effective_timeout = timeout if timeout is not None else self.command_timeout
 
-            def reader() -> None:
-                nonlocal returncode, stdout_lines, stderr_lines
-                assert self.process is not None
-                assert self.process.stdout is not None
-                while True:
-                    line = self.process.stdout.readline()
-                    if not line:
-                        # EOF
-                        logger.error("Unexpected EOF from Control Mode process")
-                        returncode = 1
-                        break
-
-                    line = line.rstrip("\n")
-
-                    if line.startswith("%begin"):
-                        # %begin time id flags
-                        parts = line.split()
-                        if len(parts) > 3:
-                            flags = int(parts[3])
-                            if flags & 1:
-                                returncode = 1
-                        continue
-                    elif line.startswith("%end"):
-                        # %end time id flags
-                        parts = line.split()
-                        if len(parts) > 3:
-                            flags = int(parts[3])
-                            if flags & 1:
-                                returncode = 1
-                        break
-                    elif line.startswith("%error"):
-                        returncode = 1
-                        stderr_lines = stdout_lines
-                        stdout_lines = []
-                        break
-                    elif line.startswith("%"):
-                        logger.debug("Control Mode Notification: %s", line)
-                        continue
-                    else:
-                        stdout_lines.append(line)
-
-            reader_exc: list[BaseException] = []
-
-            def wrapped_reader() -> None:
-                try:
-                    reader()
-                except BaseException as read_exc:  # pragma: no cover - defensive
-                    reader_exc.append(read_exc)
+            try:
+                result = self._read_response(cmd=cmd, timeout=effective_timeout)
+            except exc.ControlModeTimeout:
+                self.close()
+                raise
 
-            effective_timeout = timeout if timeout is not None else self.command_timeout
-            if effective_timeout is None:
-                wrapped_reader()
-            else:
-                reader_thread = threading.Thread(target=wrapped_reader, daemon=True)
-                reader_thread.start()
-                reader_thread.join(timeout=effective_timeout)
-                if reader_thread.is_alive():
-                    logger.error(
-                        "Control Mode command timed out waiting for response: %s",
-                        command_line,
-                    )
-                    self.close()
-                    msg = "tmux control mode command timed out"
-                    raise exc.ControlModeTimeout(msg)
-                if reader_exc:
-                    raise reader_exc[0]
-
-            # Tmux usually puts error message in stdout (captured above) for %error
-            # But we moved it to stderr_lines if %error occurred.
-
-            # If we detected failure via flags but got %end, treat stdout as potentially
-            # containing info?
-            # For now, keep stdout as is.
-
-            # Mimic subprocess.communicate output structure
-            return tmux_cmd(
-                cmd=[cmd] + (list(map(str, cmd_args)) if cmd_args else []),
-                stdout=stdout_lines,
-                stderr=stderr_lines,
-                returncode=returncode,
-            )
+            full_cmd = [self.tmux_bin or "tmux"]
+            if incoming_server_args:
+                full_cmd.extend(str(x) for x in incoming_server_args)
+            full_cmd.append(cmd)
+            if cmd_args:
+                full_cmd.extend(str(x) for x in cmd_args)
+            result.cmd = full_cmd
+            return result