feat(mcp): support MCP access group names in URL-based namespacing (#27726)
* 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 <cursoragent@cursor.com>
* 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 <cursoragent@cursor.com>
* 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 <cursoragent@cursor.com>
* 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 <cursoragent@cursor.com>
* fix(mcp): short-TTL negative cache for access-group existence lookup
An unauthenticated caller could repeatedly request /<unknown>/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 <mateo-berri@users.noreply.github.com>
* 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 <mateo-berri@users.noreply.github.com>
---------
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude <claude@anthropic.com>
Co-authored-by: mateo-berri <mateo@berri.ai>
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
baa68ebb12
commit
1294165768
@ -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 = [
|
||||
|
||||
@ -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:
|
||||
|
||||
488
tests/test_litellm/proxy/test_dynamic_mcp_route.py
Normal file
488
tests/test_litellm/proxy/test_dynamic_mcp_route.py
Normal file
@ -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
|
||||
Loading…
Reference in New Issue
Block a user