diff --git a/pyproject.toml b/pyproject.toml index f657ef6..66318cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,7 @@ dev = [ [project.scripts] noirdoc = "noirdoc.cli:main" +noirdoc-daemon = "noirdoc.daemon.server:main" [project.urls] Homepage = "https://noirdoc.de" diff --git a/src/noirdoc/cli.py b/src/noirdoc/cli.py index b0a38ec..16a8fa9 100644 --- a/src/noirdoc/cli.py +++ b/src/noirdoc/cli.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import os import sys from pathlib import Path @@ -13,6 +14,10 @@ from noirdoc.sdk import Redactor +def _daemon_disabled() -> bool: + return os.environ.get("NOIRDOC_NO_DAEMON", "").strip().lower() in {"1", "true", "yes"} + + @click.group() @click.version_option(__version__, prog_name="noirdoc") def main() -> None: @@ -39,6 +44,11 @@ def main() -> None: ) @click.option("--score-threshold", type=float, default=0.5, show_default=True) @click.option("-v", "--verbose", is_flag=True, help="Print per-file details.") +@click.option( + "--no-daemon", + is_flag=True, + help="Skip the daemon and run redaction in-process (each call reloads models).", +) def redact( inputs: tuple[Path, ...], output: Path | None, @@ -48,6 +58,7 @@ def redact( detector: str, score_threshold: float, verbose: bool, + no_daemon: bool, ) -> None: """Redact PII in one or more files.""" files = _expand_inputs(inputs) @@ -59,6 +70,145 @@ def redact( click.echo("--output only works with a single input; use --output-dir.", err=True) sys.exit(2) + use_daemon = not (no_daemon or _daemon_disabled()) + + if use_daemon: + success, last_namespace_size, switched = _redact_via_daemon( + files, + output=output, + output_dir=output_dir, + namespace=namespace, + language=language, + detector=detector, + score_threshold=score_threshold, + verbose=verbose, + ) + if switched is None: + # Daemon handled every file successfully. + if namespace and last_namespace_size is not None: + click.echo( + f"Namespace {namespace!r}: {last_namespace_size} unique entities total", + ) + sys.exit(0 if success else 1) + # Fall through with the remaining files handled in-process. + files = switched + else: + success = 0 + + success += _redact_in_process( + files, + output=output, + output_dir=output_dir, + namespace=namespace, + language=language, + detector=detector, + score_threshold=score_threshold, + verbose=verbose, + ) + sys.exit(0 if success else 1) + + +def _redact_via_daemon( + files: list[Path], + *, + output: Path | None, + output_dir: Path | None, + namespace: str | None, + language: str, + detector: str, + score_threshold: float, + verbose: bool, +) -> tuple[int, int | None, list[Path] | None]: + """Redact via the daemon. Returns (successes, last_namespace_size, fallback_remaining). + + ``fallback_remaining`` is ``None`` on full success, or a list of files + not yet processed when the daemon failed and the caller should fall + back to in-process redaction. + """ + from noirdoc.daemon.client import DaemonError, DaemonUnavailable, call_sync + + success = 0 + last_namespace_size: int | None = None + + for i, path in enumerate(files): + out_path = _choose_output_path( + path, + output=output, + output_dir=output_dir, + reconstructed=True, # daemon decides; we'll rename if it says reconstructed=False + ) + out_path.parent.mkdir(parents=True, exist_ok=True) + + params = { + "namespace": namespace, + "language": language, + "detector": detector, + "score_threshold": score_threshold, + "input": {"type": "file", "path": str(path.resolve())}, + "output_path": str(out_path.resolve()), + } + try: + result = call_sync("redact", params) + except DaemonUnavailable as exc: + click.echo( + f"noirdoc: daemon unavailable ({exc}); running locally for the rest.", + err=True, + ) + return success, last_namespace_size, list(files[i:]) + except DaemonError as exc: + click.echo( + f"noirdoc: daemon error ({exc}); running locally for the rest.", + err=True, + ) + return success, last_namespace_size, list(files[i:]) + except Exception as exc: # connection died, etc. + click.echo( + f"noirdoc: daemon failed mid-request ({exc}); running locally for the rest.", + err=True, + ) + return success, last_namespace_size, list(files[i:]) + + # If the daemon couldn't reconstruct the file, the output it wrote + # is plain text — rename to .txt to match in-process behaviour. + if not result.get("reconstructed", False): + corrected = _choose_output_path( + path, + output=output, + output_dir=output_dir, + reconstructed=False, + ) + if corrected != out_path: + Path(out_path).rename(corrected) + out_path = corrected + + last_namespace_size = result.get("namespace_size") or last_namespace_size + + if verbose: + types = ", ".join(f"{k}={v}" for k, v in sorted(result.get("entity_types", {}).items())) + click.echo( + f"{path.name}: {result['entity_count']} entities [{types}] -> {out_path}", + ) + else: + click.echo(f"{path.name}: {result['entity_count']} entities -> {out_path}") + success += 1 + + return success, last_namespace_size, None + + +def _redact_in_process( + files: list[Path], + *, + output: Path | None, + output_dir: Path | None, + namespace: str | None, + language: str, + detector: str, + score_threshold: float, + verbose: bool, +) -> int: + if not files: + return 0 + r = Redactor( namespace=namespace, language=language, @@ -93,7 +243,7 @@ def redact( if namespace: click.echo(f"Namespace {namespace!r}: {r.mapper.entity_count} unique entities total") - sys.exit(0 if success else 1) + return success @main.command() @@ -187,6 +337,85 @@ def ns_delete(namespace: str) -> None: click.echo(f"Deleted {namespace!r}") +@main.group("daemon") +def daemon_group() -> None: + """Inspect and control the long-lived noirdoc daemon.""" + + +@daemon_group.command("status") +def daemon_status() -> None: + """Print whether a daemon is running and its current stats.""" + from noirdoc.daemon import paths as daemon_paths + from noirdoc.daemon import spawn as daemon_spawn + from noirdoc.daemon.client import DaemonError, call_sync + + pid = daemon_spawn.read_pidfile() + if pid is None or not daemon_spawn.is_pid_alive(pid): + click.echo(json.dumps({"running": False}, indent=2)) + return + + try: + status = call_sync("status") + except DaemonError as exc: + click.echo( + json.dumps({"running": True, "pid": pid, "error": str(exc)}, indent=2), + err=True, + ) + sys.exit(1) + + click.echo( + json.dumps( + { + "running": True, + "pid": pid, + "socket": str(daemon_paths.socket_path()), + **status, + }, + indent=2, + ), + ) + + +@daemon_group.command("stop") +def daemon_stop() -> None: + """Send SIGTERM to the daemon and wait for it to exit.""" + from noirdoc.daemon import spawn as daemon_spawn + + pid = daemon_spawn.read_pidfile() + if pid is None or not daemon_spawn.is_pid_alive(pid): + click.echo("daemon is not running") + return + + if daemon_spawn.stop_daemon(): + click.echo(f"daemon (pid={pid}) stopped") + else: + click.echo(f"daemon (pid={pid}) did not stop in time", err=True) + sys.exit(1) + + +@daemon_group.command("restart") +@click.pass_context +def daemon_restart(ctx: click.Context) -> None: + """Stop the daemon if running. Next ``redact`` will auto-spawn a fresh one.""" + ctx.invoke(daemon_stop) + click.echo("(daemon will auto-spawn on next redact)") + + +@daemon_group.command("logs") +@click.option("-n", "--lines", type=int, default=50, show_default=True) +def daemon_logs(lines: int) -> None: + """Print the tail of the daemon log file.""" + from noirdoc.daemon import paths as daemon_paths + + log = daemon_paths.logfile_path() + if not log.exists(): + click.echo("no log file yet", err=True) + sys.exit(1) + text = log.read_text(encoding="utf-8", errors="replace") + for line in text.splitlines()[-lines:]: + click.echo(line) + + @main.group() def models() -> None: """Manage local detection models.""" diff --git a/src/noirdoc/daemon/__init__.py b/src/noirdoc/daemon/__init__.py new file mode 100644 index 0000000..5a144bb --- /dev/null +++ b/src/noirdoc/daemon/__init__.py @@ -0,0 +1,6 @@ +"""Long-lived daemon that holds detection models in memory. + +The CLI's ``redact`` command transparently spawns and talks to this daemon +over a Unix domain socket, so model loading (~10 s on cold start) is paid +once per session instead of once per invocation. +""" diff --git a/src/noirdoc/daemon/__main__.py b/src/noirdoc/daemon/__main__.py new file mode 100644 index 0000000..cb4ee8c --- /dev/null +++ b/src/noirdoc/daemon/__main__.py @@ -0,0 +1,8 @@ +"""``python -m noirdoc.daemon`` entry point.""" + +from __future__ import annotations + +from noirdoc.daemon.server import main + +if __name__ == "__main__": + main() diff --git a/src/noirdoc/daemon/client.py b/src/noirdoc/daemon/client.py new file mode 100644 index 0000000..56a2a1b --- /dev/null +++ b/src/noirdoc/daemon/client.py @@ -0,0 +1,163 @@ +"""CLI-side helpers for talking to the daemon. + +Used by ``noirdoc redact`` to transparently spawn and call the daemon. +On any daemon failure (no spawn, crash mid-request, version mismatch +that won't reconcile), the caller is expected to fall back to in-process +redaction so the user's command still succeeds. +""" + +from __future__ import annotations + +import asyncio +import json +import uuid +from pathlib import Path +from typing import Any + +from noirdoc import __version__ +from noirdoc.daemon import paths, spawn + +CONNECT_TIMEOUT = 2.0 +RPC_TIMEOUT = 600.0 # generous; covers cold-spawn warmup + slow file redaction +SHUTDOWN_DRAIN_TIMEOUT = 5.0 + + +class DaemonError(Exception): + """Daemon returned an error response or violated the protocol.""" + + +class DaemonUnavailable(DaemonError): + """Daemon could not be reached even after attempting to spawn.""" + + +async def _try_connect( + socket_path: Path, +) -> tuple[asyncio.StreamReader, asyncio.StreamWriter] | None: + try: + return await asyncio.wait_for( + asyncio.open_unix_connection(path=str(socket_path)), + timeout=CONNECT_TIMEOUT, + ) + except (FileNotFoundError, ConnectionRefusedError, TimeoutError, OSError): + return None + + +async def _spawn_and_connect( + socket_path: Path, +) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + spawn.cleanup_stale_socket() + spawn.spawn_detached() + ready = await asyncio.to_thread(spawn.wait_for_socket, 30.0) + if not ready: + raise DaemonUnavailable("daemon did not bind socket within 30s") + conn = await _try_connect(socket_path) + if conn is None: + raise DaemonUnavailable("could not connect to spawned daemon") + return conn + + +async def _send_request( + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + method: str, + params: dict[str, Any], + *, + timeout: float = RPC_TIMEOUT, +) -> dict[str, Any]: + req_id = uuid.uuid4().hex + payload = json.dumps( + {"id": req_id, "method": method, "params": params}, + ensure_ascii=False, + ) + writer.write(payload.encode("utf-8") + b"\n") + await writer.drain() + + line = await asyncio.wait_for(reader.readline(), timeout=timeout) + if not line: + raise DaemonUnavailable("daemon closed connection before sending a response") + + try: + response = json.loads(line.decode("utf-8")) + except json.JSONDecodeError as exc: + raise DaemonError(f"malformed response: {exc}") from exc + + if response.get("id") != req_id: + raise DaemonError( + f"response id mismatch: got {response.get('id')!r}, want {req_id!r}", + ) + + if "error" in response and response["error"] is not None: + err = response["error"] + raise DaemonError(f"{err.get('code', '?')}: {err.get('message', '?')}") + + result = response.get("result") + if result is None: + raise DaemonError("response missing both 'result' and 'error'") + return result + + +async def _wait_socket_gone(socket_path: Path, timeout: float) -> None: + deadline = asyncio.get_event_loop().time() + timeout + while asyncio.get_event_loop().time() < deadline: + if not socket_path.exists(): + return + await asyncio.sleep(0.1) + + +async def _close(writer: asyncio.StreamWriter) -> None: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + + +async def call(method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + """Connect (spawning if needed), handshake, send one RPC, return its result. + + On version mismatch, asks the running daemon to shut down, waits for the + socket to disappear, and tries once more with a freshly spawned daemon. + """ + socket_path = paths.socket_path() + params = params or {} + + for attempt in (0, 1): + conn = await _try_connect(socket_path) + if conn is None: + conn = await _spawn_and_connect(socket_path) + reader, writer = conn + + try: + hello = await _send_request( + reader, + writer, + "hello", + {"client_version": __version__}, + ) + if hello.get("daemon_version") != __version__: + if attempt == 1: + raise DaemonError( + f"version mismatch persists: daemon={hello.get('daemon_version')!r} " + f"client={__version__!r}", + ) + # Ask the stale daemon to exit, wait for it to release the + # socket, then loop and let _spawn_and_connect bring up a + # fresh one at the current version. + try: + await _send_request(reader, writer, "shutdown", {}) + except DaemonError: + pass + await _close(writer) + await _wait_socket_gone(socket_path, SHUTDOWN_DRAIN_TIMEOUT) + continue + + return await _send_request(reader, writer, method, params) + finally: + await _close(writer) + + raise DaemonUnavailable("daemon connect/respawn loop exhausted") + + +def call_sync(method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + """Sync wrapper used from Click commands.""" + return asyncio.run(call(method, params)) diff --git a/src/noirdoc/daemon/paths.py b/src/noirdoc/daemon/paths.py new file mode 100644 index 0000000..f007b88 --- /dev/null +++ b/src/noirdoc/daemon/paths.py @@ -0,0 +1,42 @@ +"""Filesystem paths used by the daemon (socket, pidfile, log).""" + +from __future__ import annotations + +import os +from pathlib import Path + +DEFAULT_ROOT = Path.home() / ".noirdoc" + + +def root_dir() -> Path: + """Directory holding daemon state. Honors ``NOIRDOC_DAEMON_ROOT`` for tests.""" + override = os.environ.get("NOIRDOC_DAEMON_ROOT") + if override: + return Path(override).expanduser() + return DEFAULT_ROOT + + +def socket_path() -> Path: + override = os.environ.get("NOIRDOC_DAEMON_SOCKET") + if override: + return Path(override).expanduser() + return root_dir() / "daemon.sock" + + +def pidfile_path() -> Path: + return root_dir() / "daemon.pid" + + +def logfile_path() -> Path: + return root_dir() / "daemon.log" + + +def ensure_root_dir() -> Path: + """Create the daemon root with 0o700 perms, idempotent.""" + d = root_dir() + d.mkdir(parents=True, exist_ok=True) + try: + os.chmod(d, 0o700) + except OSError: + pass + return d diff --git a/src/noirdoc/daemon/protocol.py b/src/noirdoc/daemon/protocol.py new file mode 100644 index 0000000..0768717 --- /dev/null +++ b/src/noirdoc/daemon/protocol.py @@ -0,0 +1,104 @@ +"""Wire types for the daemon JSON-lines protocol. + +One source of truth for both server (``noirdoc/daemon/server.py``) and +client (``noirdoc/daemon/client.py``) so they cannot drift. + +Message shape:: + + request: {"id": "", "method": "", "params": {...}} + response: {"id": "", "result": {...}} + {"id": "", "error": {"code": "...", "message": "..."}} + +Each message is a single JSON object terminated by ``\\n``. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal, Union + +from pydantic import BaseModel, Field + +DetectorChoice = Literal["presidio", "gliner", "ensemble"] + + +class HelloParams(BaseModel): + client_version: str + + +class HelloResult(BaseModel): + daemon_version: str + pid: int + started_at: float + + +class RedactTextInput(BaseModel): + type: Literal["text"] = "text" + value: str + + +class RedactFileInput(BaseModel): + type: Literal["file"] = "file" + path: str # absolute path on the daemon's filesystem (same user as CLI) + + +RedactInput = Annotated[ + Union[RedactTextInput, RedactFileInput], + Field(discriminator="type"), +] + + +class RedactParams(BaseModel): + namespace: str | None = None + namespace_root: str | None = None + language: str = "de" + detector: DetectorChoice = "ensemble" + score_threshold: float = 0.5 + gliner_model: str = "knowledgator/gliner-pii-edge-v1.0" + input: RedactInput + output_path: str | None = None # for file input; daemon writes here directly + + +class RedactResult(BaseModel): + redacted_text: str | None = None # populated for text input + output_path: str | None = None # populated for file input written to disk + entity_count: int + entity_types: dict[str, int] + mime_type: str | None = None + reconstructed: bool = False + namespace_size: int | None = None + + +class StatusResult(BaseModel): + uptime_s: float + models_loaded: bool + last_request_at: float | None + queue_depth: int + total_requests: int + + +class ShutdownResult(BaseModel): + ok: bool = True + + +class ErrorPayload(BaseModel): + code: str + message: str + + +class Request(BaseModel): + id: str + method: str + params: dict[str, Any] = Field(default_factory=dict) + + +class Response(BaseModel): + id: str + result: dict[str, Any] | None = None + error: ErrorPayload | None = None + + +# Error codes used in Response.error.code +ERR_BAD_REQUEST = "bad_request" +ERR_UNKNOWN_METHOD = "unknown_method" +ERR_INTERNAL = "internal" +ERR_VERSION_MISMATCH = "version_mismatch" diff --git a/src/noirdoc/daemon/server.py b/src/noirdoc/daemon/server.py new file mode 100644 index 0000000..b2269f9 --- /dev/null +++ b/src/noirdoc/daemon/server.py @@ -0,0 +1,477 @@ +"""asyncio Unix-socket daemon: serial redaction + cached models. + +One daemon, one event loop, one in-flight redact at a time. Models are +cached at the daemon level so cold-start cost is paid once per session. +Per-request namespace state is loaded from disk and saved back, never +cached, to avoid coherence issues with concurrent CLI processes. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import logging.handlers +import os +import signal +import time +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from pydantic import ValidationError + +from noirdoc import __version__ +from noirdoc.daemon import paths, spawn +from noirdoc.daemon.protocol import ( + ERR_BAD_REQUEST, + ERR_INTERNAL, + ERR_UNKNOWN_METHOD, + ErrorPayload, + HelloParams, + HelloResult, + RedactFileInput, + RedactParams, + RedactResult, + RedactTextInput, + Request, + Response, + ShutdownResult, + StatusResult, +) + +if TYPE_CHECKING: + from noirdoc.detection.base import BaseDetector + from noirdoc.detection.presidio_detector import PresidioDetector + +DEFAULT_IDLE_SECONDS = 600 # 10 minutes +IDLE_CHECK_INTERVAL = 60.0 +LOG_MAX_BYTES = 5 * 1024 * 1024 +LOG_BACKUP_COUNT = 1 +SUPPORTED_WARMUP_LANGUAGES = ("de", "en") + +log = logging.getLogger("noirdoc.daemon") + + +def _setup_logging() -> None: + paths.ensure_root_dir() + handler = logging.handlers.RotatingFileHandler( + paths.logfile_path(), + maxBytes=LOG_MAX_BYTES, + backupCount=LOG_BACKUP_COUNT, + ) + handler.setFormatter( + logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s"), + ) + logging.basicConfig(level=logging.INFO, handlers=[handler]) + + +def _idle_seconds() -> float: + env = os.environ.get("NOIRDOC_DAEMON_IDLE_SECONDS") + if env: + try: + return float(env) + except ValueError: + log.warning("ignoring invalid NOIRDOC_DAEMON_IDLE_SECONDS=%r", env) + return float(DEFAULT_IDLE_SECONDS) + + +class DaemonState: + """Mutable runtime state. One instance per daemon process.""" + + def __init__(self) -> None: + self.start_time = time.time() + self.last_request_at: float | None = None + self.total_requests = 0 + self.queue_depth = 0 # pending + in-flight redacts + self.redact_lock = asyncio.Lock() + self.shutdown_event = asyncio.Event() + + # Cached underlying detectors. EnsembleDetector wrappers are built + # per-request so each request can pick its own score threshold and + # detector subset. + self._presidio_by_lang: dict[str, PresidioDetector] = {} + # GlinerDetector when installed; typed Any so mypy permits the lazy + # instantiation without dragging in the optional dependency. + self._gliner: Any = None + self._gliner_attempted = False + self._gliner_model_name: str | None = None + self._init_lock = asyncio.Lock() # serialize lazy detector init + + self.models_loaded_event = asyncio.Event() + + @property + def models_loaded(self) -> bool: + return self.models_loaded_event.is_set() + + async def warmup(self) -> None: + """Eagerly load every detector we expect to need. + + Failures are logged but never raised — a missing GLiNER means the + ensemble degrades to Presidio-only, which is the same behaviour as + in-process ``Redactor``. + """ + from noirdoc.detection.model_manager import ensure_spacy_models + + for lang in SUPPORTED_WARMUP_LANGUAGES: + try: + await asyncio.to_thread(ensure_spacy_models, [lang]) + from noirdoc.detection.presidio_detector import PresidioDetector + + self._presidio_by_lang[lang] = PresidioDetector(languages=[lang]) + log.info("loaded presidio for language=%s", lang) + except Exception: + log.exception("presidio warmup failed for language=%s", lang) + + await self._load_gliner_if_available( + "knowledgator/gliner-pii-edge-v1.0", + ) + + self.models_loaded_event.set() + log.info("warmup complete") + + async def _load_gliner_if_available(self, model_name: str) -> None: + if self._gliner_attempted: + return + self._gliner_attempted = True + self._gliner_model_name = model_name + try: + from noirdoc.detection.gliner_detector import GlinerDetector + except ImportError: + log.info("gliner not installed; ensemble will use presidio only") + return + try: + self._gliner = await asyncio.to_thread( + GlinerDetector, + model_name=model_name, + ) + log.info("loaded gliner model=%s", model_name) + except Exception: + log.exception("gliner load failed for model=%s", model_name) + + async def get_detectors( + self, + language: str, + choice: str, + gliner_model: str, + ) -> list[BaseDetector]: + """Return cached detector instances for a request. + + Lazily loads anything warmup didn't already cover (e.g., a request + for a language not in ``SUPPORTED_WARMUP_LANGUAGES``). + """ + from noirdoc.detection.base import BaseDetector # noqa: F401 (typing) + + out: list[BaseDetector] = [] + async with self._init_lock: + if choice in ("presidio", "ensemble"): + if language not in self._presidio_by_lang: + from noirdoc.detection.model_manager import ensure_spacy_models + from noirdoc.detection.presidio_detector import PresidioDetector + + await asyncio.to_thread(ensure_spacy_models, [language]) + self._presidio_by_lang[language] = PresidioDetector( + languages=[language], + ) + out.append(self._presidio_by_lang[language]) + + if choice in ("gliner", "ensemble"): + if not self._gliner_attempted or (self._gliner_model_name != gliner_model): + # Honor a request that wants a different GLiNER model than + # the one we warmed up with. + self._gliner_attempted = False + await self._load_gliner_if_available(gliner_model) + if self._gliner is not None: + out.append(self._gliner) + elif choice == "gliner": + raise RuntimeError( + "GLiNER is not installed (pip install 'noirdoc[full]')", + ) + + return out + + +# -- request handlers -------------------------------------------------------- + + +async def handle_hello(state: DaemonState, params: dict[str, Any]) -> dict[str, Any]: + HelloParams.model_validate(params) # validates client_version present + return HelloResult( + daemon_version=__version__, + pid=os.getpid(), + started_at=state.start_time, + ).model_dump() + + +async def handle_status(state: DaemonState, params: dict[str, Any]) -> dict[str, Any]: + return StatusResult( + uptime_s=time.time() - state.start_time, + models_loaded=state.models_loaded, + last_request_at=state.last_request_at, + queue_depth=state.queue_depth, + total_requests=state.total_requests, + ).model_dump() + + +async def handle_shutdown( + state: DaemonState, + params: dict[str, Any], +) -> dict[str, Any]: + state.shutdown_event.set() + return ShutdownResult().model_dump() + + +async def handle_redact( + state: DaemonState, + params: dict[str, Any], +) -> dict[str, Any]: + from noirdoc.detection.ensemble import EnsembleDetector + from noirdoc.sdk import build_redactor + + parsed = RedactParams.model_validate(params) + + state.queue_depth += 1 + try: + async with state.redact_lock: + state.last_request_at = time.time() + state.total_requests += 1 + t0 = time.monotonic() + + detectors = await state.get_detectors( + parsed.language, + parsed.detector, + parsed.gliner_model, + ) + ensemble = EnsembleDetector( + detectors=detectors, + score_threshold=parsed.score_threshold, + ) + + redactor = build_redactor( + ensemble=ensemble, + namespace=parsed.namespace, + namespace_root=parsed.namespace_root, + language=parsed.language, + detector=parsed.detector, # type: ignore[arg-type] + score_threshold=parsed.score_threshold, + gliner_model=parsed.gliner_model, + ) + + if isinstance(parsed.input, RedactTextInput): + pseudonymized, entities = await redactor.aredact_text_detailed( + parsed.input.value, + parsed.language, + ) + entity_types: dict[str, int] = {} + for e in entities: + entity_types[e.entity_type] = entity_types.get(e.entity_type, 0) + 1 + result = RedactResult( + redacted_text=pseudonymized, + entity_count=len(entities), + entity_types=entity_types, + namespace_size=redactor.mapper.entity_count, + ) + else: + assert isinstance(parsed.input, RedactFileInput) + in_path = Path(parsed.input.path) + file_result = await redactor.aredact_file( + in_path, + language=parsed.language, + ) + if parsed.output_path: + out_path = Path(parsed.output_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_bytes(file_result.output_bytes) + output_path_str: str | None = str(out_path) + else: + output_path_str = None + result = RedactResult( + output_path=output_path_str, + entity_count=file_result.entity_count, + entity_types=file_result.entity_types, + mime_type=file_result.mime_type, + reconstructed=file_result.reconstructed, + namespace_size=redactor.mapper.entity_count, + ) + + duration_ms = int((time.monotonic() - t0) * 1000) + log.info( + "redact ns=%s lang=%s detector=%s entities=%d ms=%d", + parsed.namespace, + parsed.language, + parsed.detector, + result.entity_count, + duration_ms, + ) + return result.model_dump() + finally: + state.queue_depth = max(state.queue_depth - 1, 0) + + +HANDLERS: dict[str, Any] = { + "hello": handle_hello, + "status": handle_status, + "shutdown": handle_shutdown, + "redact": handle_redact, +} + + +# -- connection plumbing ----------------------------------------------------- + + +def _serialize(response: Response) -> bytes: + return (json.dumps(response.model_dump(exclude_none=True), ensure_ascii=False) + "\n").encode( + "utf-8", + ) + + +async def _dispatch(state: DaemonState, raw_line: bytes) -> Response: + try: + payload = json.loads(raw_line.decode("utf-8")) + request = Request.model_validate(payload) + except (json.JSONDecodeError, ValidationError) as exc: + return Response( + id="", + error=ErrorPayload(code=ERR_BAD_REQUEST, message=str(exc)), + ) + + handler = HANDLERS.get(request.method) + if handler is None: + return Response( + id=request.id, + error=ErrorPayload( + code=ERR_UNKNOWN_METHOD, + message=f"unknown method: {request.method!r}", + ), + ) + + try: + result = await handler(state, request.params) + except ValidationError as exc: + return Response( + id=request.id, + error=ErrorPayload(code=ERR_BAD_REQUEST, message=str(exc)), + ) + except Exception as exc: + log.exception("handler %s raised", request.method) + return Response( + id=request.id, + error=ErrorPayload(code=ERR_INTERNAL, message=str(exc)), + ) + + return Response(id=request.id, result=result) + + +async def _handle_connection( + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + state: DaemonState, +) -> None: + try: + while not state.shutdown_event.is_set(): + line = await reader.readline() + if not line: + return # client closed + response = await _dispatch(state, line) + writer.write(_serialize(response)) + await writer.drain() + except (ConnectionError, asyncio.IncompleteReadError, BrokenPipeError): + pass + except Exception: + log.exception("connection handler crashed") + finally: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + + +async def _idle_watcher(state: DaemonState) -> None: + idle = _idle_seconds() + while not state.shutdown_event.is_set(): + try: + await asyncio.wait_for( + state.shutdown_event.wait(), + timeout=IDLE_CHECK_INTERVAL, + ) + return + except TimeoutError: + pass + if state.last_request_at is None: + # No request yet — measure idleness from start_time. + elapsed = time.time() - state.start_time + else: + elapsed = time.time() - state.last_request_at + if elapsed > idle and not state.redact_lock.locked(): + log.info("idle shutdown after %.0fs of inactivity", elapsed) + state.shutdown_event.set() + return + + +# -- bootstrap --------------------------------------------------------------- + + +async def _async_main() -> None: + _setup_logging() + paths.ensure_root_dir() + spawn.cleanup_stale_socket() + + existing_pid = spawn.read_pidfile() + if existing_pid is not None and spawn.is_pid_alive(existing_pid): + log.info("another daemon is already running (pid=%d), exiting", existing_pid) + return + + spawn.write_pidfile(os.getpid()) + state = DaemonState() + log.info("daemon starting pid=%d version=%s", os.getpid(), __version__) + + warmup_task = asyncio.create_task(state.warmup(), name="warmup") + idle_task = asyncio.create_task(_idle_watcher(state), name="idle-watcher") + + sock_path = paths.socket_path() + server = await asyncio.start_unix_server( + lambda r, w: _handle_connection(r, w, state), + path=str(sock_path), + ) + try: + os.chmod(sock_path, 0o600) + except OSError: + log.warning("could not chmod socket %s", sock_path) + + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + try: + loop.add_signal_handler(sig, state.shutdown_event.set) + except NotImplementedError: + pass # Windows; we don't ship there but be defensive. + + log.info("listening on %s", sock_path) + try: + async with server: + await state.shutdown_event.wait() + finally: + log.info("draining and shutting down") + idle_task.cancel() + warmup_task.cancel() + for task in (idle_task, warmup_task): + try: + await task + except (asyncio.CancelledError, Exception): + pass + try: + sock_path.unlink() + except FileNotFoundError: + pass + spawn.remove_pidfile() + log.info("daemon stopped") + + +def main() -> None: + try: + asyncio.run(_async_main()) + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() diff --git a/src/noirdoc/daemon/spawn.py b/src/noirdoc/daemon/spawn.py new file mode 100644 index 0000000..cf6146e --- /dev/null +++ b/src/noirdoc/daemon/spawn.py @@ -0,0 +1,133 @@ +"""Detached daemon spawn + stale-socket recovery. + +The daemon is launched as ``python -m noirdoc.daemon`` in a new session so +it survives the parent CLI process exiting. Old sockets and stale pidfiles +left behind by a crashed daemon are cleaned up before binding. +""" + +from __future__ import annotations + +import os +import signal +import subprocess +import sys +import time +from pathlib import Path + +from noirdoc.daemon import paths + + +def is_pid_alive(pid: int) -> bool: + if pid <= 0: + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def read_pidfile() -> int | None: + pf = paths.pidfile_path() + try: + text = pf.read_text().strip() + except (FileNotFoundError, OSError): + return None + try: + return int(text) + except ValueError: + return None + + +def write_pidfile(pid: int) -> None: + paths.ensure_root_dir() + paths.pidfile_path().write_text(f"{pid}\n") + + +def remove_pidfile() -> None: + try: + paths.pidfile_path().unlink() + except FileNotFoundError: + pass + + +def cleanup_stale_socket() -> None: + """Remove ``daemon.sock`` if no live process owns it. + + A live daemon's PID lives in the pidfile; we trust that record. If the + pidfile points at a dead PID (or is missing entirely) and a socket file + is still on disk, it's a leftover from a crash and is safe to remove. + """ + sock = paths.socket_path() + if not sock.exists(): + return + pid = read_pidfile() + if pid is not None and is_pid_alive(pid): + return # Live daemon owns the socket; do not touch. + try: + sock.unlink() + except FileNotFoundError: + pass + remove_pidfile() + + +def spawn_detached() -> int: + """Fork a fully detached daemon. Returns the spawned PID.""" + paths.ensure_root_dir() + proc = subprocess.Popen( + [sys.executable, "-m", "noirdoc.daemon"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + close_fds=True, + ) + return proc.pid + + +def wait_for_socket(timeout_s: float = 30.0) -> bool: + """Poll for the daemon socket to appear. Backoff 100→500 ms.""" + sock = paths.socket_path() + deadline = time.monotonic() + timeout_s + delay = 0.1 + while time.monotonic() < deadline: + if sock.exists(): + return True + time.sleep(delay) + delay = min(delay * 1.5, 0.5) + return False + + +def stop_daemon(timeout_s: float = 5.0) -> bool: + """Send SIGTERM to the running daemon (if any) and wait for exit. + + Returns ``True`` if a daemon was running and stopped (or was already + gone), ``False`` if the PID was alive but didn't exit in time. + """ + pid = read_pidfile() + if pid is None or not is_pid_alive(pid): + cleanup_stale_socket() + return True + try: + os.kill(pid, signal.SIGTERM) + except ProcessLookupError: + cleanup_stale_socket() + return True + + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if not is_pid_alive(pid): + cleanup_stale_socket() + return True + time.sleep(0.1) + return False + + +def daemon_log_handle() -> tuple[Path, int]: + """Open ``daemon.log`` for append, returning (path, fd).""" + paths.ensure_root_dir() + log = paths.logfile_path() + fd = os.open(log, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600) + return log, fd diff --git a/src/noirdoc/sdk.py b/src/noirdoc/sdk.py index e1c8b22..9adf6cd 100644 --- a/src/noirdoc/sdk.py +++ b/src/noirdoc/sdk.py @@ -22,7 +22,7 @@ from noirdoc.pseudonymization.mapper import PseudonymMapper if TYPE_CHECKING: - from noirdoc.detection.base import BaseDetector + from noirdoc.detection.base import BaseDetector, DetectedEntity from noirdoc.detection.ensemble import EnsembleDetector Policy = Literal["pseudonymize", "extract_only"] @@ -158,15 +158,29 @@ def redact_text(self, text: str, language: str | None = None) -> str: """Detect PII in *text*, replace with pseudonyms, return the result.""" return asyncio.run(self._redact_text_async(text, language or self._language)) - async def _redact_text_async(self, text: str, language: str) -> str: + async def aredact_text(self, text: str, language: str | None = None) -> str: + """Async version of :meth:`redact_text` for callers already in an event loop.""" + return await self._redact_text_async(text, language or self._language) + + async def aredact_text_detailed( + self, + text: str, + language: str | None = None, + ) -> tuple[str, list[DetectedEntity]]: + """Like :meth:`aredact_text` but also returns the detected entities.""" if not text: - return text + return text, [] from noirdoc.pseudonymization.engine import PseudonymizationEngine + lang = language or self._language detector = await self._ensure_detector() - entities = await detector.detect(text, language) + entities = await detector.detect(text, lang) result = PseudonymizationEngine().pseudonymize(text, entities, self._mapper) self._persist() + return result, entities + + async def _redact_text_async(self, text: str, language: str) -> str: + result, _ = await self.aredact_text_detailed(text, language) return result def reveal_text(self, text: str) -> str: @@ -194,6 +208,19 @@ def redact_file( result.write(output) return result + async def aredact_file( + self, + input_path: Path | str, + *, + output: Path | str | None = None, + language: str | None = None, + ) -> RedactionResult: + """Async version of :meth:`redact_file` for callers already in an event loop.""" + result = await self._redact_file_async(Path(input_path), language or self._language) + if output is not None: + result.write(output) + return result + async def _redact_file_async(self, path: Path, language: str) -> RedactionResult: from noirdoc.file_analysis.extractor import FileTextExtractor from noirdoc.file_analysis.mime import format_for_mime @@ -286,6 +313,36 @@ def reveal_file( return revealed +def build_redactor( + *, + ensemble: EnsembleDetector | None = None, + namespace: str | None = None, + namespace_root: Path | str | None = None, + language: str = "de", + detector: DetectorChoice = "ensemble", + score_threshold: float = 0.5, + gliner_model: str = "knowledgator/gliner-pii-edge-v1.0", +) -> Redactor: + """Construct a :class:`Redactor`, optionally pre-installing a built ensemble. + + The CLI fallback path passes ``ensemble=None`` and lets the redactor + lazily build its own. The daemon passes a pre-built, cached ensemble + so model loading is paid once for the daemon's lifetime, not once per + request. + """ + r = Redactor( + namespace=namespace, + namespace_root=namespace_root, + language=language, + detector=detector, + score_threshold=score_threshold, + gliner_model=gliner_model, + ) + if ensemble is not None: + r._ensemble = ensemble + return r + + def redact( input_path: Path | str, *, diff --git a/tests/daemon/__init__.py b/tests/daemon/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/daemon/test_fallback.py b/tests/daemon/test_fallback.py new file mode 100644 index 0000000..22ac711 --- /dev/null +++ b/tests/daemon/test_fallback.py @@ -0,0 +1,136 @@ +"""When the daemon path fails, the CLI must fall back to in-process redaction.""" + +from __future__ import annotations + +from pathlib import Path + +from click.testing import CliRunner + +from noirdoc import cli as cli_module +from noirdoc import sdk as sdk_module +from noirdoc.cli import main +from noirdoc.daemon import client as daemon_client +from noirdoc.sdk import RedactionResult + + +def _make_fake_redact_file(out_bytes: bytes = b"REDACTED"): + def _fake(self, input_path, *, output=None, language=None): + return RedactionResult( + input_path=Path(input_path), + output_bytes=out_bytes, + entity_count=1, + entity_types={"PERSON": 1}, + mime_type="text/plain", + reconstructed=False, + ) + + return _fake + + +def test_cli_falls_back_when_daemon_unavailable(monkeypatch, tmp_path: Path): + """call_sync raises DaemonUnavailable → CLI prints fallback warning and runs in-process.""" + + def fake_call_sync(method, params=None): + raise daemon_client.DaemonUnavailable("simulated: daemon not running") + + # Stub the daemon RPC. + monkeypatch.setattr( + "noirdoc.daemon.client.call_sync", + fake_call_sync, + ) + # Stub the in-process Redactor.redact_file so we don't load real models. + monkeypatch.setattr( + sdk_module.Redactor, + "redact_file", + _make_fake_redact_file(b"REDACTED-LOCAL"), + ) + + inp = tmp_path / "input.txt" + inp.write_text("Max Mueller lives in Berlin.") + + out_dir = tmp_path / "out" + result = CliRunner().invoke( + main, + ["redact", str(inp), "--output-dir", str(out_dir)], + catch_exceptions=False, + ) + + assert result.exit_code == 0, result.output + assert "daemon unavailable" in result.output + assert "1 entities" in result.output + written = list(out_dir.glob("*")) + assert len(written) == 1 + assert written[0].read_bytes() == b"REDACTED-LOCAL" + + +def test_cli_skips_daemon_when_no_daemon_flag_set(monkeypatch, tmp_path: Path): + """--no-daemon must skip the daemon entirely (no DaemonUnavailable warning).""" + call_count = {"n": 0} + + def fake_call_sync(method, params=None): + call_count["n"] += 1 + raise daemon_client.DaemonUnavailable("should not be called") + + monkeypatch.setattr( + "noirdoc.daemon.client.call_sync", + fake_call_sync, + ) + monkeypatch.setattr( + sdk_module.Redactor, + "redact_file", + _make_fake_redact_file(), + ) + + inp = tmp_path / "input.txt" + inp.write_text("hello") + out_dir = tmp_path / "out" + + result = CliRunner().invoke( + main, + ["redact", "--no-daemon", str(inp), "--output-dir", str(out_dir)], + catch_exceptions=False, + ) + + assert result.exit_code == 0, result.output + assert "daemon unavailable" not in result.output.lower() + assert call_count["n"] == 0 + + +def test_cli_skips_daemon_when_env_var_set(monkeypatch, tmp_path: Path): + """NOIRDOC_NO_DAEMON=1 also skips the daemon.""" + call_count = {"n": 0} + + def fake_call_sync(method, params=None): + call_count["n"] += 1 + raise daemon_client.DaemonUnavailable("should not be called") + + monkeypatch.setenv("NOIRDOC_NO_DAEMON", "1") + monkeypatch.setattr( + "noirdoc.daemon.client.call_sync", + fake_call_sync, + ) + monkeypatch.setattr( + sdk_module.Redactor, + "redact_file", + _make_fake_redact_file(), + ) + + inp = tmp_path / "input.txt" + inp.write_text("hello") + out_dir = tmp_path / "out" + + result = CliRunner().invoke( + main, + ["redact", str(inp), "--output-dir", str(out_dir)], + catch_exceptions=False, + ) + + assert result.exit_code == 0, result.output + assert call_count["n"] == 0 + assert "daemon unavailable" not in result.output.lower() + + +def test_cli_module_uses_daemon_imports(monkeypatch): + """Smoke test: the cli module exposes the helper used by tests.""" + assert hasattr(cli_module, "_redact_via_daemon") + assert hasattr(cli_module, "_redact_in_process") diff --git a/tests/daemon/test_protocol.py b/tests/daemon/test_protocol.py new file mode 100644 index 0000000..b54fdf7 --- /dev/null +++ b/tests/daemon/test_protocol.py @@ -0,0 +1,110 @@ +"""Pydantic round-trip tests for the daemon wire protocol.""" + +from __future__ import annotations + +import json + +import pytest + +from noirdoc.daemon.protocol import ( + ErrorPayload, + HelloParams, + HelloResult, + RedactFileInput, + RedactParams, + RedactResult, + RedactTextInput, + Request, + Response, + ShutdownResult, + StatusResult, +) + + +def _roundtrip(model): + data = model.model_dump() + serialized = json.dumps(data) + return type(model).model_validate(json.loads(serialized)) + + +def test_hello_params_roundtrip(): + assert _roundtrip(HelloParams(client_version="1.2.3")).client_version == "1.2.3" + + +def test_hello_result_roundtrip(): + out = _roundtrip(HelloResult(daemon_version="1.2.3", pid=42, started_at=10.0)) + assert out.daemon_version == "1.2.3" + assert out.pid == 42 + + +def test_redact_text_input_roundtrip(): + p = RedactParams(input=RedactTextInput(value="hello world")) + out = _roundtrip(p) + assert isinstance(out.input, RedactTextInput) + assert out.input.value == "hello world" + + +def test_redact_file_input_roundtrip(): + p = RedactParams(input=RedactFileInput(path="/tmp/foo.txt"), output_path="/tmp/out.txt") + out = _roundtrip(p) + assert isinstance(out.input, RedactFileInput) + assert out.input.path == "/tmp/foo.txt" + assert out.output_path == "/tmp/out.txt" + + +def test_redact_params_defaults(): + p = RedactParams(input=RedactTextInput(value="x")) + assert p.language == "de" + assert p.detector == "ensemble" + assert p.score_threshold == 0.5 + + +def test_redact_input_discriminator_rejects_unknown_type(): + with pytest.raises(Exception): + RedactParams.model_validate({"input": {"type": "garbage", "value": "x"}}) + + +def test_redact_result_optional_fields(): + r = RedactResult(entity_count=3, entity_types={"PERSON": 2, "EMAIL": 1}) + assert r.redacted_text is None + assert r.output_path is None + assert r.reconstructed is False + + +def test_status_result_roundtrip(): + s = StatusResult( + uptime_s=10.5, + models_loaded=True, + last_request_at=None, + queue_depth=0, + total_requests=0, + ) + out = _roundtrip(s) + assert out.uptime_s == 10.5 + assert out.last_request_at is None + + +def test_shutdown_result_default(): + assert ShutdownResult().ok is True + + +def test_request_envelope_roundtrip(): + req = Request(id="abc", method="redact", params={"foo": "bar"}) + out = _roundtrip(req) + assert out.id == "abc" + assert out.method == "redact" + assert out.params == {"foo": "bar"} + + +def test_response_envelope_with_result(): + resp = Response(id="abc", result={"ok": True}) + out = _roundtrip(resp) + assert out.result == {"ok": True} + assert out.error is None + + +def test_response_envelope_with_error(): + resp = Response(id="abc", error=ErrorPayload(code="bad_request", message="x")) + out = _roundtrip(resp) + assert out.error is not None + assert out.error.code == "bad_request" diff --git a/tests/daemon/test_serial_queue.py b/tests/daemon/test_serial_queue.py new file mode 100644 index 0000000..88862e5 --- /dev/null +++ b/tests/daemon/test_serial_queue.py @@ -0,0 +1,109 @@ +"""Verify the daemon serializes redact requests via its single asyncio lock.""" + +from __future__ import annotations + +import asyncio +import time + +import pytest + +from noirdoc.daemon import server + +pytestmark = pytest.mark.asyncio + + +async def test_handle_redact_serializes(monkeypatch): + """Two concurrent ``handle_redact`` calls must run back-to-back, not in parallel.""" + state = server.DaemonState() + + sleep_for = 0.15 + + class FakeMapper: + entity_count = 0 + + class FakeRedactor: + mapper = FakeMapper() + + async def aredact_text_detailed(self, text, language): + await asyncio.sleep(sleep_for) + return text, [] + + async def fake_get_detectors(*args, **kwargs): + return [] + + monkeypatch.setattr(state, "get_detectors", fake_get_detectors) + monkeypatch.setattr( + "noirdoc.sdk.build_redactor", + lambda **kwargs: FakeRedactor(), + ) + + params = { + "namespace": None, + "language": "de", + "detector": "ensemble", + "score_threshold": 0.5, + "input": {"type": "text", "value": "hello"}, + } + + t0 = time.monotonic() + results = await asyncio.gather( + server.handle_redact(state, params), + server.handle_redact(state, params), + ) + elapsed = time.monotonic() - t0 + + assert elapsed >= sleep_for * 1.8, ( + f"expected serialized (~{sleep_for * 2}s), got {elapsed:.3f}s" + ) + assert len(results) == 2 + assert state.total_requests == 2 + assert state.queue_depth == 0 + + +async def test_queue_depth_tracks_pending(monkeypatch): + """While one request holds the lock, a queued request bumps queue_depth above 1.""" + state = server.DaemonState() + + proceed = asyncio.Event() + + class FakeMapper: + entity_count = 0 + + class FakeRedactor: + mapper = FakeMapper() + + async def aredact_text_detailed(self, text, language): + await proceed.wait() + return text, [] + + async def fake_get_detectors(*args, **kwargs): + return [] + + monkeypatch.setattr(state, "get_detectors", fake_get_detectors) + monkeypatch.setattr( + "noirdoc.sdk.build_redactor", + lambda **kwargs: FakeRedactor(), + ) + + params = { + "namespace": None, + "language": "de", + "detector": "ensemble", + "score_threshold": 0.5, + "input": {"type": "text", "value": "hello"}, + } + + task1 = asyncio.create_task(server.handle_redact(state, params)) + task2 = asyncio.create_task(server.handle_redact(state, params)) + + # Yield repeatedly so both tasks reach the lock-acquire point. + for _ in range(50): + await asyncio.sleep(0.01) + if state.queue_depth >= 2: + break + + assert state.queue_depth == 2 + + proceed.set() + await asyncio.gather(task1, task2) + assert state.queue_depth == 0 diff --git a/tests/daemon/test_spawn.py b/tests/daemon/test_spawn.py new file mode 100644 index 0000000..37f1c77 --- /dev/null +++ b/tests/daemon/test_spawn.py @@ -0,0 +1,71 @@ +"""Spawn / pidfile / stale-socket tests.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from noirdoc.daemon import paths, spawn + + +@pytest.fixture +def isolated_root(monkeypatch, tmp_path: Path) -> Path: + monkeypatch.setenv("NOIRDOC_DAEMON_ROOT", str(tmp_path)) + monkeypatch.delenv("NOIRDOC_DAEMON_SOCKET", raising=False) + return tmp_path + + +def test_pidfile_roundtrip(isolated_root: Path): + assert spawn.read_pidfile() is None + spawn.write_pidfile(12345) + assert spawn.read_pidfile() == 12345 + spawn.remove_pidfile() + assert spawn.read_pidfile() is None + + +def test_pidfile_invalid_content(isolated_root: Path): + paths.pidfile_path().write_text("not-an-int\n") + assert spawn.read_pidfile() is None + + +def test_is_pid_alive_self(): + assert spawn.is_pid_alive(os.getpid()) is True + + +def test_is_pid_alive_dead(): + # PID 0 is conventionally invalid; we treat <=0 as dead. + assert spawn.is_pid_alive(0) is False + + +def test_cleanup_stale_socket_removes_when_no_owner(isolated_root: Path): + sock = paths.socket_path() + sock.parent.mkdir(parents=True, exist_ok=True) + sock.write_bytes(b"") # leftover from a crash + # No pidfile → no live owner → socket should go. + spawn.cleanup_stale_socket() + assert not sock.exists() + + +def test_cleanup_stale_socket_removes_when_pid_dead(isolated_root: Path): + sock = paths.socket_path() + sock.parent.mkdir(parents=True, exist_ok=True) + sock.write_bytes(b"") + spawn.write_pidfile(0) # 0 is not alive in our check + spawn.cleanup_stale_socket() + assert not sock.exists() + assert spawn.read_pidfile() is None + + +def test_cleanup_stale_socket_preserves_when_owner_alive(isolated_root: Path): + sock = paths.socket_path() + sock.parent.mkdir(parents=True, exist_ok=True) + sock.write_bytes(b"") + spawn.write_pidfile(os.getpid()) # this very process is alive + spawn.cleanup_stale_socket() + assert sock.exists() # left alone + + +def test_stop_daemon_when_not_running(isolated_root: Path): + assert spawn.stop_daemon(timeout_s=0.5) is True diff --git a/tests/daemon/test_version_handshake.py b/tests/daemon/test_version_handshake.py new file mode 100644 index 0000000..20161c8 --- /dev/null +++ b/tests/daemon/test_version_handshake.py @@ -0,0 +1,152 @@ +"""Version-mismatch handling: client should ask the stale daemon to shut down and respawn. + +These tests stand up a tiny in-process Unix-socket server that pretends to be +the daemon and inspect the sequence of methods the client invokes. No real +``noirdoc-daemon`` subprocess is involved. +""" + +from __future__ import annotations + +import asyncio +import json +import tempfile +import uuid +from pathlib import Path + +import pytest + +from noirdoc.daemon import client as daemon_client +from noirdoc.daemon import paths + + +@pytest.fixture +def isolated_paths(monkeypatch): + # AF_UNIX paths are limited to ~104 bytes on macOS, so pytest tmp_path + # (which nests deeply) overflows. Use a short /tmp path instead. + short_dir = Path(tempfile.gettempdir()) / f"nd-{uuid.uuid4().hex[:8]}" + short_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("NOIRDOC_DAEMON_ROOT", str(short_dir)) + monkeypatch.setenv("NOIRDOC_DAEMON_SOCKET", str(short_dir / "d.sock")) + yield short_dir + for f in short_dir.glob("*"): + try: + f.unlink() + except OSError: + pass + try: + short_dir.rmdir() + except OSError: + pass + + +class _FakeDaemon: + """Minimal Unix-socket server scripted with daemon_version per-connection.""" + + def __init__(self, socket_path: Path, version: str): + self.socket_path = socket_path + self.version = version + self.received: list[dict] = [] + self.server: asyncio.AbstractServer | None = None + self._stop_after_shutdown = False + + async def start(self): + self.server = await asyncio.start_unix_server( + self._handle, + path=str(self.socket_path), + ) + + async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter): + try: + while True: + line = await reader.readline() + if not line: + return + req = json.loads(line.decode()) + self.received.append(req) + method = req["method"] + if method == "hello": + result = {"daemon_version": self.version, "pid": 999, "started_at": 0.0} + elif method == "shutdown": + result = {"ok": True} + self._stop_after_shutdown = True + else: + result = {"echoed": method} + resp = {"id": req["id"], "result": result} + writer.write((json.dumps(resp) + "\n").encode()) + await writer.drain() + if self._stop_after_shutdown: + return + finally: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + if self._stop_after_shutdown and self.server is not None: + self.server.close() + + async def stop(self): + if self.server is not None: + self.server.close() + await self.server.wait_closed() + try: + self.socket_path.unlink() + except FileNotFoundError: + pass + + +@pytest.mark.asyncio +async def test_matching_version_passes_through(isolated_paths): + sock = paths.socket_path() + daemon = _FakeDaemon(sock, version=daemon_client.__version__) + await daemon.start() + try: + result = await daemon_client.call("status", {}) + assert result.get("echoed") == "status" + methods = [m["method"] for m in daemon.received] + assert methods == ["hello", "status"] + finally: + await daemon.stop() + + +@pytest.mark.asyncio +async def test_version_mismatch_triggers_shutdown_and_respawn( + isolated_paths, + monkeypatch, +): + """First connection (stale version) must receive ``shutdown``; the client + then sees a fresh daemon brought up at the matching version.""" + sock = paths.socket_path() + + stale = _FakeDaemon(sock, version="0.0.0-stale") + await stale.start() + + # When the client calls _spawn_and_connect after the stale daemon dies, + # we don't actually want to fork a real daemon. Patch it to bring up a + # fresh fake at the matching version. + fresh: dict = {"daemon": None} + + async def fake_spawn_and_connect(socket_path: Path): + fresh_d = _FakeDaemon(socket_path, version=daemon_client.__version__) + await fresh_d.start() + fresh["daemon"] = fresh_d + reader, writer = await asyncio.open_unix_connection(path=str(socket_path)) + return reader, writer + + monkeypatch.setattr(daemon_client, "_spawn_and_connect", fake_spawn_and_connect) + + try: + result = await daemon_client.call("status", {}) + assert result.get("echoed") == "status" + # Stale daemon must have seen hello, then shutdown. + stale_methods = [m["method"] for m in stale.received] + assert stale_methods[0] == "hello" + assert "shutdown" in stale_methods + # Fresh daemon must have seen hello, then status. + assert fresh["daemon"] is not None + fresh_methods = [m["method"] for m in fresh["daemon"].received] + assert fresh_methods == ["hello", "status"] + finally: + await stale.stop() + if fresh["daemon"] is not None: + await fresh["daemon"].stop()