Fix MCP rate limit quota handling

This commit is contained in:
Cursor Agent 2026-06-02 09:11:23 +00:00
parent 4fc8bfb0a9
commit 73c61e5c89
No known key found for this signature in database
4 changed files with 46 additions and 6 deletions

View file

@ -949,12 +949,12 @@ def get_key_mcp_rpm_limit(
"""
if user_api_key_dict.metadata:
result = user_api_key_dict.metadata.get("mcp_rpm_limit")
if result:
if result is not None:
return result
if user_api_key_dict.team_metadata:
team_limit = user_api_key_dict.team_metadata.get("mcp_rpm_limit")
if team_limit:
if team_limit is not None:
return team_limit
return None

View file

@ -1727,10 +1727,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptors=descriptors,
)
# Per-MCP-server rate limits. Only honor mcp_server_name on actual MCP
# tool calls; otherwise a normal LLM request could inject it in its body
# to consume another server's MCP quota and 429 legitimate tool calls.
if call_type == CallTypes.call_mcp_tool.value:
# REST MCP calls pass the raw body through this hook before server
# resolution; only the later synthetic hook payload may carry this key.
if call_type == CallTypes.call_mcp_tool.value and "server_id" not in data:
mcp_server_name = data.get("mcp_server_name", None)
self._add_mcp_per_key_rate_limit_descriptor(
user_api_key_dict=user_api_key_dict,

View file

@ -14,6 +14,7 @@ from litellm.proxy.auth.auth_utils import (
abbreviate_api_key,
check_complete_credentials,
get_end_user_id_from_request_body,
get_key_mcp_rpm_limit,
get_key_model_rpm_limit,
get_key_model_tpm_limit,
get_model_from_request,
@ -92,6 +93,22 @@ class TestGetKeyModelRpmLimit:
assert result == {}
class TestGetKeyMcpRpmLimit:
def test_empty_dict_limits_are_returned(self):
key_override = UserAPIKeyAuth(
api_key="sk-123",
metadata={"mcp_rpm_limit": {}},
team_metadata={"mcp_rpm_limit": {"github": 50}},
)
assert get_key_mcp_rpm_limit(key_override) == {}
team_empty = UserAPIKeyAuth(
api_key="sk-123",
team_metadata={"mcp_rpm_limit": {}},
)
assert get_key_mcp_rpm_limit(team_empty) == {}
class TestGetKeyModelTpmLimit:
"""Tests for get_key_model_tpm_limit function."""

View file

@ -2976,6 +2976,30 @@ def test_mcp_descriptor_skipped_for_non_mcp_request_v3():
assert _find_descriptor(descriptors, "mcp_per_key") is None
def test_mcp_descriptor_skipped_for_raw_rest_body_v3():
handler, _ = _make_mcp_handler()
user_api_key_dict = UserAPIKeyAuth(
api_key=hash_token("sk-mcp-key"),
team_id="team-1",
metadata={"mcp_rpm_limit": {"github": 5}},
team_metadata={"mcp_rpm_limit": {"github": 3}},
)
descriptors = _build_mcp_descriptors(
handler,
user_api_key_dict,
{
"server_id": "slack",
"name": "demo-tool",
"arguments": {},
"mcp_server_name": "github",
},
)
assert _find_descriptor(descriptors, "mcp_per_key") is None
assert _find_descriptor(descriptors, "mcp_per_team") is None
def test_mcp_per_team_descriptor_created_from_team_metadata_v3():
handler, _ = _make_mcp_handler()
user_api_key_dict = UserAPIKeyAuth(