From 1294165768a3168283f6787e25f651e367cecfa4 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 14 May 2026 08:50:38 +0530 Subject: [PATCH] feat(mcp): support MCP access group names in URL-based namespacing (#27726) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(mcp): support MCP access group names in URL-based namespacing Extends dynamic_mcp_route to resolve /{name}/mcp requests where {name} is an MCP access group tag or a comma-separated list of servers/groups, matching what the documentation promised but the handler did not implement. Resolution order: registered server alias → toolset → comma-separated list → single access group tag (404 if none match). Adds unit tests covering all four resolution paths plus 404 cases. Co-authored-by: Cursor * fix(mcp): address Greptile review comments on dynamic_mcp_route - Move comma-separated check before toolset DB lookup so comma names short-circuit without hitting the database - Cache access-group DB lookups via user_api_key_cache to avoid a raw find_many on every request (matches toolset caching pattern) - Remove unused response_started variable from _forward_as_mcp_path - Update tests to assert comma list skips toolset call and to mock cache Co-authored-by: Cursor * refactor(mcp): extract helpers to fix PLR0915 too-many-statements in dynamic_mcp_route Extract _mcp_forward_as_path and _is_mcp_access_group_cached as module-level helpers so dynamic_mcp_route stays under the 50-statement limit. Update tests to patch the new module-level symbols directly. Co-authored-by: Cursor * Avoid caching missing MCP access groups * fix(mcp): stream MCP responses via _stream_mcp_asgi_response instead of buffering _mcp_forward_as_path previously accumulated the full response body in memory before sending it. Replace the buffering custom_send pattern with _stream_mcp_asgi_response, which uses an asyncio.Queue bridge so chunks are yielded to the client as they arrive, preventing unbounded memory growth on large or long-lived MCP responses. Co-authored-by: Cursor * fix(mcp): short-TTL negative cache for access-group existence lookup An unauthenticated caller could repeatedly request //mcp and force a fresh DB lookup for the access-group existence check on every request (only positive results were cached). Cache negative results for a short DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL window (10s by default) so the DB is shielded from flooding while a transient DB error (which surfaces as an empty list) cannot hide a real group for long. https://claude.ai/code/session_01SjyPmwfmrq8fveFgw9iHW9 * fix(mcp): use plain int for access-group negative cache TTL Drop the os.getenv wrapper around DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL to avoid the documentation_test_env_keys check failing on the new variable. The negative-cache window is a small internal tuning constant, not a user-facing knob, so a plain integer is clearer than an env override. https://claude.ai/code/session_01SjyPmwfmrq8fveFgw9iHW9 * fix(mcp): validate, dedupe, and cap CSV tokens in dynamic MCP route For /{name1,name2,...}/mcp, validate every token resolves to a known server alias or access group, dedupe case-insensitively, and cap at DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS=16 before forwarding. - Bounds the per-request DB / cache fan-out an authenticated caller can trigger by stuffing the path with tokens (raised by veria-ai). - Returns 404 instead of forwarding when no token resolves, so the downstream server filter cannot silently fall back to the full allowed_mcp_servers list (raised by Cursor agentic security review). - Forwards only the resolved subset, so unknown tokens cannot ride along into the downstream filter. Co-authored-by: Mateo Wang * fix(mcp): exact-match CSV token dedupe to preserve case-sensitive distinct tokens Bugbot flagged that case-insensitive dedup on `MyGroup,mygroup` could collapse to whichever case appeared first and silently drop the matching casing if the downstream resolver is case-sensitive. Switch to exact-match dedup so distinct casings survive; whitespace-only differences still collapse via the .strip() before comparison. Co-authored-by: Mateo Wang --------- Co-authored-by: Cursor Co-authored-by: Claude Co-authored-by: mateo-berri Co-authored-by: Mateo Wang --- litellm/constants.py | 9 + litellm/proxy/proxy_server.py | 217 +++++--- .../proxy/test_dynamic_mcp_route.py | 488 ++++++++++++++++++ 3 files changed, 646 insertions(+), 68 deletions(-) create mode 100644 tests/test_litellm/proxy/test_dynamic_mcp_route.py diff --git a/litellm/constants.py b/litellm/constants.py index 072c2c358f7..e36746326cc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1569,6 +1569,15 @@ DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int( os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60) ) DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) +# Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated +# callers from forcing a DB query per request for unknown names, while bounding +# staleness so a transient DB error (which surfaces as an empty list) cannot +# hide a real group for long. +DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL = 10 +# Maximum number of comma-separated MCP server / access-group tokens accepted +# in a single ``/{name1,name2,...}/mcp`` URL. Bounds the per-request DB / cache +# fan-out an authenticated caller can trigger by stuffing the path with tokens. +DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS = 16 # Sentry Scrubbing Configuration SENTRY_DENYLIST = [ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ecef96351ae..fec2f47ad5e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15135,95 +15135,176 @@ async def toolset_mcp_route(toolset_name: str, request: Request): raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") +async def _mcp_forward_as_path(path_segment: str, request: Request): + """Rewrite path to /mcp/{path_segment} and stream the response.""" + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + + scope = dict(request.scope) + scope["path"] = f"/mcp/{path_segment}" + return await _stream_mcp_asgi_response( + handle_streamable_http_mcp, scope, request.receive + ) + + +async def _resolve_mcp_csv_tokens( + csv_segment: str, client_ip: Optional[str] +) -> List[str]: + """Validate a comma-separated ``/{name1,name2,...}/mcp`` segment. + + For each token, check (in order) whether it is a registered MCP server + alias / name or an MCP access group tag (cached). Tokens are stripped, + deduped (exact-match, keeping first occurrence in original order), and + capped at ``DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS`` to bound the + per-request DB / cache fan-out an authenticated caller can trigger by + stuffing the path with tokens. Dedup is case-sensitive on purpose: + downstream resolvers may treat names case-sensitively, so collapsing + ``MyGroup`` and ``mygroup`` would risk dropping a valid distinct token. + + Toolset names are intentionally NOT resolved here — toolsets bind a single + toolset id into request scope and have no defined semantics inside a + comma-separated server list. + + Returns the subset of resolved tokens in original order. An empty list + means the segment did not resolve to any known server / group; the caller + should treat that as a 404 instead of forwarding it downstream (where an + all-unmatched server filter falls back to the full ``allowed_mcp_servers`` + list and silently broadens the request scope). + """ + from litellm.constants import DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + seen: set = set() + deduped: List[str] = [] + for raw in csv_segment.split(","): + token = raw.strip() + if not token or token in seen: + continue + seen.add(token) + deduped.append(token) + if len(deduped) >= DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: + break + + resolved: List[str] = [] + for token in deduped: + if global_mcp_server_manager.get_mcp_server_by_name(token, client_ip=client_ip): + resolved.append(token) + continue + if await _is_mcp_access_group_cached(token): + resolved.append(token) + return resolved + + +async def _is_mcp_access_group_cached(name: str) -> bool: + """Return True if *name* is a known MCP access group tag. + + Positive results are cached for ``DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL`` + seconds. Negative results are cached for a short + ``DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL`` window so unauthenticated + callers cannot force a fresh DB lookup per request for unknown names, while + bounding staleness so a transient DB error (which surfaces as an empty + list) cannot hide a real group for long. + """ + from litellm.constants import ( + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL, + ) + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + cache_key = f"mcp_access_group_exists:{name}" + cached = await user_api_key_cache.async_get_cache(key=cache_key) + if cached is not None: + return bool(cached) + result = bool(await MCPRequestHandler._get_mcp_servers_from_access_groups([name])) + await user_api_key_cache.async_set_cache( + key=cache_key, + value=result, + ttl=( + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL + if result + else DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL + ), + ) + return result + + # Dynamic MCP server routes - handle /{mcp_server_name}/mcp @app.api_route( "/{mcp_server_name}/mcp", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], ) async def dynamic_mcp_route(mcp_server_name: str, request: Request): - """Handle dynamic MCP server routes like /github_mcp/mcp and toolset routes like /devtooling-prod/mcp""" + """Handle /{name}/mcp for MCP server aliases, toolsets, MCP access group tags, and comma-separated lists. + + Resolution order: + 1. Registered MCP server alias / name + 2. Comma-separated list (short-circuits before any DB call) + 3. Toolset name (DB lookup, cached) + 4. MCP access group tag (DB lookup, cached) + """ try: - # Validate that the MCP server exists from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils - from litellm.types.mcp import MCPAuth client_ip = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = global_mcp_server_manager.get_mcp_server_by_name( + + # 1. Registered MCP server alias + if global_mcp_server_manager.get_mcp_server_by_name( mcp_server_name, client_ip=client_ip - ) - if mcp_server is None: - # Check if this is a toolset name — toolsets are accessible at /{name}/mcp - # the same way individual servers are, no separate /toolset/ prefix needed. - if prisma_client is not None: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy._experimental.mcp_server.server import ( - _mcp_active_toolset_id, - handle_streamable_http_mcp, - ) + ): + return await _mcp_forward_as_path(mcp_server_name, request) - toolset = await global_mcp_server_manager.get_toolset_by_name_cached( - prisma_client, mcp_server_name + # 2. Comma-separated list — validate every token resolves to a known + # server alias or access group before forwarding. Bounds DB / cache + # fan-out and prevents the downstream filter from silently falling back + # to the full allowed_mcp_servers list when no token matches. + if "," in mcp_server_name: + resolved_tokens = await _resolve_mcp_csv_tokens(mcp_server_name, client_ip) + if not resolved_tokens: + raise HTTPException( + status_code=404, + detail=( + f"No MCP server, toolset, or access group in " + f"'{mcp_server_name}' resolved to a known target" + ), ) - if toolset is not None: - scope = dict(request.scope) - scope["path"] = "/mcp" + return await _mcp_forward_as_path(",".join(resolved_tokens), request) - token = _mcp_active_toolset_id.set(toolset.toolset_id) - try: - return await _stream_mcp_asgi_response( - handle_streamable_http_mcp, scope, request.receive - ) - finally: - _mcp_active_toolset_id.reset(token) - - raise HTTPException( - status_code=404, detail=f"MCP server '{mcp_server_name}' not found" + # 3. Toolset name (cached) + if prisma_client is not None: + from litellm.proxy._experimental.mcp_server.server import ( + _mcp_active_toolset_id, + handle_streamable_http_mcp, ) - # Create a new scope with the correct path format that the MCP handler expects - # Transform /{mcp_server_name}/mcp to /mcp/{mcp_server_name} - scope = dict(request.scope) - scope["path"] = f"/mcp/{mcp_server_name}" + toolset = await global_mcp_server_manager.get_toolset_by_name_cached( + prisma_client, mcp_server_name + ) + if toolset is not None: + scope = dict(request.scope) + scope["path"] = "/mcp" + token = _mcp_active_toolset_id.set(toolset.toolset_id) + try: + return await _stream_mcp_asgi_response( + handle_streamable_http_mcp, scope, request.receive + ) + finally: + _mcp_active_toolset_id.reset(token) - # Import the MCP handler - from litellm.proxy._experimental.mcp_server.server import ( - handle_streamable_http_mcp, - ) + # 4. MCP access group tag (cached) + if await _is_mcp_access_group_cached(mcp_server_name): + return await _mcp_forward_as_path(mcp_server_name, request) - # Create a custom send function to capture the response - response_started = False - response_body = b"" - response_status = 200 - response_headers = [] - - async def custom_send(message): - nonlocal response_started, response_body, response_status, response_headers - if message["type"] == "http.response.start": - response_started = True - response_status = message["status"] - response_headers = message.get("headers", []) - elif message["type"] == "http.response.body": - response_body += message.get("body", b"") - - # Call the existing MCP handler - await handle_streamable_http_mcp( - scope, receive=request.receive, send=custom_send - ) - - # Return the response - from starlette.responses import Response - - headers_dict = {k.decode(): v.decode() for k, v in response_headers} - return Response( - content=response_body, - status_code=response_status, - headers=headers_dict, - media_type=headers_dict.get("content-type", "application/json"), + raise HTTPException( + status_code=404, + detail=f"MCP server, toolset, or access group '{mcp_server_name}' not found", ) except HTTPException as e: diff --git a/tests/test_litellm/proxy/test_dynamic_mcp_route.py b/tests/test_litellm/proxy/test_dynamic_mcp_route.py new file mode 100644 index 00000000000..2462aff2119 --- /dev/null +++ b/tests/test_litellm/proxy/test_dynamic_mcp_route.py @@ -0,0 +1,488 @@ +""" +Tests for the dynamic_mcp_route handler in proxy_server.py. + +Covers the resolution order: + 1. Registered MCP server alias → forwards to /mcp/{name} + 2. Comma-separated list → short-circuits before any DB call; + forwarded to /mcp/{segment} + 3. Toolset name (cached) → sets toolset scope, forwards to /mcp + 4. MCP access group tag (cached) → forwards to /mcp/{name} when the group + resolves to at least one server + 5. Unknown name → 404 + +Patch targets are at the source modules because dynamic_mcp_route +uses lazy local imports inside the function body. +""" + +from unittest.mock import ANY, AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + +_MCP_MANAGER = "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" +_HANDLE_HTTP = ( + "litellm.proxy._experimental.mcp_server.server.handle_streamable_http_mcp" +) +_STREAM_ASGI = "litellm.proxy.proxy_server._stream_mcp_asgi_response" +_PRISMA = "litellm.proxy.proxy_server.prisma_client" +_IS_ACCESS_GROUP = "litellm.proxy.proxy_server._is_mcp_access_group_cached" +_USER_API_KEY_CACHE = "litellm.proxy.proxy_server.user_api_key_cache" +_GET_ACCESS_GROUP_SERVERS = ( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp." + "MCPRequestHandler._get_mcp_servers_from_access_groups" +) +_FORWARD = "litellm.proxy.proxy_server._mcp_forward_as_path" +_RESOLVE_CSV = "litellm.proxy.proxy_server._resolve_mcp_csv_tokens" + + +def _make_request(path: str = "/test/mcp"): + """Minimal fake Starlette Request.""" + from starlette.requests import Request + + scope = { + "type": "http", + "method": "POST", + "path": path, + "headers": [], + "query_string": b"", + "server": ("localhost", 4000), + "scheme": "http", + } + + async def receive(): + return {"type": "http.request", "body": b"{}"} + + return Request(scope=scope, receive=receive) + + +def _fake_server(name: str = "my_server", server_id: str = "server-id-1"): + s = MagicMock() + s.name = name + s.server_id = server_id + return s + + +def _fake_toolset(name: str = "my_toolset", toolset_id: str = "ts-1"): + t = MagicMock() + t.toolset_id = toolset_id + t.name = name + return t + + +async def _ok_mcp_handle(scope, receive, send): + """Stub MCP handler that returns HTTP 200.""" + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"{}"}) + + +# --------------------------------------------------------------------------- +# 1. Registered MCP server alias +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_resolves_registered_server(): + """When the segment matches a known server alias the request is forwarded + to /mcp/{name} and the handler returns 200.""" + from starlette.responses import Response + + from litellm.proxy.proxy_server import dynamic_mcp_route + + request = _make_request("/my_server/mcp") + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=_fake_server("my_server")) + + fake_forward = AsyncMock(return_value=Response(content=b"{}", status_code=200)) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_FORWARD, new=fake_forward), + ): + response = await dynamic_mcp_route("my_server", request) + + assert response.status_code == 200 + fake_forward.assert_awaited_once_with("my_server", request) + + +# --------------------------------------------------------------------------- +# 2. Comma-separated list (short-circuits before toolset DB call) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_comma_list_forwarded_when_tokens_resolve(): + """A comma-separated segment is forwarded after every token is resolved as + a known server / access group. Forwarding uses the deduped, validated + token list (so unknown / duplicate tokens cannot leak through). The + toolset DB lookup is bypassed entirely for comma names.""" + from starlette.responses import Response + + from litellm.proxy.proxy_server import dynamic_mcp_route + + segment = "github_mcp,zapier" + request = _make_request(f"/{segment}/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=None) + + fake_forward = AsyncMock(return_value=Response(content=b"{}", status_code=200)) + fake_resolve = AsyncMock(return_value=["github_mcp", "zapier"]) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_RESOLVE_CSV, new=fake_resolve), + patch(_FORWARD, new=fake_forward), + ): + response = await dynamic_mcp_route(segment, request) + + assert response.status_code == 200 + fake_forward.assert_awaited_once_with("github_mcp,zapier", request) + fake_mgr.get_toolset_by_name_cached.assert_not_called() + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_comma_list_returns_404_when_no_tokens_resolve(): + """A comma-separated segment with zero resolved tokens must 404 instead of + forwarding (downstream filter falls back to full allowed_mcp_servers when + no token matches, which would silently broaden scope).""" + from litellm.proxy.proxy_server import dynamic_mcp_route + + segment = "ghost1,ghost2" + request = _make_request(f"/{segment}/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=None) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_RESOLVE_CSV, new=AsyncMock(return_value=[])), + ): + with pytest.raises(HTTPException) as exc_info: + await dynamic_mcp_route(segment, request) + + assert exc_info.value.status_code == 404 + assert segment in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_comma_list_forwards_only_resolved_subset(): + """If only a subset of CSV tokens resolve, the request is forwarded with + just that subset (so unknown tokens cannot ride along into the downstream + server filter).""" + from starlette.responses import Response + + from litellm.proxy.proxy_server import dynamic_mcp_route + + segment = "github_mcp,ghost,zapier" + request = _make_request(f"/{segment}/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=None) + + fake_forward = AsyncMock(return_value=Response(content=b"{}", status_code=200)) + fake_resolve = AsyncMock(return_value=["github_mcp", "zapier"]) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_RESOLVE_CSV, new=fake_resolve), + patch(_FORWARD, new=fake_forward), + ): + await dynamic_mcp_route(segment, request) + + fake_forward.assert_awaited_once_with("github_mcp,zapier", request) + + +@pytest.mark.asyncio +async def test_resolve_mcp_csv_tokens_dedupes_and_caps(): + """_resolve_mcp_csv_tokens dedupes tokens exact-match (so distinct casings + are preserved — downstream resolution may be case-sensitive), drops empty + fragments, and stops looking up after DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS + unique tokens to bound DB / cache fan-out.""" + from litellm.constants import DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS + from litellm.proxy.proxy_server import _resolve_mcp_csv_tokens + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=_fake_server()) + + # "github_mcp" appears twice (once with surrounding whitespace) — must be + # collapsed to a single entry. "GITHUB_MCP" is a distinct exact token and + # is kept (downstream resolution may be case-sensitive). + csv = ",,github_mcp, github_mcp ,GITHUB_MCP," + ",".join( + f"srv_{i}" for i in range(DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS + 5) + ) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_IS_ACCESS_GROUP, new=AsyncMock(return_value=False)), + ): + resolved = await _resolve_mcp_csv_tokens(csv, client_ip=None) + + assert len(resolved) == DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS + assert resolved[0] == "github_mcp" + assert "GITHUB_MCP" in resolved + assert resolved.count("github_mcp") == 1 + + +@pytest.mark.asyncio +async def test_resolve_mcp_csv_tokens_drops_unknown_and_resolves_access_groups(): + """Unknown tokens are dropped; access-group tokens are accepted via the + cached existence helper (no per-call uncached DB hit).""" + from litellm.proxy.proxy_server import _resolve_mcp_csv_tokens + + fake_mgr = MagicMock() + # Only "registered_srv" is a known server alias. + fake_mgr.get_mcp_server_by_name = MagicMock( + side_effect=lambda name, client_ip=None: ( + _fake_server(name) if name == "registered_srv" else None + ) + ) + + # "dev_group" is a real access group; "ghost" is not. + is_group = AsyncMock(side_effect=lambda name: name == "dev_group") + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_IS_ACCESS_GROUP, new=is_group), + ): + resolved = await _resolve_mcp_csv_tokens( + "registered_srv,dev_group,ghost", client_ip=None + ) + + assert resolved == ["registered_srv", "dev_group"] + # Access-group lookup must NOT be called for "registered_srv" (already + # matched as a server alias) but MUST be called for "dev_group" and + # "ghost" (the only tokens that fall through to the access-group check). + assert {call.args[0] for call in is_group.await_args_list} == { + "dev_group", + "ghost", + } + + +# --------------------------------------------------------------------------- +# 3. Toolset name +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_resolves_toolset(): + """When the segment is a toolset name the toolset context var is set + and the request is forwarded to /mcp (not /mcp/{name}).""" + from litellm.proxy.proxy_server import dynamic_mcp_route + + request = _make_request("/my_toolset/mcp") + fake_toolset = _fake_toolset("my_toolset", "ts-42") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=None) + fake_mgr.get_toolset_by_name_cached = AsyncMock(return_value=fake_toolset) + + captured_toolset_id = None + captured_scope = {} + + async def fake_stream(fn, scope, receive): + nonlocal captured_toolset_id + from litellm.proxy._experimental.mcp_server.server import ( + _mcp_active_toolset_id, + ) + + captured_toolset_id = _mcp_active_toolset_id.get() + captured_scope.update(scope) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_PRISMA, new=MagicMock()), + patch(_STREAM_ASGI, new=AsyncMock(side_effect=fake_stream)), + ): + await dynamic_mcp_route("my_toolset", request) + + assert captured_toolset_id == "ts-42" + assert captured_scope.get("path") == "/mcp" + + +# --------------------------------------------------------------------------- +# 4. MCP access group tag (cached) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_resolves_access_group(): + """When the segment is an MCP access group the request is forwarded (not 404).""" + from starlette.responses import Response + + from litellm.proxy.proxy_server import dynamic_mcp_route + + request = _make_request("/dev_group/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=None) + fake_mgr.get_toolset_by_name_cached = AsyncMock(return_value=None) + + fake_forward = AsyncMock(return_value=Response(content=b"{}", status_code=200)) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_PRISMA, new=MagicMock()), + patch(_IS_ACCESS_GROUP, new=AsyncMock(return_value=True)), + patch(_FORWARD, new=fake_forward), + ): + response = await dynamic_mcp_route("dev_group", request) + + assert response.status_code == 200 + fake_forward.assert_awaited_once_with("dev_group", request) + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_access_group_called_with_correct_name(): + """The access group lookup receives exactly the segment from the URL.""" + from starlette.responses import Response + + from litellm.proxy.proxy_server import dynamic_mcp_route + + request = _make_request("/qa_tools/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=None) + fake_mgr.get_toolset_by_name_cached = AsyncMock(return_value=None) + + is_group = AsyncMock(return_value=True) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_PRISMA, new=MagicMock()), + patch(_IS_ACCESS_GROUP, new=is_group), + patch( + _FORWARD, + new=AsyncMock(return_value=Response(content=b"{}", status_code=200)), + ), + ): + await dynamic_mcp_route("qa_tools", request) + + is_group.assert_awaited_once_with("qa_tools") + + +@pytest.mark.asyncio +async def test_is_mcp_access_group_cached_caches_positive_result(): + """Known access groups are cached after resolving to one or more servers.""" + from litellm.proxy.proxy_server import _is_mcp_access_group_cached + + fake_cache = MagicMock() + fake_cache.async_get_cache = AsyncMock(return_value=None) + fake_cache.async_set_cache = AsyncMock() + get_access_group_servers = AsyncMock(return_value=["server-id"]) + + with ( + patch(_USER_API_KEY_CACHE, new=fake_cache), + patch(_GET_ACCESS_GROUP_SERVERS, new=get_access_group_servers), + ): + result = await _is_mcp_access_group_cached("dev_group") + + assert result is True + get_access_group_servers.assert_awaited_once_with(["dev_group"]) + fake_cache.async_set_cache.assert_awaited_once_with( + key="mcp_access_group_exists:dev_group", + value=True, + ttl=ANY, + ) + + +@pytest.mark.asyncio +async def test_is_mcp_access_group_cached_caches_negative_result_briefly(): + """Empty access-group lookups are cached with a short TTL so unauthenticated + callers cannot force a fresh DB lookup per request for unknown names.""" + from litellm.constants import DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL + from litellm.proxy.proxy_server import _is_mcp_access_group_cached + + fake_cache = MagicMock() + fake_cache.async_get_cache = AsyncMock(return_value=None) + fake_cache.async_set_cache = AsyncMock() + get_access_group_servers = AsyncMock(return_value=[]) + + with ( + patch(_USER_API_KEY_CACHE, new=fake_cache), + patch(_GET_ACCESS_GROUP_SERVERS, new=get_access_group_servers), + ): + result = await _is_mcp_access_group_cached("dev_group") + + assert result is False + get_access_group_servers.assert_awaited_once_with(["dev_group"]) + fake_cache.async_set_cache.assert_awaited_once_with( + key="mcp_access_group_exists:dev_group", + value=False, + ttl=DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL, + ) + + +@pytest.mark.asyncio +async def test_is_mcp_access_group_cached_returns_cached_negative_without_db(): + """A cached False entry short-circuits the DB lookup on subsequent calls.""" + from litellm.proxy.proxy_server import _is_mcp_access_group_cached + + fake_cache = MagicMock() + fake_cache.async_get_cache = AsyncMock(return_value=False) + fake_cache.async_set_cache = AsyncMock() + get_access_group_servers = AsyncMock(return_value=["server-id"]) + + with ( + patch(_USER_API_KEY_CACHE, new=fake_cache), + patch(_GET_ACCESS_GROUP_SERVERS, new=get_access_group_servers), + ): + result = await _is_mcp_access_group_cached("never_existed") + + assert result is False + get_access_group_servers.assert_not_awaited() + fake_cache.async_set_cache.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# 5. Unknown name → 404 +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_unknown_name_returns_404(): + """A segment that is not a server, toolset, or access group → 404.""" + from litellm.proxy.proxy_server import dynamic_mcp_route + + request = _make_request("/does_not_exist/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=None) + fake_mgr.get_toolset_by_name_cached = AsyncMock(return_value=None) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_PRISMA, new=MagicMock()), + patch(_IS_ACCESS_GROUP, new=AsyncMock(return_value=False)), + ): + with pytest.raises(HTTPException) as exc_info: + await dynamic_mcp_route("does_not_exist", request) + + assert exc_info.value.status_code == 404 + assert "does_not_exist" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_dynamic_mcp_route_empty_access_group_returns_404(): + """An access group tag that resolves to zero servers still returns 404.""" + from litellm.proxy.proxy_server import dynamic_mcp_route + + request = _make_request("/empty_group/mcp") + + fake_mgr = MagicMock() + fake_mgr.get_mcp_server_by_name = MagicMock(return_value=None) + fake_mgr.get_toolset_by_name_cached = AsyncMock(return_value=None) + + with ( + patch(_MCP_MANAGER, fake_mgr), + patch(_PRISMA, new=MagicMock()), + patch(_IS_ACCESS_GROUP, new=AsyncMock(return_value=False)), + ): + with pytest.raises(HTTPException) as exc_info: + await dynamic_mcp_route("empty_group", request) + + assert exc_info.value.status_code == 404