refactor: delete ttyd.py, WS proxy, connect endpoint — muxterm owns terminal path
This commit is contained in:
Ken
+17
-161
@@ -725,50 +725,19 @@ async def create_session(payload: CreateSessionPayload) -> dict:
|
||||
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")
|
||||
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
|
||||
in persistent state.
|
||||
Sets active_session to None in persistent state.
|
||||
|
||||
Returns {active_session: None}.
|
||||
"""
|
||||
async with state_lock:
|
||||
state = load_state()
|
||||
active = state.get("active_session")
|
||||
state["active_session"] = None
|
||||
save_state(state)
|
||||
|
||||
if active:
|
||||
await kill_session(active)
|
||||
|
||||
return {"active_session": None}
|
||||
|
||||
|
||||
@@ -816,10 +785,6 @@ async def delete_session(name: str) -> dict:
|
||||
except Exception:
|
||||
_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}
|
||||
|
||||
|
||||
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1191,9 +1035,21 @@ async def federation_terminal_ws_proxy(websocket: WebSocket, device_id: str) ->
|
||||
Auth check uses the same cookie + bearer pattern as terminal_ws_proxy.
|
||||
Closes with code 4004 if device_id does not match any remote.
|
||||
"""
|
||||
# Auth check before accepting — same pattern as terminal_ws_proxy
|
||||
if not await _ws_auth_check(websocket):
|
||||
return
|
||||
# Auth check before accepting — BaseHTTPMiddleware doesn't cover WS scope
|
||||
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
|
||||
|
||||
# Look up remote instance by device_id
|
||||
remote = _lookup_remote_by_device_id(device_id)
|
||||
|
||||
+2
-196
@@ -356,76 +356,13 @@ def test_get_sessions_returns_empty_list_when_no_sessions(client, monkeypatch):
|
||||
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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delete_current_kills_ttyd_and_clears_active(client, monkeypatch):
|
||||
"""DELETE /api/sessions/current kills the session's ttyd and clears active_session."""
|
||||
def test_delete_current_clears_active_session(client, monkeypatch):
|
||||
"""DELETE /api/sessions/current clears active_session in state."""
|
||||
from muxplex.state import load_state, save_state
|
||||
|
||||
# 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")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["active_session"] is None
|
||||
assert kill_session_calls == ["alpha"]
|
||||
|
||||
# Verify state was persisted
|
||||
state = load_state()
|
||||
@@ -800,20 +728,6 @@ def test_api_routes_not_shadowed(client):
|
||||
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
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -967,90 +881,6 @@ def test_logout_clears_session_cookie(monkeypatch):
|
||||
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
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -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():
|
||||
"""cli.py serve() must pass log_level='info' to uvicorn.run so logs appear in journalctl."""
|
||||
import inspect
|
||||
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
)
|
||||
-471
@@ -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
|
||||
Reference in New Issue
Block a user