From 77131edc008ce8d757824a38d0b115d96ba31a4d Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:49:18 -0400 Subject: [PATCH] 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. --- fastmcp_slim/fastmcp/server/context.py | 44 ++++++++-- fastmcp_slim/fastmcp/server/sampling/run.py | 34 ++++++-- tests/server/test_protocol_eras.py | 94 ++++++++++++++++++++- 3 files changed, 157 insertions(+), 15 deletions(-) diff --git a/fastmcp_slim/fastmcp/server/context.py b/fastmcp_slim/fastmcp/server/context.py index e76c6ab6b..e2c2f1977 100644 --- a/fastmcp_slim/fastmcp/server/context.py +++ b/fastmcp_slim/fastmcp/server/context.py @@ -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 diff --git a/fastmcp_slim/fastmcp/server/sampling/run.py b/fastmcp_slim/fastmcp/server/sampling/run.py index 1fba19a5b..cb0a1e9ec 100644 --- a/fastmcp_slim/fastmcp/server/sampling/run.py +++ b/fastmcp_slim/fastmcp/server/sampling/run.py @@ -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 diff --git a/tests/server/test_protocol_eras.py b/tests/server/test_protocol_eras.py index e0df7013c..73863bd84 100644 --- a/tests/server/test_protocol_eras.py +++ b/tests/server/test_protocol_eras.py @@ -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)