From f0b5f03bb67726d8b83a817904fc21025533a3db Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 2 Jun 2026 13:15:24 +0530 Subject: [PATCH 1/6] feat(proxy): add per-MCP-server RPM rate limiting for keys and teams Adds mcp_rpm_limit, a dict keyed by MCP server name (alias if set, else the configured name) that caps requests per minute per server for a key or team. The v3 rate limiter builds a per-server descriptor only when a limit is configured for the server being called, so other servers stay uncapped and no TPM reservation is engaged. Server identity is surfaced into the request data via mcp_rate_limit_server_name so the limiter can resolve it. --- .../mcp_server/mcp_server_manager.py | 3 + litellm/proxy/_types.py | 4 + litellm/proxy/auth/auth_utils.py | 34 +++ .../hooks/parallel_request_limiter_v3.py | 86 +++++++ .../key_management_endpoints.py | 2 + .../management_endpoints/team_endpoints.py | 3 +- litellm/proxy/utils.py | 1 + scripts/test_mcp_rpm_limit.sh | 238 ++++++++++++++++++ .../mcp_server/test_mcp_hook_extra_headers.py | 84 +++++++ .../hooks/test_parallel_request_limiter_v3.py | 194 ++++++++++++++ .../management_endpoints/test_common_utils.py | 27 ++ 11 files changed, 675 insertions(+), 1 deletion(-) create mode 100755 scripts/test_mcp_rpm_limit.sh 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..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 @@ 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], @@ -1653,6 +1726,19 @@ def _create_rate_limit_descriptors( 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 @@ 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/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 @@ 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/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 @@ 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: """ From 4fc8bfb0a9e181680f833f4997f2e8ca818a1a4e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 2 Jun 2026 14:14:46 +0530 Subject: [PATCH 2/6] fix(proxy): gate MCP rpm descriptors on call_mcp_tool; document mcp_rpm_limit param Only honor mcp_server_name when the call is an actual MCP tool call. Without this, a normal LLM request could inject mcp_server_name in its body to consume a target server's MCP quota and 429 legitimate tool calls. Also adds the mcp_rpm_limit parameter docstring to update_key, new_user, and user_update so the API docs validator passes. --- .../hooks/parallel_request_limiter_v3.py | 31 +++++++++++-------- .../internal_user_endpoints.py | 2 ++ .../key_management_endpoints.py | 1 + .../hooks/test_parallel_request_limiter_v3.py | 13 ++++++-- 4 files changed, 32 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 574a5d60d3c..d81f2cd0b3e 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 @@ -1606,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. @@ -1726,18 +1727,21 @@ def _create_rate_limit_descriptors( 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, - ) + # 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: + 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 @@ -2069,6 +2073,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..4fd28b0f6eb 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} - 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} - 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 1714efd7d1a..93b552e6f4c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2424,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" 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 85fb78a95e2..c16b8955d61 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 @@ -2910,13 +2910,14 @@ 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): +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, ) @@ -2956,13 +2957,21 @@ def test_mcp_per_key_descriptor_skipped_for_non_matching_server_v3(): 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"}) + 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 From 73c61e5c8995c6b140087b205b188172608aa61e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 2 Jun 2026 09:11:23 +0000 Subject: [PATCH 3/6] Fix MCP rate limit quota handling --- litellm/proxy/auth/auth_utils.py | 4 ++-- .../hooks/parallel_request_limiter_v3.py | 7 +++--- .../proxy/auth/test_auth_utils.py | 17 +++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 24 +++++++++++++++++++ 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 4afef4c4db7..f57d1848ffd 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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 diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d81f2cd0b3e..4343747d104 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1727,10 +1727,9 @@ def _create_rate_limit_descriptors( 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, 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 c16b8955d61..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 @@ -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( From a34126afae5e804fb310a3441897ecddb9a490a6 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 2 Jun 2026 17:32:43 +0530 Subject: [PATCH 4/6] Delete scripts/test_mcp_rpm_limit.sh --- scripts/test_mcp_rpm_limit.sh | 238 ---------------------------------- 1 file changed, 238 deletions(-) delete mode 100755 scripts/test_mcp_rpm_limit.sh diff --git a/scripts/test_mcp_rpm_limit.sh b/scripts/test_mcp_rpm_limit.sh deleted file mode 100755 index fa32a88cdfb..00000000000 --- a/scripts/test_mcp_rpm_limit.sh +++ /dev/null @@ -1,238 +0,0 @@ -#!/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 From 8c74ef7d37396796dab479a8cb30ecaafe86c577 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 2 Jun 2026 13:24:11 +0000 Subject: [PATCH 5/6] docs(proxy): clarify mcp_rpm_limit is enforced for keys and teams, not per user --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 4fd28b0f6eb..7b8f0f72e13 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -386,7 +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} + - 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. @@ -1428,7 +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} + - 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. From 1d52d364296ee640b3df1be85eb156a02aa5d070 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 2 Jun 2026 13:44:08 +0000 Subject: [PATCH 6/6] fix(proxy): accept mcp_rpm_limit in generate_key_helper_fn NewUserRequest and GenerateKeyRequest inherit mcp_rpm_limit from GenerateRequestBase, so /user/new and /key/generate forwarded the field to generate_key_helper_fn, which did not accept it and returned a 500 ("unexpected keyword argument 'mcp_rpm_limit'"). Accept the param and store it in metadata, matching model_rpm_limit/model_tpm_limit, so the limit is persisted where get_key_mcp_rpm_limit reads it. --- .../proxy/management_endpoints/key_management_endpoints.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 93b552e6f4c..80ded0bdd16 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3404,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, @@ -3482,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