fix(mcp): count initialize and tools/list against the key rpm limit

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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:

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

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: