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 @@ class MCPServerManager: "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..4afef4c4db7 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: + return result + + if user_api_key_dict.team_metadata: + team_limit = user_api_key_dict.team_metadata.get("mcp_rpm_limit") + if team_limit: + 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..574a5d60d3c 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1375,6 +1375,79 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + 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], @@ -1653,6 +1726,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors=descriptors, ) + # Per-MCP-server rate limits + 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 diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0e645013b92..1714efd7d1a 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 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 @@ class ProxyLogging: "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/scripts/test_mcp_rpm_limit.sh b/scripts/test_mcp_rpm_limit.sh new file mode 100755 index 00000000000..fa32a88cdfb --- /dev/null +++ b/scripts/test_mcp_rpm_limit.sh @@ -0,0 +1,238 @@ +#!/usr/bin/env bash +# +# End-to-end manual test for per-MCP RPM rate limiting. +# +# What it does, with no other setup required: +# 1. Writes a throwaway proxy config with one stdio MCP server (the +# `uvx mcp-server-fetch` server, aliased "fetch_mcp"). +# 2. Boots the proxy in the background and waits until it is ready. +# 3. Generates two keys, both with full access to fetch_mcp: +# - "limited" key: mcp_rpm_limit caps "fetch_mcp" at 2 req/min. +# - "control" key: mcp_rpm_limit caps a DIFFERENT server name +# ("other_mcp") at 2 req/min, so calls to fetch_mcp +# are uncapped. +# 4. Fires 4 fetch_mcp calls with each key. The limited key must trip at the +# 3rd call (429); the control key must never be rate limited. This proves +# the limit is keyed per MCP server name, not globally per key. +# +# Using one physical server with two keys (rather than two servers) keeps the +# test deterministic: a single server's access resolution is exercised, and the +# only variable between the two runs is which server name the key's limit +# targets. +# +# Usage: +# ./scripts/test_mcp_rpm_limit.sh +# +# Requirements: jq, curl, uvx (for the stdio fetch MCP server), and a reachable +# DATABASE_URL (read from .env). + +set -uo pipefail + +# --- locate repo root (this script lives in /scripts) ------------------- +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +cd "${REPO_ROOT}" + +RPM_LIMIT=2 +SERVER="fetch_mcp" +OTHER_SERVER="other_mcp" +WORKDIR="$(mktemp -d)" +CONFIG="${WORKDIR}/mcp_rpm_test_config.yaml" +PROXY_LOG="${WORKDIR}/proxy.log" +PROXY_PID="" + +# --- load secrets (DATABASE_URL, provider keys, master key) ------------------- +# Parse .env line-by-line and export each KEY=VALUE verbatim. We avoid +# `source`-ing it because some values contain characters (e.g. '#') that the +# shell would try to execute. +if [[ -f .env ]]; then + while IFS= read -r line; do + [[ "${line}" =~ ^[A-Za-z_][A-Za-z0-9_]*= ]] || continue + key="${line%%=*}" + val="${line#*=}" + val="${val%\"}"; val="${val#\"}" # strip surrounding double quotes + val="${val%\'}"; val="${val#\'}" # strip surrounding single quotes + export "${key}=${val}" + done < .env +fi +MASTER_KEY="${LITELLM_MASTER_KEY:-sk-1234}" + +# --- pick a free TCP port (start at 4000) so we never collide with a proxy +# already running from a prior session ----------------------------------------- +PORT="" +for candidate in $(seq 4000 4050); do + if ! lsof -iTCP:"${candidate}" -sTCP:LISTEN -n -P >/dev/null 2>&1; then + PORT="${candidate}" + break + fi +done +if [[ -z "${PORT}" ]]; then + echo "ERROR: no free port found in 4000-4050" + exit 1 +fi +BASE="http://localhost:${PORT}" +echo ">> using port ${PORT}" + +cleanup() { + if [[ -n "${PROXY_PID}" ]] && kill -0 "${PROXY_PID}" 2>/dev/null; then + echo ">> stopping proxy (pid ${PROXY_PID})" + kill "${PROXY_PID}" 2>/dev/null + # kill the whole process group in case uvicorn spawned children + pkill -P "${PROXY_PID}" 2>/dev/null + wait "${PROXY_PID}" 2>/dev/null + fi + echo ">> logs kept at: ${PROXY_LOG}" +} +trap cleanup EXIT + +require() { command -v "$1" >/dev/null 2>&1 || { echo "ERROR: '$1' is required but not installed"; exit 1; }; } +require jq +require curl + +# --- 1. write throwaway config ------------------------------------------------ +cat > "${CONFIG}" <> config written to ${CONFIG}" + +# --- 2. start proxy ----------------------------------------------------------- +echo ">> starting proxy on :${PORT} (log: ${PROXY_LOG})" +# Put the repo root first on PYTHONPATH so the local litellm source shadows any +# stale `litellm` installed in site-packages (running the cli as a script puts +# litellm/proxy/ on sys.path instead of the repo root). +PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}" python litellm/proxy/proxy_cli.py \ + --config "${CONFIG}" \ + --port "${PORT}" \ + --detailed_debug \ + --use_v2_migration_resolver > "${PROXY_LOG}" 2>&1 & +PROXY_PID=$! + +echo -n ">> waiting for readiness" +ready=false +for _ in $(seq 1 90); do + if curl -sf "${BASE}/health/readiness" >/dev/null 2>&1; then + ready=true + break + fi + if ! kill -0 "${PROXY_PID}" 2>/dev/null; then + echo "" + echo "ERROR: proxy process died during startup. Tail of log:" + tail -n 40 "${PROXY_LOG}" + exit 1 + fi + echo -n "." + sleep 1 +done +echo "" +if [[ "${ready}" != "true" ]]; then + echo "ERROR: proxy did not become ready in time. Tail of log:" + tail -n 40 "${PROXY_LOG}" + exit 1 +fi +echo ">> proxy is ready" + +# --- 3. generate the two keys ------------------------------------------------- +# Both keys get explicit access to fetch_mcp. They differ only in which server +# name their mcp_rpm_limit targets. +generate_key() { + local rpm_target="$1" + curl -sf -X POST "${BASE}/key/generate" \ + -H "Authorization: Bearer ${MASTER_KEY}" \ + -H "Content-Type: application/json" \ + -d "{\"mcp_rpm_limit\": {\"${rpm_target}\": ${RPM_LIMIT}}, \"object_permission\": {\"mcp_servers\": [\"${SERVER}\"]}}" \ + | jq -r '.key' +} + +echo ">> generating limited key (mcp_rpm_limit {\"${SERVER}\": ${RPM_LIMIT}})" +LIMITED_KEY="$(generate_key "${SERVER}")" +echo ">> generating control key (mcp_rpm_limit {\"${OTHER_SERVER}\": ${RPM_LIMIT}})" +CONTROL_KEY="$(generate_key "${OTHER_SERVER}")" +for k in "${LIMITED_KEY}" "${CONTROL_KEY}"; do + if [[ -z "${k}" || "${k}" == "null" ]]; then + echo "ERROR: /key/generate failed. Is DATABASE_URL set and reachable?" + tail -n 40 "${PROXY_LOG}" + exit 1 + fi +done +echo ">> limited key: ${LIMITED_KEY:0:12}... control key: ${CONTROL_KEY:0:12}..." + +# --- discover a real tool name on the server ---------------------------------- +TOOL_NAME="$(curl -sf "${BASE}/mcp-rest/tools/list?server_id=${SERVER}" \ + -H "Authorization: Bearer ${LIMITED_KEY}" 2>/dev/null \ + | jq -r '.tools[0].name // empty')" +if [[ -z "${TOOL_NAME}" ]]; then + echo ">> could not auto-discover a tool name; falling back to 'fetch'" + TOOL_NAME="fetch" +fi +echo ">> using tool: ${TOOL_NAME}" + +# tool_call: the server alias is accepted directly as server_id. Point the fetch +# tool at the proxy's own health endpoint so the call is fast and always +# reachable; that way a non-429 response unambiguously means "the rate limiter +# let this through" rather than "the upstream fetch flaked". +FETCH_URL="${BASE}/health/readiness" +call_mcp() { + local key="$1" + curl -s -o /dev/null -w "%{http_code}" \ + -X POST "${BASE}/mcp-rest/tools/call" \ + -H "Authorization: Bearer ${key}" \ + -H "Content-Type: application/json" \ + -d "{\"server_id\": \"${SERVER}\", \"name\": \"${TOOL_NAME}\", \"arguments\": {\"url\": \"${FETCH_URL}\", \"max_length\": 100}}" +} + +# --- 4a. limited key: expect 429 once the cap is exceeded --------------------- +echo "" +echo "=== limited key (caps ${SERVER} at ${RPM_LIMIT}/min) ===" +limited_codes=() +for i in 1 2 3 4; do + code="$(call_mcp "${LIMITED_KEY}")" + limited_codes+=("${code}") + echo " call ${i} -> HTTP ${code}" +done + +# --- 4b. control key: caps a different server name, so fetch_mcp is uncapped -- +echo "" +echo "=== control key (caps ${OTHER_SERVER}, so ${SERVER} is uncapped) ===" +control_codes=() +for i in 1 2 3 4; do + code="$(call_mcp "${CONTROL_KEY}")" + control_codes+=("${code}") + echo " call ${i} -> HTTP ${code}" +done + +# --- evaluate ----------------------------------------------------------------- +echo "" +echo "=== result ===" +pass=true + +# limited: first two must NOT be 429, last two MUST be 429 +[[ "${limited_codes[0]}" != "429" ]] || { echo "FAIL: limited call 1 was rate limited"; pass=false; } +[[ "${limited_codes[1]}" != "429" ]] || { echo "FAIL: limited call 2 was rate limited"; pass=false; } +[[ "${limited_codes[2]}" == "429" ]] || { echo "FAIL: limited call 3 was NOT rate limited (got ${limited_codes[2]})"; pass=false; } +[[ "${limited_codes[3]}" == "429" ]] || { echo "FAIL: limited call 4 was NOT rate limited (got ${limited_codes[3]})"; pass=false; } + +# control: none may be 429 +for c in "${control_codes[@]}"; do + [[ "${c}" != "429" ]] || { echo "FAIL: control key was rate limited on ${SERVER} (got ${c})"; pass=false; } +done + +if [[ "${pass}" == "true" ]]; then + echo "PASS: ${SERVER} tripped at call 3 (429) under the limited key; the control key (which caps ${OTHER_SERVER}) was never rate limited on ${SERVER}." + exit 0 +else + echo "See proxy log for detail: ${PROXY_LOG}" + exit 1 +fi 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 @@ class TestUserAPIKeyAuthJwtClaims: 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/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 3e2eb4b02c2..85fb78a95e2 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,197 @@ 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): + 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, + ) + + +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(): + 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"}) + + assert _find_descriptor(descriptors, "mcp_per_key") 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 @@ class TestSetObjectMetadataField: _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: """