From a10ae9895e074fc31ada6a2ddfc6968d6508441f Mon Sep 17 00:00:00 2001 From: lennney <94768569+lennney@users.noreply.github.com> Date: Thu, 25 Jun 2026 01:09:49 +0800 Subject: [PATCH] fix(wrap): detect stale proxy on target port and auto-cleanup before starting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds _ensure_port_free() and helpers to _start_proxy() that: - Detect if target port is already in use (via existing _port_bind_error) - Find the owning process via /proc/net/tcp (Linux) or lsof (macOS) - Only kills processes identified as stale headroom proxies - Reports clear error for non-headroom processes - Uses zero new dependencies - macOS _is_headroom_proxy now uses ps fallback when /proc unavailable - Windows SIGKILL fallback (getattr) for _kill_process Fixes the case where a terminal is killed, leaving an orphaned headroom proxy on the port — the next wrap would wait 30s then fail with a confusing RuntimeError. Now it auto-cleans and restarts. Tests: 14 new tests for all helpers + edge cases (54 total, all pass) --- headroom/cli/wrap.py | 188 ++++++++++++++++++++++ tests/test_cli/test_wrap_helpers.py | 241 ++++++++++++++++++++++++++++ 2 files changed, 429 insertions(+) diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 216073649..39ed5f102 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -357,6 +357,188 @@ def _get_proxy_stdio_log_path() -> Path: return _get_log_path().with_name("proxy-stdio.log") +def _find_process_on_port(port: int) -> int | None: + """Find PID of process listening on a TCP port. + + Linux: parse /proc/net/tcp for the port, resolve socket inode to PID + via /proc/*/fd/. + Fallback: ``lsof -ti tcp:{port}`` for macOS / other POSIX. + + Returns ``None`` when the port is free or resolution fails. + """ + if sys.platform == "linux": + pid = _linux_find_process_on_port(port) + if pid is not None: + return pid + # Fallback: lsof (macOS, other POSIX) + try: + result = subprocess.run( + ["lsof", "-ti", f"tcp:{port}"], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0 and result.stdout.strip(): + return int(result.stdout.strip().splitlines()[0]) + except (OSError, ValueError, subprocess.TimeoutExpired): + pass + return None + + +def _linux_find_process_on_port(port: int) -> int | None: + """Parse /proc/net/tcp and /proc/net/tcp6 to find the PID owning ``port``. + + Returns ``None`` when no listening socket is found on that port. + """ + for net_path in ("/proc/net/tcp", "/proc/net/tcp6"): + try: + with open(net_path) as f: + # Skip header line + next(f, None) + for line in f: + parts = line.strip().split() + if len(parts) < 10: + continue + # local_address column: "00000000:2251" (hex) + local_addr = parts[1] + if ":" not in local_addr: + continue + # Only match listening sockets (state 0A) + st = parts[3] if len(parts) > 3 else "" + if st != "0A": + continue + _, hex_port = local_addr.rsplit(":", 1) + try: + if int(hex_port, 16) != port: + continue + except ValueError: + continue + inode = parts[9] + if inode and inode.isdigit(): + return _resolve_inode_to_pid(int(inode)) + except OSError: + continue + return None + + +def _resolve_inode_to_pid(inode: int) -> int | None: + """Search /proc/*/fd/* for a socket symlink whose inode matches. + + Returns the PID or ``None``. + """ + try: + for proc_dir in Path("/proc").iterdir(): + if not proc_dir.name.isdigit(): + continue + try: + fd_dir = proc_dir / "fd" + for fd in fd_dir.iterdir(): + try: + link = os.readlink(fd) + if link == f"socket:[{inode}]": + return int(proc_dir.name) + except OSError: + continue + except (PermissionError, OSError): + continue + except OSError: + pass + return None + + +def _is_headroom_proxy(pid: int) -> bool: + """Check whether *pid* is a headroom proxy process (by its cmdline). + + Uses ``/proc//cmdline`` (Linux) or ``ps -p -o command=`` + (macOS/BSD fallback) so ``_find_process_on_port`` via ``lsof`` can + work on macOS. + + Returns ``False`` when the cmdline can't be read. + """ + cmdline = _read_process_cmdline(pid) + if cmdline is None: + return False + return "headroom" in cmdline and "proxy" in cmdline + + +def _read_process_cmdline(pid: int) -> str | None: + """Read the command-line of *pid*, cross-platform. + + Linux: reads ``/proc//cmdline`` (null-byte separated). + macOS/BSD: ``ps -p -o command=`` (no args). + + Returns ``None`` when the cmdline can't be read. + """ + # Linux path — /proc//cmdline + try: + return Path(f"/proc/{pid}/cmdline").read_bytes().decode("utf-8", errors="replace") + except OSError: + pass + # macOS / BSD fallback — ps -p -o command= + try: + result = subprocess.run( + ["ps", "-p", str(pid), "-o", "command="], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0 and result.stdout.strip(): + return result.stdout.strip() + except (OSError, subprocess.TimeoutExpired): + pass + return None + + +def _kill_process(pid: int) -> None: + """Gracefully terminate *pid*, then escalate if it does not exit within 3s. + + Uses ``signal.SIGKILL`` on Unix; falls back to ``signal.SIGTERM`` on + platforms where ``SIGKILL`` is not available (Windows). + """ + try: + os.kill(pid, signal.SIGTERM) + for _ in range(3): + time.sleep(1) + try: + os.kill(pid, 0) # probe liveness — raises if gone + except OSError: + return + # Escalate — SIGKILL on Unix, SIGTERM (harmless re-send) on Windows + force = getattr(signal, "SIGKILL", signal.SIGTERM) + os.kill(pid, force) + except OSError: + pass + + +def _ensure_port_free(port: int) -> bool: + """Ensure *port* is available for the proxy. + + Uses the existing ``_port_bind_error`` for a fast socket.bind probe. + When the port is occupied and the owning process is a stale headroom + proxy, it is killed automatically. Other processes are left alone and + the caller is expected to raise a user-facing error. + + Returns ``True`` when the port is (now) free, ``False`` when it is + held by a non-headroom process. + """ + bind_error = _port_bind_error(port) + if bind_error is None: + return True # port is free + + pid = _find_process_on_port(port) + if pid is None: + # Port became free between check and now + return True + + if not _is_headroom_proxy(pid): + return False # non-headroom process — let caller decide + + _kill_process(pid) + # Verify port is now free + bind_error = _port_bind_error(port) + return bind_error is None + + def _start_proxy( port: int, *, @@ -377,6 +559,12 @@ def _start_proxy( `~/.headroom/logs/proxy-stdio.log`, to avoid pipe deadlock risk without competing with the rotating `proxy.log` runtime log. """ + if not _ensure_port_free(port): + raise click.ClickException( + f"Port {port} is in use by a non-headroom process. " + f"Please free it manually or use --port to specify a different port." + ) + cmd = [sys.executable, "-m", "headroom.cli", "proxy", "--port", str(port)] # Forward HEADROOM_MODE env var so the proxy respects the user's mode choice diff --git a/tests/test_cli/test_wrap_helpers.py b/tests/test_cli/test_wrap_helpers.py index 87806e5e5..e061f2abf 100644 --- a/tests/test_cli/test_wrap_helpers.py +++ b/tests/test_cli/test_wrap_helpers.py @@ -789,3 +789,244 @@ def test_resolve_1m_model_falls_back_to_default_when_unset() -> None: """With no model selected, fall back to the default Opus carrying [1m].""" assert wrap_mod._resolve_1m_model(None) == "claude-opus-4-8[1m]" assert wrap_mod._resolve_1m_model(" ") == "claude-opus-4-8[1m]" + + +class TestEnsurePortFree: + """Tests for _ensure_port_free and its helper functions.""" + + def test_ensure_port_free_port_is_free(self, monkeypatch: pytest.MonkeyPatch) -> None: + """When the port is free, _ensure_port_free returns True immediately.""" + monkeypatch.setattr(wrap_mod, "_port_bind_error", lambda port: None) + assert wrap_mod._ensure_port_free(8787) is True + + def test_ensure_port_free_stale_headroom(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A stale headroom proxy is killed and the port becomes free.""" + killed = [] + bind_calls = [] + + def _mock_bind(port): + bind_calls.append(port) + # First call: occupied; second call (after kill): free + if len(bind_calls) == 1: + return OSError(98, "Address in use") + return None + + monkeypatch.setattr(wrap_mod, "_port_bind_error", _mock_bind) + monkeypatch.setattr(wrap_mod, "_find_process_on_port", lambda port: 12345) + monkeypatch.setattr(wrap_mod, "_is_headroom_proxy", lambda pid: True) + monkeypatch.setattr(wrap_mod, "_kill_process", lambda pid: killed.append(pid)) + + assert wrap_mod._ensure_port_free(8787) is True + assert killed == [12345] + + def test_ensure_port_free_non_headroom(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A non-headroom process on the port returns False.""" + monkeypatch.setattr(wrap_mod, "_port_bind_error", lambda port: OSError(98, "Address in use")) + monkeypatch.setattr(wrap_mod, "_find_process_on_port", lambda port: 12345) + monkeypatch.setattr(wrap_mod, "_is_headroom_proxy", lambda pid: False) + + assert wrap_mod._ensure_port_free(8787) is False + + def test_linux_find_process_on_port_empty(self, monkeypatch: pytest.MonkeyPatch) -> None: + """When /proc/net/tcp has no matching entry, returns None.""" + fake_tcp = " sl local_address rem_address st tx_queue rx_queue tr tm->when retrnsmt uid timeout inode\n" + + def _mock_open(path, *args, **kwargs): + if path == "/proc/net/tcp": + from io import StringIO + return StringIO(fake_tcp) + raise FileNotFoundError(path) + + monkeypatch.setattr("builtins.open", _mock_open) + assert wrap_mod._linux_find_process_on_port(8787) is None + + def test_linux_find_process_on_port_found(self, monkeypatch: pytest.MonkeyPatch) -> None: + """When /proc/net/tcp has a matching port, resolves and returns PID.""" + # Port 8787 = 0x2253 + fake_tcp = ( + " sl local_address rem_address st tx_queue rx_queue tr tm->when retrnsmt uid timeout inode\n" + " 0: 00000000:2253 00000000:0000 0A 00000000:00000000 00:00000000 00000000 1000 0 123456 1 0000000000000000 100 0 0 10 0\n" + ) + + def _mock_open(path, *args, **kwargs): + if path == "/proc/net/tcp": + from io import StringIO + return StringIO(fake_tcp) + raise FileNotFoundError(path) + + monkeypatch.setattr("builtins.open", _mock_open) + + # Mock _resolve_inode_to_pid to return a known PID + monkeypatch.setattr(wrap_mod, "_resolve_inode_to_pid", lambda inode: 99999) + + assert wrap_mod._linux_find_process_on_port(8787) == 99999 + + def test_is_headroom_proxy_true(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A process with 'headroom' and 'proxy' in cmdline is identified as headroom.""" + monkeypatch.setattr( + "pathlib.Path.read_bytes", + lambda self: b"headroom\x00proxy\x00--port\x008787\x00", + ) + assert wrap_mod._is_headroom_proxy(12345) is True + + def test_is_headroom_proxy_false(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A process without 'headroom' is not identified as such.""" + monkeypatch.setattr( + "pathlib.Path.read_bytes", + lambda self: b"nginx\x00", + ) + assert wrap_mod._is_headroom_proxy(12345) is False + + def test_read_process_cmdline_linux_path(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Linux /proc//cmdline path returns the decoded cmdline.""" + monkeypatch.setattr( + "pathlib.Path.read_bytes", + lambda self: b"headroom\x00proxy\x00--port\x008787\x00", + ) + result = wrap_mod._read_process_cmdline(12345) + assert result is not None + assert "headroom" in result + assert "proxy" in result + + def test_read_process_cmdline_macos_fallback(self, monkeypatch: pytest.MonkeyPatch) -> None: + """When /proc is unavailable, falls back to ps command (macOS/BSD).""" + # Make /proc path fail + monkeypatch.setattr("pathlib.Path.read_bytes", lambda self: (_ for _ in ()).throw(FileNotFoundError)) + + import subprocess as _subprocess + + def _mock_run(cmd, *a, **kw): + assert cmd[0] == "ps" + assert "-p" in cmd + assert str(12345) in cmd + return _subprocess.CompletedProcess(cmd, 0, "headroom proxy --port 8787\n", "") + + monkeypatch.setattr(wrap_mod.subprocess, "run", _mock_run) + result = wrap_mod._read_process_cmdline(12345) + assert result == "headroom proxy --port 8787" + + def test_read_process_cmdline_failure(self, monkeypatch: pytest.MonkeyPatch) -> None: + """When both paths fail, returns None.""" + monkeypatch.setattr("pathlib.Path.read_bytes", lambda self: (_ for _ in ()).throw(FileNotFoundError)) + monkeypatch.setattr(wrap_mod.subprocess, "run", lambda *a, **kw: (_ for _ in ()).throw(FileNotFoundError)) + assert wrap_mod._read_process_cmdline(12345) is None + + def test_kill_process_terminates(self, monkeypatch: pytest.MonkeyPatch) -> None: + """_kill_process sends SIGTERM and returns when process dies.""" + calls = [] + + def _mock_kill(pid, sig): + calls.append((pid, sig)) + if sig == 15: # SIGTERM — die after first check + raise OSError(3, "No such process") + + monkeypatch.setattr(os, "kill", _mock_kill) + wrap_mod._kill_process(12345) + assert calls == [(12345, 15)] # Only SIGTERM, no SIGKILL needed + + def test_kill_process_force_kill(self, monkeypatch: pytest.MonkeyPatch) -> None: + """When SIGTERM doesn't work, _kill_process escalates (SIGKILL on Unix, SIGTERM on Windows).""" + import signal as _signal + ESCALATION = getattr(_signal, "SIGKILL", _signal.SIGTERM) + + calls = [] + + def _mock_kill(pid, sig): + calls.append((pid, sig)) + if sig == ESCALATION: # escalation signal — let it die + raise OSError(3, "No such process") + # SIGTERM and liveness probes succeed (process stays alive) + + monkeypatch.setattr(os, "kill", _mock_kill) + monkeypatch.setattr(wrap_mod.time, "sleep", lambda s: None) + wrap_mod._kill_process(12345) + assert len(calls) >= 2 + assert calls[0] == (12345, 15) # SIGTERM first + assert calls[-1] == (12345, ESCALATION) # escalation signal last + + def test_resolve_inode_to_pid_matches_symlink(self, monkeypatch: pytest.MonkeyPatch) -> None: + """_resolve_inode_to_pid matches the exact socket:[inode] symlink format.""" + inode = 123456 + + # Mock Path.iterdir to return a fake /proc/ structure + class FakeFd: + def __init__(self, name, link): + self.name = name + self._link = link + self.parent = self + + def is_symlink(self): + return True + + def __truediv__(self, other): + return self + + itercount = [0] + + def _patched_iterdir(self): + itercount[0] += 1 + if itercount[0] == 1: + # First call: /proc/ — return one PID dir + pid_dir = FakePath("99999", Path("/proc")) + return iter([pid_dir]) + return iter([]) + + class FakePath: + def __init__(self, name, parent=None): + self.name = name + self._parent = parent or self + + @property + def parent(self): + return self._parent + + def iterdir(self): + if self.name == "99999": + # Return fd dir + fd_dir = FakePath("fd", self) + fd_dir._links = [FakePath("3", self)] + return iter([fd_dir]) + # Return socket links + fd = FakePath("3", self) + fd._link = f"socket:[{inode}]" + return iter([fd]) + + def is_symlink(self): + return True + + def __truediv__(self, other): + return FakePath(str(other), self) + + def _mock_readlink(fd_obj): + return f"socket:[{inode}]" + + monkeypatch.setattr(Path, "iterdir", lambda self: _patched_iterdir(self)) + monkeypatch.setattr(os, "readlink", _mock_readlink) + + result = wrap_mod._resolve_inode_to_pid(inode) + assert result == 99999, f"Expected PID 99999, got {result!r}" + + def test_linux_find_process_on_port_tcp6(self, monkeypatch: pytest.MonkeyPatch) -> None: + """When port is only in /proc/net/tcp6 (IPv6), it should still be found.""" + inode = 654321 + # Port 18888 = 0x49C8, in tcp6 format (IPv6 addr) + fake_tcp6 = ( + " sl local_address rem_address " + "st tx_queue rx_queue tr tm->when retrnsmt uid timeout inode\n" + f" 0: 00000000000000000000000000000000:49C8 00000000000000000000000000000000:0000 " + f"0A 00000000:00000000 00:00000000 00000000 1000 0 {inode} 1 0000000000000000 100 0 0 10 0\n" + ) + + def _mock_open(path, *args, **kwargs): + if path == "/proc/net/tcp": + raise FileNotFoundError(path) # not in tcp + if path == "/proc/net/tcp6": + from io import StringIO + return StringIO(fake_tcp6) + raise FileNotFoundError(path) + + monkeypatch.setattr("builtins.open", _mock_open) + monkeypatch.setattr(wrap_mod, "_resolve_inode_to_pid", lambda ino: 88888) + + result = wrap_mod._linux_find_process_on_port(18888) + assert result == 88888, f"Expected PID 88888 from tcp6, got {result!r}"