session.py
python
sha256:fc04e4cae9e1774d6a21b65c45daeed0e6787eb581d13aa1b03bfe9384a34226
Merge branch 'fix/two-column-scroll-layout' into dev
Human
9 days ago
| 1 | """MCP Session Store — stateful connection management for MCP 2025-11-25. |
| 2 | |
| 3 | Manages per-client sessions required by the Streamable HTTP transport: |
| 4 | - Cryptographically secure session IDs returned in ``Mcp-Session-Id`` header |
| 5 | - Pending elicitation Futures (correlate elicitation/create ↔ client response) |
| 6 | - SSE queues for GET /mcp push channels |
| 7 | - TTL-based expiry with background cleanup |
| 8 | |
| 9 | Design: In-process only. Documented upgrade path to Redis for multi-replica |
| 10 | deployments — replace the module-level ``_SESSIONS`` dict with a Redis hash |
| 11 | and use ``asyncio.Event`` cross-process synchronisation. |
| 12 | """ |
| 13 | |
| 14 | import asyncio |
| 15 | import logging |
| 16 | import secrets |
| 17 | import time |
| 18 | from dataclasses import dataclass, field |
| 19 | from typing import AsyncIterator |
| 20 | |
| 21 | from musehub.types.json_types import JSONObject |
| 22 | |
| 23 | logger = logging.getLogger(__name__) |
| 24 | |
| 25 | # Session TTL in seconds (15 minutes); reset on each activity. |
| 26 | # 15 minutes covers the longest realistic agent workflows (multi-step tool |
| 27 | # sequences, elicitation round-trips, SSE reconnection grace) while limiting |
| 28 | # the exposure window if a session ID is stolen. |
| 29 | _SESSION_TTL_SECONDS = 900 |
| 30 | |
| 31 | # Maximum number of recent SSE events buffered per session for Last-Event-ID replay. |
| 32 | _SSE_BUFFER_SIZE = 50 |
| 33 | |
| 34 | |
| 35 | @dataclass |
| 36 | class MCPSession: |
| 37 | """Stateful MCP session for the Streamable HTTP transport. |
| 38 | |
| 39 | Attributes: |
| 40 | session_id: Cryptographically secure, globally unique session identifier. |
| 41 | user_id: Authenticated user derived from MSign handle, or ``None`` |
| 42 | for anonymous sessions. |
| 43 | client_capabilities: Capabilities advertised by the client during |
| 44 | ``initialize``, including elicitation mode support. |
| 45 | pending: Mapping from elicitation request ID → ``asyncio.Future`` whose |
| 46 | result is set when the client sends the elicitation response. |
| 47 | sse_queues: Active GET /mcp SSE consumers. Each entry is an |
| 48 | ``asyncio.Queue`` fed by :func:`push_to_session`. |
| 49 | event_buffer: Ring buffer of recent (event_id, data) pairs for |
| 50 | ``Last-Event-ID`` replay on reconnection. |
| 51 | created_at: Unix timestamp of session creation. |
| 52 | last_active: Unix timestamp of last activity (reset on each request). |
| 53 | """ |
| 54 | |
| 55 | session_id: str |
| 56 | user_id: str | None |
| 57 | client_capabilities: JSONObject |
| 58 | pending: dict[str | int, asyncio.Future[JSONObject]] = field( |
| 59 | default_factory=dict |
| 60 | ) |
| 61 | sse_queues: list[asyncio.Queue[str | None]] = field(default_factory=list) |
| 62 | event_buffer: list[tuple[str, str]] = field(default_factory=list) |
| 63 | created_at: float = field(default_factory=time.monotonic) |
| 64 | last_active: float = field(default_factory=time.monotonic) |
| 65 | repo_focus: tuple[str, str] | None = None |
| 66 | |
| 67 | def touch(self) -> None: |
| 68 | """Reset the activity timestamp to defer session expiry.""" |
| 69 | self.last_active = time.monotonic() |
| 70 | |
| 71 | def is_expired(self) -> bool: |
| 72 | """Return True if the session TTL has elapsed since last activity.""" |
| 73 | return (time.monotonic() - self.last_active) > _SESSION_TTL_SECONDS |
| 74 | |
| 75 | def supports_elicitation_form(self) -> bool: |
| 76 | """Return True if the client declared form-mode elicitation support.""" |
| 77 | elicitation = self.client_capabilities.get("elicitation") |
| 78 | if not isinstance(elicitation, dict): |
| 79 | return False |
| 80 | # Empty dict ≡ form-only per spec backwards compat. |
| 81 | return "form" in elicitation or len(elicitation) == 0 |
| 82 | |
| 83 | def supports_elicitation_url(self) -> bool: |
| 84 | """Return True if the client declared URL-mode elicitation support.""" |
| 85 | elicitation = self.client_capabilities.get("elicitation") |
| 86 | if not isinstance(elicitation, dict): |
| 87 | return False |
| 88 | return "url" in elicitation |
| 89 | |
| 90 | |
| 91 | # ── Module-level registry ───────────────────────────────────────────────────── |
| 92 | |
| 93 | type _SessionMap = dict[str, MCPSession] |
| 94 | _SESSIONS: _SessionMap = {} |
| 95 | _cleanup_task: asyncio.Task[None] | None = None |
| 96 | |
| 97 | # Hard cap on simultaneous in-memory sessions. Each session holds async queues |
| 98 | # and futures; unbounded growth would exhaust memory under a connection flood. |
| 99 | # Tune via the MUSEHUB_MAX_MCP_SESSIONS environment variable. The default of |
| 100 | # 10 000 comfortably handles tens of thousands of users when sessions are |
| 101 | # short-lived (sub-minute). Longer-lived clients should be backed by Redis. |
| 102 | _MAX_SESSIONS: int = int(__import__("os").environ.get("MUSEHUB_MAX_MCP_SESSIONS", "10000")) |
| 103 | |
| 104 | |
| 105 | class SessionCapacityError(RuntimeError): |
| 106 | """Raised when :func:`create_session` is called and the store is full. |
| 107 | |
| 108 | The MCP POST handler converts this to ``503 Service Unavailable`` with a |
| 109 | ``Retry-After: 5`` header so well-behaved clients back off automatically. |
| 110 | """ |
| 111 | |
| 112 | |
| 113 | def create_session( |
| 114 | user_id: str | None, |
| 115 | client_capabilities: JSONObject, |
| 116 | ) -> MCPSession: |
| 117 | """Create a new MCP session and register it in the store. |
| 118 | |
| 119 | Args: |
| 120 | user_id: Authenticated user ID from MSign handle, or ``None`` for anonymous. |
| 121 | client_capabilities: Client capability map from ``initialize`` params. |
| 122 | |
| 123 | Returns: |
| 124 | The newly created :class:`MCPSession`. |
| 125 | |
| 126 | Raises: |
| 127 | SessionCapacityError: When the number of active sessions already equals |
| 128 | ``_MAX_SESSIONS``. The caller should return HTTP 503. |
| 129 | """ |
| 130 | if len(_SESSIONS) >= _MAX_SESSIONS: |
| 131 | logger.warning( |
| 132 | "⚠️ MCP session store full (%d/%d) — rejecting new session for user=%s", |
| 133 | len(_SESSIONS), |
| 134 | _MAX_SESSIONS, |
| 135 | user_id, |
| 136 | ) |
| 137 | raise SessionCapacityError( |
| 138 | f"MCP session capacity ({_MAX_SESSIONS}) reached. Retry after expired " |
| 139 | "sessions are evicted (typically within 60 seconds)." |
| 140 | ) |
| 141 | session_id = secrets.token_urlsafe(32) |
| 142 | session = MCPSession( |
| 143 | session_id=session_id, |
| 144 | user_id=user_id, |
| 145 | client_capabilities=client_capabilities, |
| 146 | ) |
| 147 | _SESSIONS[session_id] = session |
| 148 | logger.info("MCP session created: %.8s... user=%s", session_id, user_id) |
| 149 | _ensure_cleanup_running() |
| 150 | return session |
| 151 | |
| 152 | |
| 153 | def get_session(session_id: str) -> MCPSession | None: |
| 154 | """Look up a session by ID and touch its activity timestamp. |
| 155 | |
| 156 | Returns ``None`` if the session does not exist or has expired. |
| 157 | """ |
| 158 | session = _SESSIONS.get(session_id) |
| 159 | if session is None: |
| 160 | return None |
| 161 | if session.is_expired(): |
| 162 | delete_session(session_id) |
| 163 | return None |
| 164 | session.touch() |
| 165 | return session |
| 166 | |
| 167 | |
| 168 | def delete_session(session_id: str) -> bool: |
| 169 | """Terminate a session, closing all open SSE streams. |
| 170 | |
| 171 | Args: |
| 172 | session_id: Session to delete. |
| 173 | |
| 174 | Returns: |
| 175 | ``True`` if the session existed and was deleted, ``False`` otherwise. |
| 176 | """ |
| 177 | session = _SESSIONS.pop(session_id, None) |
| 178 | if session is None: |
| 179 | return False |
| 180 | |
| 181 | # Signal all open GET /mcp SSE streams to close. |
| 182 | for queue in session.sse_queues: |
| 183 | queue.put_nowait(None) # None is the sentinel for stream termination. |
| 184 | |
| 185 | # Cancel any pending elicitation Futures. |
| 186 | for fut in session.pending.values(): |
| 187 | if not fut.done(): |
| 188 | fut.cancel() |
| 189 | |
| 190 | logger.info("MCP session deleted: %.8s...", session_id) |
| 191 | return True |
| 192 | |
| 193 | |
| 194 | def find_sessions_by_user(user_id: str) -> list[MCPSession]: |
| 195 | """Return all active (non-expired) sessions for the given user handle.""" |
| 196 | return [ |
| 197 | s for s in _SESSIONS.values() |
| 198 | if s.user_id == user_id and not s.is_expired() |
| 199 | ] |
| 200 | |
| 201 | |
| 202 | def find_sessions_by_repo(owner: str, slug: str) -> list[MCPSession]: |
| 203 | """Return all active sessions focused on the given owner/slug repo.""" |
| 204 | return [ |
| 205 | s for s in _SESSIONS.values() |
| 206 | if s.repo_focus == (owner, slug) and not s.is_expired() |
| 207 | ] |
| 208 | |
| 209 | |
| 210 | def list_active_sessions() -> list[MCPSession]: |
| 211 | """Return all active (non-expired) sessions.""" |
| 212 | return [s for s in _SESSIONS.values() if not s.is_expired()] |
| 213 | |
| 214 | |
| 215 | # ── SSE push ────────────────────────────────────────────────────────────────── |
| 216 | |
| 217 | |
| 218 | def push_to_session(session: MCPSession, event_text: str) -> None: |
| 219 | """Broadcast a serialised SSE event string to all open GET /mcp streams. |
| 220 | |
| 221 | Buffers the event for ``Last-Event-ID`` replay. |
| 222 | |
| 223 | Args: |
| 224 | session: Target session. |
| 225 | event_text: Fully-formatted SSE event string (including trailing ``\\n\\n``). |
| 226 | """ |
| 227 | # Append to ring buffer (drop oldest if full). |
| 228 | if len(session.event_buffer) >= _SSE_BUFFER_SIZE: |
| 229 | session.event_buffer.pop(0) |
| 230 | session.event_buffer.append(("", event_text)) |
| 231 | |
| 232 | for queue in list(session.sse_queues): |
| 233 | try: |
| 234 | queue.put_nowait(event_text) |
| 235 | except asyncio.QueueFull: |
| 236 | logger.warning("SSE queue full for session %.8s..., dropping event", session.session_id) |
| 237 | |
| 238 | |
| 239 | async def register_sse_queue( |
| 240 | session: MCPSession, |
| 241 | last_event_id: str | None = None, |
| 242 | ) -> AsyncIterator[str]: |
| 243 | """Register a new SSE consumer and yield events until disconnection. |
| 244 | |
| 245 | Replays buffered events if ``last_event_id`` is provided, then streams |
| 246 | live events from the queue. Yields ``None``-sentinel as end-of-stream. |
| 247 | |
| 248 | Args: |
| 249 | session: Active session to attach to. |
| 250 | last_event_id: ``Last-Event-ID`` header value for replay, or ``None``. |
| 251 | |
| 252 | Yields: |
| 253 | Serialised SSE event strings. |
| 254 | """ |
| 255 | queue: asyncio.Queue[str | None] = asyncio.Queue(maxsize=256) |
| 256 | session.sse_queues.append(queue) |
| 257 | |
| 258 | try: |
| 259 | # Replay buffered events after the given last_event_id. |
| 260 | if last_event_id is not None: |
| 261 | replaying = False |
| 262 | for buf_id, buf_text in session.event_buffer: |
| 263 | if replaying: |
| 264 | yield buf_text |
| 265 | elif buf_id == last_event_id: |
| 266 | replaying = True |
| 267 | |
| 268 | # Stream live events. |
| 269 | while True: |
| 270 | item = await queue.get() |
| 271 | if item is None: |
| 272 | break # Session terminated. |
| 273 | yield item |
| 274 | finally: |
| 275 | try: |
| 276 | session.sse_queues.remove(queue) |
| 277 | except ValueError: |
| 278 | pass |
| 279 | |
| 280 | |
| 281 | # ── Elicitation pending request helpers ─────────────────────────────────────── |
| 282 | |
| 283 | |
| 284 | def create_pending_elicitation( |
| 285 | session: MCPSession, |
| 286 | request_id: str | int, |
| 287 | ) -> asyncio.Future[JSONObject]: |
| 288 | """Register a pending elicitation and return its Future. |
| 289 | |
| 290 | The Future is resolved by :func:`resolve_elicitation` when the client |
| 291 | sends the elicitation response. |
| 292 | |
| 293 | Args: |
| 294 | session: Session on which the elicitation is pending. |
| 295 | request_id: JSON-RPC ID of the ``elicitation/create`` request. |
| 296 | |
| 297 | Returns: |
| 298 | ``asyncio.Future`` that will be set to the client's response dict. |
| 299 | """ |
| 300 | loop = asyncio.get_event_loop() |
| 301 | fut: asyncio.Future[JSONObject] = loop.create_future() |
| 302 | session.pending[request_id] = fut |
| 303 | return fut |
| 304 | |
| 305 | |
| 306 | def resolve_elicitation( |
| 307 | session: MCPSession, |
| 308 | request_id: str | int, |
| 309 | result: JSONObject, |
| 310 | ) -> bool: |
| 311 | """Resolve a pending elicitation Future with the client's response. |
| 312 | |
| 313 | Args: |
| 314 | session: Session owning the pending elicitation. |
| 315 | request_id: JSON-RPC ID matching the ``elicitation/create`` request. |
| 316 | result: Client elicitation result dict (``action`` + optional ``content``). |
| 317 | |
| 318 | Returns: |
| 319 | ``True`` if the pending Future was found and resolved, ``False`` otherwise. |
| 320 | """ |
| 321 | fut = session.pending.pop(request_id, None) |
| 322 | if fut is None or fut.done(): |
| 323 | return False |
| 324 | fut.set_result(result) |
| 325 | return True |
| 326 | |
| 327 | |
| 328 | def cancel_elicitation(session: MCPSession, request_id: str | int) -> bool: |
| 329 | """Cancel a pending elicitation Future (e.g. on ``notifications/cancelled``). |
| 330 | |
| 331 | Args: |
| 332 | session: Session owning the pending elicitation. |
| 333 | request_id: JSON-RPC ID to cancel. |
| 334 | |
| 335 | Returns: |
| 336 | ``True`` if the Future was found and cancelled. |
| 337 | """ |
| 338 | fut = session.pending.pop(request_id, None) |
| 339 | if fut is None or fut.done(): |
| 340 | return False |
| 341 | fut.cancel() |
| 342 | return True |
| 343 | |
| 344 | |
| 345 | # ── Background cleanup ──────────────────────────────────────────────────────── |
| 346 | |
| 347 | |
| 348 | def _ensure_cleanup_running() -> None: |
| 349 | """Start the background cleanup task if it is not already running.""" |
| 350 | global _cleanup_task |
| 351 | try: |
| 352 | loop = asyncio.get_running_loop() |
| 353 | except RuntimeError: |
| 354 | return |
| 355 | if _cleanup_task is None or _cleanup_task.done(): |
| 356 | _cleanup_task = loop.create_task(_cleanup_loop(), name="mcp-session-cleanup") |
| 357 | |
| 358 | |
| 359 | async def _cleanup_loop() -> None: |
| 360 | """Periodically expire stale sessions (runs every 5 minutes).""" |
| 361 | while True: |
| 362 | await asyncio.sleep(300) |
| 363 | expired = [sid for sid, s in list(_SESSIONS.items()) if s.is_expired()] |
| 364 | for sid in expired: |
| 365 | delete_session(sid) |
| 366 | if expired: |
| 367 | logger.info("MCP session cleanup: expired %d session(s)", len(expired)) |
File History
13 commits
sha256:fc04e4cae9e1774d6a21b65c45daeed0e6787eb581d13aa1b03bfe9384a34226
Merge branch 'fix/two-column-scroll-layout' into dev
Human
9 days ago
sha256:408916fc5973ba59c6e4eebaa80ebdcc801c0a63205651e25009d11548f79454
chore: bump version to 0.2.0.dev2 — nightly.2, matching muse
Sonnet 4.6
patch
12 days ago
sha256:d035733f21ccff27735fddebfbbe0ed24565a32a22db8de5885402262671ecd2
chore: bump version to 0.2.0rc15 for musehub#113 fix release
Sonnet 4.6
patch
15 days ago
sha256:0032d6cfa33bc3c8367436ad768e7dd0e339b4332153160247da8266cb5fa352
Merge branch 'task/version-tags-phase3-server' into dev
Human
17 days ago
sha256:4669620efda9ff41c55bdefd1f7bfe1c239d468428744c84ead9957e5a003a53
merge: rescue snapshot-recovery hardening (c00aa21d) into d…
Opus 4.8
minor
⚠
30 days ago
sha256:a59da49c4611b970fc4b6ae48678ce4943261c213a07ddbd73ce9201df869b4a
fix: remove false-positive proposal_comments index drop fro…
Sonnet 4.6
patch
34 days ago
sha256:0a240d6dbff234f07d98a28a4a9a68db702f3f9ff9260196f24219bdb1c0b6f3
feat: render markdown mists as HTML with heading anchor links
Sonnet 4.6
patch
35 days ago
sha256:24a7d47486ebc4ebd1832830580e177ec6f877b48dced8c000e198cdec4ce9d6
Merge 'task/bump-version-rc12' into 'dev' — proposal: Bump …
Human
36 days ago
sha256:b9ff931d147e0114a1f17060f415b89ed551c170a91ff226c70437aa5c85f9ee
Merge 'task/bump-version-rc12' into 'dev' — proposal: Bump …
Human
36 days ago
sha256:d1122d21e73471879b460037b22c0b50fded7c423444a176f248428f75dac39c
Merge 'task/fix-issue-pagination-cursor' into 'dev' — propo…
Human
36 days ago
sha256:01e18975e73d2b3cd5b6db7929c895bef9aa6e0d4391dc5b2adfc548b41318dd
Merge 'feat/adding-debug-logs-to-staging' into 'dev' — prop…
Human
36 days ago
sha256:6b1949fc2797ca4c1936a637a4cbfec828ef56cf52398a2e74ca3c4f494e728f
fix: use wire_bytes not mpack_bytes_raw in compute_object_b…
Sonnet 4.6
patch
48 days ago
sha256:b99f2455dc346966d040133f5203297e6e3ef5803a93728a2c30568d0a0f7583
rename: delta_add → delta_upsert across wire format, models…
Sonnet 4.6
patch
51 days ago