agent_search_gateway.doctor

  1"""Local, deterministic health diagnostics for the gateway CLI."""
  2
  3import os
  4import secrets
  5from collections.abc import Awaitable, Callable, Mapping
  6from contextlib import suppress
  7from dataclasses import dataclass
  8from enum import StrEnum
  9from pathlib import Path
 10from typing import TextIO
 11
 12from .config import load_toml, resolve_config
 13from .errors import ConfigFailure
 14from .observability import normalize_log_reason
 15from .paths import RuntimePaths
 16from .providers.academic.defaults import (
 17    build_default_academic_registry,
 18    build_default_oa_resolver_registry,
 19)
 20from .providers.defaults import build_default_registry
 21from .socket_probe import SocketProbeResult, SocketState, probe_unix_socket
 22
 23
 24class DoctorStatus(StrEnum):
 25    OK = "ok"
 26    INFO = "info"
 27    FAIL = "fail"
 28
 29
 30@dataclass(frozen=True, slots=True)
 31class DoctorCheck:
 32    status: DoctorStatus
 33    message: str
 34
 35
 36@dataclass(frozen=True, slots=True)
 37class DoctorReport:
 38    checks: tuple[DoctorCheck, ...]
 39
 40    @property
 41    def exit_code(self) -> int:
 42        return 1 if any(check.status is DoctorStatus.FAIL for check in self.checks) else 0
 43
 44
 45DirectoryProbe = Callable[[Path], DoctorCheck]
 46SocketProbe = Callable[[Path], Awaitable[SocketProbeResult]]
 47
 48
 49async def run_doctor(
 50    paths: RuntimePaths,
 51    *,
 52    environ: Mapping[str, str],
 53    directory_probe: DirectoryProbe = lambda path: probe_directory_writable(path),
 54    socket_probe: SocketProbe = probe_unix_socket,
 55) -> DoctorReport:
 56    checks: list[DoctorCheck] = []
 57    checks.extend(_configuration_checks(paths.config_file, environ))
 58    checks.append(directory_probe(paths.socket_file.parent))
 59    checks.append(directory_probe(paths.results_dir))
 60    checks.append(directory_probe(paths.logs_dir))
 61    checks.append(_socket_check(paths.socket_file, await socket_probe(paths.socket_file)))
 62    return DoctorReport(tuple(checks))
 63
 64
 65def render_doctor(report: DoctorReport, stream: TextIO) -> None:
 66    for check in report.checks:
 67        message = normalize_log_reason(check.message, max_chars=1000)
 68        stream.write(f"[{check.status.value}] {message}\n")
 69
 70
 71def probe_directory_writable(path: Path) -> DoctorCheck:
 72    candidate = path
 73    while True:
 74        try:
 75            candidate.lstat()
 76        except FileNotFoundError:
 77            parent = candidate.parent
 78            if parent == candidate:
 79                return DoctorCheck(
 80                    DoctorStatus.FAIL,
 81                    f"directory has no existing parent: {path}",
 82                )
 83            candidate = parent
 84            continue
 85        except OSError as exc:
 86            return DoctorCheck(
 87                DoctorStatus.FAIL,
 88                f"unable to inspect directory path: {path}: {_safe_os_reason(exc)}",
 89            )
 90        break
 91
 92    if not candidate.is_dir():
 93        if candidate == path:
 94            return DoctorCheck(
 95                DoctorStatus.FAIL,
 96                f"expected directory but found non-directory: {path}",
 97            )
 98        return DoctorCheck(
 99            DoctorStatus.FAIL,
100            f"directory cannot be created because parent is not a directory: {candidate}",
101        )
102    if candidate == path:
103        return _probe_existing_directory(path, target=path)
104
105    probed = _probe_existing_directory(candidate, target=path)
106    if probed.status is DoctorStatus.OK:
107        return DoctorCheck(DoctorStatus.OK, f"directory is creatable: {path}")
108    return probed
109
110
111def _configuration_checks(
112    config_file: Path,
113    environ: Mapping[str, str],
114) -> tuple[DoctorCheck, ...]:
115    if not config_file.exists():
116        return (DoctorCheck(DoctorStatus.FAIL, f"config file not found: {config_file}"),)
117    try:
118        data = load_toml(config_file)
119    except ConfigFailure as exc:
120        return (DoctorCheck(DoctorStatus.FAIL, f"config parse failed: {exc.message}"),)
121    try:
122        config = resolve_config(
123            data,
124            build_default_registry(),
125            environ,
126            academic_registry=build_default_academic_registry(),
127            oa_resolver_registry=build_default_oa_resolver_registry(),
128        )
129    except ConfigFailure as exc:
130        return (DoctorCheck(DoctorStatus.FAIL, f"configuration invalid: {exc.message}"),)
131
132    checks: list[DoctorCheck] = [DoctorCheck(DoctorStatus.OK, "configuration valid")]
133    env_names = {
134        provider.api_key_env
135        for provider in config.web.providers
136        if provider.api_key_env is not None
137    }
138    env_names.update(provider.api_key_env for provider in config.llm.providers)
139    for provider in config.academic.providers:
140        if provider.api_key_env is not None:
141            env_names.add(provider.api_key_env)
142        if provider.contact_email_env is not None:
143            env_names.add(provider.contact_email_env)
144    if config.oa_resolver is not None:
145        if config.oa_resolver.api_key_env is not None:
146            env_names.add(config.oa_resolver.api_key_env)
147        if config.oa_resolver.contact_email_env is not None:
148            env_names.add(config.oa_resolver.contact_email_env)
149    checks.extend(
150        DoctorCheck(DoctorStatus.OK, f"environment variable {name} is set")
151        for name in sorted(env_names)
152    )
153    return tuple(checks)
154
155
156def _probe_existing_directory(directory: Path, *, target: Path) -> DoctorCheck:
157    probe_file = directory / f".agent-search-gateway-doctor-{secrets.token_hex(8)}"
158    descriptor: int | None = None
159    created = False
160    try:
161        descriptor = os.open(probe_file, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
162        created = True
163        os.write(descriptor, b"probe")
164        os.close(descriptor)
165        descriptor = None
166        probe_file.unlink()
167        created = False
168    except OSError as exc:
169        if descriptor is not None:
170            with suppress(OSError):
171                os.close(descriptor)
172        if created:
173            try:
174                probe_file.unlink()
175            except OSError as cleanup_exc:
176                return DoctorCheck(
177                    DoctorStatus.FAIL,
178                    f"directory probe cleanup failed: {target}: {_safe_os_reason(cleanup_exc)}",
179                )
180        return DoctorCheck(
181            DoctorStatus.FAIL,
182            f"directory is not writable: {target}: {_safe_os_reason(exc)}",
183        )
184    return DoctorCheck(DoctorStatus.OK, f"directory writable: {target}")
185
186
187def _socket_check(path: Path, result: SocketProbeResult) -> DoctorCheck:
188    if result.state is SocketState.MISSING:
189        return DoctorCheck(DoctorStatus.INFO, "daemon not running")
190    if result.state is SocketState.LIVE:
191        return DoctorCheck(DoctorStatus.OK, "daemon running")
192    if result.state is SocketState.REFUSED:
193        return DoctorCheck(
194            DoctorStatus.FAIL,
195            f"daemon socket is stale or refusing connections: {path}",
196        )
197    if result.state is SocketState.TIMEOUT:
198        return DoctorCheck(
199            DoctorStatus.FAIL,
200            f"daemon socket did not respond in time: {path}",
201        )
202    if result.state is SocketState.NOT_SOCKET:
203        return DoctorCheck(
204            DoctorStatus.FAIL,
205            f"daemon socket path is not a Unix socket: {path}",
206        )
207    reason = normalize_log_reason(result.reason or "OS error", max_chars=300)
208    return DoctorCheck(DoctorStatus.FAIL, f"unable to inspect daemon socket: {reason}")
209
210
211def _safe_os_reason(exc: OSError) -> str:
212    return normalize_log_reason(str(exc) or type(exc).__name__, max_chars=300)