stack/dev/scripts/nb_watcher.py

"""Watch the running notebooks (marimo) service for errors and auto-file
deduplicated Gitea issues.

Two streams, one signature space (shared with nb_integration.py):

1. Container stdout/stderr via the docker socket (server errors, kernel
   tracebacks) — read with stdlib http.client over the unix socket, no
   docker CLI needed.
2. Session snapshots at <notebooks>/__marimo__/session/*.json (cell-level
   errors from live sessions and nightly integration runs alike).

Findings go through nb_issue_filer (create/comment with cooldown). This
watcher is the single closing authority: signatures unseen for
NB_CLOSE_AFTER_H hours are swept closed. State (last-seen per signature,
log cursor) persists in NB_STATE_FILE.

Runs as the `nb-watcher` compose sidecar with dev/scripts mounted at
/scripts and the notebooks dir at /notebooks (both ro).
"""

from __future__ import annotations

import http.client
import json
import os
import re
import socket
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
import nb_integration  # noqa: E402 — sibling: parse_snapshot
import nb_issue_filer  # noqa: E402 — sibling: cmd_report/cmd_sweep/signature

DOCKER_SOCK = os.environ.get("DOCKER_SOCK", "/var/run/docker.sock")
CONTAINER = os.environ.get("NB_CONTAINER", "notebooks")
NB_DIR = Path(os.environ.get("NB_DIR", "/notebooks"))
STATE_FILE = Path(os.environ.get("NB_STATE_FILE", "/state/nb-watcher-state.json"))
HEARTBEAT = Path(os.environ.get("NB_HEARTBEAT", "/tmp/heartbeat"))
POLL_S = int(os.environ.get("NB_POLL_S", "60"))
SWEEP_EVERY_S = int(os.environ.get("NB_SWEEP_EVERY_S", "3600"))
CLOSE_AFTER_H = float(os.environ.get("NB_CLOSE_AFTER_H", "24"))
SOURCE = "nb-watcher"

_TRACEBACK_START = "Traceback (most recent call last):"


# ── docker logs via unix socket ──────────────────────────────────


class _UnixConn(http.client.HTTPConnection):
    def __init__(self, path: str):
        super().__init__("localhost", timeout=30)
        self._unix_path = path

    def connect(self) -> None:
        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        s.settimeout(30)
        s.connect(self._unix_path)
        self.sock = s


def fetch_logs(since_epoch: int) -> str:
    """Return decoded log text since `since_epoch` (demuxing the docker
    multiplexed stream format when present)."""
    conn = _UnixConn(DOCKER_SOCK)
    try:
        conn.request(
            "GET",
            f"/containers/{CONTAINER}/logs?stdout=1&stderr=1&since={since_epoch}",
        )
        resp = conn.getresponse()
        raw = resp.read()
    finally:
        conn.close()
    if not raw:
        return ""
    # Multiplexed frames: [stream:1][pad:3][len:4 BE][payload]. TTY
    # containers emit raw text instead — detect by header shape.
    if raw[0] in (0, 1, 2) and raw[1:4] == b"\x00\x00\x00":
        out = bytearray()
        i = 0
        while i + 8 <= len(raw):
            size = int.from_bytes(raw[i + 4 : i + 8], "big")
            out += raw[i + 8 : i + 8 + size]
            i += 8 + size
        return out.decode(errors="replace")
    return raw.decode(errors="replace")


def parse_log_errors(text: str) -> list[dict]:
    """Extract tracebacks (and their final ExcType: message) from log text."""
    findings: list[dict] = []
    lines = text.splitlines()
    i = 0
    while i < len(lines):
        if _TRACEBACK_START in lines[i]:
            block = [lines[i]]
            i += 1
            while i < len(lines) and (
                lines[i].startswith((" ", "\t")) or not lines[i].strip()
            ):
                block.append(lines[i])
                i += 1
            # final "ExcType: message" line belongs to the traceback
            tail = lines[i] if i < len(lines) else ""
            m = re.match(
                r"^([A-Za-z_][\w.]*(?:Error|Exception|Interrupt|Warning))\b:?\s*(.*)",
                tail.strip(),
            )
            if m:
                block.append(tail)
                i += 1
                ename, evalue = m.group(1), m.group(2) or "(no message)"
            else:
                ename, evalue = "traceback", block[-1].strip()[:200] or "(unknown)"
            findings.append(
                {
                    "notebook": "(service)",
                    "ename": ename,
                    "evalue": evalue,
                    "detail": "\n".join(block)[-3500:],
                }
            )
        else:
            i += 1
    return findings


# ── session snapshots ────────────────────────────────────────────


def scan_snapshots(since_mtime: float) -> tuple[list[dict], float]:
    findings: list[dict] = []
    newest = since_mtime
    session_dir = NB_DIR / "__marimo__" / "session"
    if not session_dir.is_dir():
        return findings, newest
    for snap in session_dir.glob("*.json"):
        try:
            mtime = snap.stat().st_mtime
        except OSError:
            continue
        if mtime <= since_mtime:
            continue
        newest = max(newest, mtime)
        notebook = snap.name.removesuffix(".json")
        for err in nb_integration.parse_snapshot(snap):
            if err["ename"] in ("snapshot-missing", "snapshot-unreadable"):
                continue
            findings.append({"notebook": notebook, **err})
    return findings, newest


# ── state ────────────────────────────────────────────────────────


def load_state() -> dict:
    if STATE_FILE.exists():
        try:
            return json.loads(STATE_FILE.read_text())
        except (OSError, json.JSONDecodeError):
            pass
    return {"sigs": {}, "log_cursor": int(time.time()), "snap_cursor": 0.0}


def save_state(state: dict) -> None:
    STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
    tmp = STATE_FILE.with_suffix(".tmp")
    tmp.write_text(json.dumps(state))
    tmp.replace(STATE_FILE)


# ── main loop ────────────────────────────────────────────────────


def tick(state: dict) -> list[dict]:
    """One poll: gather findings from both streams, advance cursors."""
    findings: list[dict] = []

    cursor = state.get("log_cursor", int(time.time()))
    now = int(time.time())
    try:
        findings += parse_log_errors(fetch_logs(cursor))
        state["log_cursor"] = now
    except (OSError, http.client.HTTPException) as e:
        print(f"log fetch failed (will retry): {e}", file=sys.stderr)

    snaps, newest = scan_snapshots(state.get("snap_cursor", 0.0))
    findings += snaps
    state["snap_cursor"] = newest

    for f in findings:
        sig = nb_issue_filer.signature(f["notebook"], f["ename"], f["evalue"])
        state.setdefault("sigs", {})[sig] = time.time()
    return findings


def main() -> int:
    state = load_state()
    last_sweep = 0.0
    print(
        f"nb-watcher: container={CONTAINER} nb_dir={NB_DIR} "
        f"poll={POLL_S}s close_after={CLOSE_AFTER_H}h"
    )
    while True:
        findings = tick(state)
        if findings:
            print(f"{len(findings)} finding(s)")
            try:
                nb_issue_filer.cmd_report(findings, SOURCE)
            except Exception as e:  # noqa: BLE001 — keep the watcher alive
                print(f"filer error (continuing): {e}", file=sys.stderr)

        if time.time() - last_sweep >= SWEEP_EVERY_S:
            cutoff = time.time() - CLOSE_AFTER_H * 3600
            active = {s for s, ts in state.get("sigs", {}).items() if ts >= cutoff}
            try:
                nb_issue_filer.cmd_sweep(active, SOURCE)
            except Exception as e:  # noqa: BLE001
                print(f"sweep error (continuing): {e}", file=sys.stderr)
            last_sweep = time.time()
            # drop long-stale sigs so state doesn't grow unbounded
            state["sigs"] = {
                s: ts for s, ts in state.get("sigs", {}).items() if ts >= cutoff
            }

        save_state(state)
        HEARTBEAT.touch()
        time.sleep(POLL_S)


if __name__ == "__main__":
    raise SystemExit(main())