Spaces:
Sleeping
Sleeping
File size: 3,322 Bytes
e545bf5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Session store for the SWMM MCP server.
Each uploaded model gets a session with its own working directory under
/tmp/swmm_sessions. Sessions expire after SESSION_TTL_HOURS (default 6) and
are swept lazily on access. State is process-local: the HF Space runs a
single uvicorn worker, matching the stateful-session conclusion from the
WNTR MCP architecture work (one backend, session-ID state).
"""
from __future__ import annotations
import os
import shutil
import threading
import time
import uuid
from pathlib import Path
from typing import Any
SESSION_ROOT = Path(os.environ.get("SWMM_SESSION_ROOT", "/tmp/swmm_sessions"))
SESSION_TTL_S = float(os.environ.get("SESSION_TTL_HOURS", "6")) * 3600.0
MAX_SESSIONS = int(os.environ.get("MAX_SESSIONS", "40"))
class Session:
def __init__(self, session_id: str, workdir: Path):
self.id = session_id
self.workdir = workdir
self.created = time.time()
self.touched = time.time()
self.data: dict[str, Any] = {}
def touch(self) -> None:
self.touched = time.time()
class SessionStore:
def __init__(self) -> None:
self._sessions: dict[str, Session] = {}
self._lock = threading.Lock()
SESSION_ROOT.mkdir(parents=True, exist_ok=True)
def create(self) -> Session:
with self._lock:
self._sweep_locked()
if len(self._sessions) >= MAX_SESSIONS:
oldest = min(self._sessions.values(), key=lambda s: s.touched)
self._drop_locked(oldest.id)
sid = uuid.uuid4().hex[:12]
workdir = SESSION_ROOT / sid
workdir.mkdir(parents=True, exist_ok=True)
session = Session(sid, workdir)
self._sessions[sid] = session
return session
def get(self, session_id: str) -> Session:
with self._lock:
self._sweep_locked()
session = self._sessions.get(str(session_id))
if session is None:
raise KeyError(
f"Unknown or expired session '{session_id}'. Call upload_model first "
f"(sessions expire after {SESSION_TTL_S/3600:.0f} h of inactivity)."
)
session.touch()
return session
def list(self) -> list[dict[str, Any]]:
with self._lock:
self._sweep_locked()
return [
{
"session_id": s.id,
"model": s.data.get("filename"),
"simulated": bool(s.data.get("results")),
"age_minutes": round((time.time() - s.created) / 60.0, 1),
}
for s in self._sessions.values()
]
def drop(self, session_id: str) -> bool:
with self._lock:
return self._drop_locked(str(session_id))
# -- internal --
def _drop_locked(self, session_id: str) -> bool:
session = self._sessions.pop(session_id, None)
if session is None:
return False
shutil.rmtree(session.workdir, ignore_errors=True)
return True
def _sweep_locked(self) -> None:
now = time.time()
for sid in [s for s, v in self._sessions.items() if now - v.touched > SESSION_TTL_S]:
self._drop_locked(sid)
STORE = SessionStore()
|