diff --git a/src/context_engine/cli.py b/src/context_engine/cli.py index 507231e..b496a25 100644 --- a/src/context_engine/cli.py +++ b/src/context_engine/cli.py @@ -3505,6 +3505,17 @@ def phase_fn(msg: str) -> None: stats_path.write_text(json.dumps(stats), encoding="utf-8") +# Headless agent runs (e.g. `claude -p` workers) need the MCP tools but not +# live capture, and their host may kill the server without cleanup. With this +# set, `cce serve` starts no hook listener, so it never publishes, overwrites +# or deletes serve.port and the long-running sessions keep their rendezvous. +NO_HOOK_SERVER_ENV = "CCE_NO_HOOK_SERVER" + + +def _hook_server_disabled() -> bool: + return os.environ.get(NO_HOOK_SERVER_ENV, "").strip().lower() in ("1", "true", "yes", "on") + + async def _run_serve(config) -> None: """Start MCP server with live file watcher.""" import logging @@ -3645,13 +3656,16 @@ async def _reindex_worker(): # (capture is a non-critical feature; retrieval still works without it). hook_runner = None hook_port = None - try: - from context_engine.memory.hook_server import start_hook_server - hook_runner, hook_port = await start_hook_server( - storage_base=storage_base, project_name=project_name, - ) - except Exception as exc: - _log.warning("Memory hook server failed to start: %s", exc) + if _hook_server_disabled(): + _log.info("Memory hook server disabled by %s", NO_HOOK_SERVER_ENV) + else: + try: + from context_engine.memory.hook_server import start_hook_server + hook_runner, hook_port = await start_hook_server( + storage_base=storage_base, project_name=project_name, + ) + except Exception as exc: + _log.warning("Memory hook server failed to start: %s", exc) # Memory compression worker — drains pending_compressions in the background. # Each iteration opens a thread-local SQLite connection inside diff --git a/src/context_engine/memory/hook_server.py b/src/context_engine/memory/hook_server.py index ede1833..9558203 100644 --- a/src/context_engine/memory/hook_server.py +++ b/src/context_engine/memory/hook_server.py @@ -9,10 +9,20 @@ """ from __future__ import annotations +import contextlib import logging +import os import socket +import time +from collections.abc import Iterator from pathlib import Path +try: + import fcntl +except ImportError: # Windows + fcntl = None # type: ignore[assignment] + import msvcrt + from aiohttp import web from context_engine.memory import db as memory_db @@ -21,6 +31,142 @@ log = logging.getLogger(__name__) +# Name of the lock file, next to the authoritative serve.port, that serializes +# every publication and every compare-and-delete of the port files across the +# `cce serve` processes of one project. +PORT_LOCK_NAME = "serve.port.lock" +_PORT_LOCK_TIMEOUT_S = 2.0 +# Directory, next to serve.port, holding one entry per running server (file +# name = its port). On exit, a server whose port is still published hands the +# rendezvous over to the most recent live peer instead of leaving none. +PORT_REGISTRY_NAME = "serve.ports" +_PEER_PROBE_TIMEOUT_S = 0.2 + + +@contextlib.contextmanager +def _port_files_lock(lock_path: Path, timeout: float = _PORT_LOCK_TIMEOUT_S) -> Iterator[bool]: + """Hold an exclusive cross-process lock on `lock_path`; yields whether it is held. + + Without it, another server could publish its port between this server's + "does the file still hold my port?" read and its unlink, and lose its + rendezvous. Gives up after `timeout` (yields False) rather than hang a + startup or a shutdown on a wedged peer. + """ + lock_path.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(str(lock_path), os.O_CREAT | os.O_RDWR, 0o644) + held = False + try: + deadline = time.monotonic() + timeout + while True: + try: + if fcntl is not None: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + else: + os.lseek(fd, 0, os.SEEK_SET) + msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) + held = True + break + except OSError: + if time.monotonic() >= deadline: + log.warning("serve.port lock busy after %.1fs: %s", timeout, lock_path) + break + time.sleep(0.02) + yield held + finally: + if held: + try: + if fcntl is not None: + fcntl.flock(fd, fcntl.LOCK_UN) + else: + os.lseek(fd, 0, os.SEEK_SET) + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + except OSError: + pass + os.close(fd) + + +def _register_port(registry: Path, port: int) -> None: + try: + registry.mkdir(parents=True, exist_ok=True) + (registry / str(port)).write_text(str(os.getpid()), encoding="utf-8") + except OSError as exc: + log.warning("serve.port registry write failed for %s: %s", registry, exc) + + +def _port_is_live(port: int) -> bool: + try: + with socket.create_connection(("127.0.0.1", port), timeout=_PEER_PROBE_TIMEOUT_S): + return True + except OSError: + return False + + +def _live_successor(registry: Path, own: int | None) -> int | None: + """Most recently registered peer that still accepts connections. + + Entries of peers that are gone (killed without cleanup) are removed on + the way, so the registry does not grow across crashes. + """ + try: + entries = [e for e in registry.iterdir() if e.name.isdigit()] + except OSError: + return None + candidates = [] + for entry in entries: + port = int(entry.name) + if port == own: + continue + try: + candidates.append((entry.stat().st_mtime, port, entry)) + except OSError: + continue + for _, port, entry in sorted(candidates, reverse=True): + if _port_is_live(port): + return port + entry.unlink(missing_ok=True) + return None + + +def _prune_dead_peers(registry: Path, own: int) -> None: + """Drop every registry entry whose port no longer accepts connections. + + Servers killed without cleanup (TerminateProcess on Windows, SIGKILL) + never remove their entry, and `_live_successor` stops at the first live + peer, so without this sweep the registry only grows. Run at each start, + under the port-files lock. + """ + try: + entries = [e for e in registry.iterdir() if e.name.isdigit() and int(e.name) != own] + except OSError: + return + for entry in entries: + if not _port_is_live(int(entry.name)): + entry.unlink(missing_ok=True) + + +def _publish_port( + port_files: list[Path], port: int, lock_path: Path, registry: Path | None = None +) -> None: + """Write `port` to every port file, under the port-files lock.""" + with _port_files_lock(lock_path): + if registry is not None: + _prune_dead_peers(registry, port) + _register_port(registry, port) + written: set[Path] = set() + for f in port_files: + try: + key = f.resolve() + if key in written: + continue + f.parent.mkdir(parents=True, exist_ok=True) + f.write_text(str(port), encoding="utf-8") + written.add(key) + except OSError as exc: + # Non-fatal for the rendezvous copy — capture still works for + # users with default storage. + log.warning("serve.port write failed for %s: %s", f, exc) + + def _find_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("127.0.0.1", 0)) @@ -46,6 +192,8 @@ async def start_hook_server( # Authoritative port file lives in the project's storage_base. port_file = Path(storage_base) / "serve.port" + lock_path = Path(storage_base) / PORT_LOCK_NAME + registry = Path(storage_base) / PORT_REGISTRY_NAME # Stable rendezvous file at the *default* storage location. The hook # shell script always looks here (`${HOME}/.cce/projects//serve.port`) # because it has no way to read the user's config.yaml. When storage_path @@ -61,6 +209,10 @@ async def _close_db(app): except Exception: log.exception("memory_db close failed") + # Filled once the site is bound; the app is frozen by then, so the port + # lives in this closure rather than in app state. + bound: dict[str, int] = {} + async def _unlink_port_files(app): # Covers graceful shutdown paths only. SIGKILL bypasses # `app.on_cleanup` entirely, so the residual-port-file class of @@ -68,11 +220,43 @@ async def _unlink_port_files(app): # in the hook shell script. What this handler does cleanly cover # is the orderly SIGINT/SIGTERM/Ctrl-D path so the next session # doesn't inherit a stale serve.port from a normal exit. - for f in app.get("_port_files", []): - try: - f.unlink(missing_ok=True) - except OSError as exc: - log.warning("serve.port cleanup failed for %s: %s", f, exc) + # + # Several `cce serve` processes can share one project (one per + # agent session, or several agents on the same repo). Each start + # overwrites the port files with its own port, so by the time this + # one exits they may name another live server: only remove a file + # that still holds *this* server's port, or the survivors' hooks + # lose their rendezvous. + # + # A file that does hold this server's port is handed over to the + # most recent live peer rather than deleted: otherwise a short-lived + # server started after a long-running one (e.g. a headless agent run + # next to an interactive session) would leave the survivor without + # a rendezvous on exit. + own = bound.get("port") + # Compare-and-delete under the same lock as publication, so a peer + # cannot write its port between the read and the unlink. + with _port_files_lock(lock_path) as held: + if not held and own is not None: + # Unknown state: leave the files rather than risk deleting a + # peer's rendezvous (a stale file is caught by the hook's + # socket-liveness probe). + return + if own is not None: + (registry / str(own)).unlink(missing_ok=True) + successor = _live_successor(registry, own) + for f in app.get("_port_files", []): + try: + if own is not None and f.read_text(encoding="utf-8").strip() != str(own): + continue + if successor is not None: + f.write_text(str(successor), encoding="utf-8") + continue + f.unlink(missing_ok=True) + except FileNotFoundError: + pass + except OSError as exc: + log.warning("serve.port cleanup failed for %s: %s", f, exc) app.on_cleanup.append(_close_db) app.on_cleanup.append(_unlink_port_files) @@ -83,17 +267,9 @@ async def _unlink_port_files(app): port = _find_free_port() site = web.TCPSite(runner, host="127.0.0.1", port=port) await site.start() + bound["port"] = port - port_file.parent.mkdir(parents=True, exist_ok=True) - port_file.write_text(str(port), encoding="utf-8") - - try: - if default_rendezvous.resolve() != port_file.resolve(): - default_rendezvous.parent.mkdir(parents=True, exist_ok=True) - default_rendezvous.write_text(str(port), encoding="utf-8") - except OSError as exc: - # Non-fatal — capture still works for users with default storage. - log.warning("rendezvous port file write failed: %s", exc) + _publish_port([port_file, default_rendezvous], port, lock_path, registry) log.info("Memory hook server listening on 127.0.0.1:%d", port) return runner, port diff --git a/tests/memory/test_hook_server_cleanup.py b/tests/memory/test_hook_server_cleanup.py index 2986bfb..fa531d3 100644 --- a/tests/memory/test_hook_server_cleanup.py +++ b/tests/memory/test_hook_server_cleanup.py @@ -4,10 +4,12 @@ session's hooks would POST to a dead port. The on_cleanup hook now unlinks both the storage_base and rendezvous port files. """ +import threading from pathlib import Path import pytest +from context_engine.memory import hook_server from context_engine.memory.hook_server import start_hook_server @@ -84,3 +86,147 @@ async def test_cleanup_when_storage_equals_rendezvous(tmp_path, monkeypatch): finally: await runner.cleanup() assert not (storage_base / "serve.port").exists() + + +@pytest.mark.asyncio +async def test_cleanup_leaves_port_files_owned_by_another_server(tmp_path, monkeypatch): + """Two `cce serve` on one project: the second start overwrites the port + files. When the first one exits it must not delete them, or the hooks of + the server still running lose their rendezvous.""" + fake_home = tmp_path / "fake_home" + fake_home.mkdir() + monkeypatch.setattr(Path, "home", lambda: fake_home) + + storage_base = tmp_path / "storage" + storage_base.mkdir() + port_file = storage_base / "serve.port" + rendezvous = fake_home / ".cce" / "projects" / "proj_shared" / "serve.port" + + first, _ = await start_hook_server(storage_base=storage_base, project_name="proj_shared") + second, second_port = await start_hook_server( + storage_base=storage_base, project_name="proj_shared", + ) + try: + await first.cleanup() + assert port_file.read_text() == str(second_port) + assert rendezvous.read_text() == str(second_port) + finally: + await second.cleanup() + + assert not port_file.exists() + assert not rendezvous.exists() + + +@pytest.mark.asyncio +async def test_peer_publishing_during_cleanup_keeps_its_rendezvous(tmp_path, monkeypatch): + """A peer publishing its port between this server's ownership read and + its unlink must not lose the rendezvous: publication waits on the same + lock as the compare-and-delete, so it lands after the unlink.""" + fake_home = tmp_path / "fake_home" + fake_home.mkdir() + monkeypatch.setattr(Path, "home", lambda: fake_home) + storage_base = tmp_path / "storage" + storage_base.mkdir() + port_file = storage_base / "serve.port" + rendezvous = fake_home / ".cce" / "projects" / "proj_race" / "serve.port" + + runner, _ = await start_hook_server(storage_base=storage_base, project_name="proj_race") + + peer_port = 1 + peer = threading.Thread( + target=hook_server._publish_port, + args=([port_file, rendezvous], peer_port, storage_base / hook_server.PORT_LOCK_NAME), + ) + real_read_text = Path.read_text + + def read_then_peer_publishes(self, *args, **kwargs): + text = real_read_text(self, *args, **kwargs) + if self == port_file and not peer.is_alive() and peer.ident is None: + peer.start() + peer.join(timeout=0.3) # without the lock the peer writes right here + return text + + monkeypatch.setattr(Path, "read_text", read_then_peer_publishes) + await runner.cleanup() + monkeypatch.setattr(Path, "read_text", real_read_text) + peer.join(timeout=5) + + assert port_file.read_text() == str(peer_port) + assert rendezvous.read_text() == str(peer_port) + + +@pytest.mark.asyncio +async def test_last_started_server_exiting_hands_rendezvous_to_live_peer(tmp_path, monkeypatch): + """A short-lived server started after a long-running one (a headless agent + next to an interactive session) holds the port files when it exits: they + must move to the survivor, not disappear.""" + fake_home = tmp_path / "fake_home" + fake_home.mkdir() + monkeypatch.setattr(Path, "home", lambda: fake_home) + storage_base = tmp_path / "storage" + storage_base.mkdir() + port_file = storage_base / "serve.port" + rendezvous = fake_home / ".cce" / "projects" / "proj_handover" / "serve.port" + + first, first_port = await start_hook_server( + storage_base=storage_base, project_name="proj_handover", + ) + second, second_port = await start_hook_server( + storage_base=storage_base, project_name="proj_handover", + ) + try: + assert port_file.read_text() == str(second_port) + await second.cleanup() + assert port_file.read_text() == str(first_port) + assert rendezvous.read_text() == str(first_port) + finally: + await first.cleanup() + + assert not port_file.exists() + assert not rendezvous.exists() + assert list((storage_base / hook_server.PORT_REGISTRY_NAME).iterdir()) == [] + + +@pytest.mark.asyncio +async def test_dead_peer_in_registry_is_pruned_not_published(tmp_path, monkeypatch): + """A peer killed without cleanup leaves its registry entry: the exiting + server must not hand the rendezvous to that dead port.""" + fake_home = tmp_path / "fake_home" + fake_home.mkdir() + monkeypatch.setattr(Path, "home", lambda: fake_home) + storage_base = tmp_path / "storage" + storage_base.mkdir() + port_file = storage_base / "serve.port" + + runner, _ = await start_hook_server(storage_base=storage_base, project_name="proj_dead") + registry = storage_base / hook_server.PORT_REGISTRY_NAME + dead_port = hook_server._find_free_port() # nothing listens on it + (registry / str(dead_port)).write_text("0", encoding="utf-8") + + await runner.cleanup() + + assert not port_file.exists() + assert not (registry / str(dead_port)).exists() + + +@pytest.mark.asyncio +async def test_start_prunes_registry_entries_of_dead_peers(tmp_path, monkeypatch): + """Peers killed without cleanup never remove their entry; the next start + must sweep them, live peers excepted, so the registry does not grow.""" + fake_home = tmp_path / "fake_home" + fake_home.mkdir() + monkeypatch.setattr(Path, "home", lambda: fake_home) + storage_base = tmp_path / "storage" + registry = storage_base / hook_server.PORT_REGISTRY_NAME + registry.mkdir(parents=True) + dead = [hook_server._find_free_port() for _ in range(3)] # nothing listens on them + for port in dead: + (registry / str(port)).write_text("0", encoding="utf-8") + + first, first_port = await start_hook_server(storage_base=storage_base, project_name="proj_sweep") + second, second_port = await start_hook_server(storage_base=storage_base, project_name="proj_sweep") + try: + assert sorted(e.name for e in registry.iterdir()) == sorted([str(first_port), str(second_port)]) + finally: + await second.cleanup() + await first.cleanup() diff --git a/tests/test_cli_serve.py b/tests/test_cli_serve.py index b564838..2d0f205 100644 --- a/tests/test_cli_serve.py +++ b/tests/test_cli_serve.py @@ -56,3 +56,19 @@ async def fake_run_serve(config): "serve --project-dir loaded config from launch cwd, not the target " "project — the project's .context-engine.yaml was ignored." ) + + +@pytest.mark.parametrize( + ("value", "disabled"), + [(None, False), ("", False), ("0", False), ("1", True), ("true", True), (" Yes ", True), ("on", True)], +) +def test_hook_server_disabled_by_env(monkeypatch, value, disabled): + """Headless agent runs set CCE_NO_HOOK_SERVER so their `cce serve` never + touches serve.port (their host may kill it without cleanup).""" + from context_engine import cli + + if value is None: + monkeypatch.delenv(cli.NO_HOOK_SERVER_ENV, raising=False) + else: + monkeypatch.setenv(cli.NO_HOOK_SERVER_ENV, value) + assert cli._hook_server_disabled() is disabled