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 ac6300f380)
This commit is contained in:
yuneng-jiang 2026-10-16 15:22:16 -07:00 • committed by GitHub
parent c1d8cf165d
commit c61bb56f7a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 162 additions and 4 deletions

View file

@ -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]
)

View file

@ -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)