mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
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
This commit is contained in:
Sameer Kankute
parent
e4fd0383f3
commit
ca996e18e3
3 changed files with 216 additions and 28 deletions
|
|
@ -5,6 +5,7 @@ LiteLLM MCP Server Routes
|
|||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
|
|
@ -180,13 +181,23 @@ 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,
|
||||
)
|
||||
|
||||
# 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,
|
||||
|
|
@ -197,11 +208,12 @@ if MCP_AVAILABLE:
|
|||
|
||||
# Context managers for proper lifecycle management
|
||||
_session_manager_cm = None
|
||||
_session_manager_stateful_cm = None
|
||||
_sse_session_manager_cm = None
|
||||
|
||||
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
|
||||
|
||||
# Use async lock to prevent concurrent initialization
|
||||
async with _INITIALIZATION_LOCK:
|
||||
|
|
@ -211,11 +223,13 @@ 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__()
|
||||
|
||||
_SESSION_MANAGERS_INITIALIZED = True
|
||||
|
|
@ -225,7 +239,7 @@ 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
|
||||
|
||||
if _SESSION_MANAGERS_INITIALIZED:
|
||||
verbose_logger.info("Shutting down MCP session managers...")
|
||||
|
|
@ -233,12 +247,15 @@ if MCP_AVAILABLE:
|
|||
try:
|
||||
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
|
||||
_SESSION_MANAGERS_INITIALIZED = False
|
||||
|
||||
|
|
@ -304,7 +321,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
|
||||
|
|
@ -347,7 +364,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(
|
||||
|
|
@ -489,7 +506,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
|
||||
|
|
@ -2270,6 +2287,33 @@ 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 _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.
|
||||
|
||||
Note: Assumes the body fits in a single ASGI receive() chunk. MCP initialize
|
||||
requests are typically small; large chunked bodies may not be fully parsed.
|
||||
"""
|
||||
if not body:
|
||||
return False
|
||||
try:
|
||||
data = json.loads(body)
|
||||
return data.get("method") == "initialize"
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return False
|
||||
|
||||
async def _handle_stale_mcp_session(
|
||||
scope: Scope,
|
||||
receive: Receive,
|
||||
|
|
@ -2432,16 +2476,54 @@ if MCP_AVAILABLE:
|
|||
# Give it a moment to start up
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# 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
|
||||
first_msg = None
|
||||
|
||||
if scope.get("method") == "POST" and not session_id:
|
||||
# Peek at first chunk to detect initialize (assumes body fits in one chunk)
|
||||
first_msg = await receive()
|
||||
body = first_msg.get("body", b"") or b""
|
||||
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 "")
|
||||
)
|
||||
|
||||
# Replay first message if we consumed it for peeking
|
||||
original_receive = receive
|
||||
if first_msg is not None:
|
||||
|
||||
async def wrapped_receive():
|
||||
nonlocal first_msg
|
||||
if first_msg is not None:
|
||||
msg, first_msg = first_msg, None
|
||||
return msg
|
||||
return await original_receive()
|
||||
|
||||
receive = wrapped_receive
|
||||
|
||||
# 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
|
||||
scope, receive, send, target_manager
|
||||
)
|
||||
if handled:
|
||||
# Request was fully handled (e.g., DELETE on non-existent session)
|
||||
return
|
||||
|
||||
await session_manager.handle_request(scope, receive, send)
|
||||
await target_manager.handle_request(scope, receive, send)
|
||||
except HTTPException:
|
||||
# Re-raise HTTP exceptions to preserve status codes and details
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -720,17 +720,21 @@ 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
|
||||
|
||||
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
|
||||
|
||||
# Mock the session managers to avoid actual MCP initialization
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager"
|
||||
) as mock_session_manager, patch(
|
||||
"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, patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.verbose_logger"
|
||||
|
|
@ -740,7 +744,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
|
||||
|
|
@ -757,18 +762,21 @@ 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
|
||||
|
|
@ -777,32 +785,129 @@ async def test_concurrent_initialize_session_managers():
|
|||
# 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
|
||||
|
||||
|
||||
@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
|
||||
@pytest.mark.no_parallel
|
||||
|
|
|
|||
|
|
@ -7,9 +7,10 @@ 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
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestHandleStaleMcpSession:
|
||||
"""Unit tests for the _handle_stale_mcp_session helper."""
|
||||
|
|
@ -326,7 +327,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")
|
||||
|
|
@ -352,7 +353,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 patch(
|
||||
|
|
@ -365,11 +366,11 @@ async def test_valid_mcp_session_id_is_preserved():
|
|||
"litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
|
||||
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,
|
||||
):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue