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 datetime
import json
import logging
import os
import signal
import time
from logging.handlers import RotatingFileHandler
from pathlib import Path
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

# Persistent diagnostic state — written to /data so it survives
# container restart, OOM-kill, and Docker log rotation. The point of
# the persistence: the operator may not notice a crash for hours, and
# by then the Docker logs are gone. /data/diagnostics gives them a
# durable place to look post-mortem.
_DIAG_DIR = Path("/data/diagnostics")
_CURRENT_STATE_PATH = _DIAG_DIR / "current.json"
_LAST_SHUTDOWN_PATH = _DIAG_DIR / "last_shutdown.json"
_LAST_EXCEPTION_PATH = _DIAG_DIR / "last_exception.json"
_PREVIOUS_RUN_PATH = _DIAG_DIR / "previous_run.json"
_DIAG_LOG_PATH = _DIAG_DIR / "diag.log"

# Dedicated file logger for the heartbeat/exception/signal stream.
# Separate from the stdout logger so it can't be lost by Docker log
# rotation. Rotates at 10 MB, keeps 5 backups → 50 MB max footprint.
_file_logger: logging.Logger | None = None

_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:
    # Curator was removed in Phase 8 of the MCP-first pivot. Always False.
    return False


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


def _snapshot_dict(reason: str = "heartbeat") -> dict[str, Any]:
    """Bundle the current resource snapshot into a JSON-safe dict.

    This is the canonical 'what was the app doing right now' record —
    written to current.json every heartbeat, to last_shutdown.json on
    signal, to last_exception.json on uncaught task exception. The
    `reason` field tells you which event captured this snapshot.
    """
    return {
        "reason": reason,
        "wall_time_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "uptime_secs": _uptime_secs(),
        "rss_mb": _process_rss_mb(),
        "asyncio_tasks": _asyncio_task_count(),
        "db_pool": _db_pool_stats(),
        "curator_busy": _curator_busy(),
        "pid": os.getpid(),
    }


def _write_state_atomic(path: Path, data: dict[str, Any]) -> None:
    """Write JSON to `path` via a tmp+rename so a crash mid-write can't
    leave a half-written / unparseable file. The tmp file is on the same
    filesystem as the target so the rename is atomic.

    Silent on failure — diagnostic writes must never raise into the
    caller. If /data isn't writable, the in-memory + stdout logging
    still happens.
    """
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(path.suffix + ".tmp")
        tmp.write_text(json.dumps(data, indent=2, default=str))
        tmp.replace(path)
    except Exception:
        # Log to stdout — if this fires, /data is bad and the operator
        # needs to know. Don't propagate; diagnostics must not crash
        # the heartbeat that depends on them.
        logger.exception("Failed to write diagnostic state to %s", path)


def _setup_file_logger() -> logging.Logger:
    """Create or return the dedicated rotating file logger for diagnostics.

    Separate from the app's stdout logger so an aggressive log rotator
    or Docker log cleanup can't take it out. Writes to /data/diagnostics/diag.log
    with size-based rotation.
    """
    global _file_logger
    if _file_logger is not None:
        return _file_logger
    try:
        _DIAG_DIR.mkdir(parents=True, exist_ok=True)
        flog = logging.getLogger("fabledassistant.diagnostics.persistent")
        flog.setLevel(logging.INFO)
        flog.propagate = False  # don't double-log into the root stdout stream
        # Don't re-add handlers across module reloads / re-runs.
        if not flog.handlers:
            handler = RotatingFileHandler(
                _DIAG_LOG_PATH,
                maxBytes=10 * 1024 * 1024,
                backupCount=5,
            )
            handler.setFormatter(
                logging.Formatter(
                    "%(asctime)s %(levelname)s %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S",
                )
            )
            flog.addHandler(handler)
        _file_logger = flog
    except Exception:
        logger.exception("Failed to set up persistent diagnostic logger")
        # Fall back to the regular stdout logger so callers don't need
        # to null-check the return value.
        _file_logger = logger
    return _file_logger


def _persistent_log(level: int, msg: str, *args: Any) -> None:
    """Emit one line to BOTH the stdout logger AND the persistent file."""
    logger.log(level, msg, *args)
    try:
        _setup_file_logger().log(level, msg, *args)
    except Exception:
        pass


def _post_mortem_previous_run() -> None:
    """On startup, check the persistent state and tell the operator if
    the previous run died abruptly.

    Logic: if current.json exists and is newer than last_shutdown.json,
    then either there was never a clean shutdown (first ever run, or
    SIGKILL/OOM/abrupt termination of the previous run). If current.json
    exists AND last_shutdown.json doesn't, OR current.json's timestamp
    is newer than last_shutdown.json's → the previous run didn't get
    to write its shutdown record. Log the last-known state prominently
    so the operator notices on next visit.
    """
    try:
        if not _CURRENT_STATE_PATH.exists():
            return  # first ever run

        current_mtime = _CURRENT_STATE_PATH.stat().st_mtime
        shutdown_mtime = (
            _LAST_SHUTDOWN_PATH.stat().st_mtime
            if _LAST_SHUTDOWN_PATH.exists() else 0
        )
        if shutdown_mtime >= current_mtime:
            return  # previous run shut down cleanly after its last heartbeat

        # The previous run had a heartbeat without a subsequent clean
        # shutdown — that's a crash, OOM, or otherwise-killed run.
        try:
            last_snapshot = json.loads(_CURRENT_STATE_PATH.read_text())
        except Exception:
            last_snapshot = {"_parse_error": "current.json unreadable"}

        seconds_since = round(time.time() - current_mtime, 1)
        _persistent_log(
            logging.WARNING,
            "diag post-mortem: PREVIOUS RUN DIED ABRUPTLY. Last heartbeat "
            "was %ss before this startup. Last-known state: %s",
            seconds_since, json.dumps(last_snapshot),
        )

        # Preserve the offending snapshot for retrospection. This file
        # accumulates one rename per abrupt termination so the user can
        # diff sequences over time if it's a recurring pattern.
        try:
            _PREVIOUS_RUN_PATH.write_text(json.dumps({
                "noticed_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                "seconds_since_last_heartbeat": seconds_since,
                "had_shutdown_record": _LAST_SHUTDOWN_PATH.exists(),
                "had_exception_record": _LAST_EXCEPTION_PATH.exists(),
                "last_known_state": last_snapshot,
            }, indent=2, default=str))
        except Exception:
            logger.exception("Failed to preserve previous_run.json")
    except Exception:
        logger.exception("Post-mortem check itself failed")


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

    Each heartbeat does three things:
    1. Logs the snapshot to stdout (Docker logs — short-lived).
    2. Appends the snapshot to /data/diagnostics/diag.log (rotating;
       survives container restart).
    3. Atomically rewrites /data/diagnostics/current.json with the
       latest snapshot — so post-crash you can `cat` one file and see
       the last known good state instead of scanning the rolling log.

    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:
            snap = _snapshot_dict(reason="heartbeat")
            _persistent_log(
                logging.INFO,
                "diag heartbeat: uptime=%ss rss=%sMB asyncio_tasks=%s "
                "db_pool=%s curator_busy=%s",
                snap["uptime_secs"], snap["rss_mb"], snap["asyncio_tasks"],
                snap["db_pool"], snap["curator_busy"],
            )
            _write_state_atomic(_CURRENT_STATE_PATH, snap)
        except Exception:
            logger.exception("Heartbeat snapshot crashed (continuing)")
        try:
            await asyncio.sleep(_HEARTBEAT_INTERVAL_SECS)
        except asyncio.CancelledError:
            _persistent_log(
                logging.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__),
        )
        # Also emit to the persistent file logger so the traceback
        # survives even if Docker logs are flushed.
        try:
            _setup_file_logger().error(
                "asyncio unhandled exception in task %r (coro=%s): %s",
                task_name, coro_name, msg,
                exc_info=(type(exc), exc, exc.__traceback__),
            )
        except Exception:
            pass
        # Persist the snapshot at exception time so post-mortem can see
        # what the app's state was when it crashed.
        snap = _snapshot_dict(reason="asyncio_exception")
        snap["task_name"] = task_name
        snap["coro_name"] = coro_name
        snap["exception_type"] = type(exc).__name__
        snap["exception_message"] = str(exc)
        _write_state_atomic(_LAST_EXCEPTION_PATH, snap)
    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
    snap = _snapshot_dict(reason="signal_shutdown")
    snap["signal"] = signame
    _persistent_log(
        logging.WARNING,
        "diag shutdown: received %s, expecting graceful exit. snapshot=%s",
        signame, json.dumps(snap),
    )
    # Write before-exit snapshot so the next startup's post-mortem can
    # see this run shut down cleanly via signal (rather than crashed).
    _write_state_atomic(_LAST_SHUTDOWN_PATH, snap)


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()

    # 0. Set up the persistent file logger FIRST, then run the
    # post-mortem check. The post-mortem reads /data state from the
    # previous run and surfaces "previous run died abruptly" if
    # current.json is newer than last_shutdown.json. Operator catches
    # this on next visit even if they missed the crash window.
    _setup_file_logger()
    _post_mortem_previous_run()

    # 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
    snap = _snapshot_dict(reason="after_serving_shutdown")
    _persistent_log(
        logging.INFO,
        "diag stopped: heartbeat cancelled. Final snapshot=%s",
        json.dumps(snap),
    )
    # If we got here without the signal handler firing (e.g., Hypercorn
    # decided to shut down on its own), still record a clean exit so
    # the next startup's post-mortem sees a fresh last_shutdown.json
    # and doesn't flag this as an abrupt death.
    _write_state_atomic(_LAST_SHUTDOWN_PATH, snap)