scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(mcp): skip list_prompts/list_resources for servers that don't advertise them

Browse source 
Gate the MCP client's list_prompts, list_resources, and list_resource_templates
calls on the optional capabilities the upstream server declares in its initialize
result, so the proxy stops re-probing (and getting Method not found from) servers
that do not implement those optional methods on every aggregate /mcp request.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-01 17:08:31 +00:00
parent 23de7a15d9
commit b80578fa8f
2 changed files with 163 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
litellm/experimental_mcp_client/client.py
View file

@ -12,6 +12,7 @@ from typing import (
Dict,
Generator,
List,
Literal,
Optional,
Tuple,
TypeVar,
@ -36,6 +37,7 @@ from mcp.types import (
GetPromptResult,
Prompt,
ResourceTemplate,
ServerCapabilities,
TextContent,
)
from mcp.types import Tool as MCPTool
@ -231,6 +233,7 @@ class MCPClient:
# upstream client's auth= slot, taking precedence over the SigV4 aws_auth.
self._resolved_auth: Optional[httpx.Auth] = resolved_auth
self._last_initialize_instructions: Optional[str] = None
self._last_initialize_capabilities: Optional[ServerCapabilities] = None
self._sampling_callback: Optional[Callable] = sampling_callback
self._elicitation_callback: Optional[Callable] = elicitation_callback
self._logging_callback: Optional[Callable] = logging_callback
@ -360,10 +363,14 @@ class MCPClient:
try:
init_result = await session.initialize()
self._last_initialize_instructions = None
self._last_initialize_capabilities = None
if init_result is not None:
ins = getattr(init_result, "instructions", None)
if isinstance(ins, str) and ins.strip():
self._last_initialize_instructions = ins.strip()
capabilities = getattr(init_result, "capabilities", None)
if isinstance(capabilities, ServerCapabilities):
self._last_initialize_capabilities = capabilities
return await operation(session)
finally:
try:
@ -396,6 +403,7 @@ class MCPClient:
http_client: Optional[httpx.AsyncClient] = None
try:
self._last_initialize_instructions = None
self._last_initialize_capabilities = None
transport_ctx, http_client = self._create_transport_context()
return await self._execute_session_operation(transport_ctx, operation)
except Exception:
@ -622,15 +630,36 @@ class MCPClient:
# Return a default error result instead of raising
return self.error_tool_result(e)
def _server_declares_capability(self, capability: Literal["prompts", "resources"]) -> bool:
"""Whether the last upstream initialize advertised the given optional capability.
Per the MCP spec a server enumerates the feature groups it supports in
InitializeResult.capabilities, and a client must not send prompts/resources requests to a
server that does not declare them. When capabilities are unknown (no initialize seen yet)
we optimistically attempt the call so behavior is unchanged for servers we cannot classify.
"""
capabilities = self._last_initialize_capabilities
if capabilities is None:
return True
return getattr(capabilities, capability, None) is not None
async def list_prompts(self) -> List[Prompt]:
"""List available prompts from the server."""
verbose_logger.debug(f"MCP client listing tools from {self.server_url or 'stdio'}")
async def _list_prompts_operation(session: ClientSession):
if not self._server_declares_capability("prompts"):
return None
return await session.list_prompts()
try:
result = await self.run_with_session(_list_prompts_operation)
if result is None:
verbose_logger.debug(
f"MCP server {self.server_url or 'stdio'} does not advertise the prompts capability; "
"skipping list_prompts"
)
return []
prompt_count = len(result.prompts)
prompt_names = [prompt.name for prompt in result.prompts]
verbose_logger.info(
@ -704,10 +733,18 @@ class MCPClient:
verbose_logger.debug(f"MCP client listing resources from {self.server_url or 'stdio'}")
async def _list_resources_operation(session: ClientSession):
if not self._server_declares_capability("resources"):
return None
return await session.list_resources()
try:
result = await self.run_with_session(_list_resources_operation)
if result is None:
verbose_logger.debug(
f"MCP server {self.server_url or 'stdio'} does not advertise the resources capability; "
"skipping list_resources"
)
return []
resource_count = len(result.resources)
resource_names = [resource.name for resource in result.resources]
verbose_logger.info(
@ -740,10 +777,18 @@ class MCPClient:
verbose_logger.debug(f"MCP client listing resource templates from {self.server_url or 'stdio'}")
async def _list_resource_templates_operation(session: ClientSession):
if not self._server_declares_capability("resources"):
return None
return await session.list_resource_templates()
try:
result = await self.run_with_session(_list_resource_templates_operation)
if result is None:
verbose_logger.debug(
f"MCP server {self.server_url or 'stdio'} does not advertise the resources capability; "
"skipping list_resource_templates"
)
return []
resource_template_count = len(result.resourceTemplates)
resource_template_names = [resourceTemplate.name for resourceTemplate in result.resourceTemplates]
verbose_logger.info(

118
tests/test_litellm/experimental_mcp_client/test_mcp_client.py
View file

@ -408,6 +408,124 @@ class TestMCPClientInstructionsCapture:
assert client._last_initialize_instructions is None
# ---------------------------------------------------------------------------
# Optional-capability gating for list_prompts / list_resources / list_resource_templates
# ---------------------------------------------------------------------------
class TestOptionalCapabilityGating:
"""A server that does not advertise the prompts/resources capability in its initialize result
must not be sent the corresponding list_* request. This regresses issue #35460 where the gateway
re-probed list_prompts/list_resources on every aggregate request against servers that reject them
with "Method not found", wasting two round-trips per unsupporting server per request."""
@staticmethod
def _wire_session(mock_session_cls, mock_session, capabilities):
from mcp.types import Implementation, InitializeResult
init_result = InitializeResult(
protocolVersion="2025-06-18",
capabilities=capabilities,
serverInfo=Implementation(name="test-server", version="1.0.0"),
)
mock_session.initialize = AsyncMock(return_value=init_result)
session_ctx = MagicMock()
session_ctx.__aenter__ = AsyncMock(return_value=mock_session)
session_ctx.__aexit__ = AsyncMock(return_value=False)
mock_session_cls.return_value = session_ctx
@staticmethod
def _patch_transport(client):
transport_ctx = MagicMock()
transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock()))
transport_ctx.__aexit__ = AsyncMock(return_value=False)
return patch.object(client, "_create_transport_context", return_value=(transport_ctx, None))
@pytest.mark.asyncio
@patch("litellm.experimental_mcp_client.client.ClientSession")
async def test_list_prompts_skips_rpc_when_prompts_unsupported(self, mock_session_cls):
from mcp.types import ServerCapabilities, ToolsCapability
client = MCPClient(server_url="http://example.com/mcp", transport_type="http")
mock_session = AsyncMock()
self._wire_session(mock_session_cls, mock_session, ServerCapabilities(tools=ToolsCapability()))
with self._patch_transport(client):
result = await client.list_prompts()
assert result == []
mock_session.list_prompts.assert_not_called()
@pytest.mark.asyncio
@patch("litellm.experimental_mcp_client.client.ClientSession")
async def test_list_prompts_calls_rpc_when_prompts_supported(self, mock_session_cls):
from mcp.types import ListPromptsResult, Prompt, PromptsCapability, ServerCapabilities
client = MCPClient(server_url="http://example.com/mcp", transport_type="http")
mock_session = AsyncMock()
self._wire_session(mock_session_cls, mock_session, ServerCapabilities(prompts=PromptsCapability()))
mock_session.list_prompts = AsyncMock(return_value=ListPromptsResult(prompts=[Prompt(name="p1")]))
with self._patch_transport(client):
result = await client.list_prompts()
mock_session.list_prompts.assert_awaited_once()
assert [p.name for p in result] == ["p1"]
@pytest.mark.asyncio
@patch("litellm.experimental_mcp_client.client.ClientSession")
async def test_list_resources_skips_rpc_when_resources_unsupported(self, mock_session_cls):
from mcp.types import ServerCapabilities, ToolsCapability
client = MCPClient(server_url="http://example.com/mcp", transport_type="http")
mock_session = AsyncMock()
self._wire_session(mock_session_cls, mock_session, ServerCapabilities(tools=ToolsCapability()))
with self._patch_transport(client):
result = await client.list_resources()
assert result == []
mock_session.list_resources.assert_not_called()
@pytest.mark.asyncio
@patch("litellm.experimental_mcp_client.client.ClientSession")
async def test_list_resource_templates_skips_rpc_when_resources_unsupported(self, mock_session_cls):
from mcp.types import ServerCapabilities, ToolsCapability
client = MCPClient(server_url="http://example.com/mcp", transport_type="http")
mock_session = AsyncMock()
self._wire_session(mock_session_cls, mock_session, ServerCapabilities(tools=ToolsCapability()))
with self._patch_transport(client):
result = await client.list_resource_templates()
assert result == []
mock_session.list_resource_templates.assert_not_called()
@pytest.mark.asyncio
@patch("litellm.experimental_mcp_client.client.ClientSession")
async def test_list_resources_calls_rpc_when_resources_supported(self, mock_session_cls):
from mcp.types import ListResourcesResult, ResourcesCapability, ServerCapabilities
client = MCPClient(server_url="http://example.com/mcp", transport_type="http")
mock_session = AsyncMock()
self._wire_session(mock_session_cls, mock_session, ServerCapabilities(resources=ResourcesCapability()))
mock_session.list_resources = AsyncMock(return_value=ListResourcesResult(resources=[]))
with self._patch_transport(client):
await client.list_resources()
mock_session.list_resources.assert_awaited_once()
def test_unknown_capabilities_optimistically_attempts(self):
"""With no initialize seen yet (capabilities unknown) we do not suppress the call."""
client = MCPClient(server_url="http://example.com/mcp", transport_type="http")
assert client._last_initialize_capabilities is None
assert client._server_declares_capability("prompts") is True
assert client._server_declares_capability("resources") is True
# ---------------------------------------------------------------------------
# Transport error surfacing
# ---------------------------------------------------------------------------