mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
Forward backend capabilities in ProxyProvider (#3956)
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
parent
e3f845f558
commit
98f69bdba0
4 changed files with 258 additions and 11 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# -------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue