api/access.py

"""Isolate browser sessions and bound work on a single-instance public beta."""

from __future__ import annotations

import asyncio
import contextlib
import hashlib
import hmac
import logging
import math
import re
import time
from collections import deque
from contextvars import ContextVar
from datetime import UTC, datetime, timedelta
from typing import Any
from uuid import UUID, uuid4

import anyio
from starlette.responses import JSONResponse

from api.contracts import ContractError, public_limits
from shared.execution import Execution, ExecutionStopped, execution, register, unregister
from shared.usage_budget import (
    BudgetExceeded,
    BudgetLimits,
    BudgetUnavailable,
    DuplicateRequest,
    PostgresBudget,
    RequestBudget,
    request_budget,
)

session_id: ContextVar[str] = ContextVar("precedent_session", default="")
account_id: ContextVar[str] = ContextVar("precedent_account", default="")
_LOG = logging.getLogger(__name__)
_SESSION = re.compile(r"^[a-f0-9]{32}$")


def conversation_id(thread_id: str) -> str:
    """Keep public thread names local to the server-authenticated browser session."""
    session = session_id.get()
    if not session:
        return thread_id
    return hashlib.sha256(f"{session}:{thread_id}".encode()).hexdigest()


class PublicAccess:
    """Require the frontend's secret; serialize expensive graph turns.

    Without a configured database, limits remain process-local. Persistent
    accounting fails closed if unavailable; graph memory is still process-local.
    """

    def __init__(self, app: Any, config_provider: Any):
        self.app = app
        self.config_provider = config_provider
        self.turn = asyncio.Semaphore(1)
        self.waiting = 0
        self.users: dict[str, deque[float]] = {}
        self.day: deque[float] = deque()
        self.store: PostgresBudget | None = None

    def _prune(self) -> None:
        now = time.monotonic()
        for old_session, calls in list(self.users.items()):
            while calls and now - calls[0] >= 3600:
                calls.popleft()
            if not calls:
                self.users.pop(old_session)
        while self.day and now - self.day[0] >= 86400:
            self.day.popleft()

    def _persistent_store(self, config: Any) -> PostgresBudget | None:
        secret = getattr(config, "quota_database_url", None)
        if secret is None:
            return None
        if self.store is None:
            self.store = PostgresBudget(
                secret.get_secret_value(),
                BudgetLimits(
                    turns_hour=config.public_requests_per_hour,
                    turns_day=config.public_requests_per_day,
                    provider_calls_hour=config.public_provider_calls_per_hour,
                    provider_calls_day=config.public_provider_calls_per_day,
                    provider_tokens_hour=config.public_provider_tokens_per_hour,
                    provider_tokens_day=config.public_provider_tokens_per_day,
                    provider_calls_minute=config.public_provider_calls_per_minute,
                    provider_tokens_minute=config.public_provider_tokens_per_minute,
                    external_calls_day=config.public_external_requests_per_day,
                ),
            )
        return self.store

    def _memory_report(self, config: Any, session: str) -> dict[str, Any]:
        self._prune()
        now = time.monotonic()
        counted_at = datetime.now(UTC)
        external = {}
        if getattr(config, "enable_web_fallback", False):
            from tools.web_research import memory_limits

            external = memory_limits(config)

        def remaining(calls: deque[float], limit: int, window: int) -> dict[str, Any]:
            wait = max(1, math.ceil(calls[0] + window - now)) if calls else window
            return {
                "limit": limit,
                "used": len(calls),
                "remaining": max(0, limit - len(calls)),
                "window_seconds": window,
                "next_expiry_at": (counted_at + timedelta(seconds=wait)).isoformat()
                if calls
                else None,
                "retry_after": wait,
            }

        return {
            "persistence": "memory",
            "counted_at": counted_at.isoformat(),
            "limits": {
                **external,
                "shared_turns": remaining(self.day, config.public_requests_per_day, 86400),
                "visitor_turns": remaining(
                    self.users.get(session, deque()), config.public_requests_per_hour, 3600
                ),
            },
            "note": "These development limits reset when the backend restarts.",
        }

    @staticmethod
    async def _body(receive: Any) -> tuple[bytes, list[dict]]:
        chunks, messages, size = [], [], 0
        while True:
            message = await receive()
            messages.append(message)
            if message["type"] != "http.request":
                raise ValueError("The research request ended before its body was received.")
            part = message.get("body", b"")
            size += len(part)
            if size > 131_072:
                raise ValueError("The research request is too large.")
            chunks.append(part)
            if not message.get("more_body", False):
                return b"".join(chunks), messages

    async def __call__(self, scope: dict, receive: Any, send: Any) -> None:
        config = self.config_provider()
        if scope["type"] != "http" or not config.backend_api_key:
            await self.app(scope, receive, send)
            return
        # Liveness remains reachable by the hosting platform.
        path = scope["path"]
        root = scope.get("root_path", "")
        if root and (path == root or path.startswith(root + "/")):
            path = path[len(root) :] or "/"
        if path == "/health":
            await self.app(scope, receive, send)
            return
        headers = dict(scope.get("headers", []))
        supplied = headers.get(b"x-precedent-key", b"")
        if not hmac.compare_digest(supplied, config.backend_api_key.encode()):
            await JSONResponse({"detail": "Use the PrecedentAI website."}, 401)(
                scope, receive, send
            )
            return
        session = headers.get(b"x-precedent-session", b"").decode("ascii", errors="ignore")
        if not _SESSION.fullmatch(session):
            await JSONResponse({"detail": "Missing browser session."}, 401)(scope, receive, send)
            return
        account = headers.get(b"x-precedent-account", b"").decode("ascii", errors="ignore")
        if account and (not _SESSION.fullmatch(account) or account != session):
            await JSONResponse({"detail": "Invalid account identity."}, 401)(scope, receive, send)
            return
        token = session_id.set(session)
        account_token = account_id.set(account)
        budget_token = None
        store = self._persistent_store(config)
        ident: UUID | None = None
        response_status = 500
        locked = False
        control = None
        control_token = None
        watcher = None
        response_ended = False
        try:
            if path.startswith("/account/") and scope["method"] in {"POST", "PUT", "PATCH"}:
                try:
                    _, account_messages = await self._body(receive)
                except ValueError:
                    await JSONResponse(
                        {"detail": "The account request is too large or incomplete."}, 413
                    )(scope, receive, send)
                    return
                account_receive = receive

                async def replay_account() -> dict:
                    return account_messages.pop(0) if account_messages else await account_receive()

                receive = replay_account
            expensive = scope["method"] == "POST" and (
                path in {"/chat", "/stream", "/supervise"} or path.startswith("/review/")
            )
            if scope["method"] == "GET" and path == "/limits":
                try:
                    report = (
                        await asyncio.to_thread(store.report, session)
                        if store
                        else self._memory_report(config, session)
                    )
                except BudgetUnavailable:
                    await JSONResponse(
                        {
                            "detail": "Usage accounting is temporarily unavailable.",
                            "code": "budget_unavailable",
                        },
                        503,
                    )(scope, receive, send)
                    return
                try:
                    report = public_limits(report)
                except ContractError:
                    await JSONResponse(
                        {
                            "detail": "Allowance information is temporarily unavailable.",
                            "code": "response_contract_invalid",
                        },
                        500,
                    )(scope, receive, send)
                    return
                await JSONResponse(report)(scope, receive, send)
                return
            if expensive:
                self._prune()
                try:
                    ident = (
                        UUID(headers[b"x-precedent-request"].decode("ascii"))
                        if b"x-precedent-request" in headers
                        else uuid4()
                    )
                    body, queued = await asyncio.wait_for(self._body(receive), timeout=10)
                except (ValueError, UnicodeError, TimeoutError):
                    ident = None
                    await JSONResponse(
                        {
                            "detail": "Provide a valid request ID and complete body.",
                            "code": "invalid_request",
                        },
                        400,
                    )(scope, receive, send)
                    return
                control = Execution(getattr(config, "public_research_timeout_seconds", 120))
                if not register(session, ident, control):
                    ident = None
                    await JSONResponse(
                        {
                            "detail": "This request is already active.",
                            "code": "request_already_accepted",
                        },
                        409,
                    )(scope, receive, send)
                    return
                control_token = execution.set(control)
                upstream_receive = receive
                disconnected = asyncio.Event()

                async def watch_disconnect() -> None:
                    while True:
                        message = await upstream_receive()
                        if message["type"] == "http.disconnect":
                            if not response_ended:
                                control.stop()
                            disconnected.set()
                            return
                        await asyncio.sleep(0.05)

                async def replay_receive() -> dict:
                    if queued:
                        return queued.pop(0)
                    await disconnected.wait()
                    return {"type": "http.disconnect"}

                receive = replay_receive
                watcher = asyncio.create_task(watch_disconnect())
                if self.waiting >= 3:
                    await JSONResponse(
                        {
                            "detail": "Precedent's short research queue is full.",
                            "code": "service_busy",
                        },
                        429,
                        headers={"Retry-After": "15"},
                    )(scope, receive, send)
                    return
                self.waiting += 1
                try:
                    await asyncio.wait_for(self.turn.acquire(), timeout=5)
                except TimeoutError:
                    await JSONResponse(
                        {
                            "detail": "Precedent is helping another visitor. Please try again shortly.",
                            "code": "service_busy",
                        },
                        429,
                        headers={"Retry-After": "15"},
                    )(scope, receive, send)
                    return
                finally:
                    self.waiting -= 1
                locked = True
                control.check()
                if store:
                    operation = "research" if path in {"/chat", "/stream"} else path
                    fingerprint = hashlib.sha256(operation.encode() + b"\0" + body).hexdigest()
                    try:
                        await asyncio.to_thread(store.reserve_turn, session, ident, fingerprint)
                    except BudgetExceeded as exc:
                        await JSONResponse(
                            {
                                "detail": str(exc),
                                "code": "budget_exhausted",
                                "category": exc.category,
                            },
                            429,
                            headers={"Retry-After": str(exc.retry_after)},
                        )(scope, receive, send)
                        return
                    except DuplicateRequest as exc:
                        await JSONResponse(
                            {
                                "detail": "This request was already accepted. It will not be run again.",
                                "code": "request_id_conflict"
                                if exc.changed
                                else "request_already_accepted",
                                "request_status": exc.status,
                            },
                            409,
                        )(scope, receive, send)
                        return
                    except BudgetUnavailable:
                        await JSONResponse(
                            {
                                "detail": "Usage accounting is temporarily unavailable.",
                                "code": "budget_unavailable",
                            },
                            503,
                        )(scope, receive, send)
                        return
                    budget_token = request_budget.set(RequestBudget(store, session, ident))
                else:
                    # Recheck after waiting: another turn may have used the final slot.
                    report = self._memory_report(config, session)
                    exhausted = next(
                        (
                            category
                            for category in ("shared_turns", "visitor_turns")
                            if report["limits"][category]["remaining"] == 0
                        ),
                        None,
                    )
                    if exhausted:
                        await JSONResponse(
                            {
                                "detail": "The free beta's research limit is reached.",
                                "code": "budget_exhausted",
                                "category": exhausted,
                            },
                            429,
                            headers={
                                "Retry-After": str(report["limits"][exhausted]["retry_after"])
                            },
                        )(scope, receive, send)
                        return
                    if len(self.users) >= 4096:
                        await JSONResponse(
                            {
                                "detail": "Precedent is busy. Please try again shortly.",
                                "code": "service_busy",
                            },
                            429,
                            headers={"Retry-After": "60"},
                        )(scope, receive, send)
                        return
                    self.users.setdefault(session, deque()).append(time.monotonic())
                    self.day.append(time.monotonic())

            async def send_response(message: dict) -> None:
                nonlocal locked, response_status, response_ended
                if message["type"] == "http.response.start":
                    if path.startswith(("/account/", "/recovery/", "/review/", "/history/")):
                        message = {
                            **message,
                            "headers": [
                                *(
                                    item
                                    for item in message.get("headers", [])
                                    if item[0] != b"cache-control"
                                ),
                                (b"cache-control", b"no-store"),
                            ],
                        }
                    response_status = message["status"]
                    if ident is not None:
                        message = {
                            **message,
                            "headers": [
                                *message.get("headers", []),
                                (b"x-precedent-request", str(ident).encode()),
                            ],
                        }
                if message["type"] == "http.response.body" and not message.get("more_body", False):
                    response_ended = True
                await send(message)
                # Release when the response ends, including when a mounted
                # app still has middleware/disconnect bookkeeping to finish.
                if (
                    message["type"] == "http.response.body"
                    and not message.get("more_body", False)
                    and locked
                    and (control is None or control.idle.is_set())
                ):
                    locked = False
                    self.turn.release()

            await self.app(scope, receive, send_response)
        except ExecutionStopped as exc:
            response_status = 504 if exc.reason == "deadline" else 499
            await JSONResponse({"detail": str(exc), "code": exc.code}, response_status)(
                scope, receive, send
            )
        finally:
            # Cancelling an async wait cannot kill an already running worker.
            # Do not advertise a free slot while that worker still has side effects.
            with anyio.CancelScope(shield=True):
                if control is not None and control_token is not None:
                    if not response_ended:
                        control.stop()
                    if watcher is not None:
                        watcher.cancel()
                        with contextlib.suppress(asyncio.CancelledError):
                            await watcher
                    while not control.idle.is_set():
                        await asyncio.sleep(0.05)
                    if control.reason:
                        response_status = 504 if control.reason == "deadline" else 499
                    unregister(session, ident, control)
                if locked:
                    self.turn.release()
                if store and ident is not None and budget_token is not None:
                    try:
                        await asyncio.to_thread(store.finish_turn, session, ident, response_status)
                    except BudgetUnavailable:
                        _LOG.error("budget_settlement_unavailable")
            if control_token is not None:
                execution.reset(control_token)
            if budget_token is not None:
                request_budget.reset(budget_token)
            session_id.reset(token)
            account_id.reset(account_token)