From 05bd56bcb6b103fab5522da2b14295cf7bd8dbc1 Mon Sep 17 00:00:00 2001 From: Devanshi Vyas Date: Thu, 11 Jun 2026 17:42:43 -0700 Subject: [PATCH] fix(wrap): track shared proxy clients with markers (#877) ## Description Replace argv-based proxy client detection with per-port wrap client markers so cleanup and ephemeral restarts do not tear down a shared proxy while another wrapped session is still attached. Also prune stale markers, guard against PID reuse when process identity is available, and add coverage for the marker-based lifecycle behavior. Fixes #804 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Testing Describe the tests you ran to verify your changes: - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality --------- Co-authored-by: JerrettDavis --- headroom/cli/wrap.py | 242 ++++++++++++++++++++----- headroom/paths.py | 8 + tests/test_cli/test_wrap_helpers.py | 191 +++++++++++++++++++ tests/test_cli/test_wrap_persistent.py | 148 +++++++++++++++ tests/test_paths_backward_compat.py | 1 + 5 files changed, 543 insertions(+), 47 deletions(-) diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 2d26bcf26..050a6f642 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -1241,6 +1241,7 @@ def _run_proxy_only_watcher( """ proxy_holder: list[subprocess.Popen | None] = [None] cleanup = _make_cleanup(proxy_holder, port) + _register_proxy_client(port) signal.signal(signal.SIGINT, cleanup) signal.signal(signal.SIGTERM, cleanup) @@ -1845,13 +1846,19 @@ def _ensure_proxy( if helpers._proxy_needs_version_restart(health_payload): running_version = helpers._proxy_version(health_payload) or "unknown" active_sessions = helpers._proxy_active_session_count(health_payload) - if active_sessions > 0: + other_wrappers = helpers._live_proxy_clients(port, exclude_self=True) + if active_sessions > 0 or other_wrappers: + detail = ( + f"{active_sessions} active session(s)" + if active_sessions > 0 + else f"{len(other_wrappers)} attached wrapper(s)" + ) click.echo( f" Proxy on port {port} is running Headroom {running_version}; " f"current CLI is {_HEADROOM_VERSION}." ) click.echo( - f" Leaving it running because {active_sessions} active session(s) " + f" Leaving it running because {detail} " "are still attached; it will be restarted when idle." ) return None @@ -1885,13 +1892,23 @@ def _ensure_proxy( if helpers._proxy_needs_version_restart(health_payload): running_version = helpers._proxy_version(health_payload) or "unknown" active_sessions = helpers._proxy_active_session_count(health_payload) - if active_sessions > 0: + other_wrappers = helpers._live_proxy_clients(port, exclude_self=True) + if active_sessions > 0 or other_wrappers: + # active_sessions only counts Codex WebSocket relay; the + # marker list also covers HTTP wrap clients. Either means a + # live session is attached, so don't restart the shared + # proxy out from under it — defer until idle. + detail = ( + f"{active_sessions} active session(s)" + if active_sessions > 0 + else f"{len(other_wrappers)} attached wrapper(s)" + ) click.echo( f" Proxy on port {port} is running Headroom {running_version}; " f"current CLI is {_HEADROOM_VERSION}." ) click.echo( - f" Leaving it running because {active_sessions} active session(s) " + f" Leaving it running because {detail} " "are still attached; it will be restarted when idle." ) return None @@ -1936,35 +1953,49 @@ def _ensure_proxy( missing.append("openai-api-url") if missing: - needs_restart = True flags_str = ", ".join( f if f.startswith("--") else f"--{f.replace('_', '-')}" for f in missing ) - click.echo(f" Proxy on port {port} is missing: {flags_str}") - click.echo(" Restarting proxy with upgraded configuration...") - - # Merge: keep features the running proxy already has - memory = memory or bool(running_config.get("memory")) - learn = learn or bool(running_config.get("learn")) - code_graph = code_graph or bool(running_config.get("code_graph")) - - proxy_pid = running_config.get("pid") - if proxy_pid is not None: - if not helpers._kill_proxy_by_pid(int(proxy_pid), port): - raise click.ClickException( - f"Failed to stop existing proxy (PID {proxy_pid}) on port {port}. " - "Stop it manually and retry." - ) + other_wrappers = helpers._live_proxy_clients(port, exclude_self=True) + if other_wrappers: + # Another wrapper is attached to this proxy; restarting it + # to add flags would drop their in-flight requests. Reuse + # the running proxy as-is rather than disrupt them. + click.echo( + f" Proxy on port {port} is missing: {flags_str}, but " + f"{len(other_wrappers)} other wrapper(s) are attached." + ) + click.echo( + " Leaving it running to avoid disrupting them; this " + "session will use the existing proxy as-is." + ) else: - click.echo( - " Warning: Running proxy does not expose PID. " - "Cannot restart automatically." - ) - click.echo( - f" Please stop the proxy on port {port} manually " - f"and rerun with {flags_str}." - ) - return None + needs_restart = True + click.echo(f" Proxy on port {port} is missing: {flags_str}") + click.echo(" Restarting proxy with upgraded configuration...") + + # Merge: keep features the running proxy already has + memory = memory or bool(running_config.get("memory")) + learn = learn or bool(running_config.get("learn")) + code_graph = code_graph or bool(running_config.get("code_graph")) + + proxy_pid = running_config.get("pid") + if proxy_pid is not None: + if not helpers._kill_proxy_by_pid(int(proxy_pid), port): + raise click.ClickException( + f"Failed to stop existing proxy (PID {proxy_pid}) on port {port}. " + "Stop it manually and retry." + ) + else: + click.echo( + " Warning: Running proxy does not expose PID. " + "Cannot restart automatically." + ) + click.echo( + f" Please stop the proxy on port {port} manually " + f"and rerun with {flags_str}." + ) + return None if not needs_restart: click.echo(f" Proxy already running on port {port}") @@ -2006,35 +2037,150 @@ def _ensure_proxy( return None +def _client_marker_path(port: int) -> Path: + """Path to this process's wrap-client marker for ``port``.""" + from headroom import paths as _paths + + d = _paths.proxy_clients_dir(port) + d.mkdir(parents=True, exist_ok=True) + return d / f"{os.getpid()}.json" + + +def _proc_identity(pid: int) -> tuple[str, float] | None: + """Best-effort ``(source, start_time)`` identity for a PID. + + Used to defeat PID reuse: a marker is only trusted while the live PID is + *the same process* that wrote it. Returns ``None`` when start time can't be + determined (e.g. macOS without psutil), in which case callers fall back to + existence-only liveness — no regression, just no reuse protection there. + + The ``source`` tag ("psutil" vs "proc") guards against comparing values in + different units; we only compare like-for-like. + """ + try: + import psutil # optional dependency; portable when present + + return ("psutil", float(psutil.Process(pid).create_time())) + except Exception: + pass + # Linux fallback: field 22 of /proc//stat is starttime in clock ticks + # since boot — a stable per-process value. `comm` (field 2) may contain + # spaces/parens, so split after the final ')'. + try: + with open(f"/proc/{pid}/stat", "rb") as fh: + fields = fh.read().rpartition(b")")[2].split() + return ("proc", float(fields[19])) + except (OSError, IndexError, ValueError): + return None + + +def _register_proxy_client(port: int) -> None: + """Register this wrap process as a live client of the shared proxy. + + Best-effort: a failed write just means our marker is missing, and the + liveness pruning in :func:`_live_proxy_clients` is the real safety net. + """ + try: + payload: dict[str, Any] = {"pid": os.getpid(), "started_at": time.time()} + ident = _proc_identity(os.getpid()) + if ident is not None: + payload["start_src"], payload["start_time"] = ident + _client_marker_path(port).write_text(json.dumps(payload)) + except OSError: + pass + + +def _unregister_proxy_client(port: int) -> None: + """Remove this process's client marker (idempotent).""" + try: + _client_marker_path(port).unlink(missing_ok=True) + except OSError: + pass + + +def _pid_alive(pid: int) -> bool: + """Return True if ``pid`` names a live process.""" + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True # exists but owned by another user + except OSError: + return False + return True + + +def _marker_pid_reused(marker: Path, pid: int) -> bool: + """True only if the live ``pid`` is *provably* a different process than the + one that wrote ``marker`` (i.e. the PID was recycled after a crash). + + Conservative by design: any uncertainty (legacy marker, unknown start time, + mismatched source) returns ``False`` so a real client is never pruned. + """ + try: + rec = json.loads(marker.read_text()) + except (OSError, ValueError): + return False + src = rec.get("start_src") + recorded = rec.get("start_time") + if not isinstance(src, str) or not isinstance(recorded, (int, float)): + return False # legacy / identity-less marker — can't tell + ident = _proc_identity(pid) + if ident is None or ident[0] != src: + return False # can't compare like-for-like — don't prune + # Start times are stable per process; >1s apart means a different process. + return abs(ident[1] - float(recorded)) > 1.0 + + +def _live_proxy_clients(port: int, *, exclude_self: bool = True) -> list[int]: + """Live wrap-client PIDs for ``port``, pruning stale markers as we go.""" + from headroom import paths as _paths + + d = _paths.proxy_clients_dir(port) + if not d.exists(): + return [] + me = os.getpid() + live: list[int] = [] + for marker in d.glob("*.json"): + try: + pid = int(marker.stem) + except ValueError: + continue + # Stale if the PID is gone, or recycled by an unrelated process. + if not _pid_alive(pid) or _marker_pid_reused(marker, pid): + try: + marker.unlink(missing_ok=True) + except OSError: + pass + continue + if not (exclude_self and pid == me): + live.append(pid) + return live + + def _make_cleanup(proxy_proc_holder: list, port: int = 8787) -> Any: """Create a cleanup function that terminates the proxy on exit. - Only kills the proxy if no other headroom-wrapped clients are using it. - Checks by looking for other processes with ANTHROPIC_BASE_URL or - OPENAI_BASE_URL pointing at our port. + Only kills the proxy when no other live headroom-wrapped clients remain, + tracked via per-PID marker files in ``paths.proxy_clients_dir(port)``. """ def _other_clients_exist() -> bool: - """Check if other processes are using this proxy.""" - try: - # Count headroom wrap processes (excluding ourselves) - result = subprocess.run( - ["pgrep", "-f", f"127.0.0.1:{port}"], - capture_output=True, - text=True, - ) - pids = [p.strip() for p in result.stdout.strip().split("\n") if p.strip()] - my_pid = str(os.getpid()) - other_pids = [p for p in pids if p != my_pid] - return len(other_pids) > 0 - except Exception: - return False # If we can't check, assume no others + # Reference-count from marker files, not argv scans. Wrapped clients + # carry the proxy URL in ANTHROPIC_BASE_URL/OPENAI_BASE_URL (env, not + # argv), so `pgrep -f` could never see them — and it matched unrelated + # processes by substring. Markers are exact and OS-portable. + return len(_live_proxy_clients(port, exclude_self=True)) > 0 def cleanup(signum: int | None = None, frame: Any = None) -> None: + # Drop our own marker first so the count reflects the post-exit state; + # also covers the signal path, where the `finally` block may not run. + _unregister_proxy_client(port) proc = proxy_proc_holder[0] if proxy_proc_holder else None if proc and proc.poll() is None: if _other_clients_exist(): - # Other clients still using the proxy — leave it running + # Other clients still using the proxy — leave it running. return proc.terminate() try: @@ -2073,6 +2219,7 @@ def _launch_tool( """Common logic: start proxy, launch tool, clean up.""" proxy_holder: list[subprocess.Popen | None] = [None] cleanup = _make_cleanup(proxy_holder, port) + _register_proxy_client(port) signal.signal(signal.SIGINT, _ignore_child_sigint) signal.signal(signal.SIGTERM, cleanup) @@ -2447,6 +2594,7 @@ def claude( # Setup rtk before launching (Claude-specific) proxy_holder: list[subprocess.Popen | None] = [None] cleanup = _make_cleanup(proxy_holder, port) + _register_proxy_client(port) signal.signal(signal.SIGINT, _ignore_child_sigint) signal.signal(signal.SIGTERM, cleanup) diff --git a/headroom/paths.py b/headroom/paths.py index eb1633f8d..fda8f33af 100644 --- a/headroom/paths.py +++ b/headroom/paths.py @@ -73,6 +73,7 @@ _PROXY_LOG_FILE = "proxy.log" _DEBUG_400_DIR = "debug_400" _CODEX_WIRE_DEBUG_DIR = "codex_wire" _BIN_DIR = "bin" +_PROXY_CLIENTS_DIR = "clients" _RTK_UNIX = "rtk" _RTK_WIN = "rtk.exe" _LEAN_CTX_UNIX = "lean-ctx" @@ -269,6 +270,12 @@ def bin_dir() -> Path: return workspace_dir() / _BIN_DIR +def proxy_clients_dir(port: int) -> Path: + """Per-port dir of live wrap-client markers (one file per client PID).""" + + return workspace_dir() / _PROXY_CLIENTS_DIR / str(port) + + def rtk_path() -> Path: """Return the path to the vendored ``rtk`` binary.""" @@ -357,6 +364,7 @@ __all__ = [ "debug_400_dir", "codex_wire_debug_dir", "bin_dir", + "proxy_clients_dir", "rtk_path", "lean_ctx_path", "deploy_root", diff --git a/tests/test_cli/test_wrap_helpers.py b/tests/test_cli/test_wrap_helpers.py index fbb56460d..81cf1cf50 100644 --- a/tests/test_cli/test_wrap_helpers.py +++ b/tests/test_cli/test_wrap_helpers.py @@ -12,6 +12,10 @@ once with confusing diffs. from __future__ import annotations +import json +import os +import subprocess +import sys from pathlib import Path from typing import Any @@ -19,6 +23,7 @@ import click import pytest from click.testing import CliRunner +from headroom import paths as paths_mod from headroom.cli import wrap as wrap_mod # --------------------------------------------------------------------------- @@ -516,3 +521,189 @@ class TestApplyProjectHeaderEnv: monkeypatch.chdir(project_dir) assert wrap_mod._project_name_from_cwd() == "vibe-headroom" + + +# --------------------------------------------------------------------------- +# Proxy-client reference counting +# +# The shared proxy must only be torn down by its owner once *no* other live +# wrap clients remain. Clients carry the proxy URL in ANTHROPIC_BASE_URL / +# OPENAI_BASE_URL (env, not argv), so the old `pgrep -f "127.0.0.1:"` +# guard could neither see real clients nor reject unrelated processes that +# merely had the address in their command line. These tests pin the new +# marker-file contract: a per-PID file under paths.proxy_clients_dir(port). +# --------------------------------------------------------------------------- + + +class _FakeProxyProc: + """Minimal stand-in for the proxy ``subprocess.Popen`` handle.""" + + def __init__(self) -> None: + self.terminated = False + self.killed = False + + def poll(self) -> int | None: + return None # alive + + def terminate(self) -> None: + self.terminated = True + + def wait(self, timeout: float | None = None) -> int: + return 0 + + def kill(self) -> None: + self.killed = True + + +class TestProxyClientRefCounting: + """Proxy lifecycle is reference-counted via marker files, not pgrep.""" + + PORT = 8787 + + @pytest.fixture + def clients_dir(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """Redirect ``paths.proxy_clients_dir`` into a throwaway tmp tree.""" + base = tmp_path / "clients" + monkeypatch.setattr(paths_mod, "proxy_clients_dir", lambda port: base / str(port)) + return base + + def _write_marker( + self, + clients_dir: Path, + pid: int, + *, + identity: tuple[str, float] | None = None, + ) -> Path: + marker = clients_dir / str(self.PORT) / f"{pid}.json" + marker.parent.mkdir(parents=True, exist_ok=True) + rec: dict[str, Any] = {"pid": pid, "started_at": 0} + if identity is not None: + rec["start_src"], rec["start_time"] = identity + marker.write_text(json.dumps(rec)) + return marker + + def test_cleanup_terminates_proxy_when_only_self_registered(self, clients_dir: Path) -> None: + """The owner alone → no other clients → proxy is terminated on exit.""" + wrap_mod._register_proxy_client(self.PORT) + proc = _FakeProxyProc() + cleanup = wrap_mod._make_cleanup([proc], self.PORT) + + cleanup() + + assert proc.terminated is True + # Our own marker is removed before we count. + assert wrap_mod._live_proxy_clients(self.PORT, exclude_self=False) == [] + + def test_cleanup_leaves_proxy_running_when_other_client_alive(self, clients_dir: Path) -> None: + """A second live client (here: the test's parent) keeps the proxy up.""" + wrap_mod._register_proxy_client(self.PORT) + other_pid = os.getppid() # alive for the duration of the test run + assert other_pid != os.getpid() + self._write_marker(clients_dir, other_pid) + + proc = _FakeProxyProc() + cleanup = wrap_mod._make_cleanup([proc], self.PORT) + cleanup() + + assert proc.terminated is False + + def test_dead_client_marker_is_pruned_and_not_counted(self, clients_dir: Path) -> None: + """A marker for a dead PID is pruned from disk and never counted.""" + # Spawn and reap a child so its PID is reliably dead (not a zombie). + child = subprocess.Popen([sys.executable, "-c", "pass"]) + child.wait() + dead_pid = child.pid + marker = self._write_marker(clients_dir, dead_pid) + + live = wrap_mod._live_proxy_clients(self.PORT, exclude_self=True) + + assert dead_pid not in live + assert not marker.exists() + + def test_reused_pid_with_mismatched_identity_is_pruned( + self, clients_dir: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A *live* PID that the original client no longer owns is pruned. + + Models the PID-reuse orphan path: a wrapper crashed, the OS later + recycled its PID for an unrelated long-lived process. `os.kill(pid, 0)` + succeeds, but the recorded start time no longer matches. + """ + live_pid = os.getppid() # alive, but not the process that "registered" + marker = self._write_marker(clients_dir, live_pid, identity=("psutil", 1000.0)) + # The process currently holding that PID started much later → reuse. + monkeypatch.setattr(wrap_mod, "_proc_identity", lambda p: ("psutil", 9000.0)) + + live = wrap_mod._live_proxy_clients(self.PORT, exclude_self=True) + + assert live_pid not in live + assert not marker.exists() + + def test_matching_identity_within_tolerance_is_kept( + self, clients_dir: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Same process (start time within tolerance) is a real client — kept.""" + live_pid = os.getppid() + marker = self._write_marker(clients_dir, live_pid, identity=("psutil", 1000.0)) + monkeypatch.setattr(wrap_mod, "_proc_identity", lambda p: ("psutil", 1000.4)) + + live = wrap_mod._live_proxy_clients(self.PORT, exclude_self=True) + + assert live_pid in live + assert marker.exists() + + def test_identity_check_skipped_when_source_unavailable( + self, clients_dir: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """No reuse protection (e.g. macOS w/o psutil) → fall back to existence.""" + live_pid = os.getppid() + self._write_marker(clients_dir, live_pid, identity=("psutil", 1000.0)) + # Start time unknowable for the live PID → must not prune a real client. + monkeypatch.setattr(wrap_mod, "_proc_identity", lambda p: None) + + live = wrap_mod._live_proxy_clients(self.PORT, exclude_self=True) + + assert live_pid in live + + def test_non_marker_files_are_ignored(self, clients_dir: Path) -> None: + """Stray non-numeric / non-json files don't crash or count as clients.""" + d = clients_dir / str(self.PORT) + d.mkdir(parents=True, exist_ok=True) + (d / "not-a-pid.json").write_text("{}") + (d / "README.txt").write_text("ignore me") + + assert wrap_mod._live_proxy_clients(self.PORT, exclude_self=True) == [] + + def test_cleanup_does_not_shell_out_to_pgrep( + self, clients_dir: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Regression: liveness is never inferred from an argv scan. + + An unrelated process whose command line contains ``127.0.0.1:8787`` + used to be a false positive (orphaning the proxy). The new path never + calls ``subprocess.run`` at all, so it can't be fooled by argv. + """ + wrap_mod._register_proxy_client(self.PORT) + + def _no_subprocess(*args: Any, **kwargs: Any) -> None: + raise AssertionError("cleanup must not invoke subprocess.run (no pgrep)") + + monkeypatch.setattr(wrap_mod.subprocess, "run", _no_subprocess) + + proc = _FakeProxyProc() + cleanup = wrap_mod._make_cleanup([proc], self.PORT) + cleanup() # must not raise + + assert proc.terminated is True + + def test_register_then_unregister_is_idempotent(self, clients_dir: Path) -> None: + """Register adds exactly our marker; unregister removes it; re-call is safe.""" + wrap_mod._register_proxy_client(self.PORT) + all_clients = wrap_mod._live_proxy_clients(self.PORT, exclude_self=False) + assert all_clients == [os.getpid()] + + wrap_mod._unregister_proxy_client(self.PORT) + assert wrap_mod._live_proxy_clients(self.PORT, exclude_self=False) == [] + + # Second unregister is a no-op, not an error. + wrap_mod._unregister_proxy_client(self.PORT) diff --git a/tests/test_cli/test_wrap_persistent.py b/tests/test_cli/test_wrap_persistent.py index bddfe77f2..7911a1c77 100644 --- a/tests/test_cli/test_wrap_persistent.py +++ b/tests/test_cli/test_wrap_persistent.py @@ -1,10 +1,22 @@ from __future__ import annotations import click +import pytest import headroom.cli.wrap as wrap_cli +@pytest.fixture(autouse=True) +def _no_attached_wrappers(monkeypatch: pytest.MonkeyPatch) -> None: + """Default: no other wrap clients attached, so restart paths are hermetic. + + The ephemeral restart guards consult ``_live_proxy_clients``; without this, a + real ``headroom wrap`` session on the dev's machine could make these tests + flaky. Individual tests override this to simulate attached wrappers. + """ + monkeypatch.setattr(wrap_cli, "_live_proxy_clients", lambda *a, **kw: []) + + class _Manifest: profile = "default" preset = "persistent-service" @@ -167,6 +179,41 @@ def test_ensure_proxy_leaves_active_stale_persistent_deployment_running(monkeypa assert result is None +def test_ensure_proxy_defers_persistent_restart_when_http_wrapper_attached( + monkeypatch, +) -> None: + """A stale persistent proxy is left running while marker-tracked HTTP + wrappers are attached, even when WebSocket session count is zero.""" + health = { + "version": "0.0.1", + "runtime": {"websocket_sessions": {"active_sessions": 0, "active_relay_tasks": 0}}, + "config": {"pid": 12345}, + } + + monkeypatch.setattr(wrap_cli, "_find_persistent_manifest", lambda port: _Manifest()) + monkeypatch.setattr("headroom.install.health.probe_ready", lambda url: True) + monkeypatch.setattr(wrap_cli, "_query_proxy_health", lambda port: health) + monkeypatch.setattr(wrap_cli, "_live_proxy_clients", lambda *a, **kw: [999]) + monkeypatch.setattr( + wrap_cli, + "_restart_persistent_proxy", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("attached persistent proxy should not restart") + ), + ) + monkeypatch.setattr( + wrap_cli, + "_start_proxy", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("replacement proxy should not start") + ), + ) + + result = wrap_cli._ensure_proxy(8787, False) + + assert result is None + + def test_find_persistent_manifest_prefers_default_profile(monkeypatch) -> None: class DefaultManifest: profile = "default" @@ -377,3 +424,104 @@ def test_ensure_proxy_leaves_active_stale_ephemeral_proxy_running(monkeypatch) - result = wrap_cli._ensure_proxy(8787, False) assert result is None + + +def test_ensure_proxy_defers_version_restart_when_http_wrapper_attached(monkeypatch) -> None: + """A stale-version proxy is NOT restarted while a marker-tracked HTTP + wrapper is attached, even though the WebSocket session count is zero.""" + health = { + "version": "0.0.1", # stale → version restart wanted + # No WebSocket relay sessions — the gap that let the old code kill it. + "runtime": {"websocket_sessions": {"active_sessions": 0, "active_relay_tasks": 0}}, + "config": {"pid": "12345", "memory": False, "learn": False, "code_graph": False}, + } + + monkeypatch.setattr(wrap_cli, "_find_persistent_manifest", lambda port: None) + monkeypatch.setattr(wrap_cli, "_check_proxy", lambda port: True) + monkeypatch.setattr(wrap_cli, "_query_proxy_health", lambda port: health) + # Another HTTP wrapper (PID 999) is attached per the marker registry. + monkeypatch.setattr(wrap_cli, "_live_proxy_clients", lambda *a, **kw: [999]) + monkeypatch.setattr( + wrap_cli, + "_kill_proxy_by_pid", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("attached proxy must not be killed for a version restart") + ), + ) + monkeypatch.setattr( + wrap_cli, + "_start_proxy", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("replacement proxy must not start") + ), + ) + + result = wrap_cli._ensure_proxy(8787, False) + + assert result is None + + +def test_ensure_proxy_defers_flag_restart_when_other_wrapper_attached(monkeypatch) -> None: + """Requesting --memory must not restart the proxy out from under another + attached wrapper; reuse the running proxy as-is instead.""" + health = { + "version": wrap_cli._HEADROOM_VERSION, # same version → no version restart + "runtime": {"websocket_sessions": {"active_sessions": 0, "active_relay_tasks": 0}}, + # Running proxy lacks `memory`; this session asks for it. + "config": {"pid": "12345", "memory": False, "learn": False, "code_graph": False}, + } + + monkeypatch.setattr(wrap_cli, "_find_persistent_manifest", lambda port: None) + monkeypatch.setattr(wrap_cli, "_check_proxy", lambda port: True) + monkeypatch.setattr(wrap_cli, "_query_proxy_health", lambda port: health) + monkeypatch.setattr(wrap_cli, "_live_proxy_clients", lambda *a, **kw: [999]) + monkeypatch.setattr( + wrap_cli, + "_kill_proxy_by_pid", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("attached proxy must not be killed to add flags") + ), + ) + monkeypatch.setattr( + wrap_cli, + "_start_proxy", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("replacement proxy must not start") + ), + ) + + result = wrap_cli._ensure_proxy(8787, False, memory=True) + + assert result is None + + +def test_ensure_proxy_restarts_for_flags_when_no_other_wrapper(monkeypatch) -> None: + """Control: with no other wrapper attached, a missing-flag restart still + happens — the guard must not block the single-client upgrade path.""" + calls: list[object] = [] + health = { + "version": wrap_cli._HEADROOM_VERSION, + "runtime": {"websocket_sessions": {"active_sessions": 0, "active_relay_tasks": 0}}, + "config": {"pid": "12345", "memory": False, "learn": False, "code_graph": False}, + } + + monkeypatch.setattr(wrap_cli, "_find_persistent_manifest", lambda port: None) + monkeypatch.setattr(wrap_cli, "_check_proxy", lambda port: len(calls) == 0) + monkeypatch.setattr(wrap_cli, "_query_proxy_health", lambda port: health) + monkeypatch.setattr(wrap_cli, "_live_proxy_clients", lambda *a, **kw: []) + monkeypatch.setattr( + wrap_cli, + "_kill_proxy_by_pid", + lambda pid, port: calls.append(("kill", pid, port)) or True, + ) + monkeypatch.setattr( + wrap_cli, + "_start_proxy", + lambda *args, **kwargs: calls.append(("start", args, kwargs)), + ) + + result = wrap_cli._ensure_proxy(8787, False, memory=True) + + assert result is None + assert calls[0] == ("kill", 12345, 8787) + assert calls[1][0] == "start" diff --git a/tests/test_paths_backward_compat.py b/tests/test_paths_backward_compat.py index 4d8a1c678..0be75bca9 100644 --- a/tests/test_paths_backward_compat.py +++ b/tests/test_paths_backward_compat.py @@ -95,6 +95,7 @@ def test_canonical_only_user_workspace_bucket_relocates( assert paths.proxy_log_path() == alt_ws / "logs" / "proxy.log" assert paths.debug_400_dir() == alt_ws / "logs" / "debug_400" assert paths.bin_dir() == alt_ws / "bin" + assert paths.proxy_clients_dir(8787) == alt_ws / "clients" / "8787" assert paths.deploy_root() == alt_ws / "deploy" assert paths.beacon_lock_path(8787) == alt_ws / ".beacon_lock_8787" # Config bucket follows the derived config dir.