agent_search_gateway.socket_probe

 1"""Bounded, read-only inspection of the local daemon Unix socket."""
 2
 3import asyncio
 4import stat
 5from collections.abc import Awaitable, Callable
 6from contextlib import suppress
 7from dataclasses import dataclass
 8from enum import StrEnum
 9from pathlib import Path
10
11SOCKET_PROBE_TIMEOUT_SECONDS = 2.0
12SocketConnector = Callable[..., Awaitable[tuple[asyncio.StreamReader, asyncio.StreamWriter]]]
13
14
15class SocketState(StrEnum):
16    MISSING = "missing"
17    LIVE = "live"
18    REFUSED = "refused"
19    NOT_SOCKET = "not_socket"
20    TIMEOUT = "timeout"
21    OS_ERROR = "os_error"
22
23
24@dataclass(frozen=True, slots=True)
25class SocketProbeResult:
26    state: SocketState
27    identity: tuple[int, int] | None = None
28    reason: str = ""
29
30
31async def probe_unix_socket(
32    path: Path,
33    timeout_seconds: float = SOCKET_PROBE_TIMEOUT_SECONDS,
34    connector: SocketConnector = asyncio.open_unix_connection,
35) -> SocketProbeResult:
36    try:
37        existing = path.lstat()
38    except FileNotFoundError:
39        return SocketProbeResult(SocketState.MISSING)
40    except OSError as exc:
41        return SocketProbeResult(SocketState.OS_ERROR, reason=_safe_reason(exc))
42
43    if not stat.S_ISSOCK(existing.st_mode):
44        return SocketProbeResult(SocketState.NOT_SOCKET)
45    identity = (existing.st_dev, existing.st_ino)
46
47    try:
48        _, writer = await asyncio.wait_for(
49            connector(path=path),
50            timeout=timeout_seconds,
51        )
52    except FileNotFoundError:
53        return SocketProbeResult(SocketState.MISSING)
54    except ConnectionRefusedError as exc:
55        return SocketProbeResult(SocketState.REFUSED, identity, _safe_reason(exc))
56    except TimeoutError as exc:
57        return SocketProbeResult(SocketState.TIMEOUT, identity, _safe_reason(exc))
58    except OSError as exc:
59        return SocketProbeResult(SocketState.OS_ERROR, identity, _safe_reason(exc))
60
61    writer.close()
62    with suppress(OSError):
63        await writer.wait_closed()
64    return SocketProbeResult(SocketState.LIVE, identity)
65
66
67def _safe_reason(exc: OSError | TimeoutError) -> str:
68    text = str(exc).strip()
69    return text if text else type(exc).__name__