"""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)