fix: session switching — kill ttyd by port not just PID file (belt-and-suspenders)

kill_ttyd() now uses two strategies:
1. PID file (existing behavior — kept intact)
2. Port-based fallback: lsof -ti :{TTYD_PORT} finds and kills any orphaned
   ttyd process that wasn't tracked in the PID file

spawn_ttyd() also does a pre-spawn port-free guard: if any process is still
occupying TTYD_PORT after kill_ttyd() returns, it sends SIGKILL to force-free
the port before the new ttyd tries to bind.

Root cause: PID file became desynced (pointed to dead process). kill_ttyd()
thought it succeeded but the REAL ttyd (PID not in file) kept running on
port 7682. New spawn_ttyd() failed to bind, died silently. Old ttyd kept
serving the old session — so every session switch showed the same session.

Tests: 2 new tests in test_ttyd.py (RED → GREEN confirmed)
This commit is contained in:
Brian Krabach
2026-03-31 07:47:30 -07:00
parent 0ed03c4e9d
commit 9762098a00
2 changed files with 201 additions and 40 deletions
+94
View File
@@ -217,6 +217,100 @@ async def test_kill_orphan_ttyd_handles_pid_file_with_dead_process():
assert not pid_path.exists(), "PID file should be removed after orphan cleanup" assert not pid_path.exists(), "PID file should be removed after orphan cleanup"
async def test_kill_ttyd_kills_orphan_on_port_when_pid_file_desynced():
"""kill_ttyd() must also kill orphaned ttyd processes on TTYD_PORT.
If the PID file points to a dead process (desynced), but the REAL ttyd is
still running on TTYD_PORT (orphaned from a previous spawn), kill_ttyd()
must find and kill it via lsof -ti :<port>. This is the belt-and-suspenders
fallback that prevents the session-switching bug where a new ttyd cannot bind
the port because the old one is still running.
"""
pid_path = ttyd_mod.TTYD_PID_PATH
pid_path.parent.mkdir(parents=True, exist_ok=True)
# PID file with a dead process (desynced)
pid_path.write_text("99999")
killed_pids: list[tuple[int, int]] = []
def mock_os_kill(pid: int, sig: int) -> None:
if pid == 99999 and sig == 0:
raise ProcessLookupError # PID file process is already dead
killed_pids.append((pid, sig))
mock_lsof_result = MagicMock()
mock_lsof_result.returncode = 0
mock_lsof_result.stdout = "12345\n" # orphan PID occupying TTYD_PORT
def mock_subprocess_run(cmd, **kwargs): # noqa: ANN001
result = MagicMock()
if "lsof" in cmd and "-ti" in cmd:
return mock_lsof_result
result.returncode = 1
result.stdout = ""
return result
with (
patch("os.kill", side_effect=mock_os_kill),
patch("muxplex.ttyd._subprocess.run", side_effect=mock_subprocess_run),
):
result = await kill_ttyd()
assert result is True, "kill_ttyd must return True when orphan found on port"
orphan_killed = any(
pid == 12345 and sig == signal.SIGTERM for pid, sig in killed_pids
)
assert orphan_killed, (
"kill_ttyd must send SIGTERM to orphan process (12345) found via "
"lsof on TTYD_PORT when PID file is desynced"
)
async def test_spawn_ttyd_force_kills_process_on_port_before_binding():
"""spawn_ttyd() must force-kill any process occupying TTYD_PORT before spawning.
If kill_ttyd() completed but the port is still occupied (race condition),
spawn_ttyd() must do a final SIGKILL cleanup so the new ttyd can bind.
"""
mock_proc = _make_mock_ttyd_process(pid=22222)
# First lsof call returns an occupant; second call (after kill) returns empty
call_count = 0
def mock_subprocess_run(cmd, **kwargs): # noqa: ANN001
nonlocal call_count
result = MagicMock()
if "lsof" in cmd and "-ti" in cmd:
call_count += 1
if call_count == 1:
result.returncode = 0
result.stdout = "54321\n" # port still occupied
return result
result.returncode = 1
result.stdout = ""
return result
killed_pids: list[tuple[int, int]] = []
def mock_os_kill(pid: int, sig: int) -> None:
killed_pids.append((pid, sig))
with (
patch("asyncio.create_subprocess_exec", new=AsyncMock(return_value=mock_proc)),
patch("muxplex.ttyd._subprocess.run", side_effect=mock_subprocess_run),
patch("os.kill", side_effect=mock_os_kill),
):
await spawn_ttyd("test-session")
force_killed = any(
pid == 54321 and sig == signal.SIGKILL for pid, sig in killed_pids
)
assert force_killed, (
"spawn_ttyd must SIGKILL any process occupying TTYD_PORT before spawning "
"to prevent 'address already in use' errors"
)
async def test_kill_orphan_ttyd_handles_invalid_pid_file_content(): async def test_kill_orphan_ttyd_handles_invalid_pid_file_content():
"""kill_orphan_ttyd() gracefully handles a PID file with non-integer content.""" """kill_orphan_ttyd() gracefully handles a PID file with non-integer content."""
pid_path = ttyd_mod.TTYD_PID_PATH pid_path = ttyd_mod.TTYD_PID_PATH
+87 -20
View File
@@ -18,6 +18,7 @@ Public API:
import asyncio import asyncio
import os import os
import signal import signal
import subprocess as _subprocess
import time import time
from pathlib import Path from pathlib import Path
@@ -37,6 +38,43 @@ TTYD_PORT: int = 7682
_active_process: asyncio.subprocess.Process | None = None _active_process: asyncio.subprocess.Process | None = None
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _kill_pids_on_port(port: int, sig: int) -> bool:
"""Find and signal all processes listening on *port* via lsof.
Returns True if at least one PID was found and signalled.
Silently ignores lsof unavailability and already-dead processes.
"""
try:
result = _subprocess.run(
["lsof", "-ti", f":{port}"],
capture_output=True,
text=True,
timeout=5,
)
if result.returncode != 0 or not result.stdout.strip():
return False
sent = False
for pid_str in result.stdout.strip().split("\n"):
pid_str = pid_str.strip()
if not pid_str:
continue
try:
orphan_pid = int(pid_str)
os.kill(orphan_pid, sig)
sent = True
except (ValueError, ProcessLookupError, PermissionError):
pass
return sent
except Exception: # noqa: BLE001
# lsof not available, timed out, or other unexpected failure
return False
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Public API # Public API
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -45,45 +83,55 @@ _active_process: asyncio.subprocess.Process | None = None
async def kill_ttyd() -> bool: async def kill_ttyd() -> bool:
"""Kill the running ttyd process and clean up the PID file. """Kill the running ttyd process and clean up the PID file.
Belt-and-suspenders strategy:
Strategy 1 — PID file:
Reads the PID from TTYD_PID_PATH. If no PID file exists, returns False. Reads the PID from TTYD_PID_PATH. If no PID file exists, returns False.
If the file content is not a valid integer, removes the file and returns False. If the file content is not a valid integer, removes the file and returns
False. Checks whether the process is alive via ``os.kill(pid, 0)``. If
already gone (ProcessLookupError), cleans up and proceeds. Otherwise
sends SIGTERM and polls every 0.1 s for up to 2 s.
Checks whether the process is alive via ``os.kill(pid, 0)``. If the Strategy 2 — port-based fallback:
process is already gone (ProcessLookupError), cleans up and returns True. After the PID-file kill, finds and kills any process still listening on
Otherwise sends SIGTERM and polls every 0.1 s for up to 2 s waiting for TTYD_PORT via ``lsof -ti :<port>``. This catches orphaned ttyd processes
the process to exit. The PID file and ``_active_process`` are cleared in whose PID was never recorded in the file (e.g. after a coordinator crash).
all cases before returning. A brief 0.3 s wait is added to let the OS release the port.
Uses ``asyncio.sleep`` for polling to avoid blocking the event loop. The PID file and ``_active_process`` are cleared in all cases before
returning.
Returns: Returns:
True — process was killed or was already dead. True — a process was killed (or was already dead) via either strategy.
False — no PID file found, or PID file contained invalid content. False — no PID file found and no process was listening on the port.
""" """
global _active_process global _active_process
if not TTYD_PID_PATH.exists(): killed = False
return False
# -------------------------------------------------------------------
# Strategy 1: PID file
# -------------------------------------------------------------------
if TTYD_PID_PATH.exists():
try: try:
pid = int(TTYD_PID_PATH.read_text().strip()) pid = int(TTYD_PID_PATH.read_text().strip())
except ValueError: except ValueError:
TTYD_PID_PATH.unlink(missing_ok=True) TTYD_PID_PATH.unlink(missing_ok=True)
return False pid = None
else:
# Check whether the process is still alive. # Check whether the process is still alive.
try: try:
os.kill(pid, 0) os.kill(pid, 0)
except ProcessLookupError: except ProcessLookupError:
# Already dead — clean up and report success. # Already dead — clean up and note success.
TTYD_PID_PATH.unlink(missing_ok=True) TTYD_PID_PATH.unlink(missing_ok=True)
_active_process = None killed = True
return True pid = None
else:
# Process is alive — ask it to terminate. # Process is alive — ask it to terminate.
os.kill(pid, signal.SIGTERM) os.kill(pid, signal.SIGTERM)
# Poll up to 2 s for the process to exit, yielding to the event loop each iteration. # Poll up to 2 s for the process to exit.
deadline = time.time() + 2.0 deadline = time.time() + 2.0
while time.time() < deadline: while time.time() < deadline:
try: try:
@@ -92,10 +140,20 @@ async def kill_ttyd() -> bool:
break break
await asyncio.sleep(0.1) await asyncio.sleep(0.1)
# Always clean up regardless of whether the process exited in time.
TTYD_PID_PATH.unlink(missing_ok=True) TTYD_PID_PATH.unlink(missing_ok=True)
killed = True
pid = None # noqa: F841 (intentional)
# -------------------------------------------------------------------
# Strategy 2: port-based fallback — catch orphans not in PID file
# -------------------------------------------------------------------
if _kill_pids_on_port(TTYD_PORT, signal.SIGTERM):
killed = True
# Brief pause so the OS can release the port before the next spawn.
await asyncio.sleep(0.3)
_active_process = None _active_process = None
return True return killed
async def kill_orphan_ttyd() -> bool: async def kill_orphan_ttyd() -> bool:
@@ -122,6 +180,10 @@ async def spawn_ttyd(session_name: str) -> asyncio.subprocess.Process:
ttyd -W -m 3 -p 7682 tmux attach -t <session_name> ttyd -W -m 3 -p 7682 tmux attach -t <session_name>
Before spawning, verifies that TTYD_PORT is free. If any process is still
listening on the port (e.g. a race between kill_ttyd() and spawn_ttyd()),
it sends SIGKILL to force-free the port immediately.
stdout and stderr are discarded (DEVNULL). The PID is written to stdout and stderr are discarded (DEVNULL). The PID is written to
TTYD_PID_PATH. The process handle is stored in ``_active_process``. TTYD_PID_PATH. The process handle is stored in ``_active_process``.
@@ -133,6 +195,11 @@ async def spawn_ttyd(session_name: str) -> asyncio.subprocess.Process:
""" """
global _active_process global _active_process
# Final port-free guard — catches races where kill_ttyd() returned but
# the old ttyd hasn't fully released the socket yet.
if _kill_pids_on_port(TTYD_PORT, signal.SIGKILL):
await asyncio.sleep(0.3)
proc = await asyncio.create_subprocess_exec( proc = await asyncio.create_subprocess_exec(
"ttyd", "ttyd",
"-W", "-W",