FabledScribe/src/fabledassistant/services/diagnostics.py

"""Lightweight diagnostic instrumentation for crash investigation.

The Scribe app + its Postgres have been crashing recurrently with no
clear cause in the logs. This module adds three things designed to make
the crash class identifiable from logs alone:

1. **Heartbeat** — once per minute, log a snapshot of process resources
   (RSS memory, asyncio task count, DB pool checked-in/out, curator
   busy state). A sudden silence in heartbeats lets you bound the
   crash time to within a minute, and the last snapshot before silence
   usually rules in or out memory growth / pool exhaustion / hung
   curator pass.

2. **Signal handler** — catches SIGTERM and SIGINT, logs them with the
   sender's intent ("docker stop", "swarm restart", "manual ctrl-C")
   then lets the normal shutdown proceed. Distinguishes orderly
   shutdown from kill-9 / OOM-kill (which can't be caught and will
   show as a silent log gap followed by container exit code 137).

3. **Asyncio exception hook** — every Task that raises an uncaught
   exception logs a full traceback. Without this, `asyncio.create_task`
   exceptions are swallowed silently — the chat crash that locked us
   into 409 forever was exactly this pattern.

All three are read-only / log-only — no behavior changes. Safe to
leave running in production indefinitely; the cost is one log line
per minute and ~0.1ms of work per heartbeat.
"""
from __future__ import annotations

import asyncio
import logging
import os
import signal
import time
from typing import Any

logger = logging.getLogger(__name__)

# Heartbeat cadence. 60s is the sweet spot: short enough that a crash
# window is bounded to a useful interval, long enough that the log
# noise is negligible. If we ever need more resolution during active
# debugging, drop it temporarily to 15s.
_HEARTBEAT_INTERVAL_SECS = 60

_heartbeat_task: asyncio.Task | None = None
_shutdown_logged = False  # don't double-log shutdown if multiple signals arrive
_started_at: float | None = None


def _process_rss_mb() -> float | None:
    """Resident-set memory in MB. Read from /proc/self/status — no deps."""
    try:
        with open("/proc/self/status") as fh:
            for line in fh:
                if line.startswith("VmRSS:"):
                    # Format: 'VmRSS:   123456 kB'
                    kb = int(line.split()[1])
                    return round(kb / 1024, 1)
    except Exception:
        pass
    return None


def _db_pool_stats() -> dict[str, Any]:
    """Pool checked-in / checked-out / overflow. Direct from SQLAlchemy."""
    try:
        from fabledassistant.models import engine
        pool = engine.pool
        # Async engines wrap a sync pool; .checkedin() / .checkedout() exist
        # on the underlying QueuePool. Attribute access is documented but
        # version-fragile, so wrap in try.
        return {
            "size": getattr(pool, "size", lambda: None)(),
            "checked_in": getattr(pool, "checkedin", lambda: None)(),
            "checked_out": getattr(pool, "checkedout", lambda: None)(),
            "overflow": getattr(pool, "overflow", lambda: None)(),
        }
    except Exception:
        return {}


def _asyncio_task_count() -> int | None:
    try:
        return len([t for t in asyncio.all_tasks() if not t.done()])
    except Exception:
        return None


def _curator_busy() -> bool | None:
    try:
        from fabledassistant.services.curator import is_curator_running
        return is_curator_running()
    except Exception:
        return None


def _uptime_secs() -> float | None:
    if _started_at is None:
        return None
    return round(time.monotonic() - _started_at, 1)


async def _heartbeat_loop() -> None:
    """Forever-running task that emits one snapshot per interval.

    Exceptions inside the loop are caught and logged so the loop itself
    can't die silently — the whole point is that this thing keeps
    talking even when other things crash around it.
    """
    while True:
        try:
            rss = _process_rss_mb()
            tasks = _asyncio_task_count()
            curator = _curator_busy()
            pool = _db_pool_stats()
            uptime = _uptime_secs()
            logger.info(
                "diag heartbeat: uptime=%ss rss=%sMB asyncio_tasks=%s "
                "db_pool=%s curator_busy=%s",
                uptime, rss, tasks, pool, curator,
            )
        except Exception:
            logger.exception("Heartbeat snapshot crashed (continuing)")
        try:
            await asyncio.sleep(_HEARTBEAT_INTERVAL_SECS)
        except asyncio.CancelledError:
            logger.info("diag heartbeat: shutting down (CancelledError)")
            raise


def _asyncio_exception_handler(loop: asyncio.AbstractEventLoop, context: dict) -> None:
    """Log unhandled task exceptions instead of letting them disappear.

    Default asyncio behaviour: if a fire-and-forget task raises and
    nothing awaits the result, the exception is logged at ERROR but the
    surrounding context (which task, what coroutine) can be sparse.
    This handler enriches the log line so we can tell WHICH task crashed.
    """
    msg = context.get("message", "")
    exc = context.get("exception")
    task = context.get("task") or context.get("future")
    task_name = getattr(task, "get_name", lambda: "?")() if task else "?"
    coro = getattr(task, "get_coro", lambda: None)() if task else None
    coro_name = getattr(coro, "__qualname__", str(coro)) if coro else "?"

    if exc is not None:
        logger.error(
            "asyncio unhandled exception in task %r (coro=%s): %s",
            task_name, coro_name, msg,
            exc_info=(type(exc), exc, exc.__traceback__),
        )
    else:
        logger.error(
            "asyncio unhandled event in task %r (coro=%s): %s — context=%r",
            task_name, coro_name, msg, context,
        )


def _signal_handler(loop: asyncio.AbstractEventLoop, signame: str) -> None:
    """Catch SIGTERM/SIGINT and log them.

    Doesn't try to interfere with shutdown — Hypercorn handles its own
    graceful exit. We just want a log line so we can tell "orderly
    shutdown via signal X" apart from "silent gap then container exit"
    (which means kill-9 or OOM, neither of which is catchable).
    """
    global _shutdown_logged
    if _shutdown_logged:
        return
    _shutdown_logged = True
    logger.warning(
        "diag shutdown: received %s, expecting graceful exit "
        "(uptime=%ss, asyncio_tasks=%s, rss=%sMB, db_pool=%s)",
        signame,
        _uptime_secs(),
        _asyncio_task_count(),
        _process_rss_mb(),
        _db_pool_stats(),
    )


def start_diagnostics(loop: asyncio.AbstractEventLoop) -> None:
    """Install all three diagnostic surfaces. Idempotent."""
    global _heartbeat_task, _started_at

    if _started_at is None:
        _started_at = time.monotonic()

    # 1. Asyncio exception hook — install once.
    if loop.get_exception_handler() is None:
        loop.set_exception_handler(_asyncio_exception_handler)

    # 2. Signal handlers. add_signal_handler is Unix-only; skip on
    # Windows so dev on a non-Linux machine doesn't blow up.
    if os.name == "posix":
        for sig in (signal.SIGTERM, signal.SIGINT):
            try:
                loop.add_signal_handler(
                    sig,
                    _signal_handler,
                    loop,
                    sig.name,
                )
            except (RuntimeError, NotImplementedError, ValueError):
                # add_signal_handler can fail when not running in the main
                # thread or when the loop is already managing the signal.
                # Either case is fine — heartbeat + exception hook still work.
                pass

    # 3. Heartbeat loop. Start if not already running (idempotent on
    # restart-style reloads).
    if _heartbeat_task is None or _heartbeat_task.done():
        _heartbeat_task = asyncio.create_task(
            _heartbeat_loop(), name="diag-heartbeat",
        )
        logger.info(
            "diag started: heartbeat every %ds, signal-aware shutdown logging on, "
            "asyncio exception hook installed",
            _HEARTBEAT_INTERVAL_SECS,
        )


def stop_diagnostics() -> None:
    """Cancel the heartbeat. Called from after_serving.

    Lets the loop emit one last 'shutting down' log line so the silence
    that follows is intentional (not a crash). The exception hook stays
    installed but won't fire after the loop closes.
    """
    global _heartbeat_task
    if _heartbeat_task is not None and not _heartbeat_task.done():
        _heartbeat_task.cancel()
        _heartbeat_task = None
    logger.info(
        "diag stopped: heartbeat cancelled. Final snapshot: "
        "uptime=%ss rss=%sMB asyncio_tasks=%s db_pool=%s curator_busy=%s",
        _uptime_secs(),
        _process_rss_mb(),
        _asyncio_task_count(),
        _db_pool_stats(),
        _curator_busy(),
    )