src/services/limits.py

"""Budget pacing, rate limits, device cooldown and the kill switch (spec section 10).

State is in memory and resets on restart. Provider prepaid credit is the hard cap; these
counters only pace. One threading.Lock guards all counters: every critical section is short and
never awaits, so it is safe from async generation code and from sync admin callbacks.
"""

from __future__ import annotations

import asyncio
import hashlib
import logging
import re
import secrets
import statistics
import threading
import time
from collections import defaultdict, deque
from collections.abc import AsyncIterator, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Literal

from src.config import LimitsSection
from src.errors import BudgetReachedError, CooldownError, PausedError, RateLimitedError

logger = logging.getLogger(__name__)

MINUTE = 60.0
HOUR = 3600.0
STATS_WINDOW = 600.0
MAX_EVENTS = 5000
MAX_TRACKED_DEVICES = 5000
UNKNOWN_DEVICE = "-"

_DEVICE_ID = re.compile(r"[A-Za-z0-9_-]{16,64}")
_SALT = secrets.token_hex(8)

Kind = Literal["ok", "error", "refusal", "limited"]


def new_device_id() -> str:
    return secrets.token_urlsafe(16)


def clean_device_id(raw: object) -> str | None:
    """The stored id if it looks like one we issued. The browser controls it, so never trust it."""
    return raw if isinstance(raw, str) and _DEVICE_ID.fullmatch(raw) else None


def hash_device_id(raw: str) -> str:
    """Short salted hash for limits and logs. The salt changes on every restart."""
    return hashlib.sha256(f"{_SALT}:{raw}".encode()).hexdigest()[:12]


def device_identity(raw: object) -> tuple[str, str]:
    """(id to keep in the browser, hash to use for limits and logs)."""
    device_id = clean_device_id(raw) or new_device_id()
    return device_id, hash_device_id(device_id)


@dataclass
class Reservation:
    count: int
    cost: float
    epoch: int = 0
    settled: bool = False


@dataclass(frozen=True)
class Snapshot:
    images_used: int
    max_images: int
    spend_usd: float
    max_cost_usd: float
    in_flight: int
    paused: bool
    recent: dict[Kind, int]
    p50_seconds: dict[str, float]


class LimitService:
    def __init__(self, limits: LimitsSection, clock: Callable[[], float] = time.monotonic) -> None:
        self._limits = limits
        self._clock = clock
        self._lock = threading.Lock()
        self._slots = asyncio.Semaphore(limits.max_concurrent_generations)
        self._paused = False
        self._epoch = 0
        self._images = 0
        self._spend = 0.0
        self._in_flight = 0
        self._starts: deque[float] = deque()
        self._device_last: dict[str, float] = {}
        self._device_starts: dict[str, deque[float]] = {}
        self._events: deque[tuple[float, Kind, str, float]] = deque(maxlen=MAX_EVENTS)

    def reserve(self, *, count: int, est_cost: float, device: str | None) -> Reservation:
        """Claim budget for `count` images before calling a provider, or raise a typed error."""
        limits = self._limits
        key = device or UNKNOWN_DEVICE
        now = self._clock()
        with self._lock:
            if self._paused:
                raise PausedError
            last = self._device_last.get(key)
            if last is not None and now - last < limits.device_cooldown_seconds:
                raise CooldownError
            hour = self._device_starts.get(key) or deque()
            _drop_older(hour, now - HOUR)
            if limits.device_max_images_per_hour and (
                len(hour) + count > limits.device_max_images_per_hour
            ):
                raise RateLimitedError
            _drop_older(self._starts, now - MINUTE)
            if len(self._starts) + count > limits.max_generations_per_minute:
                raise RateLimitedError
            if (
                self._images + count > limits.max_images
                or self._spend + est_cost > limits.max_cost_usd
            ):
                raise BudgetReachedError

            self._device_last[key] = now
            self._device_starts[key] = hour
            hour.extend([now] * count)
            self._starts.extend([now] * count)
            self._images += count
            self._spend += est_cost
            self._forget_idle_devices(now)
            return Reservation(count=count, cost=est_cost, epoch=self._epoch)

    def reconcile(self, reservation: Reservation, actual_cost: float) -> None:
        with self._lock:
            if reservation.settled:
                return
            reservation.settled = True
            if reservation.epoch != self._epoch:
                return
            self._spend = max(0.0, self._spend + actual_cost - reservation.cost)

    def release(self, reservation: Reservation) -> None:
        """Return the budget of a call that produced nothing. The cooldown still counts."""
        with self._lock:
            if reservation.settled:
                return
            reservation.settled = True
            if reservation.epoch != self._epoch:
                return
            self._images = max(0, self._images - reservation.count)
            self._spend = max(0.0, self._spend - reservation.cost)

    @asynccontextmanager
    async def slot(self) -> AsyncIterator[None]:
        """Hold one of the global concurrent-generation slots."""
        async with self._slots:
            with self._lock:
                self._in_flight += 1
            try:
                yield
            finally:
                with self._lock:
                    self._in_flight -= 1

    def record(self, kind: Kind, model: str, seconds: float) -> None:
        with self._lock:
            self._events.append((self._clock(), kind, model, seconds))

    @property
    def paused(self) -> bool:
        with self._lock:
            return self._paused

    def set_paused(self, paused: bool) -> None:
        with self._lock:
            self._paused = paused
        logger.info("admin action: paused=%s", paused)

    def reset_counters(self) -> None:
        """Zero the counters. Reservations still in flight no longer count against the new ones."""
        with self._lock:
            self._epoch += 1
            self._images = 0
            self._spend = 0.0
            self._starts.clear()
            self._events.clear()
        logger.info("admin action: counters reset")

    def snapshot(self) -> Snapshot:
        now = self._clock()
        with self._lock:
            recent = [e for e in self._events if now - e[0] <= STATS_WINDOW]
            counts: dict[Kind, int] = {"ok": 0, "error": 0, "refusal": 0, "limited": 0}
            latencies: defaultdict[str, list[float]] = defaultdict(list)
            for _, kind, model, seconds in recent:
                counts[kind] += 1
                if kind == "ok":
                    latencies[model].append(seconds)
            return Snapshot(
                images_used=self._images,
                max_images=self._limits.max_images,
                spend_usd=self._spend,
                max_cost_usd=self._limits.max_cost_usd,
                in_flight=self._in_flight,
                paused=self._paused,
                recent=counts,
                p50_seconds={m: statistics.median(v) for m, v in latencies.items()},
            )

    def _forget_idle_devices(self, now: float) -> None:
        """Bound memory if many device ids are forged. Caller holds the lock."""
        if len(self._device_last) <= MAX_TRACKED_DEVICES:
            return
        stale = [k for k, t in self._device_last.items() if now - t > HOUR]
        for key in stale:
            self._device_last.pop(key, None)
        for key in [k for k, v in self._device_starts.items() if not v or now - v[-1] > HOUR]:
            del self._device_starts[key]


def _drop_older(stamps: deque[float], cutoff: float) -> None:
    while stamps and stamps[0] < cutoff:
        stamps.popleft()