ken/muxplex
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: delete ttyd.py, WS proxy, connect endpoint — muxterm owns terminal path

Browse Source
This commit is contained in:
Ken
2026-05-28 06:41:23 +00:00
parent 8c167453de
commit 266430d1c0
5 changed files with 19 additions and 1723 deletions
Download Patch File Download Diff File Expand all files Collapse all files
muxplex/main.py
+16 -160
View File
@@ -725,50 +725,19 @@ async def create_session(payload: CreateSessionPayload) -> dict:
return {"name": name, "ok": True} return {"name": name, "ok": True}
@app.post("/api/sessions/{name}/connect")
async def connect_session(name: str) -> dict:
"""Connect to a tmux session via ttyd.
Uses the process pool: if the session already has a live ttyd, returns
immediately. Otherwise spawns a new ttyd on a free port, polls until
listening, and returns the port.
Returns {active_session: name, ttyd_port: <port>}.
Raises HTTP 404 if *name* is not in the known session list (when non-empty).
"""
known = get_session_list()
if known and name not in known:
raise HTTPException(status_code=404, detail=f"Session '{name}' not found")
_log.info("Connecting to session '%s'", name)
port = await get_or_spawn(name)
async with state_lock:
state = load_state()
state["active_session"] = name
save_state(state)
return {"active_session": name, "ttyd_port": port}
@app.delete("/api/sessions/current") @app.delete("/api/sessions/current")
async def delete_current_session() -> dict: async def delete_current_session() -> dict:
"""Disconnect the current ttyd session. """Clear the active session.
Kills the pooled ttyd for the active session and clears active_session Sets active_session to None in persistent state.
in persistent state.
Returns {active_session: None}. Returns {active_session: None}.
""" """
async with state_lock: async with state_lock:
state = load_state() state = load_state()
active = state.get("active_session")
state["active_session"] = None state["active_session"] = None
save_state(state) save_state(state)
if active:
await kill_session(active)
return {"active_session": None} return {"active_session": None}
@@ -816,10 +785,6 @@ async def delete_session(name: str) -> dict:
except Exception: except Exception:
_log.warning("Delete command failed: %r", command) _log.warning("Delete command failed: %r", command)
# Clean up the pool entry (ttyd will die when tmux session dies,
# but clean up immediately so the port is freed).
await kill_session(name)
return {"ok": True, "name": name} return {"ok": True, "name": name}
@@ -1017,127 +982,6 @@ async def instance_info() -> dict:
} }
# ---------------------------------------------------------------------------
# WebSocket proxy — bridges browser to ttyd (eliminates Caddy dependency)
# ---------------------------------------------------------------------------
# _ttyd_is_listening is imported from muxplex.ttyd (pool-aware, accepts port param)
async def _ws_auth_check(websocket: WebSocket) -> bool:
"""Return True if the WebSocket caller is authorized.
Closes the WebSocket with code 4001 and returns False if the caller
is not authorized. Localhost connections (127.0.0.1 / ::1) are
unconditionally trusted. Remote callers must present a valid
``muxplex_session`` cookie OR a Bearer token matching ``_federation_key``.
"""
host = websocket.client.host if websocket.client else ""
if host in ("127.0.0.1", "::1"):
return True
session_cookie = websocket.cookies.get("muxplex_session")
cookie_ok = session_cookie and verify_session_cookie(
_auth_secret, session_cookie, _auth_ttl
)
bearer_ok = False
if _federation_key:
auth_header = websocket.headers.get("authorization", "")
if auth_header.lower().startswith("bearer "):
bearer_ok = hmac.compare_digest(auth_header[7:], _federation_key)
if not cookie_ok and not bearer_ok:
await websocket.close(code=4001)
return False
return True
@app.websocket("/terminal/ws")
async def terminal_ws_proxy(websocket: WebSocket) -> None:
"""Proxy WebSocket frames between the browser and ttyd.
Resolves the target ttyd port from the process pool BEFORE accepting the
browser WebSocket. If ttyd is not in the pool (e.g. after a service
restart), auto-spawns it using the session from query params or
active_session from state, then waits for it to bind its port.
Only after ttyd is confirmed reachable does the function call
websocket.accept() — so the browser's 'open' event only fires once a real
relay is possible. This prevents the reconnect-counter bounce bug where
the proxy accepted immediately (resetting _reconnectAttempts to 0) and
then closed as soon as it couldn't reach the dead ttyd.
"""
# Auth check before accepting — BaseHTTPMiddleware doesn't cover WebSocket scope
if not await _ws_auth_check(websocket):
return
# Determine which session (and port) to connect to.
# Priority: ?session= query param > active_session from state.
session = websocket.query_params.get("session", "")
if not session:
try:
async with state_lock:
state = load_state()
session = state.get("active_session") or ""
except Exception:
session = ""
# Look up port from pool, or auto-spawn if not listening
port = pool_port(session) if session else None
if port is None and session:
if not _ttyd_is_listening(pool_port(session) or TTYD_PORT):
_log.info(
"WS proxy: ttyd not listening, auto-spawning for '%s'",
session,
)
try:
port = await get_or_spawn(session)
except Exception as exc:
_log.warning("WS proxy: failed to auto-spawn ttyd: %s", exc)
if port is None:
# Fallback to TTYD_PORT (legacy single-process or externally managed ttyd)
port = TTYD_PORT
await websocket.accept(subprotocol="tty")
ttyd_url = f"ws://localhost:{port}/ws"
try:
async with websockets.connect(
ttyd_url, subprotocols=[Subprotocol("tty")]
) as ttyd_ws:
async def client_to_ttyd() -> None:
try:
while True:
msg = await websocket.receive()
if msg.get("bytes"):
await ttyd_ws.send(msg["bytes"])
elif msg.get("text"):
await ttyd_ws.send(msg["text"])
except Exception as exc:
_log.debug("ws relay closed (client_to_ttyd): %s", exc)
async def ttyd_to_client() -> None:
try:
async for message in ttyd_ws:
if isinstance(message, bytes):
await websocket.send_bytes(message)
else:
await websocket.send_text(message)
except Exception as exc:
_log.debug("ws relay closed (ttyd_to_client): %s", exc)
await asyncio.gather(client_to_ttyd(), ttyd_to_client())
except Exception as exc:
_log.debug("ws proxy closed: %s", exc)
finally:
try:
await websocket.close()
except Exception:
pass
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Federation helper utilities # Federation helper utilities
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -1191,8 +1035,20 @@ async def federation_terminal_ws_proxy(websocket: WebSocket, device_id: str) ->
Auth check uses the same cookie + bearer pattern as terminal_ws_proxy. Auth check uses the same cookie + bearer pattern as terminal_ws_proxy.
Closes with code 4004 if device_id does not match any remote. Closes with code 4004 if device_id does not match any remote.
""" """
# Auth check before accepting — same pattern as terminal_ws_proxy # Auth check before accepting — BaseHTTPMiddleware doesn't cover WS scope
if not await _ws_auth_check(websocket): host = websocket.client.host if websocket.client else ""
if host not in ("127.0.0.1", "::1"):
session_cookie = websocket.cookies.get("muxplex_session")
cookie_ok = session_cookie and verify_session_cookie(
_auth_secret, session_cookie, _auth_ttl
)
bearer_ok = False
if _federation_key:
auth_header = websocket.headers.get("authorization", "")
if auth_header.lower().startswith("bearer "):
bearer_ok = hmac.compare_digest(auth_header[7:], _federation_key)
if not cookie_ok and not bearer_ok:
await websocket.close(code=4001)
return return
# Look up remote instance by device_id # Look up remote instance by device_id
muxplex/tests/test_api.py
+2 -196
View File
@@ -356,76 +356,13 @@ def test_get_sessions_returns_empty_list_when_no_sessions(client, monkeypatch):
assert response.json() == [] assert response.json() == []
# ---------------------------------------------------------------------------
# POST /api/sessions/{name}/connect
# ---------------------------------------------------------------------------
def test_connect_session_returns_200(client, monkeypatch):
"""POST /api/sessions/{name}/connect returns 200 and correct body when session exists."""
monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["alpha"])
async def mock_get_or_spawn(name):
return 7682
monkeypatch.setattr("muxplex.main.get_or_spawn", mock_get_or_spawn)
response = client.post("/api/sessions/alpha/connect")
assert response.status_code == 200
data = response.json()
assert data["active_session"] == "alpha"
assert data["ttyd_port"] == 7682
def test_connect_session_sets_active_session(client, monkeypatch):
"""POST /api/sessions/{name}/connect persists active_session to state."""
from muxplex.state import load_state
monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["alpha"])
async def mock_get_or_spawn(name):
return 7682
monkeypatch.setattr("muxplex.main.get_or_spawn", mock_get_or_spawn)
client.post("/api/sessions/alpha/connect")
state = load_state()
assert state["active_session"] == "alpha"
def test_connect_session_calls_get_or_spawn(client, monkeypatch):
"""POST /api/sessions/{name}/connect calls get_or_spawn with the session name."""
monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["alpha"])
spawn_calls = []
async def mock_get_or_spawn(name):
spawn_calls.append(name)
return 7682
monkeypatch.setattr("muxplex.main.get_or_spawn", mock_get_or_spawn)
response = client.post("/api/sessions/alpha/connect")
assert response.status_code == 200
assert spawn_calls == ["alpha"]
def test_connect_nonexistent_session_returns_404(client, monkeypatch):
"""POST /api/sessions/{name}/connect returns 404 when session is not in list."""
monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["alpha", "beta"])
response = client.post("/api/sessions/gamma/connect")
assert response.status_code == 404
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# DELETE /api/sessions/current # DELETE /api/sessions/current
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def test_delete_current_kills_ttyd_and_clears_active(client, monkeypatch): def test_delete_current_clears_active_session(client, monkeypatch):
"""DELETE /api/sessions/current kills the session's ttyd and clears active_session.""" """DELETE /api/sessions/current clears active_session in state."""
from muxplex.state import load_state, save_state from muxplex.state import load_state, save_state
# Set up initial state with active session # Set up initial state with active session
@@ -438,19 +375,10 @@ def test_delete_current_kills_ttyd_and_clears_active(client, monkeypatch):
} }
) )
kill_session_calls = []
async def mock_kill_session(name):
kill_session_calls.append(name)
return True
monkeypatch.setattr("muxplex.main.kill_session", mock_kill_session)
response = client.delete("/api/sessions/current") response = client.delete("/api/sessions/current")
assert response.status_code == 200 assert response.status_code == 200
data = response.json() data = response.json()
assert data["active_session"] is None assert data["active_session"] is None
assert kill_session_calls == ["alpha"]
# Verify state was persisted # Verify state was persisted
state = load_state() state = load_state()
@@ -800,20 +728,6 @@ def test_api_routes_not_shadowed(client):
assert isinstance(response.json(), list) assert isinstance(response.json(), list)
def test_terminal_ws_route_exists():
"""The app must have a WebSocket route registered at /terminal/ws."""
from fastapi.routing import APIRoute, APIWebSocketRoute
from muxplex.main import app
ws_routes = [
r
for r in app.routes
if isinstance(r, (APIRoute, APIWebSocketRoute)) and r.path == "/terminal/ws"
]
assert len(ws_routes) == 1, "Expected exactly one /terminal/ws route"
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Auth middleware integration # Auth middleware integration
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -967,90 +881,6 @@ def test_logout_clears_session_cookie(monkeypatch):
assert "max-age=0" in set_cookie.lower() assert "max-age=0" in set_cookie.lower()
# ---------------------------------------------------------------------------
# WebSocket auth tests
# ---------------------------------------------------------------------------
def _wrap_with_client_host(wrapped_app, host: str):
"""Return an ASGI wrapper that forces websocket scope client to `host`.
This lets tests simulate a WebSocket connection appearing to originate
from a specific IP without touching Starlette internals.
"""
async def _middleware(scope, receive, send):
if scope.get("type") == "websocket":
scope = {**scope, "client": (host, 50000)}
await wrapped_app(scope, receive, send)
return _middleware
def test_ws_localhost_no_cookie_bypasses_auth():
"""WebSocket from 127.0.0.1 is accepted even without a session cookie."""
from starlette.websockets import WebSocketDisconnect
# Force scope to look like localhost so auth check is bypassed
localhost_app = _wrap_with_client_host(app, "127.0.0.1")
with TestClient(localhost_app) as c:
try:
with c.websocket_connect("/terminal/ws") as _:
pass # connection was accepted — auth bypassed for localhost
except WebSocketDisconnect as e:
# The websocket was accepted (auth bypassed); ttyd is not running so
# the proxy fails and closes with a non-4001 code.
assert e.code != 4001, (
f"Localhost WebSocket should not be rejected; got close code {e.code}"
)
def test_ws_valid_cookie_non_localhost_not_rejected_4001():
"""WebSocket from non-localhost with a valid cookie is not rejected with 4001."""
from starlette.websockets import WebSocketDisconnect
from muxplex.auth import create_session_cookie
from muxplex.main import _auth_secret, _auth_ttl
cookie = create_session_cookie(_auth_secret, _auth_ttl)
# TestClient default host is "testclient" — treated as non-localhost
with TestClient(app) as c:
c.cookies["muxplex_session"] = cookie
try:
with c.websocket_connect("/terminal/ws") as _:
pass # connection was accepted — auth passed
except WebSocketDisconnect as e:
# Auth passed; ttyd not running → proxy fails → close with code != 4001
assert e.code != 4001, (
f"Valid-cookie WebSocket should not be rejected; got close code {e.code}"
)
def test_ws_no_cookie_non_localhost_rejected_4001():
"""WebSocket from non-localhost without a cookie is closed with code 4001."""
from starlette.websockets import WebSocketDisconnect
# TestClient default host "testclient" is treated as non-localhost
with TestClient(app) as c:
with pytest.raises(WebSocketDisconnect) as exc_info:
with c.websocket_connect("/terminal/ws") as _:
pass
assert exc_info.value.code == 4001
def test_ws_invalid_cookie_non_localhost_rejected_4001():
"""WebSocket from non-localhost with a tampered cookie is closed with code 4001."""
from starlette.websockets import WebSocketDisconnect
with TestClient(app) as c:
c.cookies["muxplex_session"] = "tampered.invalid.cookie.value"
with pytest.raises(WebSocketDisconnect) as exc_info:
with c.websocket_connect("/terminal/ws") as _:
pass
assert exc_info.value.code == 4001
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# GET /api/settings # GET /api/settings
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -2464,30 +2294,6 @@ def test_create_session_logs_command(client, monkeypatch, tmp_path, caplog):
) )
def test_connect_session_logs_session_name(client, monkeypatch, caplog):
"""POST /api/sessions/{name}/connect must log the session name at INFO level."""
import logging
monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["target-session"])
async def mock_kill_ttyd():
pass
async def mock_spawn_ttyd(name):
pass
monkeypatch.setattr("muxplex.main.kill_ttyd", mock_kill_ttyd)
monkeypatch.setattr("muxplex.main.spawn_ttyd", mock_spawn_ttyd)
with caplog.at_level(logging.INFO, logger="muxplex.main"):
client.post("/api/sessions/target-session/connect")
log_messages = "\n".join(caplog.messages)
assert "target-session" in log_messages, (
f"connect_session must log session name at INFO level, got logs:\n{log_messages}"
)
def test_cli_uvicorn_log_level_is_info(): def test_cli_uvicorn_log_level_is_info():
"""cli.py serve() must pass log_level='info' to uvicorn.run so logs appear in journalctl.""" """cli.py serve() must pass log_level='info' to uvicorn.run so logs appear in journalctl."""
import inspect import inspect
muxplex/tests/test_ttyd.py
-448
View File
@@ -1,448 +0,0 @@
"""
Tests for coordinator/ttyd.py — ttyd process lifecycle management.
All 11 acceptance-criteria tests are defined here, plus 5 tests for the
_pids_on_port() lsof→fuser→ss fallback chain (the root-cause guard for systems
where lsof is not installed).
"""
import signal
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import muxplex.ttyd as ttyd_mod
from muxplex.ttyd import (
_kill_pids_on_port,
_pids_on_port,
kill_orphan_ttyd,
kill_ttyd,
spawn_ttyd,
)
# ---------------------------------------------------------------------------
# autouse fixture — redirect PID paths to tmp_path for every test
# ---------------------------------------------------------------------------
@pytest.fixture(autouse=True)
def use_tmp_pid_dir(tmp_path, monkeypatch):
"""Redirect PID file I/O to a fresh temp directory for every test."""
tmp_pid_dir = tmp_path / "ttyd"
tmp_pid_path = tmp_pid_dir / "ttyd.pid"
monkeypatch.setattr("muxplex.ttyd.TTYD_PID_DIR", tmp_pid_dir)
monkeypatch.setattr("muxplex.ttyd.TTYD_PID_PATH", tmp_pid_path)
# Clear the pool between tests
ttyd_mod._pool.clear()
# ---------------------------------------------------------------------------
# Helper for mocking asyncio.create_subprocess_exec
# ---------------------------------------------------------------------------
def _make_mock_ttyd_process(pid: int = 12345):
"""Return a mock ttyd process with the given PID."""
proc = MagicMock()
proc.pid = pid
return proc
# ---------------------------------------------------------------------------
# spawn_ttyd tests
# ---------------------------------------------------------------------------
async def test_spawn_ttyd_writes_pid_file():
"""spawn_ttyd() must write the process PID to TTYD_PID_PATH."""
mock_proc = _make_mock_ttyd_process(pid=99999)
with patch(
"asyncio.create_subprocess_exec",
new=AsyncMock(return_value=mock_proc),
):
await spawn_ttyd("my-session")
pid_path = ttyd_mod.TTYD_PID_PATH
assert pid_path.exists(), "PID file was not created"
assert pid_path.read_text().strip() == "99999"
async def test_spawn_ttyd_uses_correct_command():
"""spawn_ttyd() must call ttyd with args: -W -m 3 -p 7682 tmux attach -t <name>."""
mock_proc = _make_mock_ttyd_process(pid=54321)
with patch(
"asyncio.create_subprocess_exec",
new=AsyncMock(return_value=mock_proc),
) as mock_create:
await spawn_ttyd("test-session")
call_args = mock_create.call_args[0]
assert list(call_args) == [
"ttyd",
"-W",
"-m",
"3",
"-p",
"7682",
"tmux",
"attach",
"-t",
"test-session",
]
async def test_spawn_ttyd_returns_process_object():
"""spawn_ttyd() must return the process object from create_subprocess_exec."""
mock_proc = _make_mock_ttyd_process(pid=11111)
with patch(
"asyncio.create_subprocess_exec",
new=AsyncMock(return_value=mock_proc),
):
result = await spawn_ttyd("another-session")
assert result is mock_proc
# ---------------------------------------------------------------------------
# kill_ttyd tests
# ---------------------------------------------------------------------------
async def test_kill_ttyd_returns_false_when_no_pid_file():
"""kill_ttyd() returns False when no PID file exists."""
# autouse fixture ensures no PID file is present
result = await kill_ttyd()
assert result is False
async def test_kill_ttyd_reads_pid_file_and_sends_sigterm():
"""kill_ttyd() reads the PID file and sends SIGTERM to the running process."""
pid_path = ttyd_mod.TTYD_PID_PATH
pid_path.parent.mkdir(parents=True, exist_ok=True)
pid_path.write_text("12345")
kill_calls = []
def mock_os_kill(pid, sig):
kill_calls.append((pid, sig))
# First existence-check (sig=0) succeeds; subsequent sig=0 calls raise
if sig == 0 and sum(1 for _, s in kill_calls if s == 0) > 1:
raise ProcessLookupError
with patch("os.kill", side_effect=mock_os_kill):
result = await kill_ttyd()
assert result is True
sigterm_calls = [(pid, sig) for pid, sig in kill_calls if sig == signal.SIGTERM]
assert len(sigterm_calls) == 1
assert sigterm_calls[0][0] == 12345
async def test_kill_ttyd_removes_pid_file():
"""kill_ttyd() removes the PID file regardless of whether process was alive."""
pid_path = ttyd_mod.TTYD_PID_PATH
pid_path.parent.mkdir(parents=True, exist_ok=True)
pid_path.write_text("12345")
def mock_os_kill(pid, sig):
# All os.kill calls raise ProcessLookupError — simulates already-dead process
raise ProcessLookupError
with patch("os.kill", side_effect=mock_os_kill):
result = await kill_ttyd()
assert result is True
assert not pid_path.exists(), "PID file should be removed after kill_ttyd()"
async def test_kill_ttyd_handles_process_already_dead():
"""kill_ttyd() returns True and clears state when process is already gone."""
pid_path = ttyd_mod.TTYD_PID_PATH
pid_path.parent.mkdir(parents=True, exist_ok=True)
pid_path.write_text("99999")
# Simulate process already dead: os.kill(pid, 0) raises ProcessLookupError
with patch("os.kill", side_effect=ProcessLookupError):
result = await kill_ttyd()
assert result is True
assert not pid_path.exists(), (
"PID file should be removed when process was already dead"
)
assert ttyd_mod._active_process is None
# ---------------------------------------------------------------------------
# kill_orphan_ttyd tests
#
# kill_orphan_ttyd() is a thin delegation to kill_ttyd(). These tests verify
# both the delegation wiring and that behaviour is consistent with kill_ttyd().
# ---------------------------------------------------------------------------
async def test_kill_orphan_ttyd_returns_false_when_no_pid_file():
"""kill_orphan_ttyd() returns False when no PID file exists (no orphan)."""
# autouse fixture ensures no PID file is present
result = await kill_orphan_ttyd()
assert result is False
async def test_kill_orphan_ttyd_kills_running_process():
"""kill_orphan_ttyd() kills a running orphan process and returns True."""
pid_path = ttyd_mod.TTYD_PID_PATH
pid_path.parent.mkdir(parents=True, exist_ok=True)
pid_path.write_text("55555")
kill_calls = []
def mock_os_kill(pid, sig):
kill_calls.append((pid, sig))
# First existence-check (sig=0) succeeds; subsequent sig=0 calls raise
if sig == 0 and sum(1 for _, s in kill_calls if s == 0) > 1:
raise ProcessLookupError
with patch("os.kill", side_effect=mock_os_kill):
result = await kill_orphan_ttyd()
assert result is True
assert not pid_path.exists(), "PID file should be removed after kill_orphan_ttyd()"
sigterm_calls = [(pid, sig) for pid, sig in kill_calls if sig == signal.SIGTERM]
assert len(sigterm_calls) == 1
assert sigterm_calls[0][0] == 55555
async def test_kill_orphan_ttyd_handles_pid_file_with_dead_process():
"""kill_orphan_ttyd() handles a stale PID file whose process is already gone."""
pid_path = ttyd_mod.TTYD_PID_PATH
pid_path.parent.mkdir(parents=True, exist_ok=True)
pid_path.write_text("77777")
with patch("os.kill", side_effect=ProcessLookupError):
result = await kill_orphan_ttyd()
assert result is True
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"
)
# ---------------------------------------------------------------------------
# _pids_on_port fallback chain tests
#
# These cover the exact production failure: lsof not installed → all kill
# strategies silently returned False → stale ttyd survived restarts.
# ---------------------------------------------------------------------------
def _make_subprocess_result(returncode: int, stdout: str) -> MagicMock:
r = MagicMock()
r.returncode = returncode
r.stdout = stdout
return r
def test_pids_on_port_uses_lsof_when_available():
"""_pids_on_port() returns PID list from lsof output when lsof is available."""
def mock_run(cmd, **kwargs): # noqa: ANN001
if cmd[0] == "lsof":
return _make_subprocess_result(0, "3555095\n")
return _make_subprocess_result(1, "")
with patch("muxplex.ttyd._subprocess.run", side_effect=mock_run):
result = _pids_on_port(7682)
assert result == [3555095], "Should return PID from lsof"
def test_pids_on_port_falls_back_to_fuser_when_lsof_missing():
"""_pids_on_port() uses fuser when lsof is not installed (FileNotFoundError).
This is the exact failure mode that caused the yazi-stuck-ttyd bug:
lsof was absent, _kill_pids_on_port silently returned False, and the
stale ttyd was never killed.
"""
def mock_run(cmd, **kwargs): # noqa: ANN001
if cmd[0] == "lsof":
raise FileNotFoundError("lsof not found")
if cmd[0] == "fuser":
return _make_subprocess_result(0, " 3555095")
return _make_subprocess_result(1, "")
with patch("muxplex.ttyd._subprocess.run", side_effect=mock_run):
result = _pids_on_port(7682)
assert result == [3555095], "Should fall back to fuser when lsof is missing"
def test_pids_on_port_falls_back_to_ss_when_lsof_and_fuser_missing():
"""_pids_on_port() uses ss when both lsof and fuser are unavailable."""
ss_output = (
'LISTEN 0 128 0.0.0.0:7682 0.0.0.0:* users:(("ttyd",pid=3555095,fd=12))\n'
)
def mock_run(cmd, **kwargs): # noqa: ANN001
if cmd[0] in ("lsof", "fuser"):
raise FileNotFoundError(f"{cmd[0]} not found")
if cmd[0] == "ss":
return _make_subprocess_result(0, ss_output)
return _make_subprocess_result(1, "")
with patch("muxplex.ttyd._subprocess.run", side_effect=mock_run):
result = _pids_on_port(7682)
assert result == [3555095], "Should fall back to ss when lsof and fuser missing"
def test_pids_on_port_returns_empty_when_all_tools_fail():
"""_pids_on_port() returns [] when no discovery tool is available."""
def mock_run(cmd, **kwargs): # noqa: ANN001
raise FileNotFoundError(f"{cmd[0]} not found")
with patch("muxplex.ttyd._subprocess.run", side_effect=mock_run):
result = _pids_on_port(7682)
assert result == [], "Should return empty list when all tools fail"
def test_kill_pids_on_port_sends_signal_via_fuser_fallback():
"""_kill_pids_on_port() sends the correct signal using the fuser fallback.
End-to-end regression guard: verifies that the full chain from
_kill_pids_on_port → _pids_on_port → fuser → os.kill works correctly
so that stale ttyd processes are killed even when lsof is absent.
"""
killed: list[tuple[int, int]] = []
def mock_run(cmd, **kwargs): # noqa: ANN001
if cmd[0] == "lsof":
raise FileNotFoundError("lsof not found")
if cmd[0] == "fuser":
return _make_subprocess_result(0, " 3555095")
return _make_subprocess_result(1, "")
def mock_kill(pid: int, sig: int) -> None:
killed.append((pid, sig))
with (
patch("muxplex.ttyd._subprocess.run", side_effect=mock_run),
patch("os.kill", side_effect=mock_kill),
):
result = _kill_pids_on_port(7682, signal.SIGKILL)
assert result is True, "_kill_pids_on_port must return True when fuser finds a PID"
assert (3555095, signal.SIGKILL) in killed, (
"Must send SIGKILL to PID 3555095 discovered via fuser fallback"
)
async def test_kill_orphan_ttyd_handles_invalid_pid_file_content():
"""kill_orphan_ttyd() gracefully handles a PID file with non-integer content."""
pid_path = ttyd_mod.TTYD_PID_PATH
pid_path.parent.mkdir(parents=True, exist_ok=True)
pid_path.write_text("not-a-pid")
# Should not raise and should clean up the file.
# kill_ttyd() calls pid_path.unlink(missing_ok=True) before returning False
# on invalid content, so the file is removed even though no kill occurred.
result = await kill_orphan_ttyd()
assert result is False
assert not pid_path.exists(), "Invalid PID file should be removed"
muxplex/tests/test_ws_proxy.py
-447
View File
@@ -1,447 +0,0 @@
"""
Comprehensive tests for the WebSocket proxy in muxplex/main.py.
"""
import inspect
import threading
import time
import pytest
from fastapi.testclient import TestClient
from starlette.websockets import WebSocketDisconnect
from muxplex.auth import create_session_cookie
from muxplex.main import app, terminal_ws_proxy
# ---------------------------------------------------------------------------
# Polling helper — deterministic alternative to time.sleep() for async relay
# ---------------------------------------------------------------------------
def _wait_for(condition, timeout: float = 2.0, interval: float = 0.01) -> bool:
"""Poll *condition()* until it returns True or *timeout* seconds elapses.
Returns True if the condition was met, False on timeout.
Using a polling loop instead of a fixed sleep makes relay tests deterministic:
on fast machines the loop exits as soon as the relay completes; on slow machines
it waits up to *timeout* seconds rather than racing against a fixed 200ms budget.
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if condition():
return True
time.sleep(interval)
return False # pragma: no cover — timeout branch only on pathological machines
# ---------------------------------------------------------------------------
# autouse fixture — redirect state/PID files to tmp_path, mock startup side-effects
# ---------------------------------------------------------------------------
@pytest.fixture(autouse=True)
def patch_startup_and_state(tmp_path, monkeypatch):
"""Redirect state files to tmp_path, mock muxterm start/stop, replace _poll_loop with no-op."""
# Redirect state files
tmp_state_dir = tmp_path / "state"
tmp_state_path = tmp_state_dir / "state.json"
monkeypatch.setattr("muxplex.state.STATE_DIR", tmp_state_dir)
monkeypatch.setattr("muxplex.state.STATE_PATH", tmp_state_path)
# Mock start_muxterm/stop_muxterm so startup doesn't touch real processes
from unittest.mock import AsyncMock
monkeypatch.setattr("muxplex.main.start_muxterm", AsyncMock())
monkeypatch.setattr("muxplex.main.stop_muxterm", AsyncMock())
# Replace _poll_loop with a no-op so tests don't spin up real poll cycles
async def noop_poll_loop() -> None:
pass
monkeypatch.setattr("muxplex.main._poll_loop", noop_poll_loop)
# ---------------------------------------------------------------------------
# Helper — create TestClient with valid session cookie
# ---------------------------------------------------------------------------
def _make_authed_client():
"""Creates TestClient with valid session cookie."""
from muxplex.main import _auth_secret, _auth_ttl
cookie = create_session_cookie(_auth_secret, _auth_ttl)
client = TestClient(app)
client.cookies.set("muxplex_session", cookie)
return client
# ---------------------------------------------------------------------------
# FakeTtydWs — mock ttyd WebSocket for relay testing
# ---------------------------------------------------------------------------
class FakeTtydWs:
"""Mock ttyd WebSocket that stores sent messages and yields pre-loaded responses.
Supports send(), close(), async iterator, and async context manager.
"""
def __init__(self, responses=None):
self.sent = []
self._responses = list(responses or [])
self._closed = False
async def send(self, message):
self.sent.append(message)
async def close(self):
self._closed = True
def __aiter__(self):
return self._async_gen()
async def _async_gen(self):
for msg in self._responses:
yield msg
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.close()
return False
# ---------------------------------------------------------------------------
# Test 1: regression — proxy source must use receive(), not receive_bytes()
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# Tests: ttyd liveness check before websocket.accept
# ---------------------------------------------------------------------------
def test_ttyd_is_listening_function_exists():
"""_ttyd_is_listening() must exist in main.py (TCP probe helper)."""
# Import will fail if function doesn't exist — that IS the failing test
from muxplex.main import _ttyd_is_listening # noqa: F401
assert callable(_ttyd_is_listening)
def test_ws_proxy_checks_ttyd_before_accepting():
"""terminal_ws_proxy must check _ttyd_is_listening BEFORE websocket.accept.
Root cause of the reconnect loop: the proxy called websocket.accept() before
checking if ttyd was alive. The browser's 'open' event fired immediately,
resetting _reconnectAttempts to 0. The counter bounced 0→1→0→1 forever so
the client-side /connect POST (at >= 2 attempts) never fired.
Fix: check _ttyd_is_listening() first. If not listening, auto-spawn ttyd
THEN accept — so the browser only gets 'open' when ttyd is actually ready.
"""
source = inspect.getsource(terminal_ws_proxy)
# Use "await websocket.accept" to avoid matching the docstring mention
accept_idx = source.index("await websocket.accept")
ttyd_check_idx = source.index("_ttyd_is_listening")
assert ttyd_check_idx < accept_idx, (
"_ttyd_is_listening() must be checked BEFORE await websocket.accept() — "
"proxy must not accept the browser WS until ttyd is confirmed alive"
)
def test_ws_proxy_auto_spawns_ttyd_when_dead(monkeypatch):
"""WS proxy must call get_or_spawn when pool has no entry for the session."""
spawn_calls = []
async def mock_get_or_spawn(name: str):
spawn_calls.append(name)
return 7682
# pool_port returns None → no existing ttyd for this session
monkeypatch.setattr("muxplex.main.pool_port", lambda name: None)
# _ttyd_is_listening returns False → nothing listening on fallback port either
monkeypatch.setattr("muxplex.main._ttyd_is_listening", lambda port=7682: False)
# get_or_spawn is what the proxy should call to auto-spawn
monkeypatch.setattr("muxplex.main.get_or_spawn", mock_get_or_spawn)
# Provide a fake websockets.connect that immediately closes (no real ttyd)
fake_ws = FakeTtydWs(responses=[])
monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws)
# Patch load_state to return state with active_session
monkeypatch.setattr(
"muxplex.main.load_state",
lambda: {"active_session": "test-session", "sessions": {}, "session_order": []},
)
with _make_authed_client() as c:
with c.websocket_connect("/terminal/ws") as _:
pass
assert spawn_calls == ["test-session"], (
"get_or_spawn must be called with active_session when ttyd is not in pool"
)
def test_terminal_ws_proxy_does_not_use_receive_bytes():
"""Regression: receive_bytes() silently drops TEXT frames (like the ttyd auth token).
terminal.js sends {"AuthToken": ""} as a TEXT WebSocket frame. The original
proxy used receive_bytes() which fails on text frames, swallowed the exception,
and exited — meaning ttyd never received the auth token, never started
streaming, resulting in a permanent black screen and reconnect loop.
The proxy MUST use receive() and dispatch on message type to handle both
binary and text frames correctly.
"""
source = inspect.getsource(terminal_ws_proxy)
assert "receive_bytes" not in source, (
"client_to_ttyd must not use receive_bytes() — silently drops text frames "
'like the ttyd auth token {"AuthToken": ""}'
)
assert ".receive()" in source, (
"client_to_ttyd must use receive() to handle both text and binary frames"
)
# ---------------------------------------------------------------------------
# Tests 2–3: auth rejection
# ---------------------------------------------------------------------------
def test_ws_auth_rejection_no_cookie():
"""WebSocket from non-localhost without cookie is closed with code 4001."""
# TestClient default host is "testclient" which is treated as non-localhost
with TestClient(app) as c:
with pytest.raises(WebSocketDisconnect) as exc_info:
with c.websocket_connect("/terminal/ws") as _:
pass
assert exc_info.value.code == 4001
def test_ws_auth_rejection_invalid_cookie():
"""WebSocket from non-localhost with a tampered cookie is closed with code 4001."""
with TestClient(app) as c:
c.cookies.set("muxplex_session", "tampered.invalid.cookie.value")
with pytest.raises(WebSocketDisconnect) as exc_info:
with c.websocket_connect("/terminal/ws") as _:
pass
assert exc_info.value.code == 4001
# ---------------------------------------------------------------------------
# Test: Bearer token auth accepted
# ---------------------------------------------------------------------------
def test_ws_bearer_auth_accepted(monkeypatch):
"""WebSocket from non-localhost with valid Bearer federation key is NOT rejected with 4001.
When a valid federation key is provided as 'Authorization: Bearer <key>',
the WebSocket connection must be accepted (not closed with code 4001).
"""
fed_key = "test-federation-secret-key"
monkeypatch.setattr("muxplex.main._federation_key", fed_key)
fake_ws = FakeTtydWs(responses=[])
monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws)
# Connect without a session cookie but with a valid Bearer token.
# Should NOT raise WebSocketDisconnect with code 4001.
with TestClient(app) as c:
# If Bearer auth is not implemented, this raises WebSocketDisconnect(code=4001)
with c.websocket_connect(
"/terminal/ws",
headers={"Authorization": f"Bearer {fed_key}"},
) as _ws:
pass # Successfully connected — auth was accepted
# ---------------------------------------------------------------------------
# Tests 4–5: browser → ttyd relay
# ---------------------------------------------------------------------------
def test_browser_text_relayed_to_ttyd(monkeypatch):
"""Text message from browser is forwarded to ttyd via FakeTtydWs.send()."""
fake_ws = FakeTtydWs()
monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws)
with _make_authed_client() as c:
with c.websocket_connect("/terminal/ws") as ws:
ws.send_text("hello from browser")
_wait_for(lambda: "hello from browser" in fake_ws.sent)
assert "hello from browser" in fake_ws.sent
def test_browser_bytes_relayed_to_ttyd(monkeypatch):
"""Binary message from browser is forwarded to ttyd via FakeTtydWs.send()."""
fake_ws = FakeTtydWs()
monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws)
with _make_authed_client() as c:
with c.websocket_connect("/terminal/ws") as ws:
ws.send_bytes(b"\x00\x01\x02 binary data")
_wait_for(lambda: b"\x00\x01\x02 binary data" in fake_ws.sent)
assert b"\x00\x01\x02 binary data" in fake_ws.sent
# ---------------------------------------------------------------------------
# Tests 6–7: ttyd → browser relay
# ---------------------------------------------------------------------------
def test_ttyd_text_relayed_to_browser(monkeypatch):
"""Text message from ttyd is forwarded to browser via websocket.send_text()."""
fake_ws = FakeTtydWs(responses=["hello from ttyd"])
monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws)
with _make_authed_client() as c:
with c.websocket_connect("/terminal/ws") as ws:
msg = ws.receive_text()
assert msg == "hello from ttyd"
def test_ttyd_bytes_relayed_to_browser(monkeypatch):
"""Binary message from ttyd is forwarded to browser via websocket.send_bytes()."""
fake_ws = FakeTtydWs(responses=[b"\xde\xad\xbe\xef binary"])
monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws)
with _make_authed_client() as c:
with c.websocket_connect("/terminal/ws") as ws:
msg = ws.receive_bytes()
assert msg == b"\xde\xad\xbe\xef binary"
# ---------------------------------------------------------------------------
# Test 8: ttyd close propagates to browser
# ---------------------------------------------------------------------------
def test_ttyd_close_propagates_to_browser(monkeypatch):
"""When ttyd exhausts its messages, the proxy cleans up and closes the browser WS."""
fake_ws = FakeTtydWs(responses=[]) # no responses — exhausts immediately
monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws)
with _make_authed_client() as c:
with c.websocket_connect("/terminal/ws") as _:
# FakeTtydWs has no responses so ttyd_to_client exhausts immediately.
# Exiting the context manager closes the browser WS, which causes
# client_to_ttyd to complete, gather finishes, and the proxy
# finally-block calls fake_ws.close().
pass
# fake_ws should have been closed when the async-with block exited
assert fake_ws._closed
# ---------------------------------------------------------------------------
# Test 9: ttyd unreachable closes browser WS
# ---------------------------------------------------------------------------
def test_ttyd_unreachable_closes_browser_ws(monkeypatch):
"""OSError on ttyd connect closes the browser WebSocket (no hang, no 4001)."""
def mock_connect_raises(*args, **kwargs):
raise OSError("Connection refused — ttyd not running")
monkeypatch.setattr("muxplex.main.websockets.connect", mock_connect_raises)
with _make_authed_client() as c:
with c.websocket_connect("/terminal/ws") as ws:
# Proxy accepts, then closes after failing to reach ttyd.
# Receive the close frame — proves the proxy closed (no hang)
# and that auth was not rejected (which would use code 4001).
close_frame = ws.receive()
assert close_frame.get("type") == "websocket.close", (
"Proxy must close the WebSocket"
)
assert close_frame.get("code") != 4001, "Must not be an auth rejection (4001)"
# ---------------------------------------------------------------------------
# Test 10: concurrent sessions don't interfere
# ---------------------------------------------------------------------------
def test_concurrent_ws_sessions(monkeypatch):
"""Two simultaneous proxy sessions relay to separate FakeTtydWs instances."""
# Create two separate FakeTtydWs instances, one per connection
ws_pool = [FakeTtydWs(), FakeTtydWs()]
call_count = 0
lock = threading.Lock()
def mock_connect(*args, **kwargs):
nonlocal call_count
with lock:
idx = call_count % len(ws_pool)
call_count += 1
return ws_pool[idx]
monkeypatch.setattr("muxplex.main.websockets.connect", mock_connect)
errors = []
with _make_authed_client() as c:
def send_msg(text):
try:
with c.websocket_connect("/terminal/ws") as ws:
ws.send_text(text)
_wait_for(lambda: text in ws_pool[0].sent + ws_pool[1].sent)
except Exception as exc:
errors.append(exc)
t1 = threading.Thread(target=send_msg, args=("session_one_msg",))
t2 = threading.Thread(target=send_msg, args=("session_two_msg",))
t1.start()
t2.start()
t1.join(timeout=10)
t2.join(timeout=10)
assert not errors, f"Concurrent sessions raised errors: {errors}"
# Both messages must have been relayed (one to each fake_ws)
all_sent = ws_pool[0].sent + ws_pool[1].sent
assert "session_one_msg" in all_sent
assert "session_two_msg" in all_sent
# ---------------------------------------------------------------------------
# Task-11: federation WebSocket proxy route
# ---------------------------------------------------------------------------
def test_federation_ws_proxy_route_exists():
"""App must have a WebSocket route at /federation/{device_id}/terminal/ws."""
from starlette.routing import WebSocketRoute
ws_routes = [r for r in app.routes if isinstance(r, WebSocketRoute)]
paths = [r.path for r in ws_routes]
assert "/federation/{device_id}/terminal/ws" in paths, (
"App must have a WebSocket route at /federation/{device_id}/terminal/ws"
)
def test_federation_ws_proxy_uses_ssl_context_for_wss():
"""Federation WS proxy must pass an SSL context when connecting via wss://.
Self-signed certs on remote instances (cortex, spark-2, etc.) fail the
default SSL verification in websockets.connect(). The proxy must build an
SSLContext with CERT_NONE for wss:// URLs — the same fix already applied
to the httpx client (verify=False) but for the websockets library.
"""
from muxplex.main import federation_terminal_ws_proxy
source = inspect.getsource(federation_terminal_ws_proxy)
assert "ssl" in source and ("CERT_NONE" in source or "ssl_context" in source), (
"Federation WS proxy must configure an SSL context (CERT_NONE / ssl_context) "
"for self-signed cert support on wss:// connections"
)
muxplex/ttyd.py
-471
View File
@@ -1,471 +0,0 @@
"""
ttyd process lifecycle management for the tmux-web muxplex.
Constants:
TTYD_PID_DIR — directory for PID files (default: ~/.local/share/tmux-web/)
TTYD_PID_PATH — full path to the legacy PID file (TTYD_PID_DIR / 'ttyd.pid')
TTYD_PORT — base port for the ttyd pool (7682)
Module state:
_active_process — the currently running ttyd subprocess (or None) [backward compat]
_pool — session→PoolEntry mapping for the process pool
Pool API:
get_or_spawn(session_name) — Return port; spawn if needed, poll until listening
pool_port(session_name) — Look up port for session without spawning
kill_session(session_name) — Kill ttyd for one session, remove from pool
kill_all() — Kill all pooled ttyd processes
kill_orphans() — Startup cleanup: PID files, port range scan
Backward-compat API:
spawn_ttyd(session_name) — spawn ttyd attached to a tmux session, write PID file
kill_ttyd() — kill the running ttyd process, clean up PID file
kill_orphan_ttyd() — kill any orphaned ttyd from a previous coordinator run
"""
import asyncio
import os
import re as _re
import signal
import socket as _socket
import subprocess as _subprocess
import time
from dataclasses import dataclass
from pathlib import Path
# ---------------------------------------------------------------------------
# Paths and constants
# ---------------------------------------------------------------------------
_default_ttyd_pid_dir = Path.home() / ".local" / "share" / "tmux-web"
TTYD_PID_DIR: Path = Path(os.environ.get("TMUX_WEB_STATE_DIR", _default_ttyd_pid_dir))
TTYD_PID_PATH: Path = TTYD_PID_DIR / "ttyd.pid"
TTYD_PORT: int = 7682
_PORT_MIN: int = 7682
_PORT_MAX: int = 7701
# ---------------------------------------------------------------------------
# Module state
# ---------------------------------------------------------------------------
_active_process: asyncio.subprocess.Process | None = None
@dataclass
class PoolEntry:
"""A pooled ttyd process and its assigned port."""
process: asyncio.subprocess.Process
port: int
_pool: dict[str, PoolEntry] = {}
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _pid_path_for_port(port: int) -> Path:
"""Return the per-port PID file path: ``ttyd-{port}.pid`` in TTYD_PID_DIR."""
return TTYD_PID_DIR / f"ttyd-{port}.pid"
def _is_alive(entry: PoolEntry) -> bool:
"""Return True if the pool entry's process has not exited."""
return entry.process.returncode is None
def _ttyd_is_listening(port: int = TTYD_PORT) -> bool:
"""Return True if something is accepting TCP connections on *port*.
Uses a raw socket connect (no WebSocket handshake, no PTY spawned).
Takes < 1 ms on localhost when ttyd is running; fails immediately with
ConnectionRefusedError when it's not. OSError/TimeoutError are also
caught so the caller always gets a bool.
"""
try:
with _socket.create_connection(("127.0.0.1", port), timeout=0.5):
return True
except (ConnectionRefusedError, OSError, TimeoutError):
return False
def _allocate_port() -> int:
"""Return the next free port in the pool range [_PORT_MIN, _PORT_MAX]."""
used_ports = {e.port for e in _pool.values()}
for port in range(_PORT_MIN, _PORT_MAX + 1):
if port not in used_ports:
return port
raise RuntimeError("ttyd pool exhausted: all ports in range are in use")
def _pids_on_port(port: int) -> list[int]:
"""Return PIDs of all processes listening on *port*.
Tries three tools in order, stopping at the first that returns results:
1. ``lsof -ti :<port>`` — one PID per line, most widely available.
2. ``fuser <port>/tcp`` — space-separated PIDs on stdout (psmisc).
3. ``ss -Hnltp`` — parses ``pid=N`` from users field (iproute2).
Returns an empty list if none of the tools are available or find anything.
Silently swallows all errors so callers always get a list.
"""
pids: list[int] = []
# --- Tool 1: lsof ---
try:
result = _subprocess.run(
["lsof", "-ti", f":{port}"],
capture_output=True,
text=True,
timeout=5,
)
if result.returncode == 0 and result.stdout.strip():
for tok in result.stdout.split():
try:
pids.append(int(tok))
except ValueError:
pass
except Exception: # noqa: BLE001
pass
if pids:
return pids
# --- Tool 2: fuser (psmisc) ---
try:
result = _subprocess.run(
["fuser", f"{port}/tcp"],
capture_output=True,
text=True,
timeout=5,
)
# fuser writes the port label to stderr and PIDs to stdout.
if result.stdout.strip():
for tok in result.stdout.split():
try:
pids.append(int(tok))
except ValueError:
pass
except Exception: # noqa: BLE001
pass
if pids:
return pids
# --- Tool 3: ss (iproute2) ---
try:
result = _subprocess.run(
["ss", "-Hnltp", f"sport = :{port}"],
capture_output=True,
text=True,
timeout=5,
)
if result.returncode == 0 and result.stdout.strip():
for match in _re.finditer(r"pid=(\d+)", result.stdout):
try:
pids.append(int(match.group(1)))
except ValueError:
pass
except Exception: # noqa: BLE001
pass
return pids
def _kill_pids_on_port(port: int, sig: int) -> bool:
"""Find and signal all processes listening on *port*.
Uses :func:`_pids_on_port` (lsof → fuser → ss) to locate PIDs, then
signals each with *sig*.
Returns True if at least one PID was found and signalled.
Silently ignores unavailable tools and already-dead processes.
"""
pids = _pids_on_port(port)
if not pids:
return False
sent = False
for pid in pids:
try:
os.kill(pid, sig)
sent = True
except (ProcessLookupError, PermissionError):
pass
return sent
# ---------------------------------------------------------------------------
# Pool API
# ---------------------------------------------------------------------------
async def get_or_spawn(session_name: str) -> int:
"""Return the port for *session_name*, spawning ttyd if needed.
If the session already has a live ttyd in the pool (process alive AND
port is listening), returns immediately. Otherwise allocates the next
free port from [_PORT_MIN, _PORT_MAX], spawns ttyd, polls until
listening (up to ~1 s), and returns the port.
"""
# Check existing pool entry — require both alive process AND listening port
entry = _pool.get(session_name)
if entry is not None:
if _is_alive(entry) and _ttyd_is_listening(entry.port):
return entry.port
# Stale entry — clean up
try:
entry.process.terminate()
except ProcessLookupError:
pass
_pid_path_for_port(entry.port).unlink(missing_ok=True)
del _pool[session_name]
# Allocate next free port
port = _allocate_port()
# Force-free the port (catches races)
if _kill_pids_on_port(port, signal.SIGKILL):
await asyncio.sleep(0.3)
# Spawn ttyd
proc = await asyncio.create_subprocess_exec(
"ttyd",
"-W",
"-m",
"3",
"-p",
str(port),
"tmux",
"attach",
"-t",
session_name,
stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL,
start_new_session=True,
)
# Write per-port PID file
TTYD_PID_DIR.mkdir(parents=True, exist_ok=True)
_pid_path_for_port(port).write_text(str(proc.pid))
# Add to pool
_pool[session_name] = PoolEntry(process=proc, port=port)
# Poll until listening (up to ~1 s)
for _ in range(20):
if _ttyd_is_listening(port):
break
await asyncio.sleep(0.05)
return port
def pool_port(session_name: str) -> int | None:
"""Look up the port for *session_name* without spawning.
Returns the port if the session has a live pool entry, else None.
"""
entry = _pool.get(session_name)
if entry is not None and _is_alive(entry):
return entry.port
return None
async def kill_session(session_name: str) -> bool:
"""Kill the ttyd process for *session_name* and remove from pool.
Returns True if a pool entry was found and cleaned up, False otherwise.
"""
entry = _pool.pop(session_name, None)
if entry is None:
return False
try:
entry.process.terminate()
try:
await asyncio.wait_for(entry.process.wait(), timeout=2.0)
except asyncio.TimeoutError:
entry.process.kill()
except ProcessLookupError:
pass
_pid_path_for_port(entry.port).unlink(missing_ok=True)
return True
async def kill_all() -> None:
"""Kill all pooled ttyd processes and clear the pool."""
for name in list(_pool):
await kill_session(name)
async def kill_orphans() -> None:
"""Startup cleanup: read PID files, SIGTERM, scan port range for stragglers.
Handles both legacy ``ttyd.pid`` and per-port ``ttyd-{port}.pid`` files.
"""
await kill_orphan_ttyd()
# ---------------------------------------------------------------------------
# Backward-compat public API
# ---------------------------------------------------------------------------
async def kill_ttyd() -> bool:
"""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.
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.
Strategy 2 — port-based fallback:
After the PID-file kill, finds and kills any process still listening on
TTYD_PORT via ``_pids_on_port()`` (lsof → fuser → ss). This catches
orphaned ttyd processes whose PID was never recorded in the file (e.g.
after a coordinator crash). A brief 0.3 s wait is added to let the OS
release the port.
The PID file and ``_active_process`` are cleared in all cases before
returning.
Returns:
True — a process was killed (or was already dead) via either strategy.
False — no PID file found and no process was listening on the port.
"""
global _active_process
killed = False
# -------------------------------------------------------------------
# Strategy 1: PID file
# -------------------------------------------------------------------
if TTYD_PID_PATH.exists():
try:
pid = int(TTYD_PID_PATH.read_text().strip())
except ValueError:
TTYD_PID_PATH.unlink(missing_ok=True)
pid = None
else:
# Check whether the process is still alive.
try:
os.kill(pid, 0)
except ProcessLookupError:
# Already dead — clean up and note success.
TTYD_PID_PATH.unlink(missing_ok=True)
killed = True
pid = None
else:
# Process is alive — ask it to terminate.
os.kill(pid, signal.SIGTERM)
# Poll up to 2 s for the process to exit.
deadline = time.time() + 2.0
while time.time() < deadline:
try:
os.kill(pid, 0)
except (ProcessLookupError, PermissionError):
break
await asyncio.sleep(0.1)
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
return killed
async def kill_orphan_ttyd() -> bool:
"""Kill any orphaned ttyd process left over from a previous coordinator run.
Handles both legacy ``ttyd.pid`` and pool-era ``ttyd-{port}.pid`` files,
plus scans the full port range for stragglers.
Returns:
True — an orphan was found (process was dead or alive).
False — no PID file found, or PID file contained invalid content.
"""
# Pool-era cleanup: per-port PID files
if TTYD_PID_DIR.exists():
for pid_file in TTYD_PID_DIR.glob("ttyd-*.pid"):
try:
pid = int(pid_file.read_text().strip())
os.kill(pid, signal.SIGTERM)
except (ValueError, ProcessLookupError, PermissionError):
pass
pid_file.unlink(missing_ok=True)
# Scan additional ports beyond TTYD_PORT
for port in range(_PORT_MIN + 1, _PORT_MAX + 1):
_kill_pids_on_port(port, signal.SIGTERM)
# Original behavior: legacy PID file + TTYD_PORT cleanup
return await kill_ttyd()
async def spawn_ttyd(session_name: str) -> asyncio.subprocess.Process:
"""Spawn a ttyd process attached to *session_name* via ``tmux attach``.
Runs::
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
TTYD_PID_PATH. The process handle is stored in ``_active_process``.
Args:
session_name: The tmux session name to attach to.
Returns:
The asyncio.subprocess.Process object for the spawned ttyd.
"""
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(
"ttyd",
"-W",
"-m",
"3",
"-p",
str(TTYD_PORT),
"tmux",
"attach",
"-t",
session_name,
stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL,
start_new_session=True, # detach from parent process group so ttyd survives independently
)
# Write PID file (create parent dirs if needed)
TTYD_PID_DIR.mkdir(parents=True, exist_ok=True)
TTYD_PID_PATH.write_text(str(proc.pid))
_active_process = proc
return proc