From c61bb56f7a6820a60d979b4411613bdddf314b88 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 16 Oct 2026 15:22:16 -0700 Subject: [PATCH] fix(mcp): stop proxy-started legacy listeners from consuming server rpm (#46968) (#46982) * fix(mcp): stop proxy-started legacy listeners from consuming server rpm The listener the proxy opens on notifications/initialized for legacy MCP sessions ran the server rate-limit admission once per target, so every session spent an rpm slot before the client sent anything. Admission is now injected into the listener and applied only to client-initiated subscriptions (subscriptions/listen and resources/subscribe). Authorization and its periodic recheck still run on every path. * test(mcp): key the fake admission window by caller and server * fix(mcp): stop subscription catalog refreshes from consuming server rpm A subscription listener lists the caller's catalog to diff list_changed events, at setup and on every upstream event. Those listings went through the normal per-server admission, so against an upstream that declares listChanged each legacy session spent several rpm slots before the client asked for anything. The listener now marks its targets as already admitted for its own lifetime, so its refreshes reuse the subscription's admission. Client requests on the session are still metered. * test(mcp): exercise subscription admission through the proxy rate limiter (cherry picked from commit ac6300f3800ad9fcf12ee0f572105cf277b05a60) --- .../_experimental/mcp_server/subscriptions.py | 36 ++++- .../mcp_server/test_subscriptions.py | 130 ++++++++++++++++++ 2 files changed, 162 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/subscriptions.py b/litellm/proxy/_experimental/mcp_server/subscriptions.py index dcf24f974fe..db622ea3cbf 100644 --- a/litellm/proxy/_experimental/mcp_server/subscriptions.py +++ b/litellm/proxy/_experimental/mcp_server/subscriptions.py @@ -55,6 +55,7 @@ from litellm.proxy._experimental.mcp_server.catalog import ( ) from litellm.proxy._experimental.mcp_server.contracts import OperationContext from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.types.mcp import MCP_LEGACY_VERSIONS from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -300,6 +301,29 @@ async def _listing(event: ServerEvent, context: OperationContext) -> str | None: raise MCPError(code=-32603, message="Subscription catalog exceeds the pagination limit") +async def _enforce_rate_limit(context: OperationContext, server: MCPServer) -> None: + await operations._enforce_mcp_server_rate_limit(context.user_api_key_auth, server) # pyright: ignore[reportPrivateUsage] # same admission policy as other gateway operations + + +async def _unmetered(_: OperationContext, __: MCPServer) -> None: + return None + + +async def _admitted() -> ProxyRateLimitError | None: + return None + + +@asynccontextmanager +async def _listener_admission(targets: tuple[SubscriptionTarget, ...]) -> AsyncGenerator[None]: + scope: Final = operations._mcp_server_admission_memo.set( # pyright: ignore[reportPrivateUsage] # the listener's own catalog refreshes reuse its admission + {target.server.server_id: asyncio.create_task(_admitted()) for target in targets} + ) + try: + yield + finally: + operations._mcp_server_admission_memo.reset(scope) # pyright: ignore[reportPrivateUsage] # restore the caller's admission scope + + async def _recheck_authorization(check: Callable[[], Awaitable[OperationContext]]) -> None: while True: await asyncio.sleep(15) @@ -361,7 +385,7 @@ class GatewaySubscriptions: if ctx.protocol_version != "2026-07-28": raise MCPError(code=-32601, message="Subscriptions require the modern MCP protocol") async with self._capacity(owner): - return await self._listen(ctx, params, authorize) + return await self._listen(ctx, params, authorize, _enforce_rate_limit) @asynccontextmanager async def _capacity(self, owner: str) -> AsyncGenerator[None]: @@ -414,6 +438,7 @@ class GatewaySubscriptions: ctx, SubscriptionFilter(tools_list_changed=True, prompts_list_changed=True, resources_list_changed=True), authorization, + _unmetered, ) if (ready := self._initializing.get(key)) is not None: ready.set() @@ -427,6 +452,7 @@ class GatewaySubscriptions: ctx: ServerRequestContext, notifications: SubscriptionFilter, authorization: SubscriptionAuthorization, + admit: Callable[[OperationContext, MCPServer], Awaitable[None]], scope: anyio.CancelScope | None = None, *, task_status: TaskStatus[None], @@ -443,7 +469,7 @@ class GatewaySubscriptions: async with self._capacity(authorization.owner): await self._listen( - ctx, SubscriptionsListenRequestParams(notifications=notifications), authorization, ready + ctx, SubscriptionsListenRequestParams(notifications=notifications), authorization, admit, ready ) async def legacy_resource(self, ctx: ServerRequestContext, uri: str, *, subscribe: bool) -> None: @@ -464,6 +490,7 @@ class GatewaySubscriptions: ctx, SubscriptionFilter(resource_subscriptions=[uri]), state.authorization, + _enforce_rate_limit, scope, ) state.resources = {**state.resources, uri: scope} @@ -476,6 +503,7 @@ class GatewaySubscriptions: ctx: ServerRequestContext, params: SubscriptionsListenRequestParams, authorize: Callable[[], Awaitable[OperationContext]], + admit: Callable[[OperationContext, MCPServer], Awaitable[None]], ready: Callable[[SubscriptionFilter], None] | None = None, ) -> SubscriptionsListenResult: context: Final = await authorize() @@ -501,7 +529,7 @@ class GatewaySubscriptions: return fresh async def relay(target: SubscriptionTarget, *, task_status: TaskStatus[SubscriptionFilter]) -> None: - await operations._enforce_mcp_server_rate_limit(context.user_api_key_auth, target.server) # pyright: ignore[reportPrivateUsage] # same admission policy as other gateway operations + await admit(context, target.server) legacy: Final = LegacyNotifications() @@ -567,7 +595,7 @@ class GatewaySubscriptions: finally: self._upstream_slots.release() - async with anyio.create_task_group() as group: + async with _listener_admission(targets), anyio.create_task_group() as group: honored: Final = tuple( [SubscriptionFilter.model_validate(await group.start(relay, target)) for target in targets] ) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py index a5d42e58b99..a97ab04c506 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py @@ -1,4 +1,5 @@ import asyncio +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, Mock @@ -19,6 +20,21 @@ from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer +def allow_one_request(monkeypatch: pytest.MonkeyPatch, user_id: str, server_id: str) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + windows: Final = {(user_id, server_id): asyncio.Semaphore(1)} + + async def enforce(user: UserAPIKeyAuth | None, server: MCPServer) -> None: + window: Final = windows[(user.user_id if user is not None else None, server.server_id)] + if window.locked(): + raise ProxyRateLimitError(detail=f"{server_id} RPM exceeded") + await window.acquire() + + monkeypatch.setattr(proxy_server, "proxy_logging_obj", SimpleNamespace(enforce_mcp_server_rate_limits=enforce)) + + def context(session: ServerSession, revision: str = "2026-07-28") -> ServerRequestContext: return ServerRequestContext( session=session, @@ -356,3 +372,117 @@ async def test_proxy_route_rejects_legacy_resource_subscription_before_setup(met await getattr(server, method)(context(Mock(spec=ServerSession)), params) finally: server._mcp_proxy_mode.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("client_requests", "tools_list_admitted"), + [((), True), (("subscriptions/listen",), False), (("resources/subscribe",), False)], +) +async def test_only_client_initiated_listeners_consume_server_rate_limit( + monkeypatch: pytest.MonkeyPatch, client_requests: tuple[str, ...], tools_list_admitted: bool +) -> None: + from litellm.proxy._experimental.mcp_server import operations + + allow_one_request(monkeypatch, "alice", "upstream") + opened: Final = asyncio.Queue[None]() + + async def hold_upstream(*_: object, **__: object) -> None: + opened.put_nowait(None) + await asyncio.Event().wait() + + client: Final = Mock(spec=MCPClient) + client.run_with_session = AsyncMock(side_effect=hold_upstream) + server: Final = MCPServer(server_id="upstream", name="upstream", url="http://localhost/mcp", transport="http") + target: Final = subscriptions.SubscriptionTarget(server, client, "synthetic-fingerprint") + + async def resolve(_: OperationContext) -> tuple[subscriptions.SubscriptionTarget, ...]: + return (target,) + + service: Final = subscriptions.GatewaySubscriptions(resolve_targets=resolve) + caller: Final = OperationContext(_caller=UserAPIKeyAuth(user_id="alice")) + authorize: Final = AsyncMock(return_value=caller) + session: Final = Mock(spec=ServerSession) + session.client_params = Mock() + legacy: Final = context(session, "2025-11-25") + requests: Final = { + "subscriptions/listen": lambda: service.listen( + context(Mock(spec=ServerSession)), + SubscriptionsListenRequestParams(notifications=SubscriptionFilter(tools_list_changed=True)), + authorize, + "alice", + ), + "resources/subscribe": lambda: service.legacy_resource(legacy, "test://value", subscribe=True), + } + authorization: Final = subscriptions.SubscriptionAuthorization(caller, authorize, lambda *_: "alice") + starts: Final = (lambda: service.legacy_session(legacy, authorization), *(requests[r] for r in client_requests)) + tasks: Final = tuple(asyncio.create_task(start()) for start in starts) + try: + for _ in tasks: + await asyncio.wait_for(opened.get(), 1) + rejection: Final = await operations._mcp_server_rate_limit_rejection(server, caller.user_api_key_auth) + assert (rejection is None) is tools_list_admitted + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_legacy_listener_catalog_refreshes_do_not_consume_server_rate_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from mcp.types import ListToolsResult, ServerCapabilities, Tool, ToolsCapability + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import catalog, operations + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk + + server: Final = MCPServer(server_id="list-changed", name="list-changed", transport=MCPTransport.http) + caller: Final = UserAPIKeyAuth(api_key="list-changed-key", user_id="alice") + allow_one_request(monkeypatch, "alice", "list-changed") + listed: Final = asyncio.Queue[None]() + + async def fetch_tools(_: MCPServer, **__: object) -> tuple[ListToolsResult, ServerListOk]: + listed.put_nowait(None) + return ListToolsResult(tools=[Tool(name="add", inputSchema={"type": "object"})]), ServerListOk(tool_count=1) + + monkeypatch.setenv("LITELLM_SALT_KEY", "list-changed-test-key") + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(operations.global_mcp_server_manager, "registry", {server.server_id: server}) + monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])) + monkeypatch.setattr(catalog, "get_filtered_server_tools", fetch_tools) + + upstream: Final = Mock() + upstream.server_capabilities = ServerCapabilities(tools=ToolsCapability(list_changed=True)) + upstream.protocol_version = "2025-11-25" + upstream.send_ping = AsyncMock() + + async def run_with_session(connected, **_: object) -> None: + await connected(upstream) + + client: Final = Mock(spec=MCPClient) + client.run_with_session = AsyncMock(side_effect=run_with_session) + target: Final = subscriptions.SubscriptionTarget(server, client, "synthetic-fingerprint") + + async def resolve(_: OperationContext) -> tuple[subscriptions.SubscriptionTarget, ...]: + return (target,) + + initial: Final = operations.prepare_context(user_api_key_auth=caller, mcp_servers=[server.server_id]) + authorize: Final = AsyncMock(return_value=initial) + session: Final = Mock(spec=ServerSession) + session.client_params = Mock() + service: Final = subscriptions.GatewaySubscriptions(resolve_targets=resolve) + listener: Final = asyncio.create_task( + service.legacy_session( + context(session, "2025-11-25"), + subscriptions.SubscriptionAuthorization(initial, authorize, lambda *_: "alice"), + ) + ) + try: + for _ in range(2): + await asyncio.wait_for(listed.get(), 5) + assert await operations._mcp_server_rate_limit_rejection(server, caller) is None + finally: + listener.cancel() + await asyncio.gather(listener, return_exceptions=True)