mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
Fix MCP rate limit quota handling
This commit is contained in:
parent
4fc8bfb0a9
commit
73c61e5c89
4 changed files with 46 additions and 6 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue