Forward backend capabilities in ProxyProvider (#3956)

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Jeremiah Lowin 2026-04-17 15:22:08 -04:00 committed by GitHub
commit 98f69bdba0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 258 additions and 11 deletions

View file

@ -175,6 +175,10 @@ class LifespanMixin:
for provider in self.providers:
await stack.enter_async_context(provider.lifespan())
# After providers are up, adjust MCP handlers to reflect actual
# backend capabilities (removes handlers for unsupported methods).
self._sync_proxy_capabilities()
self._started.set()
try:
yield
@ -191,6 +195,112 @@ class LifespanMixin:
self._lifespan_result_set = False
self._lifespan_result = None
def _sync_proxy_capabilities(self: FastMCP) -> None:
"""Remove MCP handlers for capabilities the backend does not support.
After provider lifespans have run, any ProxyProvider instances have had a
chance to preload their backend's serverCapabilities. If the backend doesn't
support a capability (resources, prompts, tools) and there are no local
components of that type either, we remove the corresponding request handlers
from the low-level MCP server.
This has two effects:
1. The ``initialize`` response no longer advertises unsupported capabilities.
2. Clients that try to use an unsupported method receive a proper
``METHOD_NOT_FOUND`` (-32601) JSON-RPC error instead of an empty list.
The adjustment is conservative: if there are any providers whose capabilities
are not known (i.e. not a LocalProvider or ProxyProvider with loaded caps),
we leave the handlers untouched.
"""
import mcp.types
from fastmcp.server.providers.local_provider.local_provider import LocalProvider
from fastmcp.server.providers.proxy import ProxyProvider
from fastmcp.server.providers.wrapped_provider import _WrappedProvider
def _unwrap(p: Any) -> Any:
"""Recursively unwrap _WrappedProvider to reach the inner provider."""
while isinstance(p, _WrappedProvider):
p = p._inner
return p
# Restore handlers to the baseline that was saved at construction time
# so that a server reused across multiple lifespan cycles starts clean.
baseline = getattr(self._mcp_server, "_baseline_request_handlers", None)
if baseline is not None:
self._mcp_server.request_handlers = dict(baseline)
else:
self._mcp_server._baseline_request_handlers = dict( # type: ignore[attr-defined] # ty:ignore[unresolved-attribute]
self._mcp_server.request_handlers
)
# Unwrap all providers so we can inspect the actual provider type,
# including namespaced providers wrapped in _WrappedProvider.
unwrapped = [_unwrap(p) for p in self.providers]
all_proxy_providers = [p for p in unwrapped if isinstance(p, ProxyProvider)]
if not all_proxy_providers:
return
# If any ProxyProvider failed to preload capabilities, we can't safely
# prune: that backend's capabilities are unknown and removing handlers
# could break capabilities it can actually serve.
if any(p._backend_capabilities is None for p in all_proxy_providers):
return
# Only adjust when every provider is either a LocalProvider or a
# ProxyProvider with known capabilities. Unknown providers may have
# components we can't inspect synchronously, so we leave things alone.
if any(not isinstance(p, (LocalProvider, ProxyProvider)) for p in unwrapped):
return
# Aggregate: a capability is "supported" if ANY proxy backend supports it.
backend_caps = [
p._backend_capabilities
for p in all_proxy_providers
if p._backend_capabilities is not None
]
any_resources = any(bool(c.resources) for c in backend_caps)
any_prompts = any(bool(c.prompts) for c in backend_caps)
any_tools = any(bool(c.tools) for c in backend_caps)
# Check all LocalProvider instances for statically-registered components.
# A user may pass additional LocalProvider instances via the providers kwarg,
# so we aggregate across every LocalProvider in self.providers, not just
# the server's built-in self._local_provider.
from fastmcp.prompts.base import Prompt
from fastmcp.resources.base import Resource
from fastmcp.resources.template import ResourceTemplate
from fastmcp.tools.base import Tool
local_components = [
c
for p in unwrapped
if isinstance(p, LocalProvider)
for c in p._components.values()
]
local_has_resources = any(
isinstance(c, (Resource, ResourceTemplate)) for c in local_components
)
local_has_prompts = any(isinstance(c, Prompt) for c in local_components)
local_has_tools = any(isinstance(c, Tool) for c in local_components)
if not any_resources and not local_has_resources:
self._mcp_server.request_handlers.pop(mcp.types.ListResourcesRequest, None)
self._mcp_server.request_handlers.pop(
mcp.types.ListResourceTemplatesRequest, None
)
self._mcp_server.request_handlers.pop(mcp.types.ReadResourceRequest, None)
if not any_prompts and not local_has_prompts:
self._mcp_server.request_handlers.pop(mcp.types.ListPromptsRequest, None)
self._mcp_server.request_handlers.pop(mcp.types.GetPromptRequest, None)
if not any_tools and not local_has_tools:
self._mcp_server.request_handlers.pop(mcp.types.ListToolsRequest, None)
self._mcp_server.request_handlers.pop(mcp.types.CallToolRequest, None)
def _setup_task_protocol_handlers(self: FastMCP) -> None:
"""Register SEP-1686 task protocol handlers with SDK.

View file

@ -10,7 +10,8 @@ from __future__ import annotations
import base64
import inspect
import time
from collections.abc import Awaitable, Callable, Sequence
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any, cast
from urllib.parse import quote
@ -565,6 +566,7 @@ class ProxyProvider(Provider):
self._resources_cache: _CacheEntry[Resource] | None = None
self._templates_cache: _CacheEntry[ResourceTemplate] | None = None
self._prompts_cache: _CacheEntry[Prompt] | None = None
self._backend_capabilities: mcp.types.ServerCapabilities | None = None
async def _get_client(self) -> Client:
"""Gets a client instance by calling the sync or async factory."""
@ -720,6 +722,38 @@ class ProxyProvider(Provider):
return None
return max(matching, key=version_sort_key) # type: ignore[type-var] # ty:ignore[invalid-return-type]
# -------------------------------------------------------------------------
# Lifecycle
# -------------------------------------------------------------------------
@asynccontextmanager
async def lifespan(self) -> AsyncIterator[None]:
"""Preload backend capabilities at server startup.
Connects to the backend during lifespan to fetch its serverCapabilities
from the initialize response. These are stored on the provider and used
by the hosting FastMCP server to advertise accurate capabilities and to
remove handlers for methods the backend does not support.
"""
self._backend_capabilities = None
try:
client = await self._get_client()
async with client:
init_result = client.initialize_result
if init_result is not None:
self._backend_capabilities = init_result.capabilities
else:
logger.warning(
"ProxyProvider: backend did not return an initialize result; "
"capabilities will not be filtered"
)
except Exception as e:
logger.warning(
f"ProxyProvider: could not preload backend capabilities: {e}; "
"capabilities will not be filtered"
)
yield
# -------------------------------------------------------------------------
# Task methods
# -------------------------------------------------------------------------

View file

@ -49,6 +49,12 @@ class _WrappedProvider(Provider):
def __repr__(self) -> str:
return f"_WrappedProvider({self._inner!r}, transforms={self._transforms!r})"
@asynccontextmanager
async def lifespan(self) -> AsyncIterator[None]:
"""Delegate lifespan to the inner provider."""
async with self._inner.lifespan():
yield
# -------------------------------------------------------------------------
# Delegate to inner provider's public methods (which apply inner's transforms)
# -------------------------------------------------------------------------
@ -136,13 +142,3 @@ class _WrappedProvider(Provider):
]
if c.task_config.supports_tasks()
]
# -------------------------------------------------------------------------
# Lifecycle - combine with inner
# -------------------------------------------------------------------------
@asynccontextmanager
async def lifespan(self) -> AsyncIterator[None]:
"""Combine lifespan with inner provider."""
async with self._inner.lifespan():
yield

View file

@ -915,3 +915,110 @@ class TestProxyProviderCache:
result = await proxy.call_tool("greet", {"name": "Alice"})
mock_list.assert_not_called()
assert result.content[0].text == "Hello, Alice!" # type: ignore[union-attr] # ty:ignore[unresolved-attribute]
def _make_tools_only_backend() -> FastMCP:
"""Create a FastMCP backend that only supports tools (no resources/prompts)."""
backend = FastMCP("ToolsOnlyBackend")
@backend.tool
def my_tool() -> str:
return "result"
# Remove resource and prompt handlers to simulate a tools-only MCP server
backend._mcp_server.request_handlers.pop(mcp_types.ListResourcesRequest, None)
backend._mcp_server.request_handlers.pop(
mcp_types.ListResourceTemplatesRequest, None
)
backend._mcp_server.request_handlers.pop(mcp_types.ReadResourceRequest, None)
backend._mcp_server.request_handlers.pop(mcp_types.ListPromptsRequest, None)
backend._mcp_server.request_handlers.pop(mcp_types.GetPromptRequest, None)
return backend
class TestProxyCapabilityForwarding:
"""Proxy advertises the backend's actual capabilities (issue #3948)."""
async def test_proxy_forwards_tools_only_capabilities(self):
"""When backend only supports tools, proxy initialize response should not
include resources or prompts capabilities."""
backend = _make_tools_only_backend()
proxy = create_proxy(FastMCPTransport(backend))
async with Client(proxy) as client:
init_result = client.initialize_result
assert init_result is not None
caps = init_result.capabilities
assert caps.tools is not None
assert caps.resources is None
assert caps.prompts is None
async def test_proxy_full_backend_advertises_all_capabilities(
self, fastmcp_server: FastMCP
):
"""When backend supports tools, resources and prompts, proxy should
advertise all three."""
proxy = create_proxy(FastMCPTransport(fastmcp_server))
async with Client(proxy) as client:
init_result = client.initialize_result
assert init_result is not None
caps = init_result.capabilities
assert caps.tools is not None
assert caps.resources is not None
assert caps.prompts is not None
async def test_proxy_resources_list_returns_method_not_found(self):
"""When backend does not support resources, resources/list on the proxy
must return a METHOD_NOT_FOUND error, not an empty list."""
backend = _make_tools_only_backend()
proxy = create_proxy(FastMCPTransport(backend))
async with Client(proxy) as client:
with pytest.raises(McpError) as exc_info:
await client.list_resources()
assert exc_info.value.error.code == mcp_types.METHOD_NOT_FOUND
async def test_proxy_prompts_list_returns_method_not_found(self):
"""When backend does not support prompts, prompts/list on the proxy
must return a METHOD_NOT_FOUND error, not an empty list."""
backend = _make_tools_only_backend()
proxy = create_proxy(FastMCPTransport(backend))
async with Client(proxy) as client:
with pytest.raises(McpError) as exc_info:
await client.list_prompts()
assert exc_info.value.error.code == mcp_types.METHOD_NOT_FOUND
async def test_proxy_preserves_local_resources_even_without_backend_support(self):
"""When proxy has local resources but backend doesn't, resources capability
should still be advertised and local resources should be accessible."""
backend = _make_tools_only_backend()
proxy = create_proxy(FastMCPTransport(backend))
@proxy.resource("local://data")
def local_data() -> str:
return "local"
async with Client(proxy) as client:
init_result = client.initialize_result
assert init_result is not None
caps = init_result.capabilities
assert caps.tools is not None
assert caps.resources is not None # local resource present
assert caps.prompts is None # still no prompts
resources = await client.list_resources()
assert any(str(r.uri) == "local://data" for r in resources)
async def test_proxy_backend_capabilities_preloaded_in_lifespan(self):
"""ProxyProvider._backend_capabilities should be set after lifespan runs."""
backend = _make_tools_only_backend()
provider = ProxyProvider(lambda: ProxyClient(FastMCPTransport(backend)))
async with provider.lifespan():
caps = provider._backend_capabilities
assert caps is not None
assert caps.tools is not None
assert caps.resources is None
assert caps.prompts is None