diff --git a/fastmcp_slim/fastmcp/server/extensions.py b/fastmcp_slim/fastmcp/server/extensions.py new file mode 100644 index 000000000..5ebe466f2 --- /dev/null +++ b/fastmcp_slim/fastmcp/server/extensions.py @@ -0,0 +1,293 @@ +"""FastMCP-native server extension API (SEP-2133). + +An MCP extension is an opt-in, capability-negotiated bundle of protocol +behaviour identified by a reverse-DNS string (e.g. `io.modelcontextprotocol/tasks`). +Unlike the SDK's `mcp.server.extension.Extension`, a FastMCP `ServerExtension` +is bound to its `FastMCP` instance at registration, so its request handlers and +its `tools/call` interceptor can reach the component registry, `Context`, and +auth scope that the SDK's model withholds. + +An extension contributes any subset of four things: + +- **A negotiated capability.** `settings()` is spliced into + `ServerCapabilities.extensions[identifier]` (see `LowLevelServer.get_capabilities`). +- **New request methods.** `methods()` returns `MethodBinding`s, each wired onto + the low-level server via `add_request_handler` when the extension is registered. +- **A `tools/call` interceptor.** `intercept_tool_call()` is the last gate before + a tool body runs — it composes *after* the FastMCP middleware chain and *before* + component execution, so it can observe, short-circuit, or pass a call through. +- **A lifespan.** `lifespan()` is entered with the server's lifespan and exited on + shutdown — the hook the SDK's `Extension` lacks, needed to start backends/workers. + +The base class follows the SDK's httpx-style shape: every contribution method has +a default, so a subclass overrides only what it needs. +""" + +from __future__ import annotations + +import weakref +from collections.abc import Awaitable, Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, nullcontext +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, TypeAlias + +from mcp.server.context import ServerRequestContext +from mcp.shared.exceptions import MCPError +from mcp.shared.extension import validate_extension_identifier +from mcp_types import METHOD_NOT_FOUND, CLIENT_CAPABILITIES_META_KEY, CallToolRequestParams +from mcp_types.methods import SPEC_CLIENT_METHODS +from pydantic import BaseModel + +from fastmcp.server.dependencies import _lift_meta, bind_request_context + +if TYPE_CHECKING: + import mcp_types + + from fastmcp.server.context import Context + from fastmcp.server.server import FastMCP + from fastmcp.tools.base import ToolResult + +__all__ = [ + "MethodBinding", + "ServerExtension", + "read_client_extension_settings", +] + +# What an extension's tools/call interceptor observes and may produce: the tool +# result, or the claimed CreateTaskResult shape when the call is run as a task. +ToolCallOutcome: TypeAlias = "ToolResult | mcp_types.CreateTaskResult" + +# A method handler receives the SDK request context plus validated params and +# returns a bare result model (the runner serializes it). +ExtensionRequestHandler: TypeAlias = Callable[ + [ServerRequestContext[Any, Any], Any], + Awaitable[BaseModel | dict[str, Any] | None], +] + +# A tools/call interceptor's continuation: awaiting it runs the rest of the +# interceptor chain and, finally, the tool body. +ToolCallContinuation: TypeAlias = Callable[[], Awaitable["ToolCallOutcome"]] + + +@dataclass(frozen=True) +class MethodBinding: + """A new request method an extension serves, e.g. `tasks/get`. + + `params_type` validates incoming params before `handler` runs; it should + subclass `RequestParams` so `_meta` parses uniformly. `protocol_versions`, + when set, restricts the method to those wire versions — a request at any + other version is rejected as `METHOD_NOT_FOUND`, mirroring the spec's + `(method, version)` boundary. `None` (the default) admits every version. + + Extension methods are additive: `method` must not name a spec-defined + request method (`tools/call`, `completion/complete`, ...). Binding one would + silently shadow the server's own handler. Both constraints are enforced at + construction. + """ + + method: str + params_type: type[BaseModel] + handler: ExtensionRequestHandler + protocol_versions: frozenset[str] | None = None + + def __post_init__(self) -> None: + if self.method in SPEC_CLIENT_METHODS: + raise ValueError( + f"MethodBinding cannot bind spec method {self.method!r}; extension " + "methods are additive. Use ServerExtension.intercept_tool_call or " + "FastMCP middleware to wrap core behaviour." + ) + if self.protocol_versions is not None and not self.protocol_versions: + raise ValueError( + f"MethodBinding for {self.method!r} has an empty protocol_versions " + "set, so it could never be served; use None to admit every version." + ) + + +class ServerExtension: + """Base class for an opt-in FastMCP server extension (SEP-2133). + + Subclass, set `identifier`, and override the contribution methods that + apply. Every method has a default, so a minimal extension overrides only + `identifier` and one contribution. `identifier` is validated at + subclass-definition time when set as a class attribute, and again at + registration (which covers per-instance identifiers assigned in `__init__`). + + Register an instance with `FastMCP.add_extension(...)`, which binds the + extension to the server so `self.server`, `intercept_tool_call`, and method + handlers can reach FastMCP-level constructs. + """ + + #: Reverse-DNS extension identifier, advertised under `ServerCapabilities.extensions`. + identifier: str + + _server_ref: weakref.ref[FastMCP] | None = None + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + # A class-level identifier is validated here; a per-instance identifier + # assigned in __init__ is validated at registration instead (no class + # attribute exists to inspect at definition time). + identifier = cls.__dict__.get("identifier") + if identifier is not None: + validate_extension_identifier(identifier, owner=cls.__name__) + + def _bind(self, server: FastMCP) -> None: + """Bind this extension to its FastMCP instance (called by `add_extension`). + + A weak reference avoids a reference cycle between the server and its + extensions. Per-instance identifiers are validated here. + """ + validate_extension_identifier(self.identifier, owner=type(self).__name__) + self._server_ref = weakref.ref(server) + + @property + def server(self) -> FastMCP: + """The FastMCP server this extension is registered on. + + Handlers, interceptors, and lifespan code reach the component registry, + `Context`, and auth scope through here. Raises if the extension has not + been registered with `FastMCP.add_extension()`. + """ + ref = self._server_ref + server = ref() if ref is not None else None + if server is None: + raise RuntimeError( + f"Extension {self.identifier!r} is not bound to a FastMCP server; " + "register it with FastMCP.add_extension() before use." + ) + return server + + def settings(self) -> dict[str, Any]: + """Per-extension settings advertised at `capabilities.extensions[identifier]`. + + An empty dict (the default) advertises the extension with no settings. + """ + return {} + + def methods(self) -> Sequence[MethodBinding]: + """New request methods this extension serves (additive).""" + return () + + def lifespan(self) -> AbstractAsyncContextManager[None]: + """A context manager entered with the server's lifespan, exited on shutdown. + + Default: a no-op. Override to start and stop resources an extension owns + (a task-queue backend and worker, say). Entered once per runtime tree, at + the root — a mounted child defers to the root, as the shared Docket does. + """ + return nullcontext() + + async def intercept_tool_call( + self, + params: CallToolRequestParams, + context: Context, + call_next: ToolCallContinuation, + ) -> ToolCallOutcome: + """Wrap `tools/call`. Default: pass through unchanged. + + Runs after the FastMCP middleware chain and before the tool body, so it + is the last gate before execution. Override to observe the call, to + short-circuit (return a result without awaiting `call_next`), or to pass + it through (`return await call_next()`). `params` is the validated + `tools/call` params; `context` is the FastMCP `Context`, from which the + tool being called (`context.fastmcp.get_tool(params.name)`), auth scope, + and the server are reachable. Multiple extensions nest with the + first-registered outermost. + """ + return await call_next() + + def client_settings( + self, ctx: ServerRequestContext[Any, Any] + ) -> dict[str, Any] | None: + """This extension's per-request opt-in settings declared by the client. + + Reads the request's `_meta` client-capabilities block. Returns the + declared settings dict (possibly empty) when the client opted this + extension in for the request, or `None` when it did not. Convenience for + `read_client_extension_settings(ctx, self.identifier)`. + """ + return read_client_extension_settings(ctx, self.identifier) + + +def _extract_client_extension_settings( + meta: Mapping[str, Any] | None, identifier: str +) -> dict[str, Any] | None: + """Pull `_meta[clientCapabilities][extensions][identifier]` from a lifted meta block.""" + if not meta: + return None + client_caps = meta.get(CLIENT_CAPABILITIES_META_KEY) + if not isinstance(client_caps, Mapping): + return None + extensions = client_caps.get("extensions") + if not isinstance(extensions, Mapping): + return None + settings = extensions.get(identifier) + if isinstance(settings, Mapping): + return dict(settings) + return None + + +def read_client_extension_settings( + ctx: ServerRequestContext[Any, Any], identifier: str +) -> dict[str, Any] | None: + """Read a client's per-request extension opt-in from the request `_meta`. + + SEP-2133 extensions negotiate per request: the client repeats its extension + capabilities in each request's `_meta` under + `io.modelcontextprotocol/clientCapabilities` → `extensions` → `identifier`. + Returns the declared settings dict (possibly empty) when the extension was + opted in for this request, or `None` when it was not. + """ + return _extract_client_extension_settings(_lift_meta(ctx), identifier) + + +def build_method_handler(binding: MethodBinding) -> ExtensionRequestHandler: + """Wrap a `MethodBinding` into a low-level request handler. + + The adapter enforces `protocol_versions` gating (rejecting other versions as + `METHOD_NOT_FOUND`, since `add_request_handler` registers unconditionally) + and binds the FastMCP request context so the handler can use `get_context()`, + auth, and other request-scoped dependencies. + """ + + async def handler( + ctx: ServerRequestContext[Any, Any], params: Any + ) -> BaseModel | dict[str, Any] | None: + if ( + binding.protocol_versions is not None + and ctx.protocol_version not in binding.protocol_versions + ): + raise MCPError( + code=METHOD_NOT_FOUND, + message=( + f"Method {binding.method!r} is not available at protocol " + f"version {ctx.protocol_version!r}." + ), + ) + with bind_request_context(ctx): + return await binding.handler(ctx, params) + + return handler + + +def wrap_tool_call_interceptor( + extension: ServerExtension, + call_next: Callable[[Any], Awaitable[Any]], +) -> Callable[[Any], Awaitable[Any]]: + """Fold one extension's `intercept_tool_call` around a middleware `call_next`. + + The returned wrapper is a FastMCP `CallNext`: it hands the extension the + validated `tools/call` params, the FastMCP `Context`, and a zero-arg + continuation that runs the rest of the chain and, finally, the tool body. + """ + + async def wrapped(context: Any) -> Any: + async def cont() -> Any: + return await call_next(context) + + return await extension.intercept_tool_call( + context.message, context.fastmcp_context, cont + ) + + return wrapped diff --git a/fastmcp_slim/fastmcp/server/low_level.py b/fastmcp_slim/fastmcp/server/low_level.py index 27d9dce27..1a9362649 100644 --- a/fastmcp_slim/fastmcp/server/low_level.py +++ b/fastmcp_slim/fastmcp/server/low_level.py @@ -500,10 +500,24 @@ class LowLevelServer(_Server[LifespanResultT]): protocol_version=protocol_version, ) + # Advertise every registered extension's settings under + # capabilities.extensions[identifier]. The hand-rolled UI splice stays + # for now (MCP Apps migrates onto the extension API in a later phase); + # the two coexist. Advertisement is unconditional — the SDK's pre-2026 + # version sieve strips capabilities.extensions on legacy eras, a known + # limitation (sdk-feedback #2). existing_extensions = capabilities.extensions or {} + registered_extensions = { + extension.identifier: extension.settings() + for extension in self.fastmcp._extensions.values() + } return capabilities.model_copy( update={ "tasks": get_task_capabilities(), - "extensions": {**existing_extensions, UI_EXTENSION_ID: {}}, + "extensions": { + **existing_extensions, + UI_EXTENSION_ID: {}, + **registered_extensions, + }, } ) diff --git a/fastmcp_slim/fastmcp/server/mixins/lifespan.py b/fastmcp_slim/fastmcp/server/mixins/lifespan.py index 77374d3f7..5ea62a1fb 100644 --- a/fastmcp_slim/fastmcp/server/mixins/lifespan.py +++ b/fastmcp_slim/fastmcp/server/mixins/lifespan.py @@ -190,6 +190,30 @@ class LifespanMixin: # Reset server ContextVar _current_server.reset(server_token) + @asynccontextmanager + async def _extensions_lifespan(self: FastMCP) -> AsyncIterator[None]: + """Enter each registered extension's lifespan, exit them on shutdown. + + Extension lifespans are entered once per runtime tree, at the root. A + mounted child sees ``_lifespan_root_active`` set by its + ``FastMCPProvider`` and defers to the root, exactly as + ``_docket_lifespan`` does for the shared Docket: an extension whose + lifespan starts shared infrastructure (a task-queue backend and worker, + say) is therefore owned by the tree root, and mounted children reach it + through the same context rather than starting a second copy. + + Extensions are entered in registration order; the ``AsyncExitStack`` + exits them in reverse on teardown. + """ + if _lifespan_root_active.get() or not self._extensions: + yield + return + + async with AsyncExitStack() as stack: + for extension in self._extensions.values(): + await stack.enter_async_context(extension.lifespan()) + yield + def _capture_shared_context(self: FastMCP) -> None: """Snapshot the live ``SharedContext`` ContextVar values. @@ -238,6 +262,7 @@ class LifespanMixin: try: user_lifespan_result = await stack.enter_async_context(self._lifespan(self)) await stack.enter_async_context(self._docket_lifespan()) + await stack.enter_async_context(self._extensions_lifespan()) self._lifespan_result = user_lifespan_result self._lifespan_result_set = True diff --git a/fastmcp_slim/fastmcp/server/server.py b/fastmcp_slim/fastmcp/server/server.py index c66c9c57b..fce62550b 100644 --- a/fastmcp_slim/fastmcp/server/server.py +++ b/fastmcp_slim/fastmcp/server/server.py @@ -105,6 +105,7 @@ if TYPE_CHECKING: from fastmcp.client.client import SDKServer from fastmcp.client.sampling import SamplingHandler from fastmcp.client.transports import ClientTransport, ClientTransportT + from fastmcp.server.extensions import ServerExtension from fastmcp.server.providers.openapi import ComponentFn as OpenAPIComponentFn from fastmcp.server.providers.openapi import RouteMap from fastmcp.server.providers.openapi import RouteMapFn as OpenAPIRouteMapFn @@ -522,6 +523,12 @@ class FastMCP( self.middleware: list[Middleware] = list(middleware or []) + # Registered server extensions (SEP-2133), keyed by reverse-DNS + # identifier. Populated by add_extension; consumed by the low-level + # server (capability advertisement), the tool-call path (interception), + # and the lifespan manager (extension lifespans). + self._extensions: dict[str, ServerExtension] = {} + if dereference_schemas: from fastmcp.server.middleware.dereference import ( DereferenceRefsMiddleware, @@ -635,6 +642,64 @@ class FastMCP( def add_middleware(self, middleware: Middleware) -> None: self.middleware.append(middleware) + def add_extension(self, extension: ServerExtension) -> None: + """Register a server extension (SEP-2133). + + An extension contributes a negotiated capability, additive request + methods, a `tools/call` interceptor, and an optional lifespan — each + with access to FastMCP-level constructs (the component registry, + `Context`, auth scope). Its capability is advertised only while it is + registered. + + The extension is bound to this server (so its handlers and interceptor + can reach it), its method bindings are wired onto the low-level server, + and it is recorded for capability advertisement, interception, and + lifespan entry. Registering two extensions with the same identifier is + an error. + """ + from fastmcp.server.extensions import ( + build_method_handler, + validate_extension_identifier, + ) + + validate_extension_identifier( + extension.identifier, owner=type(extension).__name__ + ) + if extension.identifier in self._extensions: + raise ValueError( + f"An extension with identifier {extension.identifier!r} is " + "already registered." + ) + + extension._bind(self) + for binding in extension.methods(): + self._mcp_server.add_request_handler( + binding.method, + binding.params_type, + build_method_handler(binding), + ) + self._extensions[extension.identifier] = extension + + def _compose_tool_call_interceptors( + self, call_next: CallNext[Any, Any] + ) -> CallNext[Any, Any]: + """Nest every extension's `tools/call` interceptor around ``call_next``. + + Composes at the innermost point of the tool-call dispatch — after the + FastMCP middleware chain, before the tool body — so each interceptor is + the last gate before execution. First-registered extension is outermost. + A server with no extensions returns ``call_next`` unchanged, so there is + zero behaviour change. + """ + from fastmcp.server.extensions import wrap_tool_call_interceptor + + chain = call_next + for extension in reversed(list(self._extensions.values())): + chain = cast( + "CallNext[Any, Any]", wrap_tool_call_interceptor(extension, chain) + ) + return chain + def add_provider(self, provider: Provider, *, namespace: str = "") -> None: """Add a provider for dynamic tools, resources, and prompts. @@ -1347,14 +1412,21 @@ class FastMCP( method="tools/call", fastmcp_context=ctx, ) + # Extension tools/call interceptors compose here, at the + # innermost point of dispatch: the FastMCP middleware chain wraps + # the whole thing (so it observes every call), and the + # interceptors sit between it and the tool body (so each is the + # last gate before execution). return await self._dispatch_component_middleware( context=mw_context, - call_next=lambda context: self.call_tool( - context.message.name, - context.message.arguments or {}, - version=version, - run_middleware=False, - task_meta=task_meta, + call_next=self._compose_tool_call_interceptors( + lambda context: self.call_tool( + context.message.name, + context.message.arguments or {}, + version=version, + run_middleware=False, + task_meta=task_meta, + ) ), ) diff --git a/tests/server/test_extensions.py b/tests/server/test_extensions.py new file mode 100644 index 000000000..a3e65ad62 --- /dev/null +++ b/tests/server/test_extensions.py @@ -0,0 +1,533 @@ +"""Tests for the FastMCP-native server extension API (SEP-2133). + +A synthetic extension exercises every contribution kind: capability +advertisement, additive request methods (with protocol-version gating), +tools/call interception (observe and short-circuit), a lifespan hook (order and +mounted-server behaviour), and the per-request capability sniff. Registration +guards (duplicate identifier, spec-method rejection, invalid identifier) and a +zero-behaviour-change baseline round it out. +""" + +from __future__ import annotations + +from contextlib import asynccontextmanager +from typing import Any, Literal, cast + +import mcp_types +import pytest +from mcp.server.context import ServerRequestContext +from mcp.shared.exceptions import MCPError +from mcp_types import CLIENT_CAPABILITIES_META_KEY, METHOD_NOT_FOUND, RequestParams + +from fastmcp import FastMCP +from fastmcp.client import Client +from fastmcp.server.extensions import ( + MethodBinding, + ServerExtension, + read_client_extension_settings, +) +from fastmcp.tools.base import ToolResult + +EXT_ID = "com.example/synthetic" + + +class PingParams(RequestParams): + echo: str | None = None + + +class PingRequest(mcp_types.Request): + method: Literal["synthetic/ping"] = "synthetic/ping" + params: PingParams + + +class PingResult(mcp_types.Result): + pong: bool + echo: str | None = None + + +def _text(result: mcp_types.CallToolResult) -> str: + block = result.content[0] + assert isinstance(block, mcp_types.TextContent) + return block.text + + +# --------------------------------------------------------------------------- +# Capability advertisement +# --------------------------------------------------------------------------- + + +async def test_capability_advertised_to_modern_client(): + """A registered extension's settings appear under capabilities.extensions. + + Uses ``mode='auto'`` so the client negotiates the modern era via + ``server/discover`` (which reads ``get_capabilities`` directly); the SDK's + version sieve strips ``capabilities.extensions`` only on legacy eras. + """ + + class Ext(ServerExtension): + identifier = EXT_ID + + def settings(self) -> dict[str, Any]: + return {"version": "1"} + + mcp = FastMCP("t") + mcp.add_extension(Ext()) + + async with Client(mcp, mode="auto") as client: + extensions = client.server_capabilities.extensions or {} + assert extensions.get(EXT_ID) == {"version": "1"} + + +async def test_capability_absent_without_registration(): + """A server with no extensions advertises none of its own.""" + mcp = FastMCP("t") + async with Client(mcp, mode="auto") as client: + extensions = client.server_capabilities.extensions or {} + assert EXT_ID not in extensions + + +async def test_empty_settings_still_advertise(): + """The default empty-settings extension is advertised with an empty dict.""" + + class Ext(ServerExtension): + identifier = EXT_ID + + mcp = FastMCP("t") + mcp.add_extension(Ext()) + async with Client(mcp, mode="auto") as client: + extensions = client.server_capabilities.extensions or {} + assert extensions.get(EXT_ID) == {} + + +# --------------------------------------------------------------------------- +# Additive request methods +# --------------------------------------------------------------------------- + + +class _PingExtension(ServerExtension): + identifier = EXT_ID + + def methods(self) -> list[MethodBinding]: + async def handler( + ctx: ServerRequestContext[Any, Any], params: PingParams + ) -> PingResult: + return PingResult(pong=True, echo=params.echo) + + return [ + MethodBinding( + method="synthetic/ping", + params_type=PingParams, + handler=handler, + ) + ] + + +async def test_custom_method_callable_end_to_end(): + mcp = FastMCP("t") + mcp.add_extension(_PingExtension()) + + async with Client(mcp, mode="auto") as client: + result = await client.session.send_request( + request=PingRequest(params=PingParams(echo="hi")), + result_type=PingResult, + ) + assert result.pong is True + assert result.echo == "hi" + + +async def test_method_handler_reaches_server_registry(): + """A method handler can reach the FastMCP component registry via the extension.""" + + class Ext(ServerExtension): + identifier = EXT_ID + + def methods(self) -> list[MethodBinding]: + async def handler( + ctx: ServerRequestContext[Any, Any], params: PingParams + ) -> PingResult: + tools = await self.server.list_tools() + return PingResult(pong=len(tools) == 1) + + return [ + MethodBinding( + method="synthetic/ping", + params_type=PingParams, + handler=handler, + ) + ] + + mcp = FastMCP("t") + + @mcp.tool + def only_tool() -> str: + return "x" + + mcp.add_extension(Ext()) + async with Client(mcp, mode="auto") as client: + result = await client.session.send_request( + request=PingRequest(params=PingParams()), + result_type=PingResult, + ) + assert result.pong is True + + +async def test_method_protocol_version_gating(): + """A version-gated method is rejected as METHOD_NOT_FOUND off its versions.""" + + class Ext(ServerExtension): + identifier = EXT_ID + + def methods(self) -> list[MethodBinding]: + async def handler( + ctx: ServerRequestContext[Any, Any], params: PingParams + ) -> PingResult: + return PingResult(pong=True) + + return [ + MethodBinding( + method="synthetic/ping", + params_type=PingParams, + handler=handler, + protocol_versions=frozenset({"2026-07-28"}), + ) + ] + + mcp = FastMCP("t") + mcp.add_extension(Ext()) + + async with Client(mcp, mode="legacy") as client: + with pytest.raises(MCPError) as exc_info: + await client.session.send_request( + request=PingRequest(params=PingParams()), + result_type=PingResult, + ) + assert exc_info.value.error.code == METHOD_NOT_FOUND + + +# --------------------------------------------------------------------------- +# tools/call interception +# --------------------------------------------------------------------------- + + +async def test_interceptor_observes_tool_call(): + """A pass-through interceptor sees the call and the tool still runs.""" + seen: list[str] = [] + + class Ext(ServerExtension): + identifier = EXT_ID + + async def intercept_tool_call(self, params, context, call_next): + seen.append(params.name) + return await call_next() + + mcp = FastMCP("t") + + @mcp.tool + def greet() -> str: + return "hello" + + mcp.add_extension(Ext()) + async with Client(mcp, mode="auto") as client: + result = await client.call_tool("greet") + assert _text(result) == "hello" + assert seen == ["greet"] + + +async def test_interceptor_short_circuits(): + """An interceptor can return its own result without running the tool body.""" + ran = [] + + class Ext(ServerExtension): + identifier = EXT_ID + + async def intercept_tool_call(self, params, context, call_next): + return ToolResult( + content=[mcp_types.TextContent(type="text", text="intercepted")] + ) + + mcp = FastMCP("t") + + @mcp.tool + def greet(): + ran.append(True) + return "hello" + + mcp.add_extension(Ext()) + async with Client(mcp, mode="auto") as client: + result = await client.call_tool("greet") + assert _text(result) == "intercepted" + assert ran == [] + + +async def test_interceptor_reaches_tool_metadata(): + """An interceptor can resolve the tool being called through the context.""" + captured: dict[str, Any] = {} + + class Ext(ServerExtension): + identifier = EXT_ID + + async def intercept_tool_call(self, params, context, call_next): + tool = await context.fastmcp.get_tool(params.name) + captured["title"] = tool.title + return await call_next() + + mcp = FastMCP("t") + + @mcp.tool(title="A Greeting") + def greet() -> str: + return "hello" + + mcp.add_extension(Ext()) + async with Client(mcp, mode="auto") as client: + await client.call_tool("greet") + assert captured["title"] == "A Greeting" + + +async def test_interceptors_nest_first_registered_outermost(): + """Multiple interceptors nest with the first-registered extension outermost.""" + order: list[str] = [] + + def make_ext(identifier: str, label: str) -> ServerExtension: + class Ext(ServerExtension): + async def intercept_tool_call(self, params, context, call_next): + order.append(f"{label}-before") + result = await call_next() + order.append(f"{label}-after") + return result + + ext = Ext() + ext.identifier = identifier + return ext + + mcp = FastMCP("t") + + @mcp.tool + def greet() -> str: + return "hello" + + mcp.add_extension(make_ext("com.example/outer", "outer")) + mcp.add_extension(make_ext("com.example/inner", "inner")) + async with Client(mcp, mode="auto") as client: + await client.call_tool("greet") + + assert order == ["outer-before", "inner-before", "inner-after", "outer-after"] + + +async def test_no_extensions_leaves_tool_call_unchanged(): + """With no extensions registered, tools/call behaves exactly as before.""" + mcp = FastMCP("t") + + @mcp.tool + def greet() -> str: + return "hello" + + assert mcp._extensions == {} + async with Client(mcp, mode="auto") as client: + result = await client.call_tool("greet") + assert _text(result) == "hello" + + +# --------------------------------------------------------------------------- +# Lifespan hook +# --------------------------------------------------------------------------- + + +def _recording_extension(identifier: str, log: list[str]) -> ServerExtension: + class Ext(ServerExtension): + @asynccontextmanager + async def lifespan(self): + log.append(f"{identifier}:enter") + try: + yield + finally: + log.append(f"{identifier}:exit") + + ext = Ext() + ext.identifier = identifier + return ext + + +async def test_lifespan_entered_and_exited(): + log: list[str] = [] + mcp = FastMCP("t") + mcp.add_extension(_recording_extension(EXT_ID, log)) + + async with Client(mcp, mode="auto"): + assert log == [f"{EXT_ID}:enter"] + assert log == [f"{EXT_ID}:enter", f"{EXT_ID}:exit"] + + +async def test_lifespans_enter_in_order_exit_in_reverse(): + log: list[str] = [] + mcp = FastMCP("t") + mcp.add_extension(_recording_extension("com.example/a", log)) + mcp.add_extension(_recording_extension("com.example/b", log)) + + async with Client(mcp, mode="auto"): + pass + + assert log == [ + "com.example/a:enter", + "com.example/b:enter", + "com.example/b:exit", + "com.example/a:exit", + ] + + +async def test_standalone_server_enters_extension_lifespan(): + log: list[str] = [] + child = FastMCP("child") + child.add_extension(_recording_extension(EXT_ID, log)) + + async with Client(child, mode="auto"): + assert log == [f"{EXT_ID}:enter"] + assert log == [f"{EXT_ID}:enter", f"{EXT_ID}:exit"] + + +async def test_mounted_child_defers_extension_lifespan_to_root(): + """A mounted child's extension lifespan is not entered below a root. + + Mirrors the shared Docket: extension lifespans that may start shared + infrastructure are owned by the tree root, so a mounted child defers. + """ + log: list[str] = [] + child = FastMCP("child") + child.add_extension(_recording_extension(EXT_ID, log)) + + root = FastMCP("root") + root.mount(child) + + async with Client(root, mode="auto"): + pass + + assert log == [] + + +# --------------------------------------------------------------------------- +# Request-time capability sniff +# --------------------------------------------------------------------------- + + +def _ctx_with_meta(meta: dict[str, Any] | None) -> ServerRequestContext[Any, Any]: + params: dict[str, Any] = {} + if meta is not None: + params["_meta"] = meta + return ServerRequestContext( + session=cast(Any, object()), + lifespan_context={}, + protocol_version="2026-07-28", + method="synthetic/ping", + params=params, + ) + + +def test_capability_sniff_reads_declared_settings(): + meta = { + CLIENT_CAPABILITIES_META_KEY: { + "extensions": {EXT_ID: {"limit": 5}}, + } + } + assert read_client_extension_settings(_ctx_with_meta(meta), EXT_ID) == {"limit": 5} + + +def test_capability_sniff_empty_settings_is_opt_in(): + """An empty settings dict is a valid opt-in, distinct from absence (None).""" + meta = {CLIENT_CAPABILITIES_META_KEY: {"extensions": {EXT_ID: {}}}} + assert read_client_extension_settings(_ctx_with_meta(meta), EXT_ID) == {} + + +@pytest.mark.parametrize( + "meta", + [ + None, + {}, + {CLIENT_CAPABILITIES_META_KEY: {}}, + {CLIENT_CAPABILITIES_META_KEY: {"extensions": {"other/ext": {}}}}, + ], +) +def test_capability_sniff_returns_none_when_not_opted_in(meta): + assert read_client_extension_settings(_ctx_with_meta(meta), EXT_ID) is None + + +def test_client_settings_convenience_uses_own_identifier(): + class Ext(ServerExtension): + identifier = EXT_ID + + meta = {CLIENT_CAPABILITIES_META_KEY: {"extensions": {EXT_ID: {"a": 1}}}} + assert Ext().client_settings(_ctx_with_meta(meta)) == {"a": 1} + + +# --------------------------------------------------------------------------- +# Registration guards +# --------------------------------------------------------------------------- + + +async def test_duplicate_identifier_rejected(): + class Ext(ServerExtension): + identifier = EXT_ID + + mcp = FastMCP("t") + mcp.add_extension(Ext()) + with pytest.raises(ValueError, match="already registered"): + mcp.add_extension(Ext()) + + +def test_spec_method_name_rejected(): + async def handler(ctx: Any, params: Any) -> None: + return None + + with pytest.raises(ValueError, match="spec method"): + MethodBinding( + method="tools/call", + params_type=PingParams, + handler=handler, + ) + + +def test_empty_protocol_versions_rejected(): + async def handler(ctx: Any, params: Any) -> None: + return None + + with pytest.raises(ValueError, match="protocol_versions"): + MethodBinding( + method="synthetic/ping", + params_type=PingParams, + handler=handler, + protocol_versions=frozenset(), + ) + + +def test_invalid_identifier_rejected_at_class_definition(): + with pytest.raises(TypeError, match="reverse-DNS"): + + class Ext(ServerExtension): + identifier = "no-prefix" + + +async def test_per_instance_invalid_identifier_rejected_at_registration(): + class Ext(ServerExtension): + pass + + ext = Ext() + ext.identifier = "no-prefix" + mcp = FastMCP("t") + with pytest.raises(TypeError, match="reverse-DNS"): + mcp.add_extension(ext) + + +def test_bound_server_accessible_after_registration(): + class Ext(ServerExtension): + identifier = EXT_ID + + ext = Ext() + mcp = FastMCP("t") + mcp.add_extension(ext) + assert ext.server is mcp + + +def test_unbound_server_access_raises(): + class Ext(ServerExtension): + identifier = EXT_ID + + with pytest.raises(RuntimeError, match="not bound"): + _ = Ext().server