ROMFASTSQL/proxmox/lxc171-claude-agent/discord-bridge/limits.py

"""Limite: procese vii, coada per fir, timeout de tur, rate limit per user, plafon de cost.

T8. Motiv: 406 MB RSS per proces claude, pe un container cu istoric de OOM.
"""

from __future__ import annotations

import asyncio
import collections
import contextlib
import logging
import time
from typing import Callable

import config

log = logging.getLogger("discord-bridge.limits")

try:  # pragma: no cover - Lane C poate lipsi
    import alerts  # type: ignore
except ImportError:  # pragma: no cover
    class _NoAlerts:
        @staticmethod
        def alert(level: str, subject: str, body: str, dedup_key: str | None = None) -> None:
            log.warning("alerta (%s) %s: %s", level, subject, body)

    alerts = _NoAlerts()  # type: ignore


class LimitError(RuntimeError):
    """Baza pentru refuzurile de limita (mesajul e afisabil in Discord)."""


class RateLimited(LimitError):
    def __init__(self, retry_after: float):
        self.retry_after = retry_after
        super().__init__(f"prea multe mesaje; mai asteapta {retry_after:.0f}s")


def parse_cap(raw, default: float = 0.0) -> float:
    """Plafonul zilnic in USD. `off`/`none`/`nelimitat`/`0` => 0.0 (fara plafon).

    Costul raportat de CLI e pretul echivalent la API; pe abonament nu se
    factureaza, deci plafonul are sens doar cand chiar platesti per token.
    """
    if raw is None:
        return default
    text = str(raw).strip().lower()
    if text in ("", "off", "none", "no", "nu", "false", "unlimited", "nelimitat"):
        return 0.0
    try:
        val = float(text)
    except ValueError:
        return default
    return val if val > 0 else 0.0


class CostCapReached(LimitError):
    def __init__(self, spent: float, cap: float):
        self.spent, self.cap = spent, cap
        super().__init__(f"plafonul de cost pe zi a fost atins ({spent:.2f} / {cap:.2f} USD)")


class Limits:
    def __init__(
        self,
        store=None,
        *,
        max_procs: int | None = None,
        turn_timeout: float | None = None,
        rate_per_min: int | None = None,
        cost_cap: float | None = None,
        clock: Callable[[], float] = time.monotonic,
        alerter: Callable | None = None,
    ):
        self.store = store
        self.max_procs = max_procs if max_procs is not None else config.get_int("MAX_PROCS", 4)
        self.turn_timeout = (
            turn_timeout if turn_timeout is not None else config.get_float("TURN_TIMEOUT_S", 900.0)
        )
        self.rate_per_min = (
            rate_per_min if rate_per_min is not None else config.get_int("RATE_PER_USER_PER_MIN", 10)
        )
        self.cost_cap = (
            parse_cap(cost_cap) if cost_cap is not None
            else parse_cap(config.get("COST_CAP_USD_DAY"), 0.0)
        )
        self.clock = clock
        self._alert = alerter or alerts.alert
        self._slots = asyncio.Semaphore(self.max_procs)
        self._thread_locks: dict[str, asyncio.Lock] = {}
        self._hits: dict[str, collections.deque[float]] = {}
        self._local_cost = 0.0  # folosit doar daca nu avem store
        self._cap_alerted = False

    # ------------------------------------------------------------- procese
    @property
    def free_slots(self) -> int:
        return self._slots._value  # noqa: SLF001 - diagnostic pentru !status

    @contextlib.asynccontextmanager
    async def process_slot(self):
        """Maxim `max_procs` procese claude vii simultan; al 5-lea fir asteapta."""
        await self._slots.acquire()
        try:
            yield
        finally:
            self._slots.release()

    def thread_lock(self, thread_id: str) -> asyncio.Lock:
        """Coada per fir: un singur tur odata pe acelasi fir."""
        key = str(thread_id)
        lock = self._thread_locks.get(key)
        if lock is None:
            lock = self._thread_locks[key] = asyncio.Lock()
        return lock

    def queued(self, thread_id: str) -> bool:
        return self.thread_lock(thread_id).locked()

    # ---------------------------------------------------------- rate limit
    def check_rate(self, user_id: str, now: float | None = None) -> None:
        """Fereastra glisanta de 60s per utilizator. Ridica RateLimited."""
        now = self.clock() if now is None else now
        q = self._hits.setdefault(str(user_id), collections.deque())
        while q and now - q[0] >= 60.0:
            q.popleft()
        if len(q) >= self.rate_per_min:
            raise RateLimited(60.0 - (now - q[0]))
        q.append(now)

    # ---------------------------------------------------------------- cost
    def cost_today(self) -> float:
        if self.store is not None:
            with contextlib.suppress(Exception):
                return float(self.store.cost_today())
        return self._local_cost

    @property
    def cap_enabled(self) -> bool:
        """Plafon dezactivat (0 sau `off`) — cazul abonamentului."""
        return self.cost_cap > 0

    def stopped(self) -> bool:
        """Plafonul e evaluat pe ziua curenta; reset-ul zilnic vine din `cost.day`."""
        return self.cap_enabled and self.cost_today() >= self.cost_cap

    def cost_remaining(self) -> float:
        if not self.cap_enabled:
            return float("inf")
        return max(0.0, self.cost_cap - self.cost_today())

    def check_cost(self) -> None:
        if self.stopped():
            raise CostCapReached(self.cost_today(), self.cost_cap)

    def record_cost(self, usd: float, thread_id: str | None = None) -> float:
        try:
            usd = float(usd)
        except (TypeError, ValueError):
            usd = 0.0
        if self.store is not None:
            total = self.store.add_cost(thread_id, usd)
        else:
            self._local_cost = round(self._local_cost + usd, 6)
            total = self._local_cost
        if self.cap_enabled and total >= self.cost_cap and not self._cap_alerted:
            self._cap_alerted = True
            with contextlib.suppress(Exception):
                self._alert(
                    "CRITICAL",
                    "plafon de cost atins",
                    f"Cheltuit azi: {total:.2f} USD, plafon {self.cost_cap:.2f}. Puntea nu mai accepta tururi.",
                    "cost-cap",
                )
        elif not self.cap_enabled or total < self.cost_cap:
            self._cap_alerted = False
        return total

    # --------------------------------------------------------------- admit
    def admit(self, user_id: str) -> None:
        """Verificarile ieftine, inainte de a pune firul in coada. Ridica LimitError."""
        self.check_cost()
        self.check_rate(user_id)

    @contextlib.asynccontextmanager
    async def turn(self, thread_id: str, user_id: str):
        """Un tur complet: admis -> coada firului -> slot de proces."""
        self.admit(user_id)
        async with self.thread_lock(thread_id):
            async with self.process_slot():
                yield self.turn_timeout