wip: add persistent shell to ssh and local terminal backends

tools/environments/persistent_shell.py

@@ -0,0 +1,308 @@
+"""Persistent shell mixin: file-based IPC protocol for long-lived bash shells.
+
+Provides the shared logic for maintaining a persistent bash shell across
+execute() calls.  Backend-specific operations (spawning the shell, reading
+temp files, killing child processes) are implemented by subclasses via
+abstract methods.
+
+The IPC protocol writes each command's stdout/stderr/exit-code/cwd to temp
+files, then polls the status file for completion.  A daemon thread drains
+the shell's stdout to prevent pipe deadlock and detect shell death.
+"""
+
+import glob as glob_mod
+import logging
+import os
+import shlex
+import subprocess
+import threading
+import time
+import uuid
+from abc import abstractmethod
+
+from tools.interrupt import is_interrupted
+
+logger = logging.getLogger(__name__)
+
+
+class PersistentShellMixin:
+    """Mixin that adds persistent shell capability to any BaseEnvironment.
+
+    Subclasses MUST implement:
+        _spawn_shell_process() -> subprocess.Popen
+        _read_temp_files(*paths) -> list[str]
+        _kill_shell_children()
+
+    Subclasses MUST also provide ``_execute_oneshot()`` for the stdin_data
+    fallback path (commands with piped stdin cannot use the persistent shell).
+    """
+
+    # -- State (initialized by _init_persistent_shell) ---------------------
+    _shell_proc: subprocess.Popen | None = None
+    _shell_alive: bool = False
+    _shell_pid: int | None = None
+    _session_id: str = ""
+
+    # -- Abstract methods (backend-specific) -------------------------------
+
+    @abstractmethod
+    def _spawn_shell_process(self) -> subprocess.Popen:
+        """Spawn a long-lived bash shell and return the Popen handle.
+
+        Must use ``stdin=PIPE, stdout=PIPE, stderr=PIPE, text=True``.
+        """
+        ...
+
+    @abstractmethod
+    def _read_temp_files(self, *paths: str) -> list[str]:
+        """Read temp files from the execution context.
+
+        Returns contents in the same order as *paths*.  Falls back to
+        empty strings on failure.
+        """
+        ...
+
+    @abstractmethod
+    def _kill_shell_children(self):
+        """Kill the running command's processes but keep the shell alive."""
+        ...
+
+    # -- Overridable properties --------------------------------------------
+
+    @property
+    def _temp_prefix(self) -> str:
+        """Base path for temp files.  Override per backend."""
+        return f"/tmp/hermes-persistent-{self._session_id}"
+
+    # -- Shared implementation ---------------------------------------------
+
+    def _init_persistent_shell(self):
+        """Call from ``__init__`` when ``persistent=True``."""
+        self._shell_lock = threading.Lock()
+        self._session_id = ""
+        self._shell_proc = None
+        self._shell_alive = False
+        self._shell_pid = None
+        self._start_persistent_shell()
+
+    def _start_persistent_shell(self):
+        """Spawn the shell, create temp files, capture PID."""
+        self._session_id = uuid.uuid4().hex[:12]
+        p = self._temp_prefix
+        self._pshell_stdout = f"{p}-stdout"
+        self._pshell_stderr = f"{p}-stderr"
+        self._pshell_status = f"{p}-status"
+        self._pshell_cwd = f"{p}-cwd"
+        self._pshell_pid_file = f"{p}-pid"
+
+        self._shell_proc = self._spawn_shell_process()
+        self._shell_alive = True
+
+        self._drain_thread = threading.Thread(
+            target=self._drain_shell_output, daemon=True,
+        )
+        self._drain_thread.start()
+
+        # Initialize temp files and capture shell PID
+        init_script = (
+            f"touch {self._pshell_stdout} {self._pshell_stderr} "
+            f"{self._pshell_status} {self._pshell_cwd} {self._pshell_pid_file}\n"
+            f"echo $$ > {self._pshell_pid_file}\n"
+            f"pwd > {self._pshell_cwd}\n"
+        )
+        self._send_to_shell(init_script)
+
+        # Poll for PID file
+        deadline = time.monotonic() + 3.0
+        while time.monotonic() < deadline:
+            pid_str = self._read_temp_files(self._pshell_pid_file)[0].strip()
+            if pid_str.isdigit():
+                self._shell_pid = int(pid_str)
+                break
+            time.sleep(0.05)
+        else:
+            logger.warning("Could not read persistent shell PID")
+            self._shell_pid = None
+
+        if self._shell_pid:
+            logger.info(
+                "Persistent shell started (session=%s, pid=%d)",
+                self._session_id, self._shell_pid,
+            )
+
+        # Update cwd from what the shell reports
+        reported_cwd = self._read_temp_files(self._pshell_cwd)[0].strip()
+        if reported_cwd:
+            self.cwd = reported_cwd
+
+    def _drain_shell_output(self):
+        """Drain stdout to prevent pipe deadlock; detect shell death."""
+        try:
+            for _ in self._shell_proc.stdout:
+                pass  # Real output goes to temp files
+        except Exception:
+            pass
+        self._shell_alive = False
+
+    def _send_to_shell(self, text: str):
+        """Write text to the persistent shell's stdin."""
+        if not self._shell_alive or self._shell_proc is None:
+            return
+        try:
+            self._shell_proc.stdin.write(text)
+            self._shell_proc.stdin.flush()
+        except (BrokenPipeError, OSError):
+            self._shell_alive = False
+
+    def _read_persistent_output(self) -> tuple[str, int, str]:
+        """Read stdout, stderr, status, cwd.  Returns (output, exit_code, cwd)."""
+        stdout, stderr, status_raw, cwd = self._read_temp_files(
+            self._pshell_stdout, self._pshell_stderr,
+            self._pshell_status, self._pshell_cwd,
+        )
+        output = self._merge_output(stdout, stderr)
+        # Status format: "cmd_id:exit_code" — strip the ID prefix
+        status = status_raw.strip()
+        if ":" in status:
+            status = status.split(":", 1)[1]
+        try:
+            exit_code = int(status.strip())
+        except ValueError:
+            exit_code = 1
+        return output, exit_code, cwd.strip()
+
+    def _execute_persistent(self, command: str, cwd: str, *,
+                            timeout: int | None = None,
+                            stdin_data: str | None = None) -> dict:
+        """Execute a command in the persistent shell."""
+        if not self._shell_alive:
+            logger.info("Persistent shell died, restarting...")
+            self._start_persistent_shell()
+
+        exec_command, sudo_stdin = self._prepare_command(command)
+        effective_timeout = timeout or self.timeout
+
+        # Fall back to one-shot for commands needing piped stdin
+        if stdin_data or sudo_stdin:
+            return self._execute_oneshot(
+                command, cwd, timeout=timeout, stdin_data=stdin_data,
+            )
+
+        with self._shell_lock:
+            return self._execute_persistent_locked(
+                exec_command, cwd, effective_timeout,
+            )
+
+    def _execute_persistent_locked(self, command: str, cwd: str,
+                                   timeout: int) -> dict:
+        """Inner persistent execution — caller must hold ``_shell_lock``."""
+        work_dir = cwd or self.cwd
+
+        # Each command gets a unique ID written into the status file so the
+        # poll loop can distinguish the *current* command's result from a
+        # stale value left over from the previous command.  This eliminates
+        # the race where a fast local file read sees the old status before
+        # the shell has processed the truncation.
+        cmd_id = uuid.uuid4().hex[:8]
+
+        # Truncate temp files
+        truncate = (
+            f": > {self._pshell_stdout}\n"
+            f": > {self._pshell_stderr}\n"
+            f": > {self._pshell_status}\n"
+        )
+        self._send_to_shell(truncate)
+
+        # Escape command for eval
+        escaped = command.replace("'", "'\\''")
+
+        ipc_script = (
+            f"cd {shlex.quote(work_dir)}\n"
+            f"eval '{escaped}' < /dev/null > {self._pshell_stdout} 2> {self._pshell_stderr}\n"
+            f"__EC=$?\n"
+            f"pwd > {self._pshell_cwd}\n"
+            f"echo {cmd_id}:$__EC > {self._pshell_status}\n"
+        )
+        self._send_to_shell(ipc_script)
+
+        # Poll the status file for current command's ID
+        deadline = time.monotonic() + timeout
+        poll_interval = 0.15
+
+        while True:
+            if is_interrupted():
+                self._kill_shell_children()
+                output, _, _ = self._read_persistent_output()
+                return {
+                    "output": output + "\n[Command interrupted]",
+                    "returncode": 130,
+                }
+
+            if time.monotonic() > deadline:
+                self._kill_shell_children()
+                output, _, _ = self._read_persistent_output()
+                if output:
+                    return {
+                        "output": output + f"\n[Command timed out after {timeout}s]",
+                        "returncode": 124,
+                    }
+                return self._timeout_result(timeout)
+
+            if not self._shell_alive:
+                return {
+                    "output": "Persistent shell died during execution",
+                    "returncode": 1,
+                }
+
+            status_content = self._read_temp_files(self._pshell_status)[0].strip()
+            if status_content.startswith(cmd_id + ":"):
+                break
+
+            time.sleep(poll_interval)
+
+        output, exit_code, new_cwd = self._read_persistent_output()
+        if new_cwd:
+            self.cwd = new_cwd
+        return {"output": output, "returncode": exit_code}
+
+    @staticmethod
+    def _merge_output(stdout: str, stderr: str) -> str:
+        """Combine stdout and stderr into a single output string."""
+        parts = []
+        if stdout.strip():
+            parts.append(stdout.rstrip("\n"))
+        if stderr.strip():
+            parts.append(stderr.rstrip("\n"))
+        return "\n".join(parts)
+
+    def _cleanup_persistent_shell(self):
+        """Clean up persistent shell resources.  Call from ``cleanup()``."""
+        if self._shell_proc is None:
+            return
+
+        if self._session_id:
+            self._cleanup_temp_files()
+
+        try:
+            self._shell_proc.stdin.close()
+        except Exception:
+            pass
+        try:
+            self._shell_proc.terminate()
+            self._shell_proc.wait(timeout=3)
+        except subprocess.TimeoutExpired:
+            self._shell_proc.kill()
+
+        self._shell_alive = False
+        self._shell_proc = None
+
+        if hasattr(self, "_drain_thread") and self._drain_thread.is_alive():
+            self._drain_thread.join(timeout=1.0)
+
+    def _cleanup_temp_files(self):
+        """Remove local temp files.  Override for remote backends (SSH, Docker)."""
+        for f in glob_mod.glob(f"{self._temp_prefix}-*"):
+            try:
+                os.remove(f)
+            except OSError:
+                pass