backend/src/finrag/guardrails.py

"""Cost/abuse guardrails for the public deploy.

The agent endpoints call a funded Anthropic key, so a public URL without limits
is an open invitation to burn money. Two independent defenses, both enforced as
a FastAPI dependency on the paid endpoints (/answer, /agent, /agent/stream):

  1. Per-IP sliding-window rate limit  → stops one client hammering the agent.
  2. Global daily question cap          → the hard cost circuit-breaker; once the
     day's paid questions hit the cap, every further request 429s until midnight
     UTC, no matter who sends it. This is the line between "a few dollars" and
     "a surprise bill."

Both are in-process (a dict + a counter under one lock). That's correct for the
single-instance Fly deploy we target; a multi-instance/multi-worker setup would
need a shared store (Redis) instead — called out in docs/deploy.md, not hidden.

Retrieval-only /query is intentionally NOT guarded here: it makes no LLM call
(Cohere rerank only, ~$0.002) so it isn't the cost risk the cap exists for.
"""

from __future__ import annotations

import threading
import time
from collections import defaultdict, deque
from datetime import datetime, timezone

from fastapi import HTTPException, Request

from finrag.config import settings

_WINDOW_SECONDS = 60.0

_lock = threading.Lock()
# ip -> timestamps of accepted paid requests within the last _WINDOW_SECONDS
_hits: dict[str, deque[float]] = defaultdict(deque)
# (utc_date_str, count) — the global daily paid-question tally
_day: str = ""
_day_count: int = 0


def _client_ip(request: Request) -> str:
    """Best-effort real client IP. Behind Fly's proxy the socket peer is the
    proxy, so prefer the forwarded headers Fly/Vercel set; fall back to the
    socket. X-Forwarded-For is a chain — the original client is the first hop."""
    xff = request.headers.get("x-forwarded-for")
    if xff:
        return xff.split(",")[0].strip()
    fly = request.headers.get("fly-client-ip")
    if fly:
        return fly.strip()
    return request.client.host if request.client else "unknown"


def _today_utc() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def _seconds_until_utc_midnight() -> int:
    now = datetime.now(timezone.utc)
    tomorrow = now.date().toordinal() + 1
    midnight = datetime.fromordinal(tomorrow).replace(tzinfo=timezone.utc)
    return max(1, int((midnight - now).total_seconds()))


def enforce(request: Request) -> None:
    """FastAPI dependency: raise 429 if this paid request breaches either the
    per-IP rate limit or the global daily cap; otherwise record it and return.

    Recording happens here (not after the call) so an in-flight burst can't slip
    past the cap — we count on admission, which is the conservative choice for a
    cost ceiling."""
    global _day, _day_count
    now = time.monotonic()
    ip = _client_ip(request)

    with _lock:
        # ── global daily cap (reset on UTC date rollover) ──
        today = _today_utc()
        if today != _day:
            _day, _day_count = today, 0
        if _day_count >= settings.daily_question_cap:
            raise HTTPException(
                status_code=429,
                detail=(
                    "Daily demo limit reached. This is a cost-capped public "
                    "demo; it resets at 00:00 UTC. Run it locally for unlimited "
                    "use — see the repo README."
                ),
                headers={"Retry-After": str(_seconds_until_utc_midnight())},
            )

        # ── per-IP sliding window ──
        # A limit <= 0 means "no per-IP cap" (the daily cap is the real circuit
        # breaker). Guarding here also keeps bucket[0] below safe: we only reach
        # it when limit > 0 and the bucket is already at/over that limit, so it
        # can never be empty.
        limit = settings.rate_limit_per_min
        bucket = _hits[ip]
        cutoff = now - _WINDOW_SECONDS
        while bucket and bucket[0] < cutoff:
            bucket.popleft()
        if limit > 0 and len(bucket) >= limit:
            retry = max(1, int(_WINDOW_SECONDS - (now - bucket[0])))
            raise HTTPException(
                status_code=429,
                detail=(
                    f"Rate limit: max {settings.rate_limit_per_min} questions/min. "
                    f"Try again in ~{retry}s."
                ),
                headers={"Retry-After": str(retry)},
            )

        # Admitted — record against both limiters.
        bucket.append(now)
        _day_count += 1

        # Opportunistic cleanup so idle IPs don't accumulate forever.
        if len(_hits) > 4096:
            for k in [k for k, v in _hits.items() if not v]:
                del _hits[k]


def cap_status() -> dict[str, int | str]:
    """Snapshot for /health — lets the frontend show 'N questions left today'
    and makes the cap observable without reading logs."""
    with _lock:
        today = _today_utc()
        used = _day_count if today == _day else 0
        cap = settings.daily_question_cap
        return {
            "daily_cap": cap,
            "used_today": used,
            "remaining_today": max(0, cap - used),
            "resets_at": "00:00 UTC",
        }