scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

feat(mcp/v2): route list-prompts and list-resources through the egress transport

Browse source 
Override get_prompts_from_server and get_resources_from_server to use the shared _should_defer +
_v2_connection seam (resolve() + UpstreamConnection.list_prompts/list_resources), with the inherited
_create_prefixed_prompts/_create_prefixed_resources and the same degrade-to-empty-list contract as
the tools list path. _egress_list_failure now returns None (raise-on-unauthorized + log) so it is
generic across tools/prompts/resources; callers return their own typed empty list. The single-item
ops (read_resource, get_prompt) stay on v1 with call_tool (they share the raise-on-error contract).

Live-validated: prompts + resources list through v2 against an in-process FastMCP server (unit
tests); the proxy boots and tools-list is unchanged.
This commit is contained in:
Tin Chi Lo 2026-06-20 12:25:32 -07:00
parent 74f146e071
commit 382167c874
2 changed files with 105 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

78
litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py
View file

@ -3,8 +3,8 @@
``MCPServerManagerV2`` subclasses v1's ``MCPServerManager`` and overrides the per-server egress
methods to route through the v2 ``UpstreamConnection`` + ``resolve()`` instead of
``_create_mcp_client``. Registry, RBAC, cross-server aggregation, namespacing
(``_create_prefixed_tools``), and static-header resolution are inherited from v1 unchanged. It is
the egress manager, constructed at the composition root (see
(``_create_prefixed_*``), and static-header resolution are inherited from v1 unchanged. It is the
egress manager, constructed at the composition root (see
``mcp_server_manager._make_global_mcp_server_manager``); there is no opt-in flag (v2 is the egress
implementation).
@ -25,6 +25,7 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerM
from litellm.proxy._types import MCPTransport
if TYPE_CHECKING:
from mcp.types import Prompt, Resource
from mcp.types import Tool as MCPTool
from litellm.proxy._types import UserAPIKeyAuth
@ -146,20 +147,74 @@ class MCPServerManagerV2(MCPServerManager):
extra_headers=extra_headers,
)
if isinstance(conn, Error):
return self._egress_list_failure(server, conn.error)
self._egress_list_failure(server, conn.error)
return []
result = await conn.ok.list_tools()
if isinstance(result, Error):
return self._egress_list_failure(server, result.error)
self._egress_list_failure(server, result.error)
return []
return self._create_prefixed_tools(result.ok, server, add_prefix=add_prefix)
async def get_prompts_from_server(
self,
server: MCPServer,
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
add_prefix: bool = True,
raw_headers: Optional[Dict[str, str]] = None,
) -> List[Prompt]:
from litellm.proxy.gateway.mcp.result import Error
if self._should_defer(server, mcp_auth_header):
return await super().get_prompts_from_server(
server, mcp_auth_header, extra_headers, add_prefix, raw_headers
)
conn = await self._v2_connection(
server, None, raw_headers=raw_headers, extra_headers=extra_headers
)
if isinstance(conn, Error):
self._egress_list_failure(server, conn.error)
return []
result = await conn.ok.list_prompts()
if isinstance(result, Error):
self._egress_list_failure(server, result.error)
return []
return self._create_prefixed_prompts(result.ok, server, add_prefix=add_prefix)
async def get_resources_from_server(
self,
server: MCPServer,
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
add_prefix: bool = True,
raw_headers: Optional[Dict[str, str]] = None,
) -> List[Resource]:
from litellm.proxy.gateway.mcp.result import Error
if self._should_defer(server, mcp_auth_header):
return await super().get_resources_from_server(
server, mcp_auth_header, extra_headers, add_prefix, raw_headers
)
conn = await self._v2_connection(
server, None, raw_headers=raw_headers, extra_headers=extra_headers
)
if isinstance(conn, Error):
self._egress_list_failure(server, conn.error)
return []
result = await conn.ok.list_resources()
if isinstance(result, Error):
self._egress_list_failure(server, result.error)
return []
return self._create_prefixed_resources(result.ok, server, add_prefix=add_prefix)
def _egress_list_failure(
self, server: MCPServer, error: "CredError | ConnError"
) -> List[MCPTool]:
# List path: an upstream 401/403 (or a per-user mode with no usable credential) surfaces as
# MCPUpstreamAuthError so the client gets a 401 + WWW-Authenticate and starts the OAuth flow
# (this is the LIT-3795 behavior for non-delegated interactive oauth2). Any other failure
# degrades to an empty tool list (logged), so one bad server never collapses the federated
# catalog. The typed partial-failure marker is a later, separate surface.
) -> None:
# List path (tools/prompts/resources): an upstream 401/403 (or a per-user mode with no usable
# credential) surfaces as MCPUpstreamAuthError so the client gets a 401 + WWW-Authenticate and
# starts the OAuth flow (the LIT-3795 behavior for non-delegated interactive oauth2). Any
# other failure is logged and the caller degrades to an empty list, so one bad server never
# collapses the federated catalog. The typed partial-failure marker is a later surface.
if error.tag == "unauthorized":
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPUpstreamAuthError,
@ -169,9 +224,8 @@ class MCPServerManagerV2(MCPServerManager):
status_code=401, www_authenticate=None, server_name=server.name
)
verbose_logger.warning(
"v2 egress: tools unavailable for %s (%s): %s",
"v2 egress: items unavailable for %s (%s): %s",
server.name,
server.server_id,
error.summary,
)
return []

39
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py
View file

@ -41,6 +41,14 @@ def _serve_echo():
def echo(text: str) -> str:
return f"echo: {text}"
@mcp.prompt()
def greeting(name: str) -> str:
return f"Hello, {name}"
@mcp.resource("echo://info")
def info() -> str:
return "echo server info"
sock = socket.socket()
sock.bind(("127.0.0.1", 0))
port = sock.getsockname()[1]
@ -88,3 +96,34 @@ async def test_v2_override_lists_tools_via_upstream_connection(echo_server_url):
)
tools = await manager._get_tools_from_server(server, add_prefix=True)
assert any(t.name.endswith("echo") for t in tools)
@pytest.mark.asyncio
async def test_v2_override_lists_prompts_via_upstream_connection(echo_server_url):
# The none-mode prompts path goes through resolve() + UpstreamConnection.list_prompts (v2),
# namespaced via the inherited _create_prefixed_prompts.
manager = MCPServerManagerV2()
server = MCPServer(
server_id="echo1",
name="echo1",
transport=MCPTransport.http,
url=echo_server_url,
auth_type=MCPAuth.none,
)
prompts = await manager.get_prompts_from_server(server, add_prefix=True)
assert any("greeting" in p.name for p in prompts)
@pytest.mark.asyncio
async def test_v2_override_lists_resources_via_upstream_connection(echo_server_url):
# The none-mode resources path goes through resolve() + UpstreamConnection.list_resources (v2).
manager = MCPServerManagerV2()
server = MCPServer(
server_id="echo1",
name="echo1",
transport=MCPTransport.http,
url=echo_server_url,
auth_type=MCPAuth.none,
)
resources = await manager.get_resources_from_server(server, add_prefix=True)
assert len(resources) >= 1