mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(mcp): list a never-listed server before its first tools/call
The startup tool-name fill skips servers whose upstream wants the caller's own token (true_passthrough, OAuth discovery), and mcp 2 no longer runs the list handler before an uncached tools/call, so every uvicorn worker that had not served tools/list answered 404 "Tool not found" for prefixed tools/call and the REST server_id route on those servers. On a resolution miss, execute_mcp_tool now lists the prefix-matched (or server_id-requested) server once, with the caller's credentials, through the existing tools/list path, then resolves as before. Listing failures fall through to the existing 404, a worker that already listed the server never re-lists it, and a server outside the caller's allowed set is never listed.
This commit is contained in:
parent
82ddab2405
commit
92ff54f134
4 changed files with 259 additions and 9 deletions
|
|
@ -2875,6 +2875,27 @@ class MCPServerManager:
|
|||
)
|
||||
return any(owner is not None and normalize_server_name(owner) in owned for owner in mapped_owners)
|
||||
|
||||
def has_listed_tools(self, server: MCPServer) -> bool:
|
||||
"""True once this worker holds at least one tool row for ``server``."""
|
||||
owned: Final = self._owned_mapping_values(server)
|
||||
return any(
|
||||
normalize_server_name(owner) in owned for owner in self.tool_name_to_mcp_server_name_mapping.values()
|
||||
)
|
||||
|
||||
def _known_prefix_to_server(self) -> Mapping[str, MCPServer]:
|
||||
"""Every prefix form a tool name may carry, keyed to its server; a form two servers share
|
||||
stays with the one registered first."""
|
||||
return {
|
||||
normalize_server_name(known_prefix): server
|
||||
for server in reversed(tuple(self.get_registry().values()))
|
||||
for known_prefix in iter_known_server_prefixes(server)
|
||||
}
|
||||
|
||||
def server_owning_tool_name_prefix(self, tool_name: str) -> MCPServer | None:
|
||||
prefix_to_server: Final = self._known_prefix_to_server()
|
||||
matched: Final = match_known_server_prefix(tool_name, prefix_to_server.keys())
|
||||
return None if matched is None else prefix_to_server.get(matched[0])
|
||||
|
||||
def remove_server(self, mcp_server: LiteLLM_MCPServerTable):
|
||||
"""
|
||||
Remove a server from the registry
|
||||
|
|
@ -6475,15 +6496,7 @@ class MCPServerManager:
|
|||
MCPServer if found, None otherwise
|
||||
"""
|
||||
registry_servers: Final = list(self.get_registry().values())
|
||||
|
||||
# Build prefix → server lookup covering every known form a tool name
|
||||
# may take (alias / server_name / server_id / short ID). This is what
|
||||
# makes the short-prefix mode work without breaking historical names.
|
||||
prefix_to_server: Final[dict[str, MCPServer]] = {}
|
||||
for server in registry_servers:
|
||||
for known_prefix in iter_known_server_prefixes(server):
|
||||
normalised = normalize_server_name(known_prefix)
|
||||
prefix_to_server.setdefault(normalised, server)
|
||||
prefix_to_server: Final = self._known_prefix_to_server()
|
||||
|
||||
# First try with the original tool name
|
||||
if tool_name in self.tool_name_to_mcp_server_name_mapping:
|
||||
|
|
|
|||
|
|
@ -2888,6 +2888,37 @@ if MCP_AVAILABLE:
|
|||
headers={"WWW-Authenticate": get_byok_www_authenticate()},
|
||||
)
|
||||
|
||||
async def _list_tools_before_first_call(
|
||||
server: MCPServer | None,
|
||||
allowed_mcp_servers: list[MCPServer],
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
mcp_auth_header: str | None,
|
||||
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
|
||||
oauth2_headers: dict[str, str] | None,
|
||||
raw_headers: dict[str, str] | None,
|
||||
) -> None:
|
||||
"""Fill this worker's tool rows for ``server`` with the caller's own credentials.
|
||||
|
||||
The startup fill skips a server whose upstream wants the caller's token, and mcp 2 no
|
||||
longer lists before an uncached tools/call, so a worker that has not served tools/list
|
||||
would otherwise answer 404 for every tool on that server.
|
||||
"""
|
||||
if server is None or global_mcp_server_manager.has_listed_tools(server):
|
||||
return
|
||||
if all(allowed.server_id != server.server_id for allowed in allowed_mcp_servers):
|
||||
return
|
||||
try:
|
||||
await _get_tools_from_mcp_servers(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_auth_header=mcp_auth_header,
|
||||
mcp_servers=[server.server_id],
|
||||
mcp_server_auth_headers=mcp_server_auth_headers,
|
||||
oauth2_headers=oauth2_headers,
|
||||
raw_headers=raw_headers,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best effort: resolution below answers as it did before
|
||||
verbose_logger.debug("MCP tools/call: listing %s before its first call failed: %s", server.name, e)
|
||||
|
||||
async def execute_mcp_tool(
|
||||
name: str,
|
||||
arguments: dict[str, object],
|
||||
|
|
@ -2948,6 +2979,21 @@ if MCP_AVAILABLE:
|
|||
all_registry_prefixes.add(normalize_server_name(known_prefix))
|
||||
name_is_prefixed = is_tool_name_prefixed(name, known_server_prefixes=all_registry_prefixes)
|
||||
|
||||
first_call_target: Final = (
|
||||
requested_server
|
||||
if requested_server is not None and not name_is_prefixed
|
||||
else global_mcp_server_manager.server_owning_tool_name_prefix(name)
|
||||
)
|
||||
await _list_tools_before_first_call(
|
||||
server=first_call_target,
|
||||
allowed_mcp_servers=allowed_mcp_servers,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_auth_header=mcp_auth_header,
|
||||
mcp_server_auth_headers=mcp_server_auth_headers,
|
||||
oauth2_headers=oauth2_headers,
|
||||
raw_headers=raw_headers,
|
||||
)
|
||||
|
||||
if requested_server is not None and not name_is_prefixed:
|
||||
# REST callers may pass server_id with the upstream tool name (no
|
||||
# LiteLLM prefix). The first segment is not a registered server
|
||||
|
|
|
|||
|
|
@ -7437,6 +7437,162 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool
|
|||
assert captured["name"] == "echo"
|
||||
|
||||
|
||||
def _never_listed_passthrough_server() -> MCPServer:
|
||||
return MCPServer(
|
||||
server_id="lazy-map-1",
|
||||
name="lazy_map",
|
||||
server_name="lazy_map",
|
||||
url="https://up.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.true_passthrough,
|
||||
)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _worker_that_never_listed(server: MCPServer, upstream_tools: tuple[str, ...]):
|
||||
"""A worker whose tool rows for ``server`` are empty, in front of an upstream that answers
|
||||
tools/list with ``upstream_tools`` and a managed dispatch that records what reaches it."""
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
mcp_module.global_mcp_server_manager.registry[server.server_id] = server
|
||||
dispatched: dict[str, object] = {}
|
||||
|
||||
async def fake_handle_managed_mcp_tool(**kwargs):
|
||||
dispatched.update(kwargs)
|
||||
return CallToolResult(content=[TextContent(type="text", text="ok")], is_error=False)
|
||||
|
||||
async def fake_fetch_tools(client, server_name):
|
||||
return [MCPTool(name=tool_name, inputSchema={}) for tool_name in upstream_tools]
|
||||
|
||||
with (
|
||||
patch.object( # test-quality-ok: the upstream MCP session is the boundary; a real one needs an initialize handshake over a live server
|
||||
mcp_module.global_mcp_server_manager,
|
||||
"_create_mcp_client",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
) as create_client,
|
||||
patch.object( # test-quality-ok: same boundary, this is the tools/list answer the upstream would give
|
||||
mcp_module.global_mcp_server_manager,
|
||||
"_fetch_tools_with_timeout",
|
||||
side_effect=fake_fetch_tools,
|
||||
) as fetch_tools,
|
||||
patch.object( # test-quality-ok: records the resolved server and bare name the managed call would forward upstream
|
||||
mcp_module, "_handle_managed_mcp_tool", new=fake_handle_managed_mcp_tool
|
||||
),
|
||||
):
|
||||
yield SimpleNamespace(create_client=create_client, fetch_tools=fetch_tools, dispatched=dispatched)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_mcp_tool_lists_never_listed_passthrough_server_with_caller_token_first():
|
||||
"""A prefixed tools/call on a worker that has not served tools/list must list that server once
|
||||
with the caller's own credentials and then dispatch, instead of answering 404."""
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
server = _never_listed_passthrough_server()
|
||||
with _worker_that_never_listed(server, upstream_tools=("add",)) as worker:
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="lazy_map-add",
|
||||
arguments={"a": 1, "b": 2},
|
||||
allowed_mcp_servers=[server],
|
||||
start_time=datetime.now(),
|
||||
mcp_auth_header="Bearer caller-token",
|
||||
raw_headers={"authorization": "Bearer caller-token"},
|
||||
)
|
||||
|
||||
assert worker.fetch_tools.await_count == 1
|
||||
assert "caller-token" in str(worker.create_client.await_args.kwargs.get("mcp_auth_header"))
|
||||
assert worker.dispatched["server_name"] == "lazy_map"
|
||||
assert worker.dispatched["name"] == "add"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_mcp_tool_rest_server_id_lists_never_listed_server_first():
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
server = _never_listed_passthrough_server()
|
||||
with _worker_that_never_listed(server, upstream_tools=("add",)) as worker:
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="add",
|
||||
arguments={"a": 1, "b": 2},
|
||||
allowed_mcp_servers=[server],
|
||||
start_time=datetime.now(),
|
||||
mcp_auth_header="Bearer caller-token",
|
||||
requested_server_id=server.server_id,
|
||||
)
|
||||
|
||||
assert worker.fetch_tools.await_count == 1
|
||||
assert worker.dispatched["server_name"] == "lazy_map"
|
||||
assert worker.dispatched["name"] == "add"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_mcp_tool_unknown_tool_on_never_listed_server_lists_once_then_404s():
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
server = _never_listed_passthrough_server()
|
||||
with (
|
||||
_worker_that_never_listed(server, upstream_tools=("add",)) as worker,
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="lazy_map-nope",
|
||||
arguments={},
|
||||
allowed_mcp_servers=[server],
|
||||
start_time=datetime.now(),
|
||||
mcp_auth_header="Bearer caller-token",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert worker.fetch_tools.await_count == 1
|
||||
assert worker.dispatched == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_mcp_tool_does_not_relist_a_server_this_worker_already_listed():
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
server = _never_listed_passthrough_server()
|
||||
with _worker_that_never_listed(server, upstream_tools=("add",)) as worker:
|
||||
mcp_module.global_mcp_server_manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server)
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="lazy_map-add",
|
||||
arguments={"a": 1, "b": 2},
|
||||
allowed_mcp_servers=[server],
|
||||
start_time=datetime.now(),
|
||||
mcp_auth_header="Bearer caller-token",
|
||||
)
|
||||
|
||||
assert worker.fetch_tools.await_count == 0
|
||||
assert worker.dispatched["name"] == "add"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_mcp_tool_never_lists_a_server_the_caller_cannot_access():
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
server = _never_listed_passthrough_server()
|
||||
other_server = MCPServer(server_id="other-1", name="other", transport=MCPTransport.http)
|
||||
with (
|
||||
_worker_that_never_listed(server, upstream_tools=("add",)) as worker,
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="lazy_map-add",
|
||||
arguments={"a": 1, "b": 2},
|
||||
allowed_mcp_servers=[other_server],
|
||||
start_time=datetime.now(),
|
||||
mcp_auth_header="Bearer caller-token",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert worker.fetch_tools.await_count == 0
|
||||
assert worker.dispatched == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_mcp_tool_strips_a_prefix_that_contains_the_separator():
|
||||
"""A server with no alias publishes its UUID server_id as the tool prefix.
|
||||
|
|
|
|||
|
|
@ -5556,6 +5556,41 @@ class TestMCPServerManager:
|
|||
)
|
||||
mock_inject.assert_awaited_once()
|
||||
|
||||
def test_server_owning_tool_name_prefix_is_known_before_the_server_is_ever_listed(self):
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="lazy-map-1",
|
||||
name="lazy_map",
|
||||
server_name="lazy_map",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.true_passthrough,
|
||||
)
|
||||
manager.registry = {server.server_id: server}
|
||||
|
||||
assert manager._get_mcp_server_from_tool_name("lazy_map-add") is None
|
||||
assert manager.server_owning_tool_name_prefix("lazy_map-add") is server
|
||||
assert manager.server_owning_tool_name_prefix("someone_else-add") is None
|
||||
assert manager.has_listed_tools(server) is False
|
||||
|
||||
manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server)
|
||||
|
||||
assert manager.has_listed_tools(server) is True
|
||||
assert manager._get_mcp_server_from_tool_name("lazy_map-add") is server
|
||||
|
||||
def test_known_prefix_to_server_keeps_the_first_registered_owner_of_a_shared_prefix(self):
|
||||
manager = MCPServerManager()
|
||||
first = MCPServer(server_id="first-id", name="first", server_name="first", transport=MCPTransport.http)
|
||||
second = MCPServer(
|
||||
server_id="second-id", name="second", server_name="second", alias="first", transport=MCPTransport.http
|
||||
)
|
||||
manager.registry = {"first-id": first, "second-id": second}
|
||||
|
||||
prefix_to_server = manager._known_prefix_to_server()
|
||||
|
||||
assert prefix_to_server["first"] is first
|
||||
assert prefix_to_server["second"] is second
|
||||
assert manager.server_owning_tool_name_prefix("first-add") is first
|
||||
|
||||
def test_resolve_mcp_server_for_tool_call_via_prefixed_name(self):
|
||||
"""Resolution succeeds when the prefixed tool name is in the mapping."""
|
||||
manager = MCPServerManager()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue