tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 20:25:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(mcp): support stateless and stateful clients via session-id routing (#26857)

Browse Source 
* 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 <mateo-berri@users.noreply.github.com>

* 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:<sha256>.

* 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 <mateo-berri@users.noreply.github.com>

* 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 <mateo-berri@users.noreply.github.com>

* 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 <mateo-berri@users.noreply.github.com>

* 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 <cursoragent@cursor.com>

* 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 <mateo-berri@users.noreply.github.com>

* chore(mcp): trim verbose comment on lock cleanup

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* 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 <yassin@berri.ai>

* Fix MCP reinitialize session tracking

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* Fix MCP reinitialize auth context aliasing

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* Apply black formatting after merge

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Run owner-binding 403 before consuming POST body

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Harden MCP routing peek bound and stateful purge race

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Remove inadvertently committed Next.js build artifacts

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Run owner check before stale MCP session cleanup

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* 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 <mateo-berri@users.noreply.github.com>

* 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 <cursoragent@cursor.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
Co-authored-by: Sameerlite <sameerlite@users.noreply.github.com>
Co-authored-by: yuneng-jiang <yuneng@berri.ai>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
Co-authored-by: Claude Babysitter <claude@anthropic.com>
Co-authored-by: mateo-berri <mateo@berri.ai>
This commit is contained in:
Sameer Kankute
2026-05-29 20:30:36 -07:00
committed by GitHub
co-authored by Cursor Agent mateo-berri Mateo Wang Sameerlite yuneng-jiang Yassin Kortam Claude Babysitter mateo-berri
parent 70d2748d80
commit c792df64d2
6 changed files with 2163 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
+1 -1
View File
@@ -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,
litellm/proxy/_experimental/mcp_server/server.py
+597 -37
View File
@@ -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],
tests/mcp_tests/test_mcp_server.py
+15 -10
View File
@@ -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
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+5 -1
View File
@@ -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(
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+1395 -14
View File
File diff suppressed because it is too large Load Diff
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py
+150 -23
View File
@@ -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,