diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e2247e601a3..5cf94aa48c7 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -253,6 +253,11 @@ if MCP_AVAILABLE: ) _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] = {} # Keep this alias so existing references to session_manager still work session_manager = session_manager_stateless @@ -287,6 +292,7 @@ if MCP_AVAILABLE: for session_id in expired_session_ids: _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) transport = server_instances.pop(session_id, None) if transport is not None: await transport.terminate() @@ -294,6 +300,7 @@ if MCP_AVAILABLE: for session_id in list(_stateful_session_auth_context_last_seen): if session_id not in _stateful_session_auth_contexts: _stateful_session_auth_context_last_seen.pop(session_id, None) + _stateful_session_owners.pop(session_id, None) async def _cleanup_expired_stateful_session_auth_contexts() -> None: while True: @@ -2654,6 +2661,23 @@ if MCP_AVAILABLE: ) return None + def _owner_fingerprint_for( + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> str: + """ + Stable, non-reversible identifier for the caller used to bind an + mcp-session-id to its creator. ``api_key`` on UserAPIKeyAuth is + already hashed at construction time, so we can use it directly. + Falls back to user_id then to a sentinel so anonymous callers can + never share or hijack an owned session. + """ + if user_api_key_auth is not None: + if user_api_key_auth.api_key: + return f"key:{user_api_key_auth.api_key}" + if user_api_key_auth.user_id: + return f"user:{user_api_key_auth.user_id}" + return "anonymous" + def _is_initialize_request(body: bytes) -> bool: """ Check if the request body is a JSON-RPC initialize method. @@ -2953,6 +2977,27 @@ if MCP_AVAILABLE: consumed_messages, body = await _read_request_body_for_routing(receive) is_initialize = _is_initialize_request(body) + # 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. + if session_id: + expected_owner = _stateful_session_owners.get(session_id) + request_owner = _owner_fingerprint_for(user_api_key_auth) + 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 + use_stateful = bool(session_id or is_initialize) target_manager = ( session_manager_stateful if use_stateful else session_manager_stateless @@ -2986,7 +3031,11 @@ if MCP_AVAILABLE: session_id=session_id if use_stateful else None, ) if use_stateful and is_initialize: - send = _wrap_send_with_stateful_session_auth_context(send, auth_user) + send = _wrap_send_with_stateful_session_auth_context( + send, + auth_user, + _owner_fingerprint_for(user_api_key_auth), + ) async with _gateway_initialize_instructions_request_scope( user_api_key_auth, @@ -2999,6 +3048,7 @@ if MCP_AVAILABLE: if use_stateful and session_id and scope.get("method") == "DELETE": _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) except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise @@ -3006,7 +3056,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( @@ -3201,6 +3250,7 @@ if MCP_AVAILABLE: def _wrap_send_with_stateful_session_auth_context( send: Send, auth_user: MCPAuthenticatedUser, + owner_fingerprint: str, ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": @@ -3211,6 +3261,7 @@ if MCP_AVAILABLE: _stateful_session_auth_context_last_seen[session_id] = ( time.monotonic() ) + _stateful_session_owners[session_id] = owner_fingerprint break await send(message) 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 240f08b9dfb..ab6dc1752b9 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 @@ -1082,7 +1082,7 @@ async def test_streamable_http_session_manager_is_stateless(): 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) @@ -1253,37 +1253,45 @@ async def test_mcp_routing_chunked_initialize_to_stateful(): 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", - {}, + 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" - ) + assert ( + stateful_called and not stateless_called + ), "chunked initialize (no session) should route to stateful, not stateless" @pytest.mark.asyncio @@ -1398,6 +1406,7 @@ async def test_stateful_mcp_auth_contexts_expire_with_idle_sessions(): 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 @@ -1417,6 +1426,88 @@ async def test_stateful_mcp_auth_contexts_expire_with_idle_sessions(): transport.terminate.assert_awaited_once() +@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 @pytest.mark.no_parallel async def test_mcp_routing_with_conflicting_alias_and_group_name(): @@ -3277,7 +3368,7 @@ async def test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow(): """ P1 Regression: list_tools path must apply _resolve_oauth2_flow to legacy DB rows where oauth2_flow is NULL but M2M credentials are present. - + Without this fix, has_client_credentials returns False and the caller's Authorization header is forwarded upstream instead of being blocked. """ @@ -3375,7 +3466,7 @@ async def test_call_tool_empty_extra_headers_returns_none(): """ P2 Regression: When all configured extra_headers are filtered out (e.g. Authorization for M2M), the resulting extra_headers should be None, not {}. - + Downstream code that checks `if extra_headers is None` will behave differently if an empty dict is passed instead. """ @@ -3402,7 +3493,10 @@ async def test_call_tool_empty_extra_headers_returns_none(): extra_headers=["Authorization"], # Will be filtered out for M2M ) - raw_headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + raw_headers = { + "Authorization": "Bearer sk-1234", + "Content-Type": "application/json", + } captured_extra_headers = None @@ -3439,7 +3533,8 @@ async def test_call_tool_empty_extra_headers_returns_none(): pass # We only care about the captured headers # With P2 fix: extra_headers should be None (not {}) when all headers filtered - assert captured_extra_headers is None, ( - "P2 API consistency issue: expected None for empty extra_headers, got: " - + str(captured_extra_headers) + assert ( + captured_extra_headers is None + ), "P2 API consistency issue: expected None for empty extra_headers, got: " + str( + captured_extra_headers ) 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 8cba3dd2842..078c5414e59 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 @@ -228,10 +228,12 @@ async def test_stale_mcp_session_id_is_stripped(): captured_scope = {} stateful_handle_request = AsyncMock() - async def stateless_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", @@ -248,7 +250,7 @@ async def test_stale_mcp_session_id_is_stripped(): patch.object( session_manager_stateless, "handle_request", - side_effect=stateless_handle_request, + new=stateless_handle_request, ), patch.object( session_manager_stateless, @@ -515,31 +517,39 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): oauth_server.auth_type = MCPAuth.oauth2 oauth_server.needs_user_oauth_token = True - with patch( - "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", - new_callable=AsyncMock, - return_value=(user_auth, None, ["repro_oauth_server"], 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( - "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session", - new_callable=AsyncMock, - return_value=False, - ), patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", - new_callable=AsyncMock, - return_value=None, - ) as mock_get_stored_token, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", - return_value=oauth_server, - ), patch.object( - session_manager, - "handle_request", - new_callable=AsyncMock, - ) as mock_handle_request: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(user_auth, None, ["repro_oauth_server"], 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( + "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session", + new_callable=AsyncMock, + return_value=False, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new_callable=AsyncMock, + return_value=None, + ) as mock_get_stored_token, + patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + return_value=oauth_server, + ), + patch.object( + session_manager, + "handle_request", + new_callable=AsyncMock, + ) as mock_handle_request, + ): with pytest.raises(HTTPException) as exc_info: await handle_streamable_http_mcp(scope, receive, send) @@ -580,31 +590,39 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): oauth_server.auth_type = MCPAuth.oauth2 oauth_server.needs_user_oauth_token = True - with patch( - "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", - new_callable=AsyncMock, - return_value=(user_auth, None, ["repro_oauth_server"], 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( - "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session", - new_callable=AsyncMock, - return_value=False, - ), patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", - new_callable=AsyncMock, - return_value={"Authorization": "Bearer cached-token"}, - ) as mock_get_stored_token, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", - return_value=oauth_server, - ), patch.object( - session_manager, - "handle_request", - new_callable=AsyncMock, - ) as mock_handle_request: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(user_auth, None, ["repro_oauth_server"], 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( + "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session", + new_callable=AsyncMock, + return_value=False, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new_callable=AsyncMock, + return_value={"Authorization": "Bearer cached-token"}, + ) as mock_get_stored_token, + patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + return_value=oauth_server, + ), + patch.object( + session_manager, + "handle_request", + new_callable=AsyncMock, + ) as mock_handle_request, + ): await handle_streamable_http_mcp(scope, receive, send) assert mock_get_stored_token.await_count == 1