diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b3426089375..ffe9a6070af 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -115,6 +115,7 @@ if TYPE_CHECKING: _redact_mcp_resource_url: Final = redact_mcp_resource_url +_MCP_SETUP_REQUEST_METHODS: Final = frozenset({"initialize", "tools/list"}) _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS: Final = 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 @@ -135,6 +136,22 @@ _MCP_PROTOCOL_VERSION_HEADER: Final = b"mcp-protocol-version" _DECODED_JSON: Final = TypeAdapter(object) +def _mcp_setup_request_capacity( + method: str | None, + user_api_key_auth: UserAPIKeyAuth | None, +) -> contextlib.AbstractAsyncContextManager[None]: + if method not in _MCP_SETUP_REQUEST_METHODS or user_api_key_auth is None: + return contextlib.nullcontext() + + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.proxy_server import proxy_logging_obj + + limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter") + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): + return contextlib.nullcontext() + return limiter.request_capacity(user_api_key_auth, "mcp-tool-call", request_data={}) + + def reject_disallowed_mcp_origin(request: StarletteRequest) -> None: from litellm.proxy.proxy_server import origins # noqa: PLC0415 # proxy imports this module during startup @@ -1393,18 +1410,15 @@ if MCP_AVAILABLE: return b"{}" return json.dumps({"method": method}).encode() - 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. - """ + def _mcp_request_method(body: bytes) -> str | None: if not body: - return False + return None try: data: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(body))) - return data is not None and data.get("method") == "initialize" + method: Final = data.get("method") if data is not None else None + return method if isinstance(method, str) else None except (json.JSONDecodeError, UnicodeDecodeError, TypeError): - return False + return None def _jsonrpc_request_method(body: bytes) -> str | None: try: @@ -2423,7 +2437,8 @@ if MCP_AVAILABLE: ) = await extract_mcp_auth_context(scope, path) admit_mcp_request(scope, user_api_key_auth, mcp_servers) body: Final = await peek.body() if peek is not None else b"" - is_initialize: Final = _is_initialize_request(body) + mcp_method: Final = _mcp_request_method(body) + is_initialize: Final = mcp_method == "initialize" scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1 request_headers: Final = StarletteRequest(scope).headers defer_upstream_probes: Final = request_headers.get( @@ -2705,11 +2720,12 @@ if MCP_AVAILABLE: if use_stateful and session_id and scope.get("method") == "DELETE": _remove_stateful_session_tracking(session_id) + request_capacity: Final = _mcp_setup_request_capacity(mcp_method, user_api_key_auth) try: - if session_lock is not None: - async with session_lock: - await _dispatch() - else: + async with ( + session_lock if session_lock is not None else contextlib.nullcontext(), + request_capacity, + ): await _dispatch() finally: for active_request_session_id in active_request_session_ids: diff --git a/tests/integration/mcp/test_mcp_rate_limits.py b/tests/integration/mcp/test_mcp_rate_limits.py index 273419d624f..b872bd5c740 100644 --- a/tests/integration/mcp/test_mcp_rate_limits.py +++ b/tests/integration/mcp/test_mcp_rate_limits.py @@ -269,3 +269,83 @@ def test_mcp_key_team_and_server_rpm_limits_share_redis(tmp_path: Path, monkeypa server_rejected: Final = _call_tool(second_replica, key_three, server_id, "rpm-add") _assert_rate_limit(server_rejected, "mcp_server") assert tool_calls(peer.drain()) == () + + +def test_streamable_mcp_setup_requests_use_virtual_key_rpm_limit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + results_directory: Final = tmp_path / "results" + results_directory.mkdir() + monkeypatch.setenv("INTEGRATION_RESULTS_DIR", str(results_directory)) + assert os.environ.get("DATABASE_URL"), "This integration case requires disposable-database access" + + with ( + owned_redis(tmp_path) as redis, + scratch_database() as database_url, + mcp_peer() as peer, + httpx.Client() as client, + ): + monkeypatch.setenv("DATABASE_URL", database_url) + seed: Final = Gateway(client, "sk-mcp-streamable-setup-rate-limit", peer.url) + config: Final = _config_file( + tmp_path, + seed.key, + redis, + upstream=peer.url, + store_model_in_db=True, + ) + environment: Final = { + "DATABASE_URL": database_url, + "LITELLM_SALT_KEY": "mcp-streamable-setup-rate-limit", + "LITELLM_RATE_LIMIT_WINDOW_SIZE": "30", + } + options: Final = { + "config": config, + "remove_environment": ("DATABASE_URL_READ_REPLICA", "LITELLM_LICENSE", "LITELLM_LICENSE_PATH"), + } + with owned_proxy(seed, tmp_path / "proxy", environment, **options) as gateway, gateway.scenario() as scenario: + server_id: Final = register_mcp(scenario, peer, "rpm", allowed_tools=["add"]) + key: Final = scenario.key(rpm_limit=2, object_permission={"mcp_servers": [server_id]}) + + def post_mcp(payload: dict[str, object], session_id: str | None = None) -> httpx.Response: + headers: Final = { + "Authorization": f"Bearer {key}", + "Accept": "application/json, text/event-stream", + **({"mcp-session-id": session_id} if session_id is not None else {}), + } + return gateway.client.post("/mcp/", headers=headers, json=payload) + + initialize: Final = post_mcp( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-03-26", + "capabilities": {}, + "clientInfo": {"name": "litellm-test", "version": "1"}, + }, + } + ) + assert initialize.status_code == 200, initialize.text + session_id: Final = initialize.headers["mcp-session-id"] + + initialized: Final = post_mcp( + {"jsonrpc": "2.0", "method": "notifications/initialized"}, + session_id, + ) + assert initialized.status_code == 202, initialized.text + + first_list: Final = post_mcp( + {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}, + session_id, + ) + assert first_list.status_code == 200, first_list.text + + second_list: Final = post_mcp( + {"jsonrpc": "2.0", "id": 3, "method": "tools/list", "params": {}}, + session_id, + ) + assert second_list.status_code == 429, second_list.text + assert second_list.headers.get("retry-after") is not None diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py index de711d737d0..293aba4232c 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2,23 +2,21 @@ import asyncio import json import os -from typing import Final - -import pytest -from unittest.mock import ANY, AsyncMock, MagicMock, patch +from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager, nullcontext +from typing import Final +from unittest.mock import ANY, AsyncMock, MagicMock, patch +from mcp.types import CallToolResult, TextContent, Tool as MCPTool +import pytest from starlette.types import Receive, Scope, Send - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, MCPServer, MCPTransport, ) -from litellm.proxy._types import LiteLLM_ObjectPermissionTable -from mcp.types import Tool as MCPTool, CallToolResult, ListToolsResult -from mcp.types import TextContent +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth mcp_server_manager = MCPServerManager() @@ -3259,6 +3257,117 @@ async def test_mcp_routing_batch_body_is_not_stashed_for_auth(): assert await _read_request_body(_admission_request(scope)) == {} +@asynccontextmanager +async def _streamable_mcp_request_runner( + monkeypatch: pytest.MonkeyPatch, + auth: UserAPIKeyAuth, +) -> AsyncIterator[tuple[Callable[[str], Awaitable[None]], AsyncMock]]: + from litellm.caching.dual_cache import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import server + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache, ProxyLogging + + proxy_logging: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + limiter: Final = PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(dual_cache=DualCache()) + ) + proxy_logging.proxy_hook_mapping["parallel_request_limiter"] = limiter + monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging) + + dispatch: Final = AsyncMock() + session_manager: Final = MagicMock() + session_manager.handle_request = dispatch + monkeypatch.setattr(server, "_SESSION_MANAGERS_INITIALIZED", True) + monkeypatch.setattr(server, "session_manager_stateless", session_manager) + monkeypatch.setattr(server, "session_manager_stateful", session_manager) + monkeypatch.setattr( + server, + "extract_mcp_auth_context", + AsyncMock(return_value=(auth, None, ["echo_server"], {}, {}, {})), + ) + + @asynccontextmanager + async def no_initialize_scope(*_args: object, **_kwargs: object) -> AsyncIterator[None]: + yield + + monkeypatch.setattr(server, "_gateway_initialize_instructions_request_scope", no_initialize_scope) + + async def request(method: str) -> None: + body: Final = json.dumps( + { + "jsonrpc": "2.0", + "method": method, + **({"id": 1} if method != "notifications/initialized" else {}), + } + ).encode() + scope: Scope = { + "type": "http", + "method": "POST", + "path": "/mcp/echo_server", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "server": ("localhost", 4000), + "scheme": "http", + } + receive: Receive = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False}) + send: Send = AsyncMock() + await server.handle_streamable_http_mcp(scope, receive, send) + + yield request, dispatch + + +@pytest.mark.asyncio +async def test_streamable_mcp_setup_requests_obey_virtual_key_rpm_limit(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + auth: Final = UserAPIKeyAuth(api_key="mcp-setup-rpm-key", rpm_limit=2) + async with _streamable_mcp_request_runner(monkeypatch, auth) as (request, dispatch): + await request("initialize") + await request("tools/list") + with pytest.raises(ProxyRateLimitError) as rate_limit_error: + await request("tools/list") + + assert rate_limit_error.value.status_code == 429 + assert rate_limit_error.value.headers.get("retry-after") is not None + assert dispatch.await_count == 2 + + +@pytest.mark.asyncio +async def test_streamable_mcp_tools_call_is_not_charged_by_setup_limiter(monkeypatch: pytest.MonkeyPatch) -> None: + auth: Final = UserAPIKeyAuth(api_key="mcp-tools-call-rpm-key", rpm_limit=1) + async with _streamable_mcp_request_runner(monkeypatch, auth) as (request, dispatch): + await request("tools/call") + await request("initialize") + + assert dispatch.await_count == 2 + + +@pytest.mark.asyncio +async def test_streamable_mcp_initialized_notification_does_not_consume_rpm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + auth: Final = UserAPIKeyAuth(api_key="mcp-initialized-rpm-key", rpm_limit=1) + async with _streamable_mcp_request_runner(monkeypatch, auth) as (request, dispatch): + await request("notifications/initialized") + await request("initialize") + + assert dispatch.await_count == 2 + + +@pytest.mark.asyncio +async def test_streamable_mcp_setup_request_releases_parallel_request_slot( + monkeypatch: pytest.MonkeyPatch, +) -> None: + auth: Final = UserAPIKeyAuth(api_key="mcp-setup-parallel-key", max_parallel_requests=1) + async with _streamable_mcp_request_runner(monkeypatch, auth) as (request, dispatch): + await request("initialize") + await request("tools/list") + + assert dispatch.await_count == 2 + + @pytest.mark.asyncio async def test_mcp_routing_truncated_json_body_is_not_stashed_for_auth(): try: