Spaces:
Sleeping
Sleeping
| """ | |
| Per-visitor session store. | |
| Previously the app kept ONE global agent + chat history, so every visitor saw | |
| (and overwrote) everyone else's repo and conversation. Each visitor now gets an | |
| unguessable session token; their repo files, vector DB and chat history live | |
| only under that session and are deleted when it expires. | |
| """ | |
| import asyncio | |
| import hmac | |
| import logging | |
| import secrets | |
| import threading | |
| import time | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| from back_end.config import MAX_LOADED_REPOS, MAX_SESSIONS, REPO_STORAGE_ROOT, SESSION_IDLE_TTL_SECONDS | |
| from back_end.core.downloader import delete_dir | |
| from back_end.core.embeddings import drop_vector_db | |
| logger = logging.getLogger(__name__) | |
| # A loaded repo idle for this long may be evicted when the server is full | |
| EVICT_LOADED_AFTER_IDLE_SECONDS = 10 * 60 | |
| class Session: | |
| session_id: str | |
| token: str | |
| storage_dir: Path | |
| created_at: float = field(default_factory=time.time) | |
| last_seen: float = field(default_factory=time.time) | |
| holds_repo_slot: bool = False | |
| tools: list | None = None | |
| vector_db: object | None = None | |
| chat_history: list = field(default_factory=list) | |
| # Only one indexing job / one chat stream at a time per session | |
| busy: asyncio.Lock = field(default_factory=asyncio.Lock) | |
| def touch(self) -> None: | |
| self.last_seen = time.time() | |
| def idle_for(self) -> float: | |
| return time.time() - self.last_seen | |
| def reset_repo(self) -> None: | |
| drop_vector_db(self.vector_db) | |
| self.tools = None | |
| self.vector_db = None | |
| self.chat_history = [] | |
| class SessionCapacityError(RuntimeError): | |
| pass | |
| class SessionStore: | |
| def __init__(self): | |
| self._sessions: dict[str, Session] = {} | |
| self._lock = threading.Lock() | |
| self._root = Path(REPO_STORAGE_ROOT).resolve() | |
| # ── lifecycle ────────────────────────────────────────────── | |
| def create(self) -> Session: | |
| self.cleanup_expired() | |
| with self._lock: | |
| if len(self._sessions) >= MAX_SESSIONS: | |
| # Make room by dropping the least recently used session that holds nothing | |
| empty = [s for s in self._sessions.values() if not s.holds_repo_slot and not s.busy.locked()] | |
| if not empty: | |
| raise SessionCapacityError("Server is at capacity. Please try again in a few minutes.") | |
| oldest = min(empty, key=lambda s: s.last_seen) | |
| del self._sessions[oldest.session_id] | |
| session_id = secrets.token_hex(16) | |
| session = Session( | |
| session_id=session_id, | |
| token=secrets.token_urlsafe(32), | |
| storage_dir=self._root / session_id, | |
| ) | |
| self._sessions[session_id] = session | |
| return session | |
| def get(self, bearer: str | None) -> Session | None: | |
| """Bearer format: '<session_id>.<token>'. Constant-time token compare.""" | |
| if not bearer or "." not in bearer: | |
| return None | |
| session_id, _, token = bearer.partition(".") | |
| with self._lock: | |
| session = self._sessions.get(session_id) | |
| if not session or not hmac.compare_digest(session.token, token): | |
| return None | |
| if session.idle_for() > SESSION_IDLE_TTL_SECONDS and not session.busy.locked(): | |
| self._destroy(session_id) | |
| return None | |
| session.touch() | |
| return session | |
| def claim_repo_slot(self, session: Session) -> None: | |
| """Reserve capacity for a downloaded repo + vector DB, evicting a stale one if needed.""" | |
| victim = None | |
| with self._lock: | |
| if session.holds_repo_slot: | |
| return | |
| loaded = [s for s in self._sessions.values() if s.holds_repo_slot] | |
| if len(loaded) >= MAX_LOADED_REPOS: | |
| stale = [ | |
| s for s in loaded | |
| if not s.busy.locked() and s.idle_for() > EVICT_LOADED_AFTER_IDLE_SECONDS | |
| ] | |
| if not stale: | |
| raise SessionCapacityError("The server is busy analysing other repos. Please try again in a few minutes.") | |
| victim = min(stale, key=lambda s: s.last_seen) | |
| victim.holds_repo_slot = False | |
| session.holds_repo_slot = True | |
| if victim: | |
| logger.info("Evicting idle repo to make room") | |
| self.release_repo(victim) | |
| def release_repo(self, session: Session) -> None: | |
| session.reset_repo() | |
| session.holds_repo_slot = False | |
| delete_dir(session.storage_dir) | |
| def _destroy(self, session_id: str) -> None: | |
| with self._lock: | |
| session = self._sessions.pop(session_id, None) | |
| if session: | |
| self.release_repo(session) | |
| def cleanup_expired(self) -> None: | |
| with self._lock: | |
| expired = [ | |
| sid for sid, s in self._sessions.items() | |
| if s.idle_for() > SESSION_IDLE_TTL_SECONDS and not s.busy.locked() | |
| ] | |
| for sid in expired: | |
| self._destroy(sid) | |
| if expired: | |
| logger.info("Cleaned up %d expired session(s)", len(expired)) | |
| def wipe_storage_root(self) -> None: | |
| """Called at startup: remove repo copies left over from a previous run.""" | |
| delete_dir(self._root) | |
| self._root.mkdir(parents=True, exist_ok=True) | |
| store = SessionStore() | |