From c792df64d27419327eb2903c8b9d95a1cede975f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Sat, 30 May 2026 09:00:36 +0530 Subject: [PATCH] feat(mcp): support stateless and stateful clients via session-id routing (#26857) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(mcp): support stateless and stateful clients via session-id routing - Add session_manager_stateful (stateless=False) alongside stateless - Route by mcp-session-id: has ID → stateful, initialize (no ID) → stateful, else → stateless - Peek POST body to detect initialize for routing; replay via wrapped receive - Handle stale session IDs for both managers - Add test_mcp_routing_initialize_to_stateful_no_session_to_stateless - Update test_valid_mcp_session_id_is_preserved, test_concurrent_initialize_session_managers Made-with: Cursor * fix(mcp): respect stateful routing and harden initialize detection Ensure streamable MCP requests are dispatched via the computed target session manager, and guard initialize detection against non-object JSON bodies. Update stale-session test patches to target the stateful manager so routing assertions remain correct. Made-with: Cursor * test(mcp): patch stateless/stateful managers in concurrency init test Update concurrent session-manager initialization test to patch session_manager_stateless and session_manager_stateful directly, matching initialize_session_managers() behavior and preventing NameError from undefined mocks. Made-with: Cursor * Fix tests * Fix tests * Fix MCP stateful routing edge cases * Fix stateful MCP auth context refresh * Fix MCP stateful session cleanup * fix(mcp): bind stateful sessions to creator and reject hijacks Stateful mcp-session-id was usable by any authenticated proxy caller. Track the session creator's hashed API key (or user_id) when a new session is issued and reject mismatched callers with 403 before _set_or_update_auth_context overwrites the stored MCPAuthenticatedUser. Also formats nested with-statements in test_mcp_stale_session.py and fixes a pre-existing AsyncMock mismatch in test_stale_mcp_session_id_is_stripped. * fix(mcp): serialize concurrent requests on same stateful session Bugbot's 'Concurrent requests share context' finding: _update_auth_context mutates the single MCPAuthenticatedUser stored per session in place on every request, so two requests sharing one mcp-session-id can overwrite each other's mcp_servers / auth headers / oauth state / client_ip while in-flight callbacks are still reading the same object. Owner-binding alone narrows this to same-principal racing, but the in-place mutation race remains. Add a per-session asyncio.Lock around handle_request so concurrent same-session requests run sequentially. The lock is allocated on demand and torn down with the rest of the session state on DELETE / idle expiry. Co-authored-by: Mateo Wang * fix(mcp): include OAuth2 bearer in stateful session owner fingerprint UserAPIKeyAuth() for OAuth2 passthrough has no api_key/user_id, so every OAuth caller fingerprinted to "anonymous" and could hijack another OAuth caller's mcp-session-id. Hash the upstream Authorization header into the fingerprint as oauth:. * fix(mcp): don't hold stateful session lock for streaming GETs The per-session lock wraps handle_request, so a long-lived GET (SSE stream held open for the life of the session) would block every subsequent POST on the same mcp-session-id. Only POST/DELETE mutate the shared MCPAuthenticatedUser, so it's sufficient to serialize those — GETs run lock-free and stream concurrently. * fix(mcp): allow None user_api_key_auth in MCPAuthenticatedUser The set_auth_context / _set_or_update_auth_context / _update_auth_context helpers in server.py all accept Optional[UserAPIKeyAuth] and pass it straight into MCPAuthenticatedUser, but the dataclass-style constructor typed user_api_key_auth as required UserAPIKeyAuth. Mypy flagged this on the stateful-routing branch: server.py:3227: error: Incompatible types in assignment (expression has type "UserAPIKeyAuth | None", variable has type "UserAPIKeyAuth") server.py:3255: error: Argument "user_api_key_auth" to "MCPAuthenticatedUser" has incompatible type "UserAPIKeyAuth | None"; expected "UserAPIKeyAuth" Widen the parameter type to Optional[UserAPIKeyAuth] to match the call sites. Runtime behavior is unchanged. Co-authored-by: Mateo Wang * style: replace with new alias * fix(mcp): fall back to client_ip in stateful session owner fingerprint Addresses Greptile review on PR #26857: when no API key, user_id, or OAuth bearer is available (e.g. unauthenticated/passthrough callers), the owner fingerprint collapsed to a single 'anonymous' value, allowing two unrelated callers to drive each other's stateful MCP sessions. Fold client IP into the fingerprint as a fallback identity signal so distinct anonymous sources do not share an owner identity. Co-authored-by: Mateo Wang * Fix active stateful MCP session cleanup * test(mcp): cancel leaked stateful auth-context cleanup task initialize_session_managers() spawns a real asyncio.create_task running _cleanup_expired_stateful_session_auth_contexts(). The test_concurrent_initialize_session_managers test was saving and restoring the session-manager context-manager globals but did not save, cancel, or restore _stateful_auth_context_cleanup_task. Because pyproject.toml sets asyncio_default_fixture_loop_scope=session, the event loop is shared across tests in the same session, so the leaked task kept running against module-level dicts for the rest of the test run. Save and cancel the task in the finally block so the test fully cleans up after itself. Co-authored-by: Mateo Wang * Fix stateful MCP session fingerprinting * Hash MCP session user owner fingerprints * Fix stale MCP session DELETE cleanup * fix(mcp): harden owner fingerprint hashing for non-str api keys _owner_fingerprint_for assumed api_key/user_id supported .encode(); MagicMock-based tests (and any non-str truthy values) crashed with TypeError before routing. Only hash str/bytes secrets; fall through otherwise so MCP routing and session tests behave correctly. Co-authored-by: Cursor * Fix MCP stateful cleanup loop resilience * Fix stateful MCP initialize auth capture * fix(mcp): drop orphan per-session lock when auth context absent Defensive cleanup for _stateful_session_locks entries created on sessions that never enter _stateful_session_auth_contexts. The periodic cleanup loop only iterates auth_context_last_seen, so such locks would otherwise live forever. Add a test that reproduces the leak and verifies the request finalizer pops the lock. Co-authored-by: Mateo Wang * chore(mcp): trim verbose comment on lock cleanup Co-authored-by: Mateo Wang * Fix stateful MCP delete failure tracking * fix test * fix(mcp): cap routing-peek body size to bound pre-dispatch memory Authenticated clients that POST without an mcp-session-id forced the proxy to buffer the entire request body before routing, since the peek loop drained every body chunk to decide whether the JSON-RPC method was 'initialize'. Cap the peek at 4 KB (more than enough for an initialize envelope) and let the remainder stream through wrapped_receive into the downstream handler. * test: replace dall-e-3 with gpt-image-1 in health check and router tests (#27813) OpenAI returns 'The model dall-e-3 does not exist' for the test account, breaking test_openai_img_gen_health_check and test_image_generation. Switch to gpt-image-1, matching the existing TestOpenAIGPTImage1 pattern. * fix(tests): drop dall-e-only test classes; route live image tests via gpt-image-1 Second wave of failures from the 2026-05-12 DALL-E shutdown: - tests/image_gen_tests/test_image_edits.py::TestOpenAIImageEditDallE2 and tests/image_gen_tests/test_image_generation.py::TestOpenAIDalle3 are explicitly named for the deprecated models and can't pass; remove. gpt-image-1 coverage already exists in sibling classes. - tests/local_testing/test_router.py image gen tests use dall-e-3 only as a routing example; swap to gpt-image-1. - tests/local_testing/test_custom_callback_input.py image_generation success/failure paths swapped to gpt-image-1. * Fix MCP initialize session active tracking Co-authored-by: Yassin Kortam * Fix MCP reinitialize session tracking Co-authored-by: Yassin Kortam * Fix MCP reinitialize auth context aliasing Co-authored-by: Yassin Kortam * Apply black formatting after merge Co-authored-by: Mateo Wang * Run owner-binding 403 before consuming POST body Co-authored-by: Mateo Wang * Harden MCP routing peek bound and stateful purge race Co-authored-by: Mateo Wang * Remove inadvertently committed Next.js build artifacts Co-authored-by: Mateo Wang * Run owner check before stale MCP session cleanup Co-authored-by: Mateo Wang * fix(mcp): reverse cleanup ordering to terminate transport before clearing owner Reverses _purge_expired_stateful_session_auth_contexts so the transport is popped from server_instances and terminated BEFORE owner/auth tracking is cleared. The previous order left a window where _stateful_session_owners was already empty but server_instances still served the session, so a concurrent request would observe expected_owner is None and bypass the owner-binding check. Addresses Greptile review on PR #26857. Co-authored-by: Mateo Wang * test(mcp): fully reset stateful session tracking in auth-context refresh test Use _remove_stateful_session_tracking in teardown so the test no longer leaks _stateful_session_auth_context_last_seen and _stateful_session_locks between tests, matching the cleanup used by the sibling stateful tests. * fix(mcp): cap concurrent stateful sessions per caller to bound memory --------- Co-authored-by: Cursor Agent Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Mateo Wang Co-authored-by: Sameerlite Co-authored-by: yuneng-jiang Co-authored-by: Yassin Kortam Co-authored-by: Claude Babysitter Co-authored-by: mateo-berri --- .../mcp_server/auth/litellm_auth_handler.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 634 +++++++- tests/mcp_tests/test_mcp_server.py | 25 +- .../auth/test_user_api_key_auth_mcp.py | 6 +- .../mcp_server/test_mcp_server.py | 1409 ++++++++++++++++- .../mcp_server/test_mcp_stale_session.py | 173 +- 6 files changed, 2163 insertions(+), 86 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py index 75b75d3ba4..7122c64ec6 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -20,7 +20,7 @@ class MCPAuthenticatedUser(AuthenticatedUser): def __init__( self, - user_api_key_auth: UserAPIKeyAuth, + user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f31005be0c..a05ce3f741 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,6 +6,8 @@ LiteLLM MCP Server Routes import asyncio import contextlib +import hashlib +import json import time import types import traceback @@ -28,7 +30,7 @@ from fastapi import FastAPI, HTTPException from pydantic import AnyUrl, ConfigDict from starlette.requests import Request as StarletteRequest from starlette.responses import JSONResponse -from starlette.types import Receive, Scope, Send +from starlette.types import Message, Receive, Scope, Send from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG @@ -74,6 +76,19 @@ from litellm.utils import Rules, client, function_setup _byok_cred_cache: Dict[Tuple[str, str], Tuple[Optional[str], float]] = {} _BYOK_CRED_CACHE_TTL = 60 # seconds _BYOK_CRED_CACHE_MAX_SIZE = 4096 # cap to prevent unbounded growth +_STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS = 30 * 60 +# Upper bound on concurrent stateful sessions a single caller may hold. Each +# `initialize` creates a session that survives until the idle timeout, so +# without a cap an authenticated client could spam `initialize` and exhaust +# memory. The caller's own oldest idle sessions are evicted to make room; if +# the cap is still hit (every session in flight), the new `initialize` is +# rejected with 429. +_MAX_STATEFUL_SESSIONS_PER_OWNER = 100 +# Maximum bytes to peek when sniffing the JSON-RPC method on a POST. +# An `initialize` envelope is a few hundred bytes; capping the peek +# prevents an authenticated client from forcing the proxy to buffer an +# arbitrarily large body just to make a routing decision. +_MCP_ROUTING_PEEK_MAX_BYTES = 4096 def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: @@ -242,13 +257,45 @@ if MCP_AVAILABLE: sse: SseServerTransport = SseServerTransport("/mcp/sse/messages") # Create session managers - session_manager = StreamableHTTPSessionManager( + session_manager_stateless = StreamableHTTPSessionManager( app=server, event_store=None, json_response=False, # enables SSE streaming stateless=True, ) + session_manager_stateful = StreamableHTTPSessionManager( + app=server, + event_store=None, # TODO: Add EventStore for reconnection/event replay if needed + json_response=False, # enables SSE streaming + stateless=False, + ) + _stateful_session_auth_contexts: Dict[str, MCPAuthenticatedUser] = {} + _stateful_session_auth_context_last_seen: Dict[str, float] = {} + # Maps session_id -> owner identifier (hashed API key/token) so we can + # reject requests that supply a session_id created by a different caller. + # Without this, a leaked mcp-session-id could be driven (or terminated) + # by any other authenticated proxy user. + _stateful_session_owners: Dict[str, str] = {} + # Per-session lock that serializes ``handle_request`` for the same + # mcp-session-id. The stored ``MCPAuthenticatedUser`` is mutated in place + # by ``_update_auth_context`` each request; without this lock, two + # concurrent requests on the same session would clobber each other's + # auth headers / mcp_servers / oauth state while in-flight callbacks are + # still reading the shared object. + _stateful_session_locks: Dict[str, asyncio.Lock] = {} + _stateful_session_active_request_counts: Dict[str, int] = {} + + def _remove_stateful_session_tracking(session_id: str) -> None: + _stateful_session_auth_contexts.pop(session_id, None) + _stateful_session_auth_context_last_seen.pop(session_id, None) + _stateful_session_owners.pop(session_id, None) + _stateful_session_locks.pop(session_id, None) + _stateful_session_active_request_counts.pop(session_id, None) + + # Keep this alias so existing references to session_manager still work + session_manager = session_manager_stateless + # Create SSE session manager sse_session_manager = StreamableHTTPSessionManager( app=server, @@ -259,11 +306,100 @@ if MCP_AVAILABLE: # Context managers for proper lifecycle management _session_manager_cm = None + _session_manager_stateful_cm = None _sse_session_manager_cm = None + _stateful_auth_context_cleanup_task: Optional[asyncio.Task] = None + + async def _purge_expired_stateful_session_auth_contexts( + now: Optional[float] = None, + ) -> None: + """Terminate expired stateful sessions and drop their auth contexts.""" + now = time.monotonic() if now is None else now + server_instances = getattr(session_manager_stateful, "_server_instances", {}) + expired_session_ids = [] + for session_id, last_seen in _stateful_session_auth_context_last_seen.items(): + if _stateful_session_active_request_counts.get(session_id, 0) > 0: + continue + if ( + now - last_seen >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + or session_id not in server_instances + ): + expired_session_ids.append(session_id) + + for session_id in expired_session_ids: + # Re-check the active-request count immediately before tearing + # the session down. ``await transport.terminate()`` yields to + # the event loop, so a request that started after the first + # collection pass could otherwise observe its transport being + # ripped out from under it mid-flight. + if _stateful_session_active_request_counts.get(session_id, 0) > 0: + continue + # Pop transport + terminate BEFORE removing owner/auth tracking. + # Reversing the order avoids a window where ``_stateful_session_owners`` + # is empty but ``server_instances`` still serves the session — a + # concurrent request in that window would observe ``expected_owner + # is None`` and bypass the owner-binding check. + transport = server_instances.pop(session_id, None) + if transport is not None: + await transport.terminate() + _remove_stateful_session_tracking(session_id) + + for session_id in list(_stateful_session_auth_context_last_seen): + if session_id not in _stateful_session_auth_contexts: + _remove_stateful_session_tracking(session_id) + + async def _enforce_stateful_session_cap_for_owner(owner: str) -> bool: + """ + Bound the number of concurrent stateful sessions a single caller holds + before routing a new ``initialize`` to the stateful manager. + + Evicts the caller's *own* oldest idle sessions (no in-flight requests) + to make room, so a busy-but-legitimate client keeps its newest sessions + and other callers are never affected. Returns ``True`` if the new + session may proceed, or ``False`` when the caller is already at the cap + with every session in flight (the new ``initialize`` should be rejected). + """ + server_instances = getattr(session_manager_stateful, "_server_instances", {}) + + def _owned_live_session_ids() -> List[str]: + return [ + session_id + for session_id, session_owner in _stateful_session_owners.items() + if session_owner == owner and session_id in server_instances + ] + + owned = _owned_live_session_ids() + if len(owned) < _MAX_STATEFUL_SESSIONS_PER_OWNER: + return True + + for session_id in sorted( + owned, + key=lambda sid: _stateful_session_auth_context_last_seen.get(sid, 0.0), + ): + if len(_owned_live_session_ids()) < _MAX_STATEFUL_SESSIONS_PER_OWNER: + break + if _stateful_session_active_request_counts.get(session_id, 0) > 0: + continue + transport = server_instances.pop(session_id, None) + if transport is not None: + await transport.terminate() + _remove_stateful_session_tracking(session_id) + + return len(_owned_live_session_ids()) < _MAX_STATEFUL_SESSIONS_PER_OWNER + + async def _cleanup_expired_stateful_session_auth_contexts() -> None: + while True: + await asyncio.sleep(_STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS) + try: + await _purge_expired_stateful_session_auth_contexts() + except Exception as e: + verbose_logger.exception( + f"Error cleaning up expired MCP stateful sessions: {e}" + ) async def initialize_session_managers(): """Initialize the session managers. Can be called from main app lifespan.""" - global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _sse_session_manager_cm + global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _session_manager_stateful_cm, _sse_session_manager_cm, _stateful_auth_context_cleanup_task # Use async lock to prevent concurrent initialization async with _INITIALIZATION_LOCK: @@ -273,12 +409,17 @@ if MCP_AVAILABLE: verbose_logger.info("Initializing MCP session managers...") # Start the session managers with context managers - _session_manager_cm = session_manager.run() + _session_manager_cm = session_manager_stateless.run() + _session_manager_stateful_cm = session_manager_stateful.run() _sse_session_manager_cm = sse_session_manager.run() # Enter the context managers await _session_manager_cm.__aenter__() + await _session_manager_stateful_cm.__aenter__() await _sse_session_manager_cm.__aenter__() + _stateful_auth_context_cleanup_task = asyncio.create_task( + _cleanup_expired_stateful_session_auth_contexts() + ) _SESSION_MANAGERS_INITIALIZED = True verbose_logger.info( @@ -287,21 +428,29 @@ if MCP_AVAILABLE: async def shutdown_session_managers(): """Shutdown the session managers.""" - global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _sse_session_manager_cm + global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _session_manager_stateful_cm, _sse_session_manager_cm, _stateful_auth_context_cleanup_task if _SESSION_MANAGERS_INITIALIZED: verbose_logger.info("Shutting down MCP session managers...") try: + if _stateful_auth_context_cleanup_task: + _stateful_auth_context_cleanup_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await _stateful_auth_context_cleanup_task if _session_manager_cm: await _session_manager_cm.__aexit__(None, None, None) + if _session_manager_stateful_cm: + await _session_manager_stateful_cm.__aexit__(None, None, None) if _sse_session_manager_cm: await _sse_session_manager_cm.__aexit__(None, None, None) except Exception as e: verbose_logger.exception(f"Error during session manager shutdown: {e}") _session_manager_cm = None + _session_manager_stateful_cm = None _sse_session_manager_cm = None + _stateful_auth_context_cleanup_task = None _SESSION_MANAGERS_INITIALIZED = False @contextlib.asynccontextmanager @@ -366,7 +515,7 @@ if MCP_AVAILABLE: @server.call_tool() async def mcp_server_tool_call( - name: str, arguments: Dict[str, Any] | None + name: str, arguments: Optional[Dict[str, Any]] ) -> CallToolResult: """ Call a specific tool with the provided arguments @@ -409,7 +558,7 @@ if MCP_AVAILABLE: if host_token and hasattr(host_ctx, "session") and host_ctx.session: host_session = host_ctx.session - async def forward_progress(progress: float, total: float | None): + async def forward_progress(progress: float, total: Optional[float]): """Forward progress notifications from external MCP to Host""" try: await host_session.send_progress_notification( @@ -551,7 +700,7 @@ if MCP_AVAILABLE: @server.get_prompt() async def get_prompt( - name: str, arguments: dict[str, str] | None + name: str, arguments: Optional[Dict[str, str]] ) -> GetPromptResult: """ Get a specific prompt with the provided arguments @@ -2697,6 +2846,144 @@ if MCP_AVAILABLE: raw_headers, ) + def _get_session_id_from_scope(scope: Scope) -> Optional[str]: + """ + Extract mcp-session-id from ASGI scope headers. + Returns None if not present. + """ + for header_name, header_value in scope.get("headers", []): + name = ( + header_name if isinstance(header_name, bytes) else header_name.encode() + ) + if name.lower() == b"mcp-session-id": + return ( + header_value.decode() + if isinstance(header_value, bytes) + else str(header_value) + ) + return None + + def _owner_fingerprint_for( + user_api_key_auth: Optional[UserAPIKeyAuth], + oauth2_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, + ) -> str: + """ + Stable, non-reversible identifier for the caller used to bind an + mcp-session-id to its creator. Hash the resolved credential before + using it so custom key formats are never stored in cleartext. + + For OAuth2 passthrough (``UserAPIKeyAuth()`` with no key/user_id), + the caller's identity is the upstream OAuth bearer; hash it so two + OAuth callers with different tokens don't both fingerprint to + ``anonymous`` and end up sharing a session. + + When no caller-identifying credentials are available at all + (e.g. proxy running without master key, or an unauthenticated + passthrough path), fall back to the client IP so two unrelated + anonymous callers from different sources do not collapse to a + single ``anonymous`` owner and end up able to drive each other's + stateful sessions. Note: when even client IP is unavailable + (exotic deployments without trusted X-Forwarded-For and direct + socket info), the fingerprint degrades to the ``anonymous`` + sentinel and cannot meaningfully protect against another + unauthenticated caller who learns the session id — owner-binding + is best-effort in that mode. + """ + + def _bytes_for_hash(value: Any) -> Optional[bytes]: + """Only hash str/bytes secrets; skip mocks and other unexpected types.""" + if value is None: + return None + if isinstance(value, (bytes, bytearray)): + return bytes(value) + if isinstance(value, str): + return value.encode("utf-8") + return None + + if user_api_key_auth is not None: + key_material = _bytes_for_hash(getattr(user_api_key_auth, "api_key", None)) + if key_material: + api_key_hash = hashlib.sha256(key_material).hexdigest() + return f"key:{api_key_hash}" + uid_material = _bytes_for_hash(getattr(user_api_key_auth, "user_id", None)) + if uid_material: + user_id_hash = hashlib.sha256(uid_material).hexdigest() + return f"user:{user_id_hash}" + if oauth2_headers: + authz = oauth2_headers.get("Authorization") or oauth2_headers.get( + "authorization" + ) + authz_bytes = _bytes_for_hash(authz) + if authz_bytes: + return f"oauth:{hashlib.sha256(authz_bytes).hexdigest()}" + if client_ip and isinstance(client_ip, str): + return f"ip:{hashlib.sha256(client_ip.encode('utf-8')).hexdigest()}" + return "anonymous" + + def _is_initialize_request(body: bytes) -> bool: + """ + Check if the request body is a JSON-RPC initialize method. + Returns True if method is "initialize", False otherwise or on parse error. + """ + if not body: + return False + try: + data = json.loads(body) + return isinstance(data, dict) and data.get("method") == "initialize" + except (json.JSONDecodeError, TypeError): + return False + + async def _read_request_body_for_routing( + receive: Receive, + ) -> Tuple[List[Message], bytes]: + """ + Read just enough of the request body to decide whether this is a + JSON-RPC ``initialize`` call. Returns the consumed ASGI messages so + the caller can replay them faithfully to the downstream handler, and + the peeked body bytes (capped at ``_MCP_ROUTING_PEEK_MAX_BYTES``). + + Stops reading from the wire as soon as either (a) we have peeked + ``_MCP_ROUTING_PEEK_MAX_BYTES`` of body, or (b) the body is complete. + The remainder of an oversized body is streamed lazily through + ``wrapped_receive`` in the caller — so an authenticated client cannot + force the proxy to buffer an arbitrarily large payload just to make a + routing decision. + """ + consumed_messages: List[Message] = [] + body_chunks: List[bytes] = [] + peeked_bytes = 0 + + while True: + message = await receive() + consumed_messages.append(message) + + if message.get("type") != "http.request": + break + + body = message.get("body", b"") or b"" + if body: + # Only retain up to the remaining peek budget for sniffing. + # The full ``message`` is already in memory (delivered by + # the ASGI server) and must round-trip to the downstream + # handler via ``consumed_messages``, but ``body_chunks`` is + # purely for the JSON-RPC method check — there is no reason + # to copy a large body frame into a second buffer. + remaining = _MCP_ROUTING_PEEK_MAX_BYTES - peeked_bytes + if remaining > 0: + body_chunks.append(body[:remaining]) + peeked_bytes += min(len(body), remaining) + + if not message.get("more_body", False): + break + + if peeked_bytes >= _MCP_ROUTING_PEEK_MAX_BYTES: + # Stop draining; downstream replay will pull remaining chunks + # directly from the original `receive` via wrapped_receive. + break + + return consumed_messages, b"".join(body_chunks) + async def _handle_stale_mcp_session( scope: Scope, receive: Receive, @@ -2760,6 +3047,7 @@ if MCP_AVAILABLE: method = scope.get("method", "").upper() if method == "DELETE": + _remove_stateful_session_tracking(_session_id) verbose_logger.info( "DELETE request for non-existent MCP session '%s'. " "Returning success (idempotent DELETE).", @@ -2993,7 +3281,7 @@ if MCP_AVAILABLE: detail="Forbidden", ) - async def handle_streamable_http_mcp( + async def handle_streamable_http_mcp( # noqa: PLR0915 scope: Scope, receive: Receive, send: Send ) -> None: """Handle MCP requests through StreamableHTTP.""" @@ -3086,38 +3374,215 @@ if MCP_AVAILABLE: if _debug_headers: send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers) - # Set the auth context variable for easy access in MCP functions - set_auth_context( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - client_ip=_client_ip, - ) - # Ensure session managers are initialized if not _SESSION_MANAGERS_INITIALIZED: await initialize_session_managers() # Give it a moment to start up await asyncio.sleep(0.1) - # Handle stale session IDs - either strip them for reconnection - # or return success for idempotent DELETE operations - handled = await _handle_stale_mcp_session( - scope, receive, send, session_manager - ) - if handled: - # Request was fully handled (e.g., DELETE on non-existent session) - return + # Route based on mcp-session-id and request method: + # - Has session ID → stateful (Claude Code, Cursor, VSCode) + # - No session ID + initialize → stateful (so client gets mcp-session-id) + # - No session ID + other → stateless (curl, Inspector, Notion) + session_id = _get_session_id_from_scope(scope) + is_initialize = False + consumed_messages: List[Message] = [] - async with _gateway_initialize_instructions_request_scope( - user_api_key_auth, - mcp_servers, - _client_ip, - ): - await session_manager.handle_request(scope, receive, send) + # Owner-binding: a live stateful session may only be driven by the + # caller that created it. Reject mismatches with 403 so a leaked + # mcp-session-id cannot be hijacked by another authenticated user. + # + # Run before ``_handle_stale_mcp_session`` so a non-owner cannot + # force-clean another caller's residual tracking entries via a + # stale DELETE, and before peeking the request body so the 403 + # response sees a pristine ``receive`` channel. + if session_id: + expected_owner = _stateful_session_owners.get(session_id) + request_owner = _owner_fingerprint_for( + user_api_key_auth, oauth2_headers, _client_ip + ) + if expected_owner is not None and expected_owner != request_owner: + verbose_logger.warning( + "Rejecting MCP request: session '%s' owner mismatch.", + session_id, + ) + forbidden_response = JSONResponse( + status_code=403, + content={ + "error": "Forbidden", + "details": "mcp-session-id is bound to a different caller.", + }, + ) + await forbidden_response(scope, receive, send) + return + + # Handle stale session IDs before choosing a target manager. Stale + # non-DELETE requests have their session header stripped and should + # be routed as no-session requests. + if session_id: + handled = await _handle_stale_mcp_session( + scope, receive, send, session_manager_stateful + ) + if handled: + # Request was fully handled (e.g., DELETE on non-existent session) + return + session_id = _get_session_id_from_scope(scope) + + if scope.get("method") == "POST": + consumed_messages, body = await _read_request_body_for_routing(receive) + is_initialize = _is_initialize_request(body) + + use_stateful = bool(session_id or is_initialize) + target_manager = ( + session_manager_stateful if use_stateful else session_manager_stateless + ) + + verbose_logger.debug( + f"MCP routing to {'stateful' if use_stateful else 'stateless'} manager" + + (f" (session={session_id[:8]}...)" if session_id else "") + + (" (initialize)" if is_initialize else "") + ) + + # A new `initialize` (no session id) is about to create a stateful + # session. Cap how many a single caller can hold so an authenticated + # client cannot spam `initialize` and exhaust memory. + if is_initialize and not session_id: + request_owner = _owner_fingerprint_for( + user_api_key_auth, oauth2_headers, _client_ip + ) + if not await _enforce_stateful_session_cap_for_owner(request_owner): + verbose_logger.warning( + "Rejecting MCP initialize: caller already holds the maximum " + "number of active stateful sessions." + ) + too_many_response = JSONResponse( + status_code=429, + content={ + "error": "Too Many Requests", + "details": "Too many active MCP sessions for this caller.", + }, + ) + await too_many_response(scope, receive, send) + return + + # Replay body messages if we consumed them for peeking + original_receive = receive + if consumed_messages: + + async def wrapped_receive(): + if consumed_messages: + return consumed_messages.pop(0) + return await original_receive() + + receive = wrapped_receive + + # Serialize requests on the same stateful session so concurrent + # callers don't clobber each other's auth context mid-flight. + # + # Skip the lock for streaming GETs (SSE channels held open for the + # life of the session): holding a per-session lock for a long-lived + # stream would block every subsequent POST on the same session. + # POST/DELETE are the methods that actually mutate the shared + # auth context, so serializing those is sufficient for the + # clobbering race between concurrent JSON-RPC calls. + session_lock: Optional[asyncio.Lock] = None + request_method = (scope.get("method") or "").upper() + if use_stateful and session_id and request_method in ("POST", "DELETE"): + session_lock = _stateful_session_locks.setdefault( + session_id, asyncio.Lock() + ) + + active_request_session_ids: List[str] = [] + + def _increment_active_request_session(session_id_to_track: str) -> None: + if session_id_to_track in active_request_session_ids: + return + active_request_session_ids.append(session_id_to_track) + _stateful_session_active_request_counts[session_id_to_track] = ( + _stateful_session_active_request_counts.get(session_id_to_track, 0) + + 1 + ) + + if use_stateful and session_id: + _increment_active_request_session(session_id) + + def _track_initialized_stateful_session( + initialized_session_id: str, + ) -> None: + _increment_active_request_session(initialized_session_id) + + async def _dispatch() -> None: + auth_user = _set_or_update_auth_context( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + session_id=session_id if use_stateful else None, + touch_last_seen=(scope.get("method") or "").upper() != "DELETE", + copy_existing_session_auth_context=is_initialize, + ) + local_send = send + if use_stateful and is_initialize: + local_send = _wrap_send_with_stateful_session_auth_context( + local_send, + auth_user, + _owner_fingerprint_for( + user_api_key_auth, oauth2_headers, _client_ip + ), + _track_initialized_stateful_session, + ) + + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth, + mcp_servers, + _client_ip, + ): + await target_manager.handle_request(scope, receive, local_send) + if use_stateful and session_id and scope.get("method") == "DELETE": + _remove_stateful_session_tracking(session_id) + + try: + if session_lock is not None: + async with session_lock: + await _dispatch() + else: + await _dispatch() + finally: + for active_request_session_id in active_request_session_ids: + active_request_count = ( + _stateful_session_active_request_counts.get( + active_request_session_id, 0 + ) + - 1 + ) + if active_request_count > 0: + _stateful_session_active_request_counts[ + active_request_session_id + ] = active_request_count + else: + _stateful_session_active_request_counts.pop( + active_request_session_id, None + ) + + if ( + scope.get("method") != "DELETE" + and active_request_session_id in _stateful_session_auth_contexts + ): + _stateful_session_auth_context_last_seen[ + active_request_session_id + ] = time.monotonic() + + # Periodic cleanup iterates _stateful_session_auth_context_last_seen, + # so locks for untracked sessions must be dropped here. + if ( + active_request_count <= 0 + and active_request_session_id + not in _stateful_session_auth_contexts + ): + _stateful_session_locks.pop(active_request_session_id, None) except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise @@ -3125,7 +3590,6 @@ if MCP_AVAILABLE: verbose_logger.exception(f"Error handling MCP request: {e}") # Try to send a graceful error response for non-HTTP exceptions try: - from starlette.responses import JSONResponse from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR error_response = JSONResponse( @@ -3231,8 +3695,9 @@ if MCP_AVAILABLE: ############ Auth Context Functions #################### ######################################################## - def set_auth_context( - user_api_key_auth: UserAPIKeyAuth, + def _update_auth_context( + auth_user: MCPAuthenticatedUser, + user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, @@ -3240,6 +3705,23 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, ) -> None: + auth_user.user_api_key_auth = user_api_key_auth + auth_user.mcp_auth_header = mcp_auth_header + auth_user.mcp_servers = mcp_servers + auth_user.mcp_server_auth_headers = mcp_server_auth_headers or {} + auth_user.oauth2_headers = oauth2_headers + auth_user.raw_headers = raw_headers + auth_user.client_ip = client_ip + + def set_auth_context( + user_api_key_auth: Optional[UserAPIKeyAuth], + mcp_auth_header: Optional[str] = None, + mcp_servers: Optional[List[str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, + oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, + ) -> MCPAuthenticatedUser: """ Set the UserAPIKeyAuth in the auth context variable. @@ -3260,6 +3742,84 @@ if MCP_AVAILABLE: client_ip=client_ip, ) auth_context_var.set(auth_user) + return auth_user + + def _set_or_update_auth_context( + user_api_key_auth: Optional[UserAPIKeyAuth], + mcp_auth_header: Optional[str] = None, + mcp_servers: Optional[List[str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, + oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, + session_id: Optional[str] = None, + touch_last_seen: bool = True, + copy_existing_session_auth_context: bool = False, + ) -> MCPAuthenticatedUser: + auth_user = ( + _stateful_session_auth_contexts.get(session_id) if session_id else None + ) + if auth_user is not None and session_id is not None: + if touch_last_seen: + _stateful_session_auth_context_last_seen[session_id] = time.monotonic() + if copy_existing_session_auth_context: + return set_auth_context( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + _update_auth_context( + auth_user=auth_user, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + auth_context_var.set(auth_user) + return auth_user + return set_auth_context( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + + def _wrap_send_with_stateful_session_auth_context( + send: Send, + auth_user: MCPAuthenticatedUser, + owner_fingerprint: str, + on_session_registered: Optional[Callable[[str], None]] = None, + ) -> Send: + async def wrapped_send(message: Message) -> None: + if message.get("type") == "http.response.start": + for key, value in message.get("headers", []): + header_name = key if isinstance(key, bytes) else str(key).encode() + if header_name.lower() == b"mcp-session-id": + session_id = ( + value.decode() if isinstance(value, bytes) else str(value) + ) + if on_session_registered is not None: + on_session_registered(session_id) + auth_context_var.set(auth_user) + _stateful_session_auth_contexts[session_id] = auth_user + _stateful_session_auth_context_last_seen[session_id] = ( + time.monotonic() + ) + _stateful_session_owners[session_id] = owner_fingerprint + break + await send(message) + + return wrapped_send def get_auth_context() -> Tuple[ Optional[UserAPIKeyAuth], diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index c20fb09eeb..27df05225e 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -400,11 +400,11 @@ async def test_mcp_http_transport_tool_not_found(): @pytest.mark.asyncio async def test_streamable_http_mcp_handler_mock(): """Test the streamable HTTP MCP handler functionality""" - from litellm.proxy._types import UserAPIKeyAuth - - # Mock the session manager and its methods - mock_session_manager = AsyncMock() - mock_session_manager.handle_request = AsyncMock() + # Mock streamable HTTP session managers and their methods + mock_session_manager_stateless = AsyncMock() + mock_session_manager_stateless.handle_request = AsyncMock() + mock_session_manager_stateful = AsyncMock() + mock_session_manager_stateful.handle_request = AsyncMock() # Mock scope, receive, send with proper ASGI scope format mock_scope = { @@ -416,7 +416,7 @@ async def test_streamable_http_mcp_handler_mock(): "server": ("localhost", 8000), "scheme": "http", } - mock_receive = AsyncMock() + mock_receive = AsyncMock(return_value={"body": b"{}", "more_body": False}) mock_send = AsyncMock() # Mock extract_mcp_auth_context to bypass auth checks in the handler @@ -428,8 +428,12 @@ async def test_streamable_http_mcp_handler_mock(): True, ), patch( - "litellm.proxy._experimental.mcp_server.server.session_manager", - mock_session_manager, + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + mock_session_manager_stateless, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + mock_session_manager_stateful, ), patch( "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", @@ -446,8 +450,9 @@ async def test_streamable_http_mcp_handler_mock(): # Call the handler await handle_streamable_http_mcp(mock_scope, mock_receive, mock_send) - # Verify session manager handle_request was called - mock_session_manager.handle_request.assert_called_once() + # Verify stateless session manager handle_request was called + mock_session_manager_stateless.handle_request.assert_called_once() + mock_session_manager_stateful.handle_request.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index ceefc41052..a7d9ce64f8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -2269,7 +2269,11 @@ def test_mcp_path_based_server_segregation(monkeypatch): ) monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.server.session_manager", + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + MagicMock(handle_request=dummy_handle_request), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", MagicMock(handle_request=dummy_handle_request), ) monkeypatch.setattr( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index d62720ed36..fb21e4ee11 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,4 +1,6 @@ import asyncio +import contextlib +import contextvars from datetime import datetime, timedelta from unittest.mock import AsyncMock, MagicMock, patch @@ -1002,18 +1004,25 @@ async def test_concurrent_initialize_session_managers(): # Reset state before test original_initialized = mcp_server._SESSION_MANAGERS_INITIALIZED original_session_cm = mcp_server._session_manager_cm + original_session_stateful_cm = mcp_server._session_manager_stateful_cm original_sse_session_cm = mcp_server._sse_session_manager_cm + original_cleanup_task = mcp_server._stateful_auth_context_cleanup_task try: mcp_server._SESSION_MANAGERS_INITIALIZED = False mcp_server._session_manager_cm = None + mcp_server._session_manager_stateful_cm = None mcp_server._sse_session_manager_cm = None + mcp_server._stateful_auth_context_cleanup_task = None # Mock the session managers to avoid actual MCP initialization with ( patch( - "litellm.proxy._experimental.mcp_server.server.session_manager" - ) as mock_session_manager, + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless" + ) as mock_session_manager_stateless, + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful" + ) as mock_session_manager_stateful, patch( "litellm.proxy._experimental.mcp_server.server.sse_session_manager" ) as mock_sse_session_manager, @@ -1024,7 +1033,8 @@ async def test_concurrent_initialize_session_managers(): mock_cm.__aenter__ = AsyncMock() mock_cm.__aexit__ = AsyncMock() - mock_session_manager.run.return_value = mock_cm + mock_session_manager_stateless.run.return_value = mock_cm + mock_session_manager_stateful.run.return_value = mock_cm mock_sse_session_manager.run.return_value = mock_cm # Create multiple concurrent tasks that call initialize_session_managers @@ -1041,52 +1051,1423 @@ async def test_concurrent_initialize_session_managers(): result == "success" for result in results ), f"Some tasks failed: {results}" - # session_manager.run() should only be called once due to the lock + # Each session manager.run() should only be called once due to the lock assert ( - mock_session_manager.run.call_count == 1 - ), f"Expected 1 call to session_manager.run(), got {mock_session_manager.run.call_count}" + mock_session_manager_stateless.run.call_count == 1 + ), f"Expected 1 call to session_manager_stateless.run(), got {mock_session_manager_stateless.run.call_count}" + assert ( + mock_session_manager_stateful.run.call_count == 1 + ), f"Expected 1 call to session_manager_stateful.run(), got {mock_session_manager_stateful.run.call_count}" assert ( mock_sse_session_manager.run.call_count == 1 ), f"Expected 1 call to sse_session_manager.run(), got {mock_sse_session_manager.run.call_count}" - # The context managers should only be entered once each + # The context managers should only be entered once each (3 managers) assert ( - mock_cm.__aenter__.call_count == 2 - ), f"Expected 2 calls to __aenter__ (one for each session manager), got {mock_cm.__aenter__.call_count}" + mock_cm.__aenter__.call_count == 3 + ), f"Expected 3 calls to __aenter__ (one per session manager), got {mock_cm.__aenter__.call_count}" # State should be properly set assert mcp_server._SESSION_MANAGERS_INITIALIZED is True finally: + # Cancel the background cleanup task that initialize_session_managers() + # spawned. Otherwise it keeps running against module-level dicts for the + # rest of the test session (asyncio_default_fixture_loop_scope=session). + leaked_task = mcp_server._stateful_auth_context_cleanup_task + if leaked_task is not None and leaked_task is not original_cleanup_task: + leaked_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await leaked_task + # Restore original state mcp_server._SESSION_MANAGERS_INITIALIZED = original_initialized mcp_server._session_manager_cm = original_session_cm + mcp_server._session_manager_stateful_cm = original_session_stateful_cm mcp_server._sse_session_manager_cm = original_sse_session_cm + mcp_server._stateful_auth_context_cleanup_task = original_cleanup_task @pytest.mark.asyncio async def test_streamable_http_session_manager_is_stateless(): """ - Test that the StreamableHTTPSessionManager is initialized with stateless=True. + Test that the StreamableHTTPSessionManager is initialized with both stateless and stateful managers. Regression test for GitHub issue #20242 / PR #19809. When stateless=False, the mcp library rejects non-initialize requests that lack an mcp-session-id header, breaking clients like MCP Inspector, curl, and any HTTP client without automatic session management. + + Now we support both: + - stateless manager for clients without session IDs (curl, Inspector) + - stateful manager for clients with session IDs (Claude Code, Cursor, VSCode) """ try: - from litellm.proxy._experimental.mcp_server.server import session_manager + from litellm.proxy._experimental.mcp_server.server import ( + session_manager_stateful, + session_manager_stateless, + ) except ImportError: pytest.skip("MCP server not available") - # The session manager must be stateless to avoid requiring mcp-session-id + # The stateless session manager must be stateless to avoid requiring mcp-session-id # on every request. This was regressed by PR #19809 (stateless=True -> False). - assert session_manager.stateless is True, ( - "StreamableHTTPSessionManager must be initialized with stateless=True. " + assert session_manager_stateless.stateless is True, ( + "session_manager_stateless must be initialized with stateless=True. " "stateless=False breaks MCP clients that don't manage session IDs. " "See: https://github.com/BerriAI/litellm/issues/20242" ) + # The stateful session manager must be stateful to support progress notifications + assert session_manager_stateful.stateless is False, ( + "session_manager_stateful must be initialized with stateless=False. " + "stateless=True breaks progress notifications for clients that manage session IDs." + ) + + +@pytest.mark.asyncio +async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(): + """ + Test that routing correctly sends: + - initialize (no mcp-session-id) → stateful manager (so client gets mcp-session-id) + - tools/list (no mcp-session-id) → stateless manager (curl, Inspector) + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + async def make_request(method_body: bytes, path: str = "/mcp/progress_test"): + scope = { + "type": "http", + "method": "POST", + "path": path, + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": method_body, + "more_body": False, + } + ) + send = AsyncMock() + + stateless_called = [] + stateful_called = [] + + async def stateless_handle(s, r, se): + stateless_called.append(1) + + async def stateful_handle(s, r, se): + stateful_called.append(1) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, ["progress_test"], None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.set_auth_context", + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateless, + "handle_request", + side_effect=stateless_handle, + ), + patch.object( + session_manager_stateful, + "handle_request", + side_effect=stateful_handle, + ), + patch.object( + session_manager_stateless, + "_server_instances", + {}, + ), + patch.object( + session_manager_stateful, + "_server_instances", + {}, + ), + ): + await handle_streamable_http_mcp(scope, receive, send) + + return bool(stateless_called), bool(stateful_called) + + # initialize → stateful + init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}' + stateless_called, stateful_called = await make_request(init_body) + assert ( + stateful_called and not stateless_called + ), "initialize (no session) should route to stateful, not stateless" + + # tools/list → stateless + tools_body = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}' + stateless_called, stateful_called = await make_request(tools_body) + assert ( + stateless_called and not stateful_called + ), "tools/list (no session) should route to stateless, not stateful" + + +@pytest.mark.asyncio +async def test_mcp_routing_chunked_initialize_to_stateful(): + """ + Test that chunked initialize requests route to the stateful manager. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/progress_test", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + messages = [ + { + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":1,', + "more_body": True, + }, + { + "type": "http.request", + "body": b'"method":"initialize","params":{}}', + "more_body": False, + }, + ] + receive = AsyncMock(side_effect=messages) + send = AsyncMock() + stateless_called = [] + stateful_called = [] + + async def stateless_handle(s, r, se): + stateless_called.append(1) + + async def stateful_handle(s, r, se): + stateful_called.append(1) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, ["progress_test"], None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.set_auth_context", + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateless, + "handle_request", + side_effect=stateless_handle, + ), + patch.object( + session_manager_stateful, + "handle_request", + side_effect=stateful_handle, + ), + patch.object( + session_manager_stateless, + "_server_instances", + {}, + ), + patch.object( + session_manager_stateful, + "_server_instances", + {}, + ), + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert ( + stateful_called and not stateless_called + ), "chunked initialize (no session) should route to stateful, not stateless" + + +@pytest.mark.asyncio +async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body(): + """ + A no-session-id POST with a very large chunked body should not force + the proxy to buffer the entire body just to decide routing — the peek + should stop once ``_MCP_ROUTING_PEEK_MAX_BYTES`` worth of body has been + consumed, and the remaining chunks should stream through the original + receive into the downstream handler. + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + peek_cap = mcp_server._MCP_ROUTING_PEEK_MAX_BYTES + # First chunk fills the peek budget; subsequent chunks are oversized payload. + first_chunk = b"x" * peek_cap + oversized_tail = [b"y" * 65536 for _ in range(4)] + + messages = [ + {"type": "http.request", "body": first_chunk, "more_body": True}, + *[ + {"type": "http.request", "body": chunk, "more_body": True} + for chunk in oversized_tail + ], + {"type": "http.request", "body": b"", "more_body": False}, + ] + receive_calls = {"count": 0} + + async def receive(): + idx = receive_calls["count"] + receive_calls["count"] += 1 + return messages[idx] + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/progress_test", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + send = AsyncMock() + + stateless_received_chunks = [] + receive_count_at_dispatch = {"value": -1} + + async def stateless_handle(s, r, se): + # Snapshot how many wire reads happened BEFORE dispatch — the cap + # check is meaningful only against pre-dispatch consumption. + receive_count_at_dispatch["value"] = receive_calls["count"] + # Drain the wrapped receive the same way the SDK would. + while True: + msg = await r() + if msg.get("type") != "http.request": + break + stateless_received_chunks.append(msg.get("body", b"") or b"") + if not msg.get("more_body", False): + break + + async def stateful_handle(s, r, se): + raise AssertionError("non-initialize POST should not reach stateful manager") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, ["progress_test"], None, None, None), + ), + patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateless, "handle_request", side_effect=stateless_handle + ), + patch.object( + session_manager_stateful, "handle_request", side_effect=stateful_handle + ), + patch.object(session_manager_stateless, "_server_instances", {}), + patch.object(session_manager_stateful, "_server_instances", {}), + ): + await handle_streamable_http_mcp(scope, receive, send) + + # The routing peek must stop pulling from the wire once the cap is reached. + # Without the cap fix, every chunk would have been pulled before dispatch, + # so this assertion guards against unbounded pre-dispatch buffering. + assert receive_count_at_dispatch["value"] == 1, ( + "routing should stop reading after the peek cap is filled, " + f"but consumed {receive_count_at_dispatch['value']} chunks before dispatching" + ) + # All chunks must still reach the downstream handler via replay+stream. + total_streamed = sum(len(b) for b in stateless_received_chunks) + assert total_streamed == len(first_chunk) + sum(len(b) for b in oversized_tail) + + +@pytest.mark.asyncio +async def test_enforce_stateful_session_cap_evicts_oldest_idle_then_rejects(): + """ + A caller at the per-owner session cap should have its own oldest *idle* + session evicted to make room for a new one, but be rejected outright when + every one of its sessions is in flight (nothing safe to evict). + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + terminated = [] + + class FakeTransport: + def __init__(self, session_id): + self.session_id = session_id + + async def terminate(self): + terminated.append(self.session_id) + + instances = {f"s{i}": FakeTransport(f"s{i}") for i in range(3)} + owners = {f"s{i}": "owner-A" for i in range(3)} + last_seen = {"s0": 1.0, "s1": 2.0, "s2": 3.0} + contexts = {f"s{i}": MagicMock() for i in range(3)} + + with ( + patch.object(session_manager_stateful, "_server_instances", instances), + patch.object(mcp_server, "_MAX_STATEFUL_SESSIONS_PER_OWNER", 3), + patch.dict(mcp_server._stateful_session_owners, owners, clear=True), + patch.dict( + mcp_server._stateful_session_auth_context_last_seen, last_seen, clear=True + ), + patch.dict(mcp_server._stateful_session_auth_contexts, contexts, clear=True), + patch.dict(mcp_server._stateful_session_active_request_counts, {}, clear=True), + ): + # All idle -> oldest (s0) is evicted, request may proceed. + allowed = await mcp_server._enforce_stateful_session_cap_for_owner("owner-A") + assert allowed is True + assert terminated == ["s0"] + assert "s0" not in instances + assert "s0" not in mcp_server._stateful_session_owners + + # A different owner at the cap is unaffected by owner-A's sessions. + terminated.clear() + allowed_other = await mcp_server._enforce_stateful_session_cap_for_owner( + "owner-B" + ) + assert allowed_other is True + assert terminated == [] + + # Now every session is in flight -> nothing evictable -> reject. + terminated.clear() + instances = {f"s{i}": FakeTransport(f"s{i}") for i in range(3)} + owners = {f"s{i}": "owner-A" for i in range(3)} + active = {f"s{i}": 1 for i in range(3)} + + with ( + patch.object(session_manager_stateful, "_server_instances", instances), + patch.object(mcp_server, "_MAX_STATEFUL_SESSIONS_PER_OWNER", 3), + patch.dict(mcp_server._stateful_session_owners, owners, clear=True), + patch.dict( + mcp_server._stateful_session_auth_context_last_seen, + {f"s{i}": float(i) for i in range(3)}, + clear=True, + ), + patch.dict( + mcp_server._stateful_session_active_request_counts, active, clear=True + ), + ): + rejected = await mcp_server._enforce_stateful_session_cap_for_owner("owner-A") + assert rejected is False + assert terminated == [] + assert len(instances) == 3 + + +@pytest.mark.asyncio +async def test_mcp_routing_initialize_rejected_when_owner_at_session_cap(): + """ + A new ``initialize`` (no session id) must be rejected with 429 when the + caller already holds the maximum number of in-flight stateful sessions, + and must not reach the stateful session manager. + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + cap = 2 + + class FakeTransport: + async def terminate(self): + pass + + instances = {f"s{i}": FakeTransport() for i in range(cap)} + owners = {f"s{i}": "owner-X" for i in range(cap)} + active = {f"s{i}": 1 for i in range(cap)} # all in flight -> cannot evict + contexts = {f"s{i}": MagicMock() for i in range(cap)} + + init_body = ( + b'{"jsonrpc":"2.0","id":1,"method":"initialize",' + b'"params":{"protocolVersion":"2024-11-05"}}' + ) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/progress_test", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + receive = AsyncMock( + return_value={"type": "http.request", "body": init_body, "more_body": False} + ) + send = AsyncMock() + + stateful_called = [] + + async def stateful_handle(s, r, se): + stateful_called.append(1) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, ["progress_test"], None, None, None), + ), + patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object(mcp_server, "_owner_fingerprint_for", return_value="owner-X"), + patch.object(mcp_server, "_MAX_STATEFUL_SESSIONS_PER_OWNER", cap), + patch.object( + session_manager_stateful, "handle_request", side_effect=stateful_handle + ), + patch.object(session_manager_stateful, "_server_instances", instances), + patch.object(session_manager_stateless, "_server_instances", {}), + patch.dict(mcp_server._stateful_session_owners, owners, clear=True), + patch.dict( + mcp_server._stateful_session_auth_context_last_seen, + {f"s{i}": float(i) for i in range(cap)}, + clear=True, + ), + patch.dict( + mcp_server._stateful_session_active_request_counts, active, clear=True + ), + patch.dict(mcp_server._stateful_session_auth_contexts, contexts, clear=True), + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert not stateful_called, "initialize at session cap must not reach the manager" + start_messages = [ + call.args[0] + for call in send.call_args_list + if call.args and call.args[0].get("type") == "http.response.start" + ] + assert start_messages, "a response should have been sent" + assert start_messages[0]["status"] == 429 + + +@pytest.mark.asyncio +async def test_stateful_mcp_requests_refresh_session_auth_context(): + """ + Stateful MCP sessions run callbacks in the initialize task's context; the + stored auth object must be refreshed for each mcp-session-id request. + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + get_auth_context, + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "stateful-session-1" + initialize_auth = UserAPIKeyAuth(api_key="initialize-key", user_id="user-a") + current_auth = UserAPIKeyAuth(api_key="current-key", user_id="user-b") + callback_context = contextvars.copy_context() + callback_context.run( + mcp_server.set_auth_context, + initialize_auth, + None, + ["old-server"], + None, + None, + None, + "1.1.1.1", + ) + mcp_server._stateful_session_auth_contexts[session_id] = callback_context.run( + mcp_server.auth_context_var.get + ) + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/current-server", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer current-key"), + (b"mcp-session-id", session_id.encode()), + ], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', + "more_body": False, + } + ) + send = AsyncMock() + + captured_context = None + + async def stateful_handle(s, r, se): + nonlocal captured_context + captured_context = callback_context.run(get_auth_context) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=( + current_auth, + "current-mcp-auth", + ["current-server"], + {"current-server": {"Authorization": "Bearer server-key"}}, + {"Authorization": "Bearer oauth-key"}, + {"mcp-session-id": session_id}, + ), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, + "handle_request", + side_effect=stateful_handle, + ), + patch.object( + session_manager_stateful, + "_server_instances", + {session_id: MagicMock()}, + ), + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert captured_context == ( + current_auth, + "current-mcp-auth", + ["current-server"], + {"current-server": {"Authorization": "Bearer server-key"}}, + {"Authorization": "Bearer oauth-key"}, + {"mcp-session-id": session_id}, + "", + ) + mcp_server._remove_stateful_session_tracking(session_id) + + +@pytest.mark.asyncio +async def test_initialize_response_capture_accepts_str_headers_and_sets_auth_context(): + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + session_id = "initialize-session-1" + auth_user = mcp_server.MCPAuthenticatedUser( + user_api_key_auth=UserAPIKeyAuth(api_key="initialize-key", user_id="user-a") + ) + previous_auth_user = mcp_server.MCPAuthenticatedUser( + user_api_key_auth=UserAPIKeyAuth(api_key="previous-key", user_id="user-b") + ) + sent_messages = [] + + async def send(message): + sent_messages.append(message) + + wrapped_send = mcp_server._wrap_send_with_stateful_session_auth_context( + send, + auth_user, + "owner-fingerprint", + ) + token = mcp_server.auth_context_var.set(previous_auth_user) + try: + await wrapped_send( + { + "type": "http.response.start", + "headers": [("mcp-session-id", session_id)], + } + ) + + assert mcp_server.auth_context_var.get() is auth_user + assert mcp_server._stateful_session_auth_contexts[session_id] is auth_user + assert mcp_server._stateful_session_owners[session_id] == "owner-fingerprint" + assert session_id in mcp_server._stateful_session_auth_context_last_seen + assert sent_messages == [ + { + "type": "http.response.start", + "headers": [("mcp-session-id", session_id)], + } + ] + finally: + mcp_server.auth_context_var.reset(token) + mcp_server._remove_stateful_session_tracking(session_id) + + +@pytest.mark.asyncio +async def test_initialize_request_tracks_active_session_after_response_header(): + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "initialize-active-session-1" + owner_auth = UserAPIKeyAuth(api_key="initialize-key", user_id="user-a") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer initialize-key"), + ], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', + "more_body": False, + } + ) + + async def stateful_handle(s, r, se): + await se( + { + "type": "http.response.start", + "headers": [(b"mcp-session-id", session_id.encode())], + } + ) + assert mcp_server._stateful_session_active_request_counts[session_id] == 1 + now = ( + mcp_server._stateful_session_auth_context_last_seen[session_id] + + mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + ) + await mcp_server._purge_expired_stateful_session_auth_contexts(now=now) + assert session_id in mcp_server._stateful_session_auth_contexts + + async def stateless_handle(s, r, se): + raise AssertionError("initialize request should use stateful manager") + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(owner_auth, None, None, None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, + "handle_request", + side_effect=stateful_handle, + ), + patch.object( + session_manager_stateless, + "handle_request", + side_effect=stateless_handle, + ), + patch.object(session_manager_stateful, "_server_instances", {}), + ): + await handle_streamable_http_mcp(scope, receive, AsyncMock()) + + assert session_id not in mcp_server._stateful_session_active_request_counts + assert session_id in mcp_server._stateful_session_auth_contexts + finally: + mcp_server._remove_stateful_session_tracking(session_id) + + +@pytest.mark.asyncio +async def test_initialize_request_with_existing_session_tracks_new_session(): + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + existing_session_id = "existing-initialize-session" + new_session_id = "reinitialized-session" + owner_auth = UserAPIKeyAuth(api_key="initialize-key", user_id="user-a") + owner_fingerprint = mcp_server._owner_fingerprint_for(owner_auth) + existing_auth_user = mcp_server.MCPAuthenticatedUser( + user_api_key_auth=owner_auth, + mcp_auth_header="old-mcp-auth", + mcp_servers=["old-server"], + mcp_server_auth_headers={"old-server": {"Authorization": "Bearer old-key"}}, + oauth2_headers={"Authorization": "Bearer old-oauth"}, + raw_headers={"x-old-header": "old"}, + client_ip="old-client-ip", + ) + initialize_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}' + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer initialize-key"), + (b"mcp-session-id", existing_session_id.encode()), + ], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": initialize_body, + "more_body": False, + } + ) + stateful_called = [] + + async def stateful_handle(s, r, se): + stateful_called.append(1) + message = await r() + assert message["body"] == initialize_body + await se( + { + "type": "http.response.start", + "headers": [(b"mcp-session-id", new_session_id.encode())], + } + ) + assert mcp_server._stateful_session_auth_contexts[new_session_id] + assert mcp_server._stateful_session_owners[new_session_id] == owner_fingerprint + assert mcp_server._stateful_session_active_request_counts[new_session_id] == 1 + now = ( + mcp_server._stateful_session_auth_context_last_seen[new_session_id] + + mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + ) + await mcp_server._purge_expired_stateful_session_auth_contexts(now=now) + assert new_session_id in mcp_server._stateful_session_auth_contexts + assert ( + mcp_server._stateful_session_auth_contexts[new_session_id] + is not existing_auth_user + ) + assert ( + mcp_server._stateful_session_auth_contexts[new_session_id].mcp_auth_header + == "new-mcp-auth" + ) + + async def stateless_handle(s, r, se): + raise AssertionError( + "initialize request with session should use stateful manager" + ) + + try: + mcp_server._stateful_session_auth_contexts[existing_session_id] = ( + existing_auth_user + ) + mcp_server._stateful_session_auth_context_last_seen[existing_session_id] = 1.0 + mcp_server._stateful_session_owners[existing_session_id] = owner_fingerprint + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=( + owner_auth, + "new-mcp-auth", + ["new-server"], + {"new-server": {"Authorization": "Bearer new-key"}}, + {"Authorization": "Bearer new-oauth"}, + {"x-new-header": "new"}, + ), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, + "handle_request", + side_effect=stateful_handle, + ), + patch.object( + session_manager_stateless, + "handle_request", + side_effect=stateless_handle, + ), + patch.object( + session_manager_stateful, + "_server_instances", + {existing_session_id: MagicMock()}, + ), + ): + await handle_streamable_http_mcp(scope, receive, AsyncMock()) + + assert stateful_called + assert new_session_id not in mcp_server._stateful_session_active_request_counts + assert new_session_id in mcp_server._stateful_session_auth_contexts + assert ( + mcp_server._stateful_session_auth_contexts[existing_session_id] + is existing_auth_user + ) + assert existing_auth_user.mcp_auth_header == "old-mcp-auth" + assert existing_auth_user.mcp_servers == ["old-server"] + finally: + mcp_server._remove_stateful_session_tracking(existing_session_id) + mcp_server._remove_stateful_session_tracking(new_session_id) + + +@pytest.mark.asyncio +async def test_stateful_mcp_auth_contexts_expire_with_idle_sessions(): + """Expired session auth contexts should not remain in memory indefinitely.""" + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + session_id = "expired-stateful-session" + auth_user = UserAPIKeyAuth(api_key="expired-key", user_id="expired-user") + transport = MagicMock() + transport.terminate = AsyncMock() + now = 1000.0 + + mcp_server._stateful_session_auth_contexts[session_id] = auth_user + mcp_server._stateful_session_auth_context_last_seen[session_id] = ( + now - mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + ) + + with patch.object( + mcp_server.session_manager_stateful, + "_server_instances", + {session_id: transport}, + ): + await mcp_server._purge_expired_stateful_session_auth_contexts(now=now) + + assert session_id not in mcp_server._stateful_session_auth_contexts + assert session_id not in mcp_server._stateful_session_auth_context_last_seen + transport.terminate.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_stateful_mcp_auth_contexts_do_not_expire_active_sessions(): + """Active stateful sessions should not be terminated by idle cleanup.""" + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + session_id = "active-stateful-session" + auth_user = UserAPIKeyAuth(api_key="active-key", user_id="active-user") + transport = MagicMock() + transport.terminate = AsyncMock() + now = 1000.0 + + mcp_server._stateful_session_auth_contexts[session_id] = auth_user + mcp_server._stateful_session_auth_context_last_seen[session_id] = ( + now - mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + ) + mcp_server._stateful_session_active_request_counts[session_id] = 1 + + try: + with patch.object( + mcp_server.session_manager_stateful, + "_server_instances", + {session_id: transport}, + ): + await mcp_server._purge_expired_stateful_session_auth_contexts(now=now) + + assert session_id in mcp_server._stateful_session_auth_contexts + assert session_id in mcp_server._stateful_session_auth_context_last_seen + transport.terminate.assert_not_awaited() + finally: + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_auth_context_last_seen.pop(session_id, None) + mcp_server._stateful_session_active_request_counts.pop(session_id, None) + + +@pytest.mark.asyncio +async def test_stateful_mcp_auth_context_cleanup_respects_zero_now(): + """Explicit now=0 should be used as-is instead of falling back to monotonic.""" + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + session_id = "zero-now-stateful-session" + auth_user = UserAPIKeyAuth(api_key="zero-now-key", user_id="zero-now-user") + transport = MagicMock() + transport.terminate = AsyncMock() + + mcp_server._stateful_session_auth_contexts[session_id] = auth_user + mcp_server._stateful_session_auth_context_last_seen[session_id] = 0.0 + + try: + with ( + patch.object( + mcp_server.session_manager_stateful, + "_server_instances", + {session_id: transport}, + ), + patch.object( + mcp_server.time, + "monotonic", + return_value=mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + 1, + ), + ): + await mcp_server._purge_expired_stateful_session_auth_contexts(now=0.0) + + assert session_id in mcp_server._stateful_session_auth_contexts + assert session_id in mcp_server._stateful_session_auth_context_last_seen + transport.terminate.assert_not_awaited() + finally: + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_auth_context_last_seen.pop(session_id, None) + + +@pytest.mark.asyncio +async def test_stateful_mcp_cleanup_loop_survives_purge_errors(): + """Cleanup loop should keep running after one purge attempt fails.""" + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + purge = AsyncMock( + side_effect=[RuntimeError("terminate failed"), asyncio.CancelledError()] + ) + + with ( + patch.object(mcp_server.asyncio, "sleep", AsyncMock(return_value=None)), + patch.object( + mcp_server, "_purge_expired_stateful_session_auth_contexts", purge + ), + ): + with pytest.raises(asyncio.CancelledError): + await mcp_server._cleanup_expired_stateful_session_auth_contexts() + + assert purge.await_count == 2 + + +@pytest.mark.asyncio +async def test_owner_fingerprint_distinguishes_oauth_callers(): + """ + OAuth2 passthrough callers all share `UserAPIKeyAuth()` with no api_key + or user_id. Without folding the upstream bearer into the fingerprint + they would all collapse to a single 'anonymous' owner and one OAuth + user could hijack another's mcp-session-id. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + _owner_fingerprint_for, + ) + except ImportError: + pytest.skip("MCP server not available") + + anon_auth = UserAPIKeyAuth() + fp_a = _owner_fingerprint_for(anon_auth, {"Authorization": "Bearer token-A"}) + fp_b = _owner_fingerprint_for(anon_auth, {"Authorization": "Bearer token-B"}) + fp_a_again = _owner_fingerprint_for(anon_auth, {"authorization": "Bearer token-A"}) + fp_no_oauth = _owner_fingerprint_for(anon_auth, None) + + assert fp_a != fp_b + assert fp_a == fp_a_again + assert fp_a.startswith("oauth:") + assert fp_no_oauth == "anonymous" + assert "Bearer token-A" not in fp_a + + # When no API key, user_id, or OAuth bearer is available, fall back to + # client IP so two unrelated unauthenticated callers from different + # sources don't collapse to a single 'anonymous' owner and end up able + # to drive each other's stateful sessions. + fp_ip_a = _owner_fingerprint_for(anon_auth, None, "10.0.0.1") + fp_ip_b = _owner_fingerprint_for(anon_auth, None, "10.0.0.2") + fp_ip_a_again = _owner_fingerprint_for(anon_auth, None, "10.0.0.1") + + assert fp_ip_a != fp_ip_b + assert fp_ip_a == fp_ip_a_again + assert fp_ip_a.startswith("ip:") + assert "10.0.0.1" not in fp_ip_a + + +@pytest.mark.asyncio +async def test_owner_fingerprint_hashes_custom_api_keys(): + """Custom API key formats should not appear in owner fingerprints.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _owner_fingerprint_for, + ) + except ImportError: + pytest.skip("MCP server not available") + + auth = UserAPIKeyAuth(api_key="custom-master-key") + fp = _owner_fingerprint_for(auth) + fp_again = _owner_fingerprint_for(auth) + + assert fp == fp_again + assert fp.startswith("key:") + assert "custom-master-key" not in fp + assert fp != "key:custom-master-key" + + +@pytest.mark.asyncio +async def test_stateful_mcp_session_owner_mismatch_returns_403(): + """ + A stateful mcp-session-id is bound to its creator. A different + authenticated caller presenting the same session_id must be rejected + with 403, and the stateful manager must never be invoked. + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "owned-session-1" + owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner") + intruder_auth = UserAPIKeyAuth(api_key="intruder-key", user_id="intruder") + + mcp_server._stateful_session_auth_contexts[session_id] = MagicMock() + mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for( + owner_auth + ) + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer intruder-key"), + (b"mcp-session-id", session_id.encode()), + ], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', + "more_body": False, + } + ) + sent_messages: list = [] + + async def capture_send(message): + sent_messages.append(message) + + handle_request_mock = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(intruder_auth, None, None, None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, + "handle_request", + side_effect=handle_request_mock, + ), + patch.object( + session_manager_stateful, + "_server_instances", + {session_id: MagicMock()}, + ), + ): + await handle_streamable_http_mcp(scope, receive, capture_send) + + handle_request_mock.assert_not_awaited() + statuses = [ + m["status"] for m in sent_messages if m.get("type") == "http.response.start" + ] + assert statuses == [403] + + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_owners.pop(session_id, None) + + +@pytest.mark.asyncio +async def test_stateful_mcp_session_serializes_concurrent_requests(): + """ + Concurrent requests on the same stateful mcp-session-id must be + serialized so they cannot observe each other's mutation of the shared + MCPAuthenticatedUser while in-flight callbacks are still running. + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "serialized-session-1" + owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner") + mcp_server._stateful_session_auth_contexts[session_id] = ( + mcp_server.MCPAuthenticatedUser(user_api_key_auth=owner_auth) + ) + mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for( + owner_auth + ) + + inside = 0 + max_inside = 0 + gate = asyncio.Event() + + async def slow_handle(s, r, se): + nonlocal inside, max_inside + inside += 1 + max_inside = max(max_inside, inside) + await gate.wait() + inside -= 1 + + async def make_request(): + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"mcp-session-id", session_id.encode())], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', + "more_body": False, + } + ) + send = AsyncMock() + await handle_streamable_http_mcp(scope, receive, send) + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(owner_auth, None, None, None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, "handle_request", side_effect=slow_handle + ), + patch.object( + session_manager_stateful, + "_server_instances", + {session_id: MagicMock()}, + ), + ): + tasks = [asyncio.create_task(make_request()) for _ in range(3)] + await asyncio.sleep(0.05) + gate.set() + await asyncio.gather(*tasks) + finally: + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_owners.pop(session_id, None) + mcp_server._stateful_session_locks.pop(session_id, None) + + assert ( + max_inside == 1 + ), "concurrent requests on same stateful session must be serialized" + + +@pytest.mark.asyncio +async def test_stateful_mcp_lock_does_not_leak_when_auth_context_missing(): + """ + If a per-session lock is created for a session_id that is not tracked in + ``_stateful_session_auth_contexts`` (e.g., a defensive path), the request + finalizer must drop the lock so it isn't orphaned. The periodic cleanup + loop only iterates ``_stateful_session_auth_context_last_seen``, so a + leaked lock would otherwise live forever. + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "untracked-session-1" + owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner") + + async def handle(s, r, se): + return None + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"mcp-session-id", session_id.encode())], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', + "more_body": False, + } + ) + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(owner_auth, None, None, None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, "handle_request", side_effect=handle + ), + patch.object( + session_manager_stateful, + "_server_instances", + {session_id: MagicMock()}, + ), + ): + assert session_id not in mcp_server._stateful_session_auth_contexts + await handle_streamable_http_mcp(scope, receive, AsyncMock()) + + assert ( + session_id not in mcp_server._stateful_session_locks + ), "lock entry must be cleaned up for untracked stateful session" + finally: + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_owners.pop(session_id, None) + mcp_server._stateful_session_locks.pop(session_id, None) + mcp_server._stateful_session_active_request_counts.pop(session_id, None) + + +@pytest.mark.asyncio +async def test_stateful_mcp_get_stream_does_not_block_post(): + """ + A long-lived GET (server-to-client SSE stream) on a stateful session + must NOT hold the per-session lock — otherwise subsequent POSTs on the + same mcp-session-id hang for the lifetime of the stream. + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "stream-session-1" + owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner") + mcp_server._stateful_session_auth_contexts[session_id] = ( + mcp_server.MCPAuthenticatedUser(user_api_key_auth=owner_auth) + ) + mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for( + owner_auth + ) + + stream_release = asyncio.Event() + post_finished = asyncio.Event() + + async def handle(s, r, se): + if s.get("method") == "GET": + await stream_release.wait() + else: + post_finished.set() + + async def call(method: str, body: bytes = b""): + scope = { + "type": "http", + "method": method, + "path": "/mcp", + "headers": [(b"mcp-session-id", session_id.encode())], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": body, + "more_body": False, + } + ) + await handle_streamable_http_mcp(scope, receive, AsyncMock()) + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(owner_auth, None, None, None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, "handle_request", side_effect=handle + ), + patch.object( + session_manager_stateful, + "_server_instances", + {session_id: MagicMock()}, + ), + ): + stream_task = asyncio.create_task(call("GET")) + await asyncio.sleep(0.05) + assert not stream_task.done(), "GET stream should still be open" + + post_task = asyncio.create_task( + call( + "POST", + body=b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', + ) + ) + await asyncio.wait_for(post_finished.wait(), timeout=1.0) + await post_task + + stream_release.set() + await stream_task + finally: + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_owners.pop(session_id, None) + mcp_server._stateful_session_locks.pop(session_id, None) + @pytest.mark.asyncio @pytest.mark.no_parallel diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 7b3bb81e04..549afd774b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -7,11 +7,12 @@ they may send a stale `mcp-session-id` header. This test verifies that: 2. For DELETE requests: idempotent behavior returns success even if session doesn't exist """ -import pytest +import asyncio from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException from litellm.types.mcp import MCPAuth +import pytest class TestHandleStaleMcpSession: @@ -53,32 +54,55 @@ class TestHandleStaleMcpSession: try: from litellm.proxy._experimental.mcp_server.server import ( _handle_stale_mcp_session, + _stateful_session_active_request_counts, + _stateful_session_auth_context_last_seen, + _stateful_session_auth_contexts, + _stateful_session_locks, + _stateful_session_owners, ) except ImportError: pytest.skip("MCP server not available") + stale_session_id = "stale-id" scope = { "type": "http", "method": "DELETE", "headers": [ (b"content-type", b"application/json"), - (b"mcp-session-id", b"stale-id"), + (b"mcp-session-id", stale_session_id.encode()), ], } receive = AsyncMock() send = AsyncMock() mgr = MagicMock() mgr._server_instances = {} # no active sessions + _stateful_session_auth_contexts[stale_session_id] = MagicMock() + _stateful_session_auth_context_last_seen[stale_session_id] = 1.0 + _stateful_session_owners[stale_session_id] = "owner" + _stateful_session_locks[stale_session_id] = MagicMock() + _stateful_session_active_request_counts[stale_session_id] = 1 - handled = await _handle_stale_mcp_session(scope, receive, send, mgr) + try: + handled = await _handle_stale_mcp_session(scope, receive, send, mgr) - # Should be fully handled (returns True) - assert handled is True - # Should have sent a success response - assert send.called - # Header should NOT be stripped (DELETE needs the session ID) - header_names = [k for k, _ in scope["headers"]] - assert b"mcp-session-id" in header_names + # Should be fully handled (returns True) + assert handled is True + # Should have sent a success response + assert send.called + # Header should NOT be stripped (DELETE needs the session ID) + header_names = [k for k, _ in scope["headers"]] + assert b"mcp-session-id" in header_names + assert stale_session_id not in _stateful_session_auth_contexts + assert stale_session_id not in _stateful_session_auth_context_last_seen + assert stale_session_id not in _stateful_session_owners + assert stale_session_id not in _stateful_session_locks + assert stale_session_id not in _stateful_session_active_request_counts + finally: + _stateful_session_auth_contexts.pop(stale_session_id, None) + _stateful_session_auth_context_last_seen.pop(stale_session_id, None) + _stateful_session_owners.pop(stale_session_id, None) + _stateful_session_locks.pop(stale_session_id, None) + _stateful_session_active_request_counts.pop(stale_session_id, None) @pytest.mark.asyncio async def test_preserves_valid_session_id(self): @@ -202,7 +226,8 @@ async def test_stale_mcp_session_id_is_stripped(): try: from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, - session_manager, + session_manager_stateful, + session_manager_stateless, ) except ImportError: pytest.skip("MCP server not available") @@ -225,11 +250,14 @@ async def test_stale_mcp_session_id_is_stripped(): # Simulate: session manager has NO sessions (the stale one was cleaned up) captured_scope = {} + stateful_handle_request = AsyncMock() - async def mock_handle_request(s, r, se): + async def _stateless_capture(s, r, se): # Capture the scope that was actually passed captured_scope.update(s) + stateless_handle_request = AsyncMock(side_effect=_stateless_capture) + with ( patch( "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", @@ -244,12 +272,22 @@ async def test_stale_mcp_session_id_is_stripped(): True, ), patch.object( - session_manager, + session_manager_stateless, "handle_request", - side_effect=mock_handle_request, + new=stateless_handle_request, ), patch.object( - session_manager, + session_manager_stateless, + "_server_instances", + {}, + ), + patch.object( + session_manager_stateful, + "handle_request", + side_effect=stateful_handle_request, + ), + patch.object( + session_manager_stateful, "_server_instances", {}, # Empty dict = no active sessions ), @@ -261,6 +299,12 @@ async def test_stale_mcp_session_id_is_stripped(): assert ( b"mcp-session-id" not in header_names ), "Stale mcp-session-id header should have been stripped from the scope" + assert ( + stateless_handle_request.called + ), "Stale non-initialize requests should route stateless" + assert ( + not stateful_handle_request.called + ), "Stale non-initialize requests should not route stateful" @pytest.mark.asyncio @@ -332,6 +376,89 @@ async def test_delete_stale_mcp_session_returns_success(): assert send.called, "A response should have been sent" +@pytest.mark.asyncio +async def test_failed_delete_preserves_stateful_session_tracking(): + """ + When the SDK fails to terminate an existing stateful session, keep the + owner/auth tracking so the session cannot be hijacked or hidden from cleanup. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + _owner_fingerprint_for, + _stateful_session_auth_context_last_seen, + _stateful_session_auth_contexts, + _stateful_session_locks, + _stateful_session_owners, + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "delete-failure-session" + user_auth = MagicMock() + user_auth.api_key = "sk-test" + user_auth.user_id = "test-user" + auth_context = MagicMock() + session_lock = asyncio.Lock() + mock_instances = {session_id: MagicMock()} + + scope = { + "type": "http", + "method": "DELETE", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"mcp-session-id", session_id.encode()), + (b"authorization", b"Bearer sk-test"), + ], + } + receive = AsyncMock() + send = AsyncMock() + + _stateful_session_auth_contexts[session_id] = auth_context + _stateful_session_auth_context_last_seen[session_id] = 1.0 + _stateful_session_owners[session_id] = _owner_fingerprint_for(user_auth) + _stateful_session_locks[session_id] = session_lock + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(user_auth, None, None, None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, + "handle_request", + new_callable=AsyncMock, + side_effect=RuntimeError("delete failed"), + ) as mock_handle_request, + patch.object( + session_manager_stateful, + "_server_instances", + mock_instances, + ), + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert mock_handle_request.await_count == 1 + assert _stateful_session_auth_contexts[session_id] is auth_context + assert _stateful_session_auth_context_last_seen[session_id] == 1.0 + assert _stateful_session_owners[session_id] == _owner_fingerprint_for(user_auth) + assert _stateful_session_locks[session_id] is session_lock + assert session_id in mock_instances + finally: + _stateful_session_auth_contexts.pop(session_id, None) + _stateful_session_auth_context_last_seen.pop(session_id, None) + _stateful_session_owners.pop(session_id, None) + _stateful_session_locks.pop(session_id, None) + + @pytest.mark.asyncio async def test_valid_mcp_session_id_is_preserved(): """ @@ -341,7 +468,7 @@ async def test_valid_mcp_session_id_is_preserved(): try: from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, - session_manager, + session_manager_stateful, ) except ImportError: pytest.skip("MCP server not available") @@ -367,7 +494,7 @@ async def test_valid_mcp_session_id_is_preserved(): async def mock_handle_request(s, r, se): captured_scope.update(s) - # Session manager HAS this session + # Stateful session manager HAS this session (requests with mcp-session-id route there) mock_instances = {valid_session_id: MagicMock()} with ( @@ -384,12 +511,12 @@ async def test_valid_mcp_session_id_is_preserved(): True, ), patch.object( - session_manager, + session_manager_stateful, "handle_request", side_effect=mock_handle_request, ), patch.object( - session_manager, + session_manager_stateful, "_server_instances", mock_instances, ), @@ -476,7 +603,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): try: from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, - session_manager, + session_manager_stateless, ) except ImportError: pytest.skip("MCP server not available") @@ -525,7 +652,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): return_value=oauth_server, ), patch.object( - session_manager, + session_manager_stateless, "handle_request", new_callable=AsyncMock, ) as mock_handle_request, @@ -549,7 +676,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): try: from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, - session_manager, + session_manager_stateless, ) except ImportError: pytest.skip("MCP server not available") @@ -598,7 +725,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): return_value=oauth_server, ), patch.object( - session_manager, + session_manager_stateless, "handle_request", new_callable=AsyncMock, ) as mock_handle_request,