mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
* 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:
parent
c1d8cf165d
commit
c61bb56f7a
2 changed files with 162 additions and 4 deletions
|
|
@ -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]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue