mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
feat(mcp/v2): route get_resource_templates through the egress transport
Add UpstreamConnection.list_resource_templates (resources/templates/list) and override get_resource_templates_from_server in the v2 manager, mirroring the resources list path: _should_defer -> _v2_connection -> conn.ok.list_resource_templates -> _create_prefixed_resource_templates, degrading to [] on failure. This was the last MCP verb still unconditionally on v1 (no v2 override existed); every protocol op is now on v2 for the mapped modes. Validated: 9 unit tests (the templates path exercised against an in-process FastMCP server with a resource template); proxy boots clean.
This commit is contained in:
parent
e868947c80
commit
b28a636fdc
3 changed files with 62 additions and 1 deletions
|
|
@ -34,6 +34,7 @@ if TYPE_CHECKING:
|
|||
Prompt,
|
||||
ReadResourceResult,
|
||||
Resource,
|
||||
ResourceTemplate,
|
||||
)
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import AnyUrl
|
||||
|
|
@ -270,6 +271,34 @@ class MCPServerManagerV2(MCPServerManager):
|
|||
return []
|
||||
return self._create_prefixed_resources(result.ok, server, add_prefix=add_prefix)
|
||||
|
||||
async def get_resource_templates_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[ResourceTemplate]:
|
||||
from litellm.proxy.gateway.mcp.result import Error
|
||||
|
||||
if self._should_defer(server, mcp_auth_header):
|
||||
return await super().get_resource_templates_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_resource_templates()
|
||||
if isinstance(result, Error):
|
||||
self._egress_list_failure(server, result.error)
|
||||
return []
|
||||
return self._create_prefixed_resource_templates(
|
||||
result.ok, server, add_prefix=add_prefix
|
||||
)
|
||||
|
||||
async def read_resource_from_server(
|
||||
self,
|
||||
server: MCPServer,
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ if TYPE_CHECKING:
|
|||
from mcp import ReadResourceResult, Resource
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import ProgressFnT
|
||||
from mcp.types import CallToolResult, GetPromptResult, Prompt
|
||||
from mcp.types import CallToolResult, GetPromptResult, Prompt, ResourceTemplate
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import AnyUrl
|
||||
|
||||
|
|
@ -231,3 +231,11 @@ class UpstreamConnection:
|
|||
return await session.read_resource(uri)
|
||||
|
||||
return await self._run(op)
|
||||
|
||||
async def list_resource_templates(
|
||||
self,
|
||||
) -> Result[List[ResourceTemplate], ConnError]:
|
||||
async def op(session: ClientSession) -> List[ResourceTemplate]:
|
||||
return (await session.list_resource_templates()).resourceTemplates
|
||||
|
||||
return await self._run(op)
|
||||
|
|
|
|||
|
|
@ -49,6 +49,10 @@ def _serve_echo():
|
|||
def info() -> str:
|
||||
return "echo server info"
|
||||
|
||||
@mcp.resource("echo://item/{item_id}")
|
||||
def item(item_id: str) -> str:
|
||||
return f"item {item_id}"
|
||||
|
||||
sock = socket.socket()
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
port = sock.getsockname()[1]
|
||||
|
|
@ -129,6 +133,26 @@ async def test_v2_override_lists_resources_via_upstream_connection(echo_server_u
|
|||
assert len(resources) >= 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v2_override_lists_resource_templates_via_upstream_connection(
|
||||
echo_server_url,
|
||||
):
|
||||
# resources/templates/list path: resolve() + UpstreamConnection.list_resource_templates (v2),
|
||||
# namespaced via the inherited _create_prefixed_resource_templates.
|
||||
manager = MCPServerManagerV2()
|
||||
server = MCPServer(
|
||||
server_id="echo1",
|
||||
name="echo1",
|
||||
transport=MCPTransport.http,
|
||||
url=echo_server_url,
|
||||
auth_type=MCPAuth.none,
|
||||
)
|
||||
templates = await manager.get_resource_templates_from_server(
|
||||
server, add_prefix=True
|
||||
)
|
||||
assert len(templates) >= 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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue