tracer-cloud--opensre/tools/system/fleet_monitoring/probe.py

"""Per-PID process helpers for the monitor-local-agents fleet view.

Pure collectors: no background loop, no caching, no UI wiring. The
wiring layer (#1490) batches calls in a REPL background task; the
registry layer (#1487) decides which PIDs to ask about.

``psutil`` is intentionally confined to this module. The helpers
exposed here (:func:`cwd_for_pid`, :func:`started_at_for_pid`,
:func:`env_value_for_pid`) let other layers query process state
without importing psutil themselves.
"""

from __future__ import annotations

from collections.abc import Iterable, Iterator
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path

import psutil

# 1 MiB exactly. The dataclass field below is named ``rss_mb`` because
# every monitoring tool in this space (htop, top, ps, k8s metrics,
# Datadog, Grafana) labels the same 1024² unit as "MB"; the constant
# stays precise so the unit math is unambiguous.
_BYTES_PER_MIB = 1024 * 1024

PROCESS_NOT_FOUND: tuple[type[BaseException], ...] = (psutil.NoSuchProcess,)
PROCESS_INACCESSIBLE_OR_GONE: tuple[type[BaseException], ...] = (
    psutil.NoSuchProcess,
    psutil.AccessDenied,
)
PROCESS_ERROR: tuple[type[BaseException], ...] = (psutil.Error,)


@dataclass(frozen=True)
class ProcessSnapshot:
    """Single-instant resource snapshot for a process.

    Fields not available on the current platform or for the current
    user are ``None`` rather than raising — file descriptors are
    POSIX-only, and the connection count requires elevated privileges
    on some systems.
    """

    pid: int
    cpu_percent: float
    rss_mb: float
    num_fds: int | None
    num_connections: int | None
    status: str
    started_at: datetime
    last_output_at: datetime | None = None


def pid_exists(pid: int) -> bool:
    """Return whether ``pid`` corresponds to a process the OS knows about.

    Thin wrapper over ``psutil.pid_exists`` exposed here so ``psutil``
    stays confined to this module per the issue #1489 acceptance
    criterion. Unlike ``probe()``, this returns ``True`` for processes
    we can't introspect (cross-user on macOS, restricted ``/proc``,
    etc.) — the OS-level existence check doesn't traverse the access
    boundary, which is exactly what the boot sweep (#1501) needs to
    avoid pruning live foreign-user agents.

    Returns ``False`` for PIDs outside the platform's valid range
    (e.g. an int that overflows the kernel's PID type) — psutil
    raises ``OverflowError`` or ``ValueError`` on those, which we
    treat as "not a real PID" rather than propagating.
    """
    try:
        return psutil.pid_exists(pid)
    except (OverflowError, ValueError):
        return False


def process(pid: int) -> psutil.Process:
    """Return a handle for ``pid`` while keeping psutil access local."""
    return psutil.Process(pid)


def process_iter(attrs: Iterable[str]) -> Iterator[psutil.Process]:
    """Yield process handles with preloaded attrs from the local probe module."""
    return psutil.process_iter(list(attrs))


def process_has_open_codex_rollout(pid: int) -> bool:
    """Return whether ``pid`` has an open Codex ``rollout-*.jsonl`` file."""
    try:
        proc = psutil.Process(pid)
        open_files = proc.open_files()
    # The stdlib exceptions cover invalid/raced PIDs and platform-specific
    # ``open_files()`` failures that psutil may surface directly.
    except PROCESS_ERROR + (
        ProcessLookupError,
        OSError,
        ValueError,
        OverflowError,
    ):
        return False

    for open_file in open_files:
        path = getattr(open_file, "path", None)
        if isinstance(path, str):
            name = Path(path).name
            if name.startswith("rollout-") and name.endswith(".jsonl"):
                return True
    return False


def probe(pid: int, *, cpu_interval: float = 0.1) -> ProcessSnapshot | None:
    """Return a one-shot resource snapshot for ``pid``.

    ``cpu_interval`` blocks for that many seconds to compute an
    accurate CPU percentage. Pass ``0.0`` for a non-blocking sample —
    the first such call returns ``0.0`` because psutil needs a delta
    baseline; callers that want accuracy without blocking should
    manage their own ``psutil.Process`` instances and call this
    function with ``cpu_interval=0.0`` on subsequent samples.

    Returns ``None`` for PIDs that don't exist, are zombies, or whose
    fields are inaccessible (typically processes owned by another user
    on macOS or Linux setups with restricted ``/proc``). Never raises
    ``psutil.NoSuchProcess``, ``psutil.ZombieProcess``, or
    ``psutil.AccessDenied``.
    """
    try:
        proc = psutil.Process(pid)
    except (psutil.NoSuchProcess, ProcessLookupError):
        return None

    try:
        cpu = proc.cpu_percent(interval=cpu_interval)

        with proc.oneshot():
            rss_mb = proc.memory_info().rss / _BYTES_PER_MIB
            num_fds = _safe_num_fds(proc)
            num_connections = _safe_num_connections(proc)
            status = proc.status()
            started_at = datetime.fromtimestamp(proc.create_time(), tz=UTC)
    except (psutil.NoSuchProcess, psutil.ZombieProcess, psutil.AccessDenied):
        # Process exited or its core fields (memory, status, create
        # time) are inaccessible to this user. The wiring layer in
        # #1490 treats both as "no snapshot this tick" and renders an
        # empty cell rather than tearing the REPL background task.
        return None

    return ProcessSnapshot(
        pid=pid,
        cpu_percent=cpu,
        rss_mb=rss_mb,
        num_fds=num_fds,
        num_connections=num_connections,
        status=status,
        started_at=started_at,
    )


def cwd_for_pid(pid: int) -> Path | None:
    """Return the cwd of ``pid``, or ``None`` on any psutil failure.

    macOS hardened-runtime processes can deny ``cwd()`` even
    cross-user-same-user. Callers treat ``None`` as "unobservable
    this tick" and retry next tick rather than poison a cache.
    """
    try:
        proc = psutil.Process(pid)
        return Path(proc.cwd())
    except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess, OSError):
        return None


def started_at_for_pid(pid: int) -> float | None:
    """Return the POSIX-epoch start time of ``pid``, or ``None``."""
    try:
        proc = psutil.Process(pid)
        return float(proc.create_time())
    except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess, OSError):
        return None


def env_value_for_pid(pid: int, key: str) -> str | None:
    """Return env var ``key`` for ``pid``, or ``None``.

    ``environ()`` denies on macOS hardened-runtime processes even for
    the same user; callers must accept ``None`` and fall back to
    other resolution sources.
    """
    try:
        proc = psutil.Process(pid)
        environ = proc.environ()
    except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess, OSError):
        return None
    value = environ.get(key)
    if isinstance(value, str) and value:
        return value
    return None


def open_files_for_pid(pid: int) -> tuple[Path, ...]:
    """Return open file paths for ``pid``, or an empty tuple on psutil failure."""
    try:
        proc = psutil.Process(pid)
        files = proc.open_files()
    except (
        psutil.NoSuchProcess,
        psutil.AccessDenied,
        psutil.ZombieProcess,
        OSError,
        ValueError,
    ):
        return ()
    return tuple(Path(item.path) for item in files if item.path)


def _safe_num_fds(proc: psutil.Process) -> int | None:
    """File-descriptor count is POSIX-only; ``None`` on Windows.

    ``Process.num_fds`` is missing from the Windows-facing typeshed shape
    and is absent at runtime; use ``getattr`` so ``mypy --platform win32``
    stays clean.
    """
    num_fds_fn = getattr(proc, "num_fds", None)
    if num_fds_fn is None:
        return None
    try:
        n = num_fds_fn()
    except (psutil.AccessDenied, NotImplementedError, TypeError, ValueError):
        return None
    return int(n)


def _safe_num_connections(proc: psutil.Process) -> int | None:
    """Connection count requires elevated privileges on some platforms.

    Lazy fallback via ``hasattr`` rather than ``getattr(..., default)``
    because the latter eagerly evaluates ``proc.connections`` even when
    ``net_connections`` exists; a future psutil release that drops the
    deprecated ``connections`` method would then raise ``AttributeError``
    on the working code path.
    """
    method = proc.net_connections if hasattr(proc, "net_connections") else proc.connections
    try:
        connections = method()
    except (psutil.AccessDenied, NotImplementedError):
        return None
    return len(connections)