CyberGril/deploy/hel835_guard.py

"""HEL-835 deploy-layer call guard v2: durable caps + precise settlement.

Deployment-layer protection only; no product code is modified.

v2 fixes two reviewed defects of v1:
  1. settle() matched (second, pid, category) so same-second concurrency could
     overwrite the wrong entry. Every reserve now persists a unique call_id and
     settlement matches exactly that id; cancelled/closed wrappers explicitly
     close the upstream stream and settle.
  2. classify() matched feature strings inside the whole system prompt, so a
     user-editable relation boundary containing a feature string could re-route
     a real chat turn into an uncapped category. Categories now bind to the
     router Scene at route() time (the product's own scene table), never to
     prompt text; unknown scenes are refused with a diagnostic instead of
     silently falling into an uncapped, outward-calling category.

Ledger validity: the three capped counters must exist as non-negative integers
in the authoritative file; missing, empty or corrupt ledgers refuse startup and
every call (no silent reset to zero).

Caps (cumulative across runs/restarts, including failures and auto retries):
  memory_gist <= 8, dialogue_main (real chat turns) <= 6, persona_generation <= 1.
Run: python -m deploy.hel835_guard   (instead of python -m cybergril.api)
"""
from __future__ import annotations

import asyncio
import contextlib
import dataclasses
import fcntl
import json
import logging
import os
import secrets
import time
from pathlib import Path

import uvicorn

from cybergril.api.app import create_app
from cybergril.config import ModelsConfig
from cybergril.think.budget import BudgetConfig, BudgetGuard
from cybergril.think.router import ModelRouter
from cybergril.think.scenes import Scene
from cybergril.think.wiring import build_routing, build_tier_providers

log = logging.getLogger("hel835.guard")

CAPS = {
    "memory_gist": 8,
    "dialogue_main": 6,
    "persona_generation": 1,
}

# Scene -> ledger category. This is the product's own scene usage table
# (app.py DIALOGUE_MAIN, episodic.py MEMORY_GIST, persona.py DIALOGUE_FOLLOWUP
# at its only call site and FACT_EXTRACT for growth, runtime.py LIFE_EVENT,
# sharing.py OUTPUT_GUARD); prompt text is never consulted.
SCENE_CATEGORY = {
    Scene.DIALOGUE_MAIN: "dialogue_main",
    Scene.MEMORY_GIST: "memory_gist",
    Scene.DIALOGUE_FOLLOWUP: "persona_generation",
    Scene.FACT_EXTRACT: "growth_extract",
    Scene.LIFE_EVENT: "life_event",
    Scene.OUTPUT_GUARD: "sharing_judge",
}


class LedgerInvalid(RuntimeError):
    pass


class GuardLedger:
    """File-locked durable ledger; the single authoritative call account."""

    def __init__(self, path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        if not self.path.exists():
            raise SystemExit(
                f"hel835 guard: refusing to start without an authoritative ledger at {self.path}; "
                "create it from the audited call history first (never zero)."
            )
        try:
            self.counts()
        except LedgerInvalid as exc:
            raise SystemExit(f"hel835 guard: ledger at {self.path} invalid: {exc}") from exc

    def _locked(self):
        fh = open(self.path, "a+", encoding="utf-8")
        fcntl.flock(fh, fcntl.LOCK_EX)
        return fh

    @staticmethod
    def _validated(data):
        counts = data.get("counts")
        if not isinstance(counts, dict):
            raise LedgerInvalid("counts missing or not an object")
        for key in CAPS:
            value = counts.get(key)
            if not isinstance(value, int) or isinstance(value, bool) or value < 0:
                raise LedgerInvalid(f"capped counter {key} missing/invalid: {value!r}")
        if not isinstance(data.get("calls"), list):
            raise LedgerInvalid("calls missing or not a list")
        return data

    @staticmethod
    def _save(fh, data):
        fh.seek(0)
        fh.truncate(0)
        json.dump(data, fh, ensure_ascii=False, indent=1)
        fh.flush()

    def counts(self):
        with self._locked() as fh:
            fh.seek(0)
            data = self._validated(json.load(fh))
        return {c: int(data["counts"][c]) for c in CAPS}, data

    def reserve(self, category):
        """Atomically count one attempt; refuse capped categories at limit."""
        with self._locked() as fh:
            fh.seek(0)
            data = self._validated(json.load(fh))
            counts = data["counts"]
            used = int(counts.get(category, 0))
            cap = CAPS.get(category)
            if cap is not None and used >= cap:
                self._save(fh, data)
                raise RuntimeError(f"hel835_guard_cap:{category}:{used}/{cap}")
            counts[category] = used + 1
            call_id = f"{time.time_ns()}-{os.getpid()}-{secrets.token_hex(4)}"
            entry = {"call_id": call_id, "t": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                     "category": category, "pid": os.getpid(), "outcome": "attempt"}
            calls = data["calls"]
            calls.append(entry)
            data["calls"] = calls[-2000:]
            self._save(fh, data)
            return entry

    def settle(self, entry, outcome, detail=""):
        """Settle exactly the reserved entry by its unique call_id, once."""
        with self._locked() as fh:
            fh.seek(0)
            data = self._validated(json.load(fh))
            target = entry["call_id"]
            for call in reversed(data["calls"]):
                if call.get("call_id") == target:
                    if call["outcome"] != "attempt":
                        raise LedgerInvalid(
                            f"call_id {target!r} already settled as {call['outcome']!r}"
                        )
                    call["outcome"] = outcome
                    if detail:
                        call["detail"] = str(detail)[:300]
                    break
            else:
                raise LedgerInvalid(f"unknown call_id {target!r} to settle")
            self._save(fh, data)


class SceneBoundProvider:
    """Provider wrapper bound to a router Scene; category never comes from text."""

    def __init__(self, real, ledger, scene):
        self.real = real
        self.ledger = ledger
        self.scene = scene

    async def stream(self, messages, **kwargs):
        """Async generator like the wrapped provider; reserves on first pull.

        Settlement semantics: seeing a done delta (or stop finish) before the
        consumer closes means the upstream call actually completed, so an early
        close settles "ok", not "cancelled".
        """
        category = SCENE_CATEGORY.get(self.scene)
        entry = self.ledger.reserve(category)
        log.info("hel835 guard call scene=%s category=%s", self.scene.value, category)
        upstream = self.real.stream(messages, **kwargs)
        saw_done = False
        try:
            while True:
                delta = await upstream.__anext__()
                if delta.done or getattr(delta, "finish_reason", None) == "stop":
                    saw_done = True
                yield delta
            # pragma: no cover - while True never falls through
        except StopAsyncIteration:
            self.ledger.settle(entry, "ok")
        except BaseException as exc:
            if isinstance(exc, asyncio.CancelledError):
                self.ledger.settle(entry, "ok" if saw_done else "cancelled", repr(exc))
            elif isinstance(exc, GeneratorExit):
                self.ledger.settle(entry, "ok" if saw_done else "cancelled")
            else:
                self.ledger.settle(entry, "fail", repr(exc))
            raise
        finally:
            with contextlib.suppress(Exception):
                await upstream.aclose()

    def __getattr__(self, name):
        return getattr(self.real, name)


class GuardedRouter(ModelRouter):
    """route() wraps the chosen provider with the scene-bound guard."""

    def __init__(self, *, providers, routing, guard, ledger):
        super().__init__(providers=providers, routing=routing, guard=guard)
        self.guard_ledger = ledger

    def route(self, scene, **kwargs):
        if scene not in SCENE_CATEGORY:
            # Refuse regardless of tier/provider, leaving a diagnostic; never
            # silently map an unmapped scene onto an uncapped caller.
            raise RuntimeError(f"hel835_guard_unknown_scene:{scene.value}")
        decision = super().route(scene, **kwargs)
        if decision.provider is None:
            return decision
        return dataclasses.replace(
            decision, provider=SceneBoundProvider(decision.provider, self.guard_ledger, scene)
        )


def build_guarded_router(models, ledger_path):
    providers = build_tier_providers(models)
    ledger = GuardLedger(ledger_path)
    return GuardedRouter(
        providers=providers,
        routing=build_routing(models),
        guard=BudgetGuard(BudgetConfig(
            daily_tokens=models.budget.daily_tokens,
            warn_ratio=models.budget.warn_ratio,
            hard_stop_ratio=models.budget.hard_stop_ratio,
        )),
        ledger=ledger,
    )


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
    import yaml
    config = Path(os.environ.get("CYBERGRIL_MODELS", "config/models.yaml"))
    models = ModelsConfig(**yaml.safe_load(config.read_text()))
    ledger_path = os.environ.get(
        "CYBERGRIL_GUARD_LEDGER",
        str(Path(os.environ["CYBERGRIL_DB"]).parent / "guard-ledger.json"),
    )
    router = build_guarded_router(models, ledger_path)
    app = create_app(router=router)
    uvicorn.run(app,  # noqa: PLC0415 - deploy entry, single worker, no reload
                host=os.environ.get("CYBERGRIL_HOST", "127.0.0.1"),
                port=int(os.environ.get("CYBERGRIL_PORT", "8350")),
                workers=1)


if __name__ == "__main__":
    main()