mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
fix(mcp): recheck cached LazyMCP catalog visibility
Filter cached LazyMCP catalogs through the current allowed-server set before returning list/status metadata, and mark mcp_client_ip as an optional SecretFields key.
This commit is contained in:
parent
03d67173c3
commit
9ce8a2d956
3 changed files with 116 additions and 3 deletions
|
|
@ -2080,7 +2080,27 @@ if MCP_AVAILABLE:
|
|||
cache_key = f"lazymcp:catalog:{scope_hash}"
|
||||
cached = await _lazymcp_cache_get(cache_key)
|
||||
if isinstance(cached, dict):
|
||||
return cached
|
||||
allowed_servers = await _get_lazymcp_allowed_servers(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_servers=mcp_servers,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
allowed_names = {
|
||||
_get_lazymcp_server_label(server).lower() for server in allowed_servers
|
||||
}
|
||||
filtered_servers = [
|
||||
server
|
||||
for server in cached.get("servers", [])
|
||||
if str(server.get("name", "")).lower() in allowed_names
|
||||
]
|
||||
return {
|
||||
**cached,
|
||||
"servers": filtered_servers,
|
||||
"server_count": len(filtered_servers),
|
||||
"tool_count": sum(
|
||||
server.get("tool_count", 0) for server in filtered_servers
|
||||
),
|
||||
}
|
||||
|
||||
allowed_servers = await _get_lazymcp_allowed_servers(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from typing import Optional
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
|
||||
class RedactedDict(dict):
|
||||
|
|
@ -24,4 +24,4 @@ class SecretFields(TypedDict):
|
|||
"""
|
||||
|
||||
raw_headers: dict
|
||||
mcp_client_ip: Optional[str]
|
||||
mcp_client_ip: NotRequired[Optional[str]]
|
||||
|
|
|
|||
|
|
@ -427,6 +427,99 @@ def test_lazymcp_cache_helpers_tolerate_cache_errors(monkeypatch):
|
|||
fake_cache.async_set_cache.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazymcp_cached_catalog_rechecks_current_visibility(monkeypatch):
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server_module
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
cached_catalog = {
|
||||
"description": "cached catalog",
|
||||
"server_count": 2,
|
||||
"tool_count": 3,
|
||||
"servers": [
|
||||
{"name": "visible", "description": "ok", "tool_count": 1, "tools": []},
|
||||
{
|
||||
"name": "revoked",
|
||||
"description": "must not leak",
|
||||
"tool_count": 2,
|
||||
"tools": [],
|
||||
},
|
||||
],
|
||||
}
|
||||
allowed_server = types.SimpleNamespace(
|
||||
server_id="visible-id", alias="visible", server_name=None, name="visible"
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
mcp_server_module,
|
||||
"_lazymcp_cache_get",
|
||||
AsyncMock(return_value=cached_catalog),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mcp_server_module,
|
||||
"_get_lazymcp_allowed_servers",
|
||||
AsyncMock(return_value=[allowed_server]),
|
||||
)
|
||||
|
||||
catalog = await mcp_server_module._get_lazymcp_catalog(
|
||||
user_api_key_auth=None,
|
||||
mcp_auth_header=None,
|
||||
mcp_servers=None,
|
||||
mcp_server_auth_headers=None,
|
||||
oauth2_headers=None,
|
||||
raw_headers=None,
|
||||
client_ip="10.0.0.1",
|
||||
)
|
||||
|
||||
assert catalog["server_count"] == 1
|
||||
assert catalog["tool_count"] == 1
|
||||
assert [server["name"] for server in catalog["servers"]] == ["visible"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazymcp_cached_catalog_hides_all_revoked_servers(monkeypatch):
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server_module
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
monkeypatch.setattr(
|
||||
mcp_server_module,
|
||||
"_lazymcp_cache_get",
|
||||
AsyncMock(
|
||||
return_value={
|
||||
"description": "cached catalog",
|
||||
"server_count": 1,
|
||||
"tool_count": 1,
|
||||
"servers": [
|
||||
{"name": "revoked", "description": "must not leak", "tools": []}
|
||||
],
|
||||
}
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mcp_server_module,
|
||||
"_get_lazymcp_allowed_servers",
|
||||
AsyncMock(return_value=[]),
|
||||
)
|
||||
|
||||
catalog = await mcp_server_module._get_lazymcp_catalog(
|
||||
user_api_key_auth=None,
|
||||
mcp_auth_header=None,
|
||||
mcp_servers=None,
|
||||
mcp_server_auth_headers=None,
|
||||
oauth2_headers=None,
|
||||
raw_headers=None,
|
||||
client_ip="10.0.0.1",
|
||||
)
|
||||
|
||||
assert catalog["server_count"] == 0
|
||||
assert catalog["tool_count"] == 0
|
||||
assert catalog["servers"] == []
|
||||
|
||||
|
||||
def test_invalidating_toolset_cache_tolerates_lazymcp_invalidation_error():
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue