feat: muxterm process supervision (start, restart on crash, stop)

muxplex/muxterm.py

@@ -0,0 +1,165 @@
+"""
+muxterm process supervision — start, monitor, restart on crash, stop.
+
+Module state:
+    _muxterm_process — the currently running muxterm subprocess (or None)
+    _restart_task    — the background monitor/restart task (or None)
+
+Public API:
+    start_muxterm(secret, port, binary_path, auto_restart)
+    stop_muxterm()
+
+Internal:
+    _find_muxterm_binary(binary_path)
+    _monitor_and_restart(secret, port, binary_path)
+"""
+
+import asyncio
+import logging
+import shutil
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Module state
+# ---------------------------------------------------------------------------
+
+_muxterm_process: asyncio.subprocess.Process | None = None
+_restart_task: asyncio.Task | None = None  # type: ignore[type-arg]
+
+# ---------------------------------------------------------------------------
+# Binary discovery
+# ---------------------------------------------------------------------------
+
+
+def _find_muxterm_binary(binary_path: str | None = None) -> str:
+    """Locate the muxterm binary.
+
+    Checks explicit *binary_path* first, then falls back to
+    ``shutil.which('muxterm')``.  Raises :exc:`FileNotFoundError` if
+    the binary cannot be found.
+    """
+    if binary_path is not None:
+        return binary_path
+    found = shutil.which("muxterm")
+    if found is None:
+        raise FileNotFoundError(
+            "muxterm binary not found on PATH; pass binary_path explicitly"
+        )
+    return found
+
+
+# ---------------------------------------------------------------------------
+# Start
+# ---------------------------------------------------------------------------
+
+
+async def start_muxterm(
+    secret: str,
+    port: int = 7682,
+    binary_path: str | None = None,
+    auto_restart: bool = True,
+) -> None:
+    """Start the muxterm process.
+
+    Finds the binary via :func:`_find_muxterm_binary`, spawns it with
+    ``--addr 127.0.0.1:<port> --secret <secret>``, and optionally starts
+    the background monitor/restart loop.
+    """
+    global _muxterm_process, _restart_task
+
+    binary = _find_muxterm_binary(binary_path)
+
+    _muxterm_process = await asyncio.create_subprocess_exec(
+        binary,
+        "--addr",
+        f"127.0.0.1:{port}",
+        "--secret",
+        secret,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+    )
+    logger.info("muxterm started (pid=%s, port=%s)", _muxterm_process.pid, port)
+
+    if auto_restart:
+        _restart_task = asyncio.create_task(
+            _monitor_and_restart(secret, port, binary_path)
+        )
+
+
+# ---------------------------------------------------------------------------
+# Monitor / restart
+# ---------------------------------------------------------------------------
+
+
+async def _monitor_and_restart(
+    secret: str,
+    port: int,
+    binary_path: str | None,
+) -> None:
+    """Wait for process exit; restart on non-zero exit codes.
+
+    Normal exit (rc=0) ends the loop.  On crash, waits 1 s before
+    restarting, with 5 s backoff on repeated failures.
+    """
+    global _muxterm_process
+
+    backoff = 1.0
+    while True:
+        if _muxterm_process is None:
+            return
+
+        rc = await _muxterm_process.wait()
+
+        if rc == 0:
+            logger.info("muxterm exited cleanly (rc=0)")
+            return
+
+        logger.warning("muxterm exited with rc=%s, restarting in %.0fs", rc, backoff)
+        await asyncio.sleep(backoff)
+
+        try:
+            binary = _find_muxterm_binary(binary_path)
+            _muxterm_process = await asyncio.create_subprocess_exec(
+                binary,
+                "--addr",
+                f"127.0.0.1:{port}",
+                "--secret",
+                secret,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            logger.info(
+                "muxterm restarted (pid=%s, port=%s)", _muxterm_process.pid, port
+            )
+            # Reset backoff on successful spawn, increase on repeated failures
+            backoff = min(backoff + 1.0, 5.0)
+        except Exception:
+            logger.exception("failed to restart muxterm, retrying in 5s")
+            backoff = 5.0
+            await asyncio.sleep(backoff)
+
+
+# ---------------------------------------------------------------------------
+# Stop
+# ---------------------------------------------------------------------------
+
+
+async def stop_muxterm() -> None:
+    """Stop the muxterm process and cancel the restart task.
+
+    Terminates with a 5 s grace period, then kills if still alive.
+    """
+    global _muxterm_process, _restart_task
+
+    if _restart_task is not None:
+        _restart_task.cancel()
+        _restart_task = None
+
+    if _muxterm_process is not None:
+        _muxterm_process.terminate()
+        try:
+            await asyncio.wait_for(_muxterm_process.wait(), timeout=5.0)
+        except (asyncio.TimeoutError, ProcessLookupError):
+            _muxterm_process.kill()
+        _muxterm_process = None

muxplex/tests/test_muxterm.py

@@ -0,0 +1,150 @@
+"""
+Tests for muxplex/muxterm.py — muxterm process supervision (start, restart, stop).
+"""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+import muxplex.muxterm as muxterm_mod
+from muxplex.muxterm import (
+    _find_muxterm_binary,
+    start_muxterm,
+    stop_muxterm,
+)
+
+
+# ---------------------------------------------------------------------------
+# autouse fixture — reset module-level state between tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(autouse=True)
+def _reset_muxterm_state():
+    """Reset module-level muxterm state before every test."""
+    muxterm_mod._muxterm_process = None
+    muxterm_mod._restart_task = None
+    yield
+    # Teardown: clean up again
+    muxterm_mod._muxterm_process = None
+    muxterm_mod._restart_task = None
+
+
+# ---------------------------------------------------------------------------
+# _find_muxterm_binary tests
+# ---------------------------------------------------------------------------
+
+
+def test_find_muxterm_binary_explicit_path(tmp_path):
+    """Explicit binary path is returned as-is."""
+    binary = tmp_path / "muxterm"
+    binary.touch()
+    binary.chmod(0o755)
+    result = _find_muxterm_binary(str(binary))
+    assert result == str(binary)
+
+
+def test_find_muxterm_binary_uses_which():
+    """Falls back to shutil.which('muxterm') when no explicit path."""
+    with patch("shutil.which", return_value="/usr/bin/muxterm") as mock_which:
+        result = _find_muxterm_binary()
+        mock_which.assert_called_once_with("muxterm")
+        assert result == "/usr/bin/muxterm"
+
+
+def test_find_muxterm_binary_raises_when_not_found():
+    """Raises FileNotFoundError if muxterm is not found."""
+    with patch("shutil.which", return_value=None):
+        with pytest.raises(FileNotFoundError):
+            _find_muxterm_binary()
+
+
+# ---------------------------------------------------------------------------
+# start_muxterm tests
+# ---------------------------------------------------------------------------
+
+
+async def test_start_muxterm_spawns_process():
+    """start_muxterm calls create_subprocess_exec with correct args and stores process."""
+    mock_proc = MagicMock()
+    mock_proc.returncode = None
+
+    with (
+        patch(
+            "muxplex.muxterm._find_muxterm_binary",
+            return_value="/usr/local/bin/muxterm",
+        ),
+        patch(
+            "asyncio.create_subprocess_exec",
+            new=AsyncMock(return_value=mock_proc),
+        ) as mock_create,
+    ):
+        await start_muxterm("test-secret", port=7682, auto_restart=False)
+
+    # Verify the binary path was passed as first arg
+    mock_create.assert_called_once()
+    call_args = mock_create.call_args
+    assert call_args[0][0] == "/usr/local/bin/muxterm"
+    assert "--addr" in call_args[0]
+    assert "127.0.0.1:7682" in call_args[0]
+    assert "--secret" in call_args[0]
+    assert "test-secret" in call_args[0]
+
+    # Verify process stored in module state
+    assert muxterm_mod._muxterm_process is mock_proc
+
+
+# ---------------------------------------------------------------------------
+# stop_muxterm tests
+# ---------------------------------------------------------------------------
+
+
+async def test_stop_muxterm_terminates_process():
+    """stop_muxterm terminates the process and sets state to None."""
+    mock_proc = MagicMock()
+    mock_proc.terminate = MagicMock()
+    mock_proc.kill = MagicMock()
+    mock_proc.wait = AsyncMock()
+
+    muxterm_mod._muxterm_process = mock_proc
+
+    await stop_muxterm()
+
+    mock_proc.terminate.assert_called_once()
+    assert muxterm_mod._muxterm_process is None
+
+
+async def test_stop_muxterm_cancels_restart_task():
+    """stop_muxterm cancels the restart task if active."""
+    mock_proc = MagicMock()
+    mock_proc.terminate = MagicMock()
+    mock_proc.kill = MagicMock()
+    mock_proc.wait = AsyncMock()
+
+    mock_task = MagicMock()
+    mock_task.cancel = MagicMock()
+
+    muxterm_mod._muxterm_process = mock_proc
+    muxterm_mod._restart_task = mock_task
+
+    await stop_muxterm()
+
+    mock_task.cancel.assert_called_once()
+    assert muxterm_mod._restart_task is None
+
+
+async def test_stop_muxterm_kills_on_timeout():
+    """stop_muxterm kills the process if terminate doesn't finish in time."""
+    mock_proc = MagicMock()
+    mock_proc.terminate = MagicMock()
+    mock_proc.kill = MagicMock()
+    mock_proc.wait = AsyncMock(side_effect=asyncio.TimeoutError)
+
+    muxterm_mod._muxterm_process = mock_proc
+
+    await stop_muxterm()
+
+    mock_proc.terminate.assert_called_once()
+    mock_proc.kill.assert_called_once()
+    assert muxterm_mod._muxterm_process is None