engine/sessions.py

"""Bounded, browser-session workspaces. No shared corpus and no user-selected server paths."""
from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
import re
import secrets
import shutil
import threading
import time

from .config import Settings, SETTINGS


class SessionError(ValueError):
    pass


@dataclass
class Session:
    token: str
    owner: str
    path: Path
    touched: float = field(default_factory=time.monotonic)
    lock: threading.Lock = field(default_factory=threading.Lock)
    index: object = None
    version: str = ""
    rubric: object = None
    last_result: object = None
    last_key: str = ""
    history: list[str] = field(default_factory=list)
    delete_pending: bool = False
    course_id: str = ""
    transcript: str = ""
    transcript_question: str = ""


class SessionStore:
    def __init__(self, settings: Settings = SETTINGS):
        self.settings = settings
        self.root = settings.root
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        self.guard = threading.RLock()
        self.sessions: dict[str, Session] = {}

    @staticmethod
    def new_token() -> str:
        return secrets.token_hex(24)

    def _get(self, token: str, owner: str) -> Session:
        if not isinstance(token, str) or not re.fullmatch(r"[0-9a-f]{48}", token) or not owner:
            raise SessionError("This browser session is unavailable. Refresh the page.")
        with self.guard:
            session = self.sessions.get(token)
            if session is None:
                if len(self.sessions) >= self.settings.max_sessions:
                    raise SessionError("The server has reached its active-session limit. Please try again later.")
                path = self.root / token
                path.mkdir(mode=0o700, exist_ok=True)
                session = Session(token, owner, path)
                self.sessions[token] = session
            if session.owner != owner:
                raise SessionError("The requested workspace does not belong to this browser session.")
            return session

    @contextmanager
    def lease(self, token: str, owner: str):
        session = self._get(token, owner)
        if not session.lock.acquire(blocking=False):
            raise SessionError("Another operation is running in this session. Please wait for it to finish.")
        try:
            with self.guard:
                if self.sessions.get(token) is not session or session.delete_pending:
                    raise SessionError("This session has expired. Refresh the page and upload again.")
            session.touched = time.monotonic()
            yield session
        finally:
            session.touched = time.monotonic()
            session.lock.release()
            if session.delete_pending:
                self.drop(token)

    def clear_locked(self, session: Session) -> None:
        # Called only while holding the session's lease.
        shutil.rmtree(session.path, ignore_errors=True)
        session.path.mkdir(mode=0o700, exist_ok=True)
        session.index = session.rubric = session.last_result = None
        session.version = session.last_key = ""
        session.history.clear()
        session.course_id = session.transcript = session.transcript_question = ""

    def drop(self, token: str) -> None:
        with self.guard:
            session = self.sessions.get(token)
            if session is None:
                return
            if not session.lock.acquire(blocking=False):
                session.delete_pending = True
                return
            try:
                self.sessions.pop(token, None)
                shutil.rmtree(session.path, ignore_errors=True)
            finally:
                session.lock.release()

    def cleanup(self) -> None:
        now = time.monotonic()
        with self.guard:
            expired = [t for t, s in self.sessions.items() if now - s.touched > self.settings.session_ttl]
            active = set(self.sessions)
        for token in expired:
            self.drop(token)
        # Clean orphan workspaces from a previous process. Never touch arbitrary
        # names outside this private root. Busy/live sessions are excluded.
        for path in self.root.iterdir():
            if path.is_dir() and re.fullmatch(r"[0-9a-f]{48}", path.name) and path.name not in active:
                if time.time() - path.stat().st_mtime > self.settings.session_ttl:
                    shutil.rmtree(path, ignore_errors=True)

    def start_reaper(self) -> None:
        def run():
            while True:
                time.sleep(60)
                try:
                    self.cleanup()
                except OSError:
                    pass
        threading.Thread(target=run, daemon=True, name="session-cleanup").start()