mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
feat(mcp/v2): route read_resource and get_prompt through the egress transport
Override read_resource_from_server and get_prompt_from_server to use the shared _should_defer + _v2_connection seam with UpstreamConnection.read_resource/get_prompt. These are single-result ops, so on failure they raise via a new _egress_item_failure: an exhaustive match over the CredError|ConnError tags (assert_never keeps it total) mapping unauthorized -> MCPUpstreamAuthError (401 re-auth) and the rest -> a new typed MCPUpstreamError with a semantic status (502 unreachable, 428 precondition, 500 misconfigured) instead of a bare exception. MCPUpstreamError currently falls through to the endpoints' generic handler; wiring it into the three to_http_exception conversion sites so the semantic status reaches the client is a small follow-up. Validated: read_resource + get_prompt via v2 against an in-process FastMCP server (7 unit tests pass); proxy boots and tools-list unchanged.
This commit is contained in:
parent
382167c874
commit
fcdda93e19
3 changed files with 140 additions and 2 deletions
|
|
@ -79,3 +79,25 @@ class MCPUpstreamAuthError(Exception):
|
|||
detail=detail,
|
||||
headers={"www-authenticate": challenge} if challenge else None,
|
||||
)
|
||||
|
||||
|
||||
class MCPUpstreamError(Exception):
|
||||
"""Raised when an egress call to an upstream MCP server fails for a non-auth reason: the upstream
|
||||
is unreachable, or the server's credential config is invalid/unsupported.
|
||||
|
||||
Single-result routes (``call_tool``, ``read_resource``, ``get_prompt``) raise this since there is
|
||||
no list to degrade to; auth failures use :class:`MCPUpstreamAuthError` instead so clients can
|
||||
trigger re-auth. Carries a ``status_code`` so the gateway returns a semantically correct response
|
||||
(e.g. 502 unreachable, 428 precondition, 500 misconfigured) rather than a blanket 500. Unlike the
|
||||
auth error there is no ``WWW-Authenticate`` challenge to fabricate, so it converts to an
|
||||
``HTTPException`` with no request context.
|
||||
"""
|
||||
|
||||
def __init__(self, status_code: int, server_name: str, detail: str) -> None:
|
||||
self.status_code = status_code
|
||||
self.server_name = server_name
|
||||
self.detail = detail
|
||||
super().__init__(f"Upstream MCP server {server_name!r}: {detail}")
|
||||
|
||||
def to_http_exception(self) -> HTTPException:
|
||||
return HTTPException(status_code=self.status_code, detail=self.detail)
|
||||
|
|
|
|||
|
|
@ -18,15 +18,18 @@ override/inbound-token path, and the JWT-signer guardrail.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Union
|
||||
from typing import TYPE_CHECKING, Dict, List, NoReturn, Optional, Union
|
||||
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
|
||||
from litellm.proxy._types import MCPTransport
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.types import Prompt, Resource
|
||||
from mcp.types import GetPromptResult, Prompt, ReadResourceResult, Resource
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import AnyUrl
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.gateway.mcp.outbound_credentials.types import CredError
|
||||
|
|
@ -207,6 +210,62 @@ class MCPServerManagerV2(MCPServerManager):
|
|||
return []
|
||||
return self._create_prefixed_resources(result.ok, server, add_prefix=add_prefix)
|
||||
|
||||
async def read_resource_from_server(
|
||||
self,
|
||||
server: MCPServer,
|
||||
url: AnyUrl,
|
||||
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
raw_headers: Optional[Dict[str, str]] = None,
|
||||
) -> ReadResourceResult:
|
||||
from litellm.proxy.gateway.mcp.result import Error
|
||||
|
||||
if self._should_defer(server, mcp_auth_header):
|
||||
return await super().read_resource_from_server(
|
||||
server, url, mcp_auth_header, extra_headers, raw_headers
|
||||
)
|
||||
conn = await self._v2_connection(
|
||||
server, None, raw_headers=raw_headers, extra_headers=extra_headers
|
||||
)
|
||||
if isinstance(conn, Error):
|
||||
self._egress_item_failure(server, conn.error)
|
||||
result = await conn.ok.read_resource(url)
|
||||
if isinstance(result, Error):
|
||||
self._egress_item_failure(server, result.error)
|
||||
return result.ok
|
||||
|
||||
async def get_prompt_from_server(
|
||||
self,
|
||||
server: MCPServer,
|
||||
prompt_name: str,
|
||||
arguments: Optional[Dict[str, object]] = None,
|
||||
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
raw_headers: Optional[Dict[str, str]] = None,
|
||||
) -> GetPromptResult:
|
||||
from litellm.proxy.gateway.mcp.result import Error
|
||||
|
||||
if self._should_defer(server, mcp_auth_header):
|
||||
return await super().get_prompt_from_server(
|
||||
server,
|
||||
prompt_name,
|
||||
arguments,
|
||||
mcp_auth_header,
|
||||
extra_headers,
|
||||
raw_headers,
|
||||
)
|
||||
conn = await self._v2_connection(
|
||||
server, None, raw_headers=raw_headers, extra_headers=extra_headers
|
||||
)
|
||||
if isinstance(conn, Error):
|
||||
self._egress_item_failure(server, conn.error)
|
||||
# MCP prompt arguments are strings on the wire; coerce the loose dict to match the op's type.
|
||||
str_args = {k: str(v) for k, v in arguments.items()} if arguments else None
|
||||
result = await conn.ok.get_prompt(prompt_name, str_args)
|
||||
if isinstance(result, Error):
|
||||
self._egress_item_failure(server, result.error)
|
||||
return result.ok
|
||||
|
||||
def _egress_list_failure(
|
||||
self, server: MCPServer, error: "CredError | ConnError"
|
||||
) -> None:
|
||||
|
|
@ -229,3 +288,28 @@ class MCPServerManagerV2(MCPServerManager):
|
|||
server.server_id,
|
||||
error.summary,
|
||||
)
|
||||
|
||||
def _egress_item_failure(
|
||||
self, server: MCPServer, error: "CredError | ConnError"
|
||||
) -> NoReturn:
|
||||
# Single-result path (call_tool / read_resource / get_prompt): no list to degrade to, so
|
||||
# every failure raises. unauthorized -> MCPUpstreamAuthError (401 + re-auth); the rest map to
|
||||
# a semantically correct MCPUpstreamError. assert_never keeps the map total -- a new error
|
||||
# variant fails the build here until it is given a mapping.
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import (
|
||||
MCPUpstreamAuthError,
|
||||
MCPUpstreamError,
|
||||
)
|
||||
|
||||
match error.tag:
|
||||
case "unauthorized":
|
||||
raise MCPUpstreamAuthError(
|
||||
status_code=401, www_authenticate=None, server_name=server.name
|
||||
)
|
||||
case "upstream_unavailable":
|
||||
raise MCPUpstreamError(502, server.name, error.summary)
|
||||
case "precondition_required":
|
||||
raise MCPUpstreamError(428, server.name, error.summary)
|
||||
case "misconfigured" | "unsupported_mode" | "not_implemented":
|
||||
raise MCPUpstreamError(500, server.name, error.summary)
|
||||
assert_never(error.tag)
|
||||
|
|
|
|||
|
|
@ -127,3 +127,35 @@ async def test_v2_override_lists_resources_via_upstream_connection(echo_server_u
|
|||
)
|
||||
resources = await manager.get_resources_from_server(server, add_prefix=True)
|
||||
assert len(resources) >= 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v2_override_reads_resource_via_upstream_connection(echo_server_url):
|
||||
# Single-result read path: resolve() + UpstreamConnection.read_resource (v2), raises on failure.
|
||||
from pydantic import AnyUrl
|
||||
|
||||
manager = MCPServerManagerV2()
|
||||
server = MCPServer(
|
||||
server_id="echo1",
|
||||
name="echo1",
|
||||
transport=MCPTransport.http,
|
||||
url=echo_server_url,
|
||||
auth_type=MCPAuth.none,
|
||||
)
|
||||
result = await manager.read_resource_from_server(server, AnyUrl("echo://info"))
|
||||
assert result.contents
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v2_override_gets_prompt_via_upstream_connection(echo_server_url):
|
||||
# Single-result prompt path: resolve() + UpstreamConnection.get_prompt (v2), raises on failure.
|
||||
manager = MCPServerManagerV2()
|
||||
server = MCPServer(
|
||||
server_id="echo1",
|
||||
name="echo1",
|
||||
transport=MCPTransport.http,
|
||||
url=echo_server_url,
|
||||
auth_type=MCPAuth.none,
|
||||
)
|
||||
result = await manager.get_prompt_from_server(server, "greeting", {"name": "Tin"})
|
||||
assert result.messages
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue