diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py index d5ca5bc6a32..4d71e79da73 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/v2_egress.py b/litellm/proxy/_experimental/mcp_server/v2_egress.py index f2d4b1c536b..f1430849793 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_egress.py +++ b/litellm/proxy/_experimental/mcp_server/v2_egress.py @@ -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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py index 9b5ba53ce05..05bae0ba56b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py @@ -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.