mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-26 23:44:17 +02:00
Don't era-gate server-side sampling handlers on modern connections
The era-gate blocked every ctx.sample/sample_step on a 2026-07-28 connection, but a server-configured sampling handler answers server-side without the client back-channel. Gate only when the request would hit the removed client path; force the handler path (client_available=False) on modern so "fallback" goes straight to the handler instead of a bare client-attempt failure.
This commit is contained in:
Jeremiah Lowin
parent
b023aae51a
commit
77131edc00
3 changed files with 157 additions and 15 deletions
|
|
@ -80,10 +80,12 @@ ResultT = TypeVar("ResultT", default=str)
|
|||
# Import ToolChoiceOption from sampling module (after other imports)
|
||||
from fastmcp.server.sampling.run import ToolChoiceOption # noqa: E402
|
||||
|
||||
# Warn-once guards for the sampling deprecation. Server-initiated createMessage
|
||||
# Warn-once guard for the sampling deprecation. Server-initiated createMessage
|
||||
# was removed from MCP as of 2026-07-28 (SEP-2577); the warning fires a single
|
||||
# time per process to flag that ctx.sample/ctx.sample_step are on their way out.
|
||||
_sample_deprecation_warned = False
|
||||
# A mutable set (mutated in place, never rebound) rather than a `global` boolean
|
||||
# so the warn-once state is unambiguously read and written from the module.
|
||||
_sample_deprecation_warned: set[bool] = set()
|
||||
|
||||
_SAMPLING_DEPRECATION_MESSAGE = (
|
||||
"ctx.sample() and ctx.sample_step() are deprecated and will be removed in a "
|
||||
|
|
@ -110,10 +112,9 @@ def _warn_sampling_deprecated() -> None:
|
|||
Gated on ``settings.deprecation_warnings`` like every other FastMCP
|
||||
deprecation; fires a single time (module-level flag) rather than per call.
|
||||
"""
|
||||
global _sample_deprecation_warned
|
||||
if _sample_deprecation_warned or not fastmcp.settings.deprecation_warnings:
|
||||
return
|
||||
_sample_deprecation_warned = True
|
||||
_sample_deprecation_warned.add(True)
|
||||
warnings.warn(
|
||||
_SAMPLING_DEPRECATION_MESSAGE,
|
||||
FastMCPDeprecationWarning,
|
||||
|
|
@ -918,6 +919,23 @@ class Context:
|
|||
return False
|
||||
return rc.protocol_version in MODERN_PROTOCOL_VERSIONS
|
||||
|
||||
def _server_can_sample(self) -> bool:
|
||||
"""True when a server-configured sampling handler can serve the request
|
||||
without the client back-channel.
|
||||
|
||||
FastMCP supports a server-side sampling handler (``FastMCP(sampling_handler=...)``).
|
||||
With ``sampling_handler_behavior="always"`` the handler always answers;
|
||||
with ``"fallback"`` it answers whenever the client cannot. On modern
|
||||
connections the client back-channel is gone, so either configuration lets
|
||||
the server answer entirely server-side as long as a handler is set. (For
|
||||
``"always"`` without a handler the sampling implementation raises its own
|
||||
clear "no handler configured" error, which is not an era concern.)
|
||||
"""
|
||||
fastmcp = self.fastmcp
|
||||
if fastmcp.sampling_handler_behavior == "always":
|
||||
return True
|
||||
return fastmcp.sampling_handler is not None
|
||||
|
||||
async def sample_step(
|
||||
self,
|
||||
messages: str | Sequence[str | SamplingMessage],
|
||||
|
|
@ -983,7 +1001,13 @@ class Context:
|
|||
messages = step.history
|
||||
"""
|
||||
_warn_sampling_deprecated()
|
||||
if self._is_modern_protocol():
|
||||
# On modern (2026-07-28) connections the client back-channel is gone
|
||||
# (SEP-2577). A server-configured sampling handler can still answer
|
||||
# entirely server-side; only raise the era error when nothing can serve
|
||||
# the request. When modern, force the handler path (never attempt the
|
||||
# dead client) by passing client_available=False.
|
||||
client_available = not self._is_modern_protocol()
|
||||
if not client_available and not self._server_can_sample():
|
||||
raise ToolError(_SAMPLING_MODERN_ERROR)
|
||||
return await sample_step_impl(
|
||||
self,
|
||||
|
|
@ -997,6 +1021,7 @@ class Context:
|
|||
auto_execute_tools=execute_tools,
|
||||
mask_error_details=mask_error_details,
|
||||
tool_concurrency=tool_concurrency,
|
||||
client_available=client_available,
|
||||
)
|
||||
|
||||
@overload
|
||||
|
|
@ -1093,7 +1118,13 @@ class Context:
|
|||
future FastMCP release. Call an LLM directly from your server instead.
|
||||
"""
|
||||
_warn_sampling_deprecated()
|
||||
if self._is_modern_protocol():
|
||||
# On modern (2026-07-28) connections the client back-channel is gone
|
||||
# (SEP-2577). A server-configured sampling handler can still answer
|
||||
# entirely server-side; only raise the era error when nothing can serve
|
||||
# the request. When modern, force the handler path (never attempt the
|
||||
# dead client) by passing client_available=False.
|
||||
client_available = not self._is_modern_protocol()
|
||||
if not client_available and not self._server_can_sample():
|
||||
raise ToolError(_SAMPLING_MODERN_ERROR)
|
||||
return await sample_impl( # ty: ignore[invalid-return-type]
|
||||
self,
|
||||
|
|
@ -1106,6 +1137,7 @@ class Context:
|
|||
result_type=result_type,
|
||||
mask_error_details=mask_error_details,
|
||||
tool_concurrency=tool_concurrency,
|
||||
client_available=client_available,
|
||||
)
|
||||
|
||||
@overload
|
||||
|
|
|
|||
|
|
@ -138,12 +138,19 @@ def _parse_model_preferences(
|
|||
# --- Standalone functions for sample_step() ---
|
||||
|
||||
|
||||
def determine_handler_mode(context: Context, needs_tools: bool) -> bool:
|
||||
def determine_handler_mode(
|
||||
context: Context, needs_tools: bool, *, client_available: bool = True
|
||||
) -> bool:
|
||||
"""Determine whether to use fallback handler or client for sampling.
|
||||
|
||||
Args:
|
||||
context: The MCP context.
|
||||
needs_tools: Whether the sampling request requires tool support.
|
||||
client_available: Whether the client back-channel can be reached at all.
|
||||
On modern (2026-07-28) connections the server-initiated createMessage
|
||||
back-channel was removed (SEP-2577), so the client can never serve a
|
||||
sampling request; pass False there to force the server-side handler
|
||||
path (``"fallback"`` behaves like ``"always"`` when a handler exists).
|
||||
|
||||
Returns:
|
||||
True if fallback handler should be used, False to use client.
|
||||
|
|
@ -154,11 +161,13 @@ def determine_handler_mode(context: Context, needs_tools: bool) -> bool:
|
|||
fastmcp = context.fastmcp
|
||||
session = context.session
|
||||
|
||||
# Check what capabilities the client has
|
||||
has_sampling = session.check_client_capability(
|
||||
# Check what capabilities the client has. On connections without a
|
||||
# back-channel the client can never serve the request regardless of the
|
||||
# capabilities it advertised, so treat both as unavailable.
|
||||
has_sampling = client_available and session.check_client_capability(
|
||||
capability=ClientCapabilities(sampling=SamplingCapability())
|
||||
)
|
||||
has_tools_capability = session.check_client_capability(
|
||||
has_tools_capability = client_available and session.check_client_capability(
|
||||
capability=ClientCapabilities(
|
||||
sampling=SamplingCapability(tools=SamplingToolsCapability())
|
||||
)
|
||||
|
|
@ -498,11 +507,17 @@ async def sample_step_impl(
|
|||
auto_execute_tools: bool = True,
|
||||
mask_error_details: bool | None = None,
|
||||
tool_concurrency: int | None = None,
|
||||
client_available: bool = True,
|
||||
) -> SampleStep:
|
||||
"""Implementation of Context.sample_step().
|
||||
|
||||
Make a single LLM sampling call. This is a stateless function that makes
|
||||
exactly one LLM call and optionally executes any requested tools.
|
||||
|
||||
When ``client_available`` is False (e.g. a modern 2026-07-28 connection with
|
||||
no back-channel), the client is never used and a configured sampling handler
|
||||
serves the request; the caller is responsible for raising a clear era error
|
||||
when no handler can serve it.
|
||||
"""
|
||||
# Convert messages to SamplingMessage objects
|
||||
current_messages = prepare_messages(messages)
|
||||
|
|
@ -517,7 +532,9 @@ async def sample_step_impl(
|
|||
)
|
||||
|
||||
# Determine whether to use fallback handler or client
|
||||
use_fallback = determine_handler_mode(context, bool(sampling_tools))
|
||||
use_fallback = determine_handler_mode(
|
||||
context, bool(sampling_tools), client_available=client_available
|
||||
)
|
||||
|
||||
# Build tool choice
|
||||
effective_tool_choice: ToolChoice | None = None
|
||||
|
|
@ -633,12 +650,18 @@ async def sample_impl(
|
|||
result_type: type[ResultT] | None = None,
|
||||
mask_error_details: bool | None = None,
|
||||
tool_concurrency: int | None = None,
|
||||
client_available: bool = True,
|
||||
) -> SamplingResult[ResultT]:
|
||||
"""Implementation of Context.sample().
|
||||
|
||||
Send a sampling request to the client and await the response. This method
|
||||
runs to completion automatically, executing a tool loop until the LLM
|
||||
provides a final text response.
|
||||
|
||||
When ``client_available`` is False (e.g. a modern 2026-07-28 connection with
|
||||
no back-channel), the client is never used and a configured sampling handler
|
||||
serves the request; the caller is responsible for raising a clear era error
|
||||
when no handler can serve it.
|
||||
"""
|
||||
# Safety limit to prevent infinite loops
|
||||
max_iterations = 100
|
||||
|
|
@ -675,6 +698,7 @@ async def sample_impl(
|
|||
tool_choice=tool_choice,
|
||||
mask_error_details=mask_error_details,
|
||||
tool_concurrency=tool_concurrency,
|
||||
client_available=client_available,
|
||||
)
|
||||
|
||||
# Check for final_response tool call for structured output
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from mcp_types.version import (
|
|||
)
|
||||
from pydantic import FileUrl
|
||||
|
||||
import fastmcp
|
||||
from fastmcp import Client as FastMCPClient
|
||||
from fastmcp import Context, FastMCP
|
||||
from fastmcp.server.elicitation import AcceptedElicitation
|
||||
|
|
@ -340,6 +341,91 @@ async def test_elicit_sample_degradation_message_is_clear_on_modern(push_server,
|
|||
assert "server-initiated" in message
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3a-bis. Server-configured sampling handler answers WITHOUT the client
|
||||
# back-channel, so ctx.sample()/ctx.sample_step() must keep working on modern
|
||||
# connections. The era-gate only fires when nothing can serve the request.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _handler_server(behavior) -> FastMCP:
|
||||
"""A server whose sampling is answered by a server-side handler."""
|
||||
|
||||
def sampling_handler(messages, params, ctx) -> str:
|
||||
return "handler-answer"
|
||||
|
||||
mcp = FastMCP("handler", sampling_handler=sampling_handler)
|
||||
if behavior is not None:
|
||||
mcp.sampling_handler_behavior = behavior
|
||||
|
||||
@mcp.tool
|
||||
async def do_sample(ctx: Context) -> str:
|
||||
result = await ctx.sample("hello")
|
||||
return f"sampled {result.text}"
|
||||
|
||||
@mcp.tool
|
||||
async def do_sample_step(ctx: Context) -> str:
|
||||
step = await ctx.sample_step("hello")
|
||||
return f"stepped {step.text}"
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", MODERN_MODES)
|
||||
@pytest.mark.parametrize("behavior", ["always", "fallback"])
|
||||
@pytest.mark.parametrize("method", ["do_sample", "do_sample_step"])
|
||||
async def test_server_sampling_handler_works_on_modern(mode, behavior, method):
|
||||
"""A server-side sampling handler answers entirely server-side, so it works
|
||||
on modern (2026-07-28) connections regardless of behavior. The era-gate must
|
||||
NOT block these — nothing touches the removed client back-channel. Crucially,
|
||||
'fallback' must go straight to the handler (no bare client-attempt failure)."""
|
||||
server = _handler_server(behavior)
|
||||
async with SDKClient(_server(server), mode=mode) as client:
|
||||
result = await client.call_tool(method, {})
|
||||
assert result.is_error is False
|
||||
assert "handler-answer" in " ".join(_texts(result.content))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("behavior", ["always", "fallback"])
|
||||
@pytest.mark.parametrize("method", ["do_sample", "do_sample_step"])
|
||||
async def test_server_sampling_handler_works_on_legacy(behavior, method):
|
||||
"""Handshake-era behavior is unchanged: the server-side handler still answers
|
||||
on legacy connections."""
|
||||
server = _handler_server(behavior)
|
||||
async with SDKClient(_server(server), mode="legacy") as client:
|
||||
result = await client.call_tool(method, {})
|
||||
assert result.is_error is False
|
||||
assert "handler-answer" in " ".join(_texts(result.content))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", MODERN_MODES)
|
||||
@pytest.mark.parametrize("method", ["do_sample", "do_sample_step"])
|
||||
async def test_sampling_without_handler_still_era_gated_on_modern(
|
||||
push_server, mode, method
|
||||
):
|
||||
"""With no server-side handler configured, the request would hit the removed
|
||||
client back-channel, so the clear era error still fires on modern."""
|
||||
# push_server only defines do_sample; add a do_sample_step twin inline.
|
||||
mcp = FastMCP("no-handler")
|
||||
|
||||
@mcp.tool
|
||||
async def do_sample(ctx: Context) -> str:
|
||||
result = await ctx.sample("hello")
|
||||
return f"sampled {result.text}"
|
||||
|
||||
@mcp.tool
|
||||
async def do_sample_step(ctx: Context) -> str:
|
||||
step = await ctx.sample_step("hello")
|
||||
return f"stepped {step.text}"
|
||||
|
||||
async with SDKClient(
|
||||
_server(mcp), mode=mode, sampling_callback=_sampling_cb
|
||||
) as client:
|
||||
result = await client.call_tool(method, {})
|
||||
assert result.is_error is True
|
||||
assert "server-initiated" in " ".join(_texts(result.content)).lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3b. Sampling deprecation warning (SEP-2577): ctx.sample/ctx.sample_step warn
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -350,12 +436,13 @@ def reset_sample_warn_flag():
|
|||
"""Reset the process-wide warn-once flag so a warning can be observed."""
|
||||
import fastmcp.server.context as context_module
|
||||
|
||||
original = context_module._sample_deprecation_warned
|
||||
context_module._sample_deprecation_warned = False
|
||||
original = set(context_module._sample_deprecation_warned)
|
||||
context_module._sample_deprecation_warned.clear()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
context_module._sample_deprecation_warned = original
|
||||
context_module._sample_deprecation_warned.clear()
|
||||
context_module._sample_deprecation_warned.update(original)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["do_sample", "do_sample_step"])
|
||||
|
|
@ -421,7 +508,6 @@ async def test_sampling_deprecation_warning_suppressible_via_settings(
|
|||
matching the house pattern for every other FastMCP deprecation."""
|
||||
import warnings as _warnings
|
||||
|
||||
import fastmcp
|
||||
from fastmcp.exceptions import FastMCPDeprecationWarning
|
||||
|
||||
monkeypatch.setattr(fastmcp.settings, "deprecation_warnings", False)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue