Mujoco_WASM/decision_server/web_sessions.py

"""Bounded anonymous sessions; no credentials or session metadata on disk."""

import asyncio
import secrets
import shutil
import time
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path

from .protocol import DecisionError
from .server import Service
from .web_config import MemoryConnections


@dataclass
class Limits:
    sessions: int = 128
    idle: int = 1800
    lifetime: int = 28800
    inference: int = 8
    codex: int = 2
    ip_sessions: int = 30
    ip_logins: int = 12
    ip_calls: int = 600


@dataclass
class Visitor:
    ident: str
    csrf: str
    service: Service
    created: float
    touched: float
    codex_reserved: bool = False
    login_deadline: float = 0
    closed: bool = False
    account_lock: asyncio.Lock = field(default_factory=asyncio.Lock)


class Sessions:
    def __init__(self, directory, origin, limits=None):
        self.directory = Path(directory)
        self.origin = origin
        self.limits = limits or Limits()
        self.values = {}
        self.ip_buckets = {}
        self.inference = 0
        self.http = None

    def rate(self, ip, kind, maximum):
        now = time.monotonic()
        self.ip_buckets = {k: v for k, v in self.ip_buckets.items() if v and now - v[-1] < 3600}
        key = (ip, kind)
        if key not in self.ip_buckets:
            if len(self.ip_buckets) >= 4096:
                raise DecisionError("server_capacity", 429)
            self.ip_buckets[key] = deque()
        bucket = self.ip_buckets[key]
        while bucket and now - bucket[0] >= 3600:
            bucket.popleft()
        if len(bucket) >= maximum:
            raise DecisionError("ip_rate_limit", 429)
        bucket.append(now)

    def create(self, ip):
        self.rate(ip, "sessions", self.limits.ip_sessions)
        if len(self.values) >= self.limits.sessions:
            raise DecisionError("session_capacity", 429)
        ident = secrets.token_urlsafe(32)
        service = Service(
            self.directory / ident, "", {self.origin}, 8768, connections=MemoryConnections()
        )
        service.session = self.http
        now = time.monotonic()
        visitor = Visitor(ident, secrets.token_urlsafe(32), service, now, now)
        self.values[ident] = visitor
        return visitor

    def reserve_codex(self, visitor):
        if not visitor.codex_reserved:
            if sum(v.codex_reserved for v in self.values.values()) >= self.limits.codex:
                raise DecisionError("subscription_capacity", 429)
            visitor.codex_reserved = True

    async def close_codex(self, visitor):
        visitor.service.invalidate()
        await visitor.service.codex.close()
        visitor.codex_reserved = False
        visitor.login_deadline = 0

    async def destroy(self, visitor):
        if visitor.closed:
            return
        visitor.closed = True
        self.values.pop(visitor.ident, None)
        service = visitor.service
        service.invalidate()
        await asyncio.gather(*list(service.active.values()), return_exceptions=True)
        async with visitor.account_lock:
            await self.close_codex(visitor)
        service.connections.values.clear()
        service.records.clear()
        service.runs.clear()
        service.cancelled.clear()
        # Only internally generated session directories; never accept paths from HTTP.
        shutil.rmtree(self.directory / visitor.ident, ignore_errors=True)

    async def expire(self):
        now = time.monotonic()
        for visitor in list(self.values.values()):
            if (
                now - visitor.touched >= self.limits.idle
                or now - visitor.created >= self.limits.lifetime
            ):
                await self.destroy(visitor)
            elif visitor.login_deadline and now >= visitor.login_deadline:
                async with visitor.account_lock:
                    if visitor.login_deadline and now >= visitor.login_deadline:
                        await self.close_codex(visitor)

    async def close(self):
        await asyncio.gather(*(self.destroy(v) for v in list(self.values.values())))