diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b4678a50b2c..98bc5ae7a90 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2650,6 +2650,9 @@ async def pre_call_tool_check( "name": name, "arguments": arguments, "server_name": server_name, + "mcp_rate_limit_server_name": server.alias + or server.server_name + or server.name, "user_api_key_auth": user_api_key_auth, "user_api_key_user_id": ( getattr(user_api_key_auth, "user_id", None) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 98a17e4be95..b44a4a62995 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1050,6 +1050,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None model_tpm_limit: Optional[dict] = None + mcp_rpm_limit: Optional[Dict[str, int]] = None guardrails: Optional[List[str]] = None policies: Optional[List[str]] = None prompts: Optional[List[str]] = None @@ -1851,6 +1852,7 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None + mcp_rpm_limit: Optional[Dict[str, int]] = None team_member_budget: Optional[float] = ( None # allow user to set a budget for all team members ) @@ -1920,6 +1922,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): prompts: Optional[List[str]] = None model_rpm_limit: Optional[Dict[str, int]] = None model_tpm_limit: Optional[Dict[str, int]] = None + mcp_rpm_limit: Optional[Dict[str, int]] = None allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None enforced_batch_output_expires_after: Optional[dict] = None enforced_file_expires_after: Optional[dict] = None @@ -4285,6 +4288,7 @@ class PassThroughEndpointLoggingTypedDict(TypedDict): LiteLLM_ManagementEndpoint_MetadataFields = [ "model_rpm_limit", "model_tpm_limit", + "mcp_rpm_limit", "rpm_limit_type", "tpm_limit_type", "enforced_params", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 86265270357..f57d1848ffd 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -934,6 +934,40 @@ def get_team_model_tpm_limit( return None +def get_key_mcp_rpm_limit( + user_api_key_dict: UserAPIKeyAuth, +) -> Optional[Dict[str, int]]: + """ + Get the per-MCP-server rpm limit for a given api key. + + Priority order (returns first found): + 1. Key metadata (mcp_rpm_limit) + 2. Team metadata (mcp_rpm_limit) + + The returned dict is keyed by MCP server name (alias if set, else the + configured server name). + """ + if user_api_key_dict.metadata: + result = user_api_key_dict.metadata.get("mcp_rpm_limit") + 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 is not None: + return team_limit + + return None + + +def get_team_mcp_rpm_limit( + user_api_key_dict: UserAPIKeyAuth, +) -> Optional[Dict[str, int]]: + if user_api_key_dict.team_metadata: + return user_api_key_dict.team_metadata.get("mcp_rpm_limit") + return None + + def get_project_model_rpm_limit( user_api_key_dict: UserAPIKeyAuth, ) -> Optional[Dict[str, int]]: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d03ad70562a..4343747d104 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -36,7 +36,7 @@ from litellm.proxy.auth.auth_utils import get_model_rate_limit_from_metadata from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject -from litellm.types.utils import ModelResponse, Usage +from litellm.types.utils import CallTypes, ModelResponse, Usage if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -1375,6 +1375,79 @@ def _add_model_per_key_rate_limit_descriptor( ) ) + def _add_mcp_per_key_rate_limit_descriptor( + self, + user_api_key_dict: UserAPIKeyAuth, + mcp_server_name: Optional[str], + descriptors: List[RateLimitDescriptor], + ) -> None: + """ + Add a per-MCP-server rpm descriptor for the API key, if a limit is + configured for the server being called. + + MCP tool calls have no token usage, so only requests_per_unit is set; + tokens_per_unit stays None so the TPM reservation path is never engaged. + """ + from litellm.proxy.auth.auth_utils import get_key_mcp_rpm_limit + + if not mcp_server_name or not user_api_key_dict.api_key: + return + + mcp_rpm_limit = get_key_mcp_rpm_limit(user_api_key_dict) + if not mcp_rpm_limit: + return + + server_rpm_limit = mcp_rpm_limit.get(mcp_server_name) + if server_rpm_limit is None: + return + + descriptors.append( + RateLimitDescriptor( + key="mcp_per_key", + value=f"{user_api_key_dict.api_key}:{mcp_server_name}", + rate_limit={ + "requests_per_unit": server_rpm_limit, + "tokens_per_unit": None, + "window_size": self.window_size, + }, + ) + ) + + def _add_mcp_per_team_rate_limit_descriptor( + self, + user_api_key_dict: UserAPIKeyAuth, + mcp_server_name: Optional[str], + descriptors: List[RateLimitDescriptor], + ) -> None: + """ + Add a per-MCP-server rpm descriptor for the team, if a limit is + configured for the server being called. + """ + from litellm.proxy.auth.auth_utils import get_team_mcp_rpm_limit + + if not mcp_server_name or not user_api_key_dict.team_id: + return + + mcp_rpm_limit = get_team_mcp_rpm_limit(user_api_key_dict) + if not mcp_rpm_limit: + return + + server_rpm_limit = mcp_rpm_limit.get(mcp_server_name) + if server_rpm_limit is None: + return + + descriptors.append( + RateLimitDescriptor( + key="mcp_per_team", + value=f"{user_api_key_dict.team_id}:{mcp_server_name}", + rate_limit={ + "requests_per_unit": server_rpm_limit, + "tokens_per_unit": None, + "window_size": self.window_size, + }, + ) + ) + def _should_enforce_rate_limit( self, limit_type: Optional[str], @@ -1533,6 +1606,7 @@ def _create_rate_limit_descriptors( rpm_limit_type: Optional[str], tpm_limit_type: Optional[str], model_has_failures: bool, + call_type: Optional[str] = None, ) -> List[RateLimitDescriptor]: """ Create all rate limit descriptors for the request. @@ -1653,6 +1727,21 @@ def _create_rate_limit_descriptors( descriptors=descriptors, ) + # 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, + mcp_server_name=mcp_server_name, + descriptors=descriptors, + ) + self._add_mcp_per_team_rate_limit_descriptor( + user_api_key_dict=user_api_key_dict, + mcp_server_name=mcp_server_name, + descriptors=descriptors, + ) + if ( get_team_model_rpm_limit(user_api_key_dict) is not None or get_team_model_tpm_limit(user_api_key_dict) is not None @@ -1983,6 +2072,7 @@ async def async_pre_call_hook( rpm_limit_type=rpm_limit_type, tpm_limit_type=tpm_limit_type, model_has_failures=model_has_failures, + call_type=call_type, ) # Add team model rate limits from team_metadata diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 75eb5cd55ef..7b8f0f72e13 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -386,6 +386,7 @@ async def new_user( - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) + - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user. - model_tpm_limit: Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) - spend: Optional[float] - Amount spent by user. Default is 0. Will be updated by proxy whenever user is used. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). - agent_id: Optional[str] - The agent id associated with the user. @@ -1427,6 +1428,7 @@ async def user_update( - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) + - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user. - model_tpm_limit: Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) - spend: Optional[float] - Amount spent by user. Default is 0. Will be updated by proxy whenever user is used. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). - agent_id: Optional[str] - The agent id associated with the user. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0e645013b92..80ded0bdd16 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1388,6 +1388,7 @@ async def generate_key_fn( - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. + - mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit. - tpm_limit_type: Optional[str] - Type of tpm limit. Options: "best_effort_throughput" (no error if we're overallocating tpm), "guaranteed_throughput" (raise an error if we're overallocating tpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput". - rpm_limit_type: Optional[str] - Type of rpm limit. Options: "best_effort_throughput" (no error if we're overallocating rpm), "guaranteed_throughput" (raise an error if we're overallocating rpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput". - allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request @@ -1606,6 +1607,7 @@ async def generate_service_account_key_fn( - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. + - mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit. - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request @@ -2422,6 +2424,7 @@ async def update_key_fn( # noqa: PLR0915 - tpm_limit: Optional[int] - Tokens per minute limit - rpm_limit: Optional[int] - Requests per minute limit - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} + - mcp_rpm_limit: Optional[dict] - Per-MCP-server RPM limits, keyed by MCP server name {"github": 100, "slack": 200} - model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000} - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" @@ -3401,6 +3404,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 model_max_budget: Optional[dict] = {}, model_rpm_limit: Optional[dict] = None, model_tpm_limit: Optional[dict] = None, + mcp_rpm_limit: Optional[dict] = None, guardrails: Optional[list] = None, policies: Optional[list] = None, prompts: Optional[list] = None, @@ -3479,6 +3483,9 @@ async def generate_key_helper_fn( # noqa: PLR0915 if model_tpm_limit is not None: metadata = metadata or {} metadata["model_tpm_limit"] = model_tpm_limit + if mcp_rpm_limit is not None: + metadata = metadata or {} + metadata["mcp_rpm_limit"] = mcp_rpm_limit if guardrails is not None: metadata = metadata or {} metadata["guardrails"] = guardrails diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 8a8e703831b..ae7da0d29f2 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -863,8 +863,9 @@ async def new_team( # noqa: PLR0915 - members_with_roles: List[{"role": "admin" or "user", "user_id": ""}] - A list of users and their roles in the team. Get user_id when making a new user via `/user/new`. - team_member_permissions: Optional[List[str]] - A list of routes that non-admin team members can access. example: ["/key/generate", "/key/update", "/key/delete"] - metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"extra_info": "some info"} - - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit for this team - applied across all keys for this team. + - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit for this team - applied across all keys for this team. - model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit for this team - applied across all keys for this team. + - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. - tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement. diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0e72f47e224..8bd50a50a38 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -643,6 +643,7 @@ def _convert_mcp_to_llm_format(self, request_obj, kwargs: dict) -> dict: "user_api_key_request_route": kwargs.get("user_api_key_request_route"), "mcp_tool_name": request_obj.tool_name, # Keep original for reference "mcp_arguments": request_obj.arguments, # Keep original for reference + "mcp_server_name": kwargs.get("mcp_rate_limit_server_name"), # Raw Bearer token from the original HTTP request — allows guardrails # (e.g. MCPJWTSigner) to independently verify the caller's identity # before re-signing an outbound token (FR-5 verify+re-sign). diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index cbea386a69c..04ff1e4be20 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -826,3 +826,87 @@ def test_jwt_claims_set_after_construction(self): auth.jwt_claims = claims assert auth.jwt_claims == claims assert auth.jwt_claims["groups"] == ["admin"] + + +class TestMcpRateLimitServerNameSurfacing: + """ + The per-MCP-server rate limiter only sees the request `data` dict, so the + server identity must be surfaced into it. These tests pin the contract + between pre_call_tool_check, _convert_mcp_to_llm_format, and the limiter. + """ + + def setup_method(self): + self.proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + + def test_convert_mcp_to_llm_format_surfaces_rate_limit_server_name(self): + request_obj = MagicMock() + request_obj.tool_name = "list_repos" + request_obj.arguments = {"org": "acme"} + + result = self.proxy_logging._convert_mcp_to_llm_format( + request_obj, {"mcp_rate_limit_server_name": "github"} + ) + + assert result["mcp_server_name"] == "github" + + def test_convert_mcp_to_llm_format_server_name_none_when_absent(self): + request_obj = MagicMock() + request_obj.tool_name = "list_repos" + request_obj.arguments = {} + + result = self.proxy_logging._convert_mcp_to_llm_format(request_obj, {}) + + assert result["mcp_server_name"] is None + + @pytest.mark.asyncio + async def test_pre_call_tool_check_resolves_alias_for_rate_limit(self): + """ + The rate-limit server key must be the alias when set (falling back to + server_name), matching how an admin keys mcp_rpm_limit in config. + """ + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="gh", + alias="gh", + server_name="github_full_name", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + captured = {} + + def capture_convert(request_obj, kwargs): + captured["kwargs"] = kwargs + return {"model": "fake"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + side_effect=capture_convert + ) + proxy_logging.pre_call_hook = AsyncMock(return_value=None) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": {}} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + await manager.pre_call_tool_check( + name="list_repos", + arguments={}, + server_name="github_full_name", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert captured["kwargs"]["mcp_rate_limit_server_name"] == "gh" diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 2d40db9017e..60cf50efc75 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -14,6 +14,7 @@ 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 @@ def test_team_metadata_empty_rpm_dict_falls_through_to_deployment_default(self): 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.""" diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 3e2eb4b02c2..676f623a5dd 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -2893,3 +2893,230 @@ async def test_pre_call_hook_rejects_caller_supplied_stash_values(): ): leaked = [k for k in _LITELLM_STASH_KEYS if k in channel] assert not leaked, f"caller-supplied stash survived in {channel!r}: {leaked}" + + +# ----------------------- Per-MCP-server rate limiting (v3) ----------------------- + + +def _make_mcp_handler(): + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + return handler, local_cache + + +def _find_descriptor(descriptors, key): + return next((d for d in descriptors if d["key"] == key), None) + + +def _build_mcp_descriptors(handler, user_api_key_dict, data, call_type="call_mcp_tool"): + return handler._create_rate_limit_descriptors( + user_api_key_dict=user_api_key_dict, + data=data, + rpm_limit_type=None, + tpm_limit_type=None, + model_has_failures=False, + call_type=call_type, + ) + + +def test_mcp_per_key_descriptor_created_for_matching_server_v3(): + handler, _ = _make_mcp_handler() + api_key = hash_token("sk-mcp-key") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + metadata={"mcp_rpm_limit": {"github": 5}}, + ) + + descriptors = _build_mcp_descriptors( + handler, user_api_key_dict, {"mcp_server_name": "github"} + ) + + descriptor = _find_descriptor(descriptors, "mcp_per_key") + assert descriptor is not None + assert descriptor["value"] == f"{api_key}:github" + assert descriptor["rate_limit"]["requests_per_unit"] == 5 + # MCP tool calls have no token usage; tokens_per_unit must stay None so the + # TPM reservation path is never engaged (otherwise budget would leak). + assert descriptor["rate_limit"]["tokens_per_unit"] is None + + +def test_mcp_per_key_descriptor_skipped_for_non_matching_server_v3(): + handler, _ = _make_mcp_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-mcp-key"), + metadata={"mcp_rpm_limit": {"github": 5}}, + ) + + descriptors = _build_mcp_descriptors( + handler, user_api_key_dict, {"mcp_server_name": "slack"} + ) + + assert _find_descriptor(descriptors, "mcp_per_key") is None + + +def test_mcp_descriptor_skipped_for_non_mcp_request_v3(): + """A non-MCP request must not create an MCP descriptor even if the caller + injects mcp_server_name in the body; otherwise an LLM call could consume a + target server's MCP quota and 429 legitimate tool calls.""" + handler, _ = _make_mcp_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-mcp-key"), + metadata={"mcp_rpm_limit": {"github": 5}}, + ) + + descriptors = _build_mcp_descriptors( + handler, + user_api_key_dict, + {"model": "gpt-4", "mcp_server_name": "github"}, + call_type="completion", + ) + + 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( + api_key=hash_token("sk-mcp-key"), + team_id="team-1", + team_metadata={"mcp_rpm_limit": {"github": 3}}, + ) + + descriptors = _build_mcp_descriptors( + handler, user_api_key_dict, {"mcp_server_name": "github"} + ) + + descriptor = _find_descriptor(descriptors, "mcp_per_team") + assert descriptor is not None + assert descriptor["value"] == "team-1:github" + assert descriptor["rate_limit"]["requests_per_unit"] == 3 + assert descriptor["rate_limit"]["tokens_per_unit"] is None + + +@pytest.mark.asyncio +async def test_mcp_per_key_rpm_enforced_v3(monkeypatch): + """ + A key configured with mcp_rpm_limit={"github": 2} must allow 2 calls to the + github MCP server within the window and reject the 3rd with a 429, while + calls to a different MCP server are unaffected. + """ + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") + api_key = hash_token("sk-mcp-enforce") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + window_starts: Dict[str, int] = {} + request_counts: Dict[str, int] = {} + + async def mock_batch_rate_limiter(*args, **kwargs): + keys = kwargs.get("keys") if kwargs else args[0] + args_list = kwargs.get("args") if kwargs else args[1] + now = args_list[0] + window_size = args_list[1] + results = [] + for i in range(0, len(keys), 2): + window_key = keys[i] + counter_key = keys[i + 1] + prev_window = window_starts.get(window_key) + prev_counter = request_counts.get(counter_key, 0) + if prev_window is None or (now - prev_window) >= window_size: + window_starts[window_key] = now + new_counter = 1 + else: + new_counter = prev_counter + 1 + request_counts[counter_key] = new_counter + results.append(now) + results.append(new_counter) + return results + + handler.batch_rate_limiter_script = mock_batch_rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + metadata={"mcp_rpm_limit": {"github": 2}}, + ) + + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"mcp_server_name": "github"}, + call_type="call_mcp_tool", + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"mcp_server_name": "github"}, + call_type="call_mcp_tool", + ) + assert exc_info.value.status_code == 429 + + # A different server has no configured limit -> not rate limited. + for _ in range(5): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"mcp_server_name": "slack"}, + call_type="call_mcp_tool", + ) + + # The TPM counter must never be created for an MCP descriptor. + assert not any(":tokens" in key and "github" in key for key in request_counts) + + +def test_get_key_mcp_rpm_limit_precedence(): + from litellm.proxy.auth.auth_utils import ( + get_key_mcp_rpm_limit, + get_team_mcp_rpm_limit, + ) + + # Key metadata takes precedence over team metadata. + key_first = UserAPIKeyAuth( + api_key=hash_token("sk-mcp-key"), + metadata={"mcp_rpm_limit": {"github": 10}}, + team_metadata={"mcp_rpm_limit": {"github": 99}}, + ) + assert get_key_mcp_rpm_limit(key_first) == {"github": 10} + + # Falls back to team metadata when key has none. + team_only = UserAPIKeyAuth( + api_key=hash_token("sk-mcp-key"), + team_metadata={"mcp_rpm_limit": {"github": 7}}, + ) + assert get_key_mcp_rpm_limit(team_only) == {"github": 7} + assert get_team_mcp_rpm_limit(team_only) == {"github": 7} + + # No configuration anywhere. + none_set = UserAPIKeyAuth(api_key=hash_token("sk-mcp-key")) + assert get_key_mcp_rpm_limit(none_set) is None + assert get_team_mcp_rpm_limit(none_set) is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index f898763d2cb..d53ea6fa34d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -482,6 +482,33 @@ def test_set_object_metadata_field_initializes_metadata_if_none(self): _set_object_metadata_field(team, "model_rpm_limit", {"x": 1}) assert team.metadata == {"model_rpm_limit": {"x": 1}} + def test_mcp_rpm_limit_is_hoisted_into_metadata(self): + """ + Per-MCP-server rpm limits are stored in the metadata JSON column, not a + dedicated DB column. The key/team management endpoints rely on + LiteLLM_ManagementEndpoint_MetadataFields to move the request field into + metadata; this regression guards that mcp_rpm_limit is in that list and + round-trips through the same loop the endpoints use. + """ + from litellm.proxy._types import LiteLLM_ManagementEndpoint_MetadataFields + + assert "mcp_rpm_limit" in LiteLLM_ManagementEndpoint_MetadataFields + + from types import SimpleNamespace + + team = LiteLLM_TeamTable(team_id="t1", metadata={}) + mcp_rpm_limit = {"github": 100} + data = SimpleNamespace(mcp_rpm_limit=mcp_rpm_limit) + + with patch( + "litellm.proxy.management_endpoints.common_utils._premium_user_check" + ): + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if getattr(data, field, None) is not None: + _set_object_metadata_field(team, field, getattr(data, field)) + + assert team.metadata["mcp_rpm_limit"] == mcp_rpm_limit + class TestRequireCallerUserIdForNonAdmin: """