chore: import upstream snapshot with attribution

tests/server/mcpserver/test_extension.py

@@ -0,0 +1,441 @@
+"""Tests for the core SEP-2133 extension API (`Extension`, `MCPServer` wiring).
+
+These exercise the closed set of extension contribution kinds - tools,
+resources, request methods, and the single `tools/call` interceptor - through
+the highest-level public surface (in-memory `Client`), plus the
+`compose_tool_call_interceptor` helper directly.
+"""
+
+from typing import Any, Literal, cast
+
+import mcp_types as types
+import pytest
+from inline_snapshot import snapshot
+from mcp_types import (
+    METHOD_NOT_FOUND,
+    MISSING_REQUIRED_CLIENT_CAPABILITY,
+    CallToolResult,
+    TextContent,
+)
+
+from mcp.client import advertise
+from mcp.client.client import Client
+from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
+from mcp.server.extension import (
+    Extension,
+    MethodBinding,
+    ResourceBinding,
+    ToolBinding,
+    compose_tool_call_interceptor,
+)
+from mcp.server.mcpserver import Context, MCPServer, require_client_extension
+from mcp.server.mcpserver.resources import TextResource
+from mcp.shared.exceptions import MCPError
+
+pytestmark = pytest.mark.anyio
+
+_TOOL_META: dict[str, Any] = {"com.example/marker": {"v": 1}}
+
+
+class _AdditiveExt(Extension):
+    """Override `tools()`/`resources()` only - a purely additive extension."""
+
+    identifier = "com.example/additive"
+
+    def tools(self):
+        def ping() -> str:
+            """Reply with pong."""
+            return "pong"
+
+        return [ToolBinding(fn=ping, meta=_TOOL_META)]
+
+    def resources(self):
+        return [ResourceBinding(resource=TextResource(uri="ext://greeting", name="greeting", text="hello"))]
+
+
+class _SettingsExt(Extension):
+    """Override `settings()` so the extension advertises a non-empty settings map."""
+
+    identifier = "com.example/settings"
+
+    def settings(self) -> dict[str, Any]:
+        return {"feature": {"enabled": True}}
+
+
+class _PingParams(types.RequestParams):
+    pass
+
+
+class _PingResult(types.Result):
+    pong: bool
+
+
+class _PingRequest(types.Request[_PingParams, Literal["com.example/ping"]]):
+    method: Literal["com.example/ping"] = "com.example/ping"
+    params: _PingParams
+
+
+async def _pong_handler(ctx: ServerRequestContext[Any, Any], params: _PingParams) -> _PingResult:
+    """The shared `com.example/ping` handler (dispatched by the reachability test)."""
+    return _PingResult(pong=True)
+
+
+class _MethodExt(Extension):
+    """Override `methods()` to serve a new vendor request verb."""
+
+    identifier = "com.example/method"
+
+    def methods(self) -> list[MethodBinding]:
+        return [MethodBinding("com.example/ping", _PingParams, _pong_handler)]
+
+
+class _ReplacingExt(Extension):
+    """Override `intercept_tool_call()` to short-circuit with a fixed result."""
+
+    identifier = "com.example/replacing"
+
+    async def intercept_tool_call(
+        self, params: types.CallToolRequestParams, ctx: ServerRequestContext[Any, Any], call_next: CallNext
+    ) -> HandlerResult:
+        return CallToolResult(content=[TextContent(type="text", text="intercepted")])
+
+
+class _PassThroughExt(Extension):
+    """Override `intercept_tool_call()` but always delegate to `call_next` unchanged."""
+
+    identifier = "com.example/passthrough"
+
+    async def intercept_tool_call(
+        self, params: types.CallToolRequestParams, ctx: ServerRequestContext[Any, Any], call_next: CallNext
+    ) -> HandlerResult:
+        return await call_next(ctx)
+
+
+class _DefaultExt(Extension):
+    """Override nothing - relies on the base `intercept_tool_call` default (pass through)."""
+
+    identifier = "com.example/default"
+
+
+class _RecordingExt(Extension):
+    """Override `intercept_tool_call()` to record `(identifier, tool_name)` then pass through."""
+
+    def __init__(self, identifier: str, log: list[tuple[str, str]]) -> None:
+        self.identifier = identifier
+        self._log = log
+
+    async def intercept_tool_call(
+        self, params: types.CallToolRequestParams, ctx: ServerRequestContext[Any, Any], call_next: CallNext
+    ) -> HandlerResult:
+        self._log.append((self.identifier, params.name))
+        return await call_next(ctx)
+
+
+def _echo(value: str) -> str:
+    """Echo the input value (shared tool body across interceptor tests)."""
+    return value
+
+
+async def test_additive_extension_registers_its_tool_and_resource() -> None:
+    """SDK-defined: an `Extension` overriding `tools()`/`resources()` surfaces both
+    through `MCPServer`'s normal `list_tools`/`list_resources`, and the tool's
+    `_meta` round-trips equal to the exact dict the binding carried (identity can't
+    hold - the value is JSON-serialized over the transport)."""
+    server = MCPServer("test", extensions=[_AdditiveExt()])
+
+    async with Client(server) as client:
+        tools = await client.list_tools()
+        resources = await client.list_resources()
+        called = await client.call_tool("ping", {})
+
+    assert [t.name for t in tools.tools] == ["ping"]
+    assert tools.tools[0].meta == _TOOL_META
+    assert called == snapshot(CallToolResult(content=[TextContent(text="pong")], structured_content={"result": "pong"}))
+    assert resources == snapshot(
+        types.ListResourcesResult(
+            resources=[types.Resource(name="greeting", uri="ext://greeting", mime_type="text/plain")]
+        )
+    )
+
+
+async def test_extension_settings_advertised_under_server_capabilities() -> None:
+    """SDK-defined: `settings()` rides `server/discover` and lands under
+    `server_capabilities.extensions[identifier]` on the modern (`auto`) path."""
+    server = MCPServer("test", extensions=[_SettingsExt()])
+
+    async with Client(server, mode="auto") as client:
+        extensions = client.server_capabilities.extensions
+
+    assert extensions == snapshot({"com.example/settings": {"feature": {"enabled": True}}})
+
+
+async def test_extension_settings_dropped_on_legacy_handshake() -> None:
+    """Pinned gap: the 2025 `ServerCapabilities` wire schema has no `extensions`
+    field, so a legacy `initialize` handshake drops the advertised extension even
+    though the modern `auto` path carries it."""
+    server = MCPServer("test", extensions=[_SettingsExt()])
+
+    async with Client(server, mode="legacy") as client:
+        assert client.server_capabilities.extensions is None
+
+
+def test_duplicate_extension_identifier_raises() -> None:
+    """SDK-defined: registering two extensions with the same `identifier` is a
+    construction error."""
+    with pytest.raises(ValueError):
+        MCPServer("test", extensions=[_SettingsExt(), _SettingsExt()])
+
+
+async def test_extension_method_reachable_via_session_send_request() -> None:
+    """SDK-defined: an `Extension` overriding `methods()` wires a new request verb
+    onto the low-level server, reachable through `client.session.send_request`."""
+    server = MCPServer("test", extensions=[_MethodExt()])
+
+    async with Client(server) as client:
+        request = _PingRequest(params=_PingParams())
+        result = await client.session.send_request(request, _PingResult)
+
+    assert result == snapshot(_PingResult(pong=True))
+
+
+async def test_pass_through_interceptor_leaves_tool_result_unchanged() -> None:
+    """SDK-defined: an extension whose `intercept_tool_call` delegates to
+    `call_next` does not alter the underlying tool's `CallToolResult`."""
+    server = MCPServer("test", extensions=[_PassThroughExt()])
+    server.tool(name="echo")(_echo)
+
+    async with Client(server) as client:
+        result = await client.call_tool("echo", {"value": "hi"})
+
+    assert result == snapshot(CallToolResult(content=[TextContent(text="hi")], structured_content={"result": "hi"}))
+
+
+async def test_short_circuiting_interceptor_replaces_tool_result() -> None:
+    """SDK-defined: an extension that returns from `intercept_tool_call` without
+    calling `call_next` replaces the tool's result wholesale (the tool never runs)."""
+    server = MCPServer("test", extensions=[_ReplacingExt()])
+    server.tool(name="echo", structured_output=False)(_echo)
+
+    async with Client(server) as client:
+        result = await client.call_tool("echo", {"value": "hi"})
+
+    assert result == snapshot(CallToolResult(content=[TextContent(text="intercepted")]))
+
+
+def test_plain_extension_installs_no_tool_call_interceptor() -> None:
+    """SDK-defined: an extension that does not override `intercept_tool_call` adds no
+    middleware - the composed interceptor exists only when at least one extension
+    overrides it."""
+    baseline = len(MCPServer("test")._lowlevel_server.middleware)
+    server = MCPServer("test", extensions=[_AdditiveExt()])
+
+    assert len(server._lowlevel_server.middleware) == baseline
+
+
+def test_overriding_extension_installs_one_tool_call_interceptor() -> None:
+    """SDK-defined: an extension that overrides `intercept_tool_call` composes exactly
+    one additional `tools/call` middleware."""
+    baseline = len(MCPServer("test")._lowlevel_server.middleware)
+    server = MCPServer("test", extensions=[_ReplacingExt()])
+
+    assert len(server._lowlevel_server.middleware) == baseline + 1
+
+
+async def test_default_interceptor_passes_through_alongside_an_overriding_one() -> None:
+    """SDK-defined: an extension that does not override `intercept_tool_call` runs the
+    base-class default (pass through) when another extension forces the composed
+    middleware to exist, leaving the tool result untouched."""
+    server = MCPServer("test", extensions=[_DefaultExt(), _PassThroughExt()])
+    server.tool(name="echo")(_echo)
+
+    async with Client(server) as client:
+        result = await client.call_tool("echo", {"value": "hi"})
+
+    assert result == snapshot(CallToolResult(content=[TextContent(text="hi")], structured_content={"result": "hi"}))
+
+
+async def test_interceptors_run_in_registration_order_with_threaded_params() -> None:
+    """SDK-defined: `compose_tool_call_interceptor` nests extensions first-outermost, so
+    two passing-through interceptors record in registration order, each seeing the
+    validated `tools/call` params (the real tool name)."""
+    log: list[tuple[str, str]] = []
+    server = MCPServer(
+        "test",
+        extensions=[_RecordingExt("com.example/first", log), _RecordingExt("com.example/second", log)],
+    )
+    server.tool(name="echo")(_echo)
+
+    async with Client(server) as client:
+        await client.call_tool("echo", {"value": "hi"})
+
+    assert log == [("com.example/first", "echo"), ("com.example/second", "echo")]
+
+
+async def test_compose_tool_call_interceptor_passes_through_non_tools_call() -> None:
+    """SDK-defined: the composed middleware is a no-op for any method other than
+    `tools/call` - it forwards to `call_next` without touching the interceptors."""
+    sentinel = types.EmptyResult()
+
+    async def call_next(ctx: ServerRequestContext[Any, Any]) -> HandlerResult:
+        return sentinel
+
+    middleware = compose_tool_call_interceptor([_ReplacingExt()])
+    ctx = ServerRequestContext(
+        session=cast("Any", None),
+        lifespan_context={},
+        protocol_version="2026-07-28",
+        method="tasks/get",
+        params={"taskId": "t-1"},
+    )
+
+    result = await middleware(ctx, call_next)
+
+    assert result is sentinel
+
+
+def test_extension_subclass_without_prefixed_identifier_is_rejected_at_definition() -> None:
+    """SDK-defined: SEP-2133 requires a `vendor-prefix/name` identifier, enforced when the
+    subclass is defined (a bare name with no prefix is a TypeError)."""
+    with pytest.raises(TypeError):
+        type("_BadExt", (Extension,), {"identifier": "noprefix"})
+
+
+def test_extension_without_identifier_is_rejected_at_registration() -> None:
+    """SDK-defined: a subclass that never sets `identifier` (neither class-level nor in
+    `__init__`) is rejected when the server applies it."""
+
+    class _NoIdExt(Extension):
+        pass
+
+    with pytest.raises(TypeError):
+        MCPServer("test", extensions=[_NoIdExt()])
+
+
+class _VersionPinnedParams(types.RequestParams):
+    pass
+
+
+class _VersionPinnedResult(types.Result):
+    ok: bool
+
+
+class _VersionPinnedRequest(types.Request[_VersionPinnedParams, Literal["com.example/pinned"]]):
+    method: Literal["com.example/pinned"] = "com.example/pinned"
+    params: _VersionPinnedParams
+
+
+class _VersionPinnedExt(Extension):
+    """A method scoped to 2026-07-28 only via `MethodBinding.protocol_versions`."""
+
+    identifier = "com.example/pinned"
+
+    def methods(self):
+        async def handler(ctx: ServerRequestContext[Any, Any], params: _VersionPinnedParams) -> _VersionPinnedResult:
+            return _VersionPinnedResult(ok=True)
+
+        return [MethodBinding("com.example/pinned", _VersionPinnedParams, handler, frozenset({"2026-07-28"}))]
+
+
+async def test_version_pinned_method_is_served_at_an_allowed_version() -> None:
+    """SDK-defined: a `MethodBinding` with `protocol_versions` serves the method at a version
+    in the set."""
+    server = MCPServer("test", extensions=[_VersionPinnedExt()])
+
+    async with Client(server, mode="2026-07-28") as client:
+        request = _VersionPinnedRequest(params=_VersionPinnedParams())
+        result = await client.session.send_request(request, _VersionPinnedResult)
+
+    assert result == snapshot(_VersionPinnedResult(ok=True))
+
+
+async def test_version_pinned_method_is_method_not_found_at_a_disallowed_version() -> None:
+    """SDK-defined: the same method at a version outside `protocol_versions` is rejected with
+    METHOD_NOT_FOUND, mirroring the spec's per-version boundary."""
+    server = MCPServer("test", extensions=[_VersionPinnedExt()])
+
+    async with Client(server, mode="legacy") as client:
+        request = _VersionPinnedRequest(params=_VersionPinnedParams())
+        with pytest.raises(MCPError) as exc_info:
+            await client.session.send_request(request, _VersionPinnedResult)
+
+    assert exc_info.value.code == METHOD_NOT_FOUND
+    assert exc_info.value.error.data == "com.example/pinned"
+
+
+@pytest.mark.parametrize("method", ["tools/list", "completion/complete"])
+def test_method_binding_rejects_spec_methods(method: str) -> None:
+    """SDK-defined: extension methods are additive — binding a spec-defined request method
+    would silently shadow (or be shadowed by) the server's own handler, so it is rejected
+    when the binding is constructed."""
+    with pytest.raises(ValueError):
+        MethodBinding(method, _PingParams, _pong_handler)
+
+
+def test_method_binding_rejects_empty_protocol_versions() -> None:
+    """SDK-defined: an empty `protocol_versions` set would make the method unreachable at
+    every version; `None` is the universal-version spelling."""
+    with pytest.raises(ValueError) as exc_info:
+        MethodBinding("com.example/dead", _PingParams, _pong_handler, frozenset())
+    assert str(exc_info.value) == snapshot(
+        "MethodBinding for 'com.example/dead' has an empty protocol_versions set, so it could "
+        "never be served; use None to admit every version"
+    )
+
+
+class _OtherMethodExt(Extension):
+    """A second extension binding the same verb as `_MethodExt`."""
+
+    identifier = "com.example/other-method"
+
+    def methods(self) -> list[MethodBinding]:
+        return [MethodBinding("com.example/ping", _PingParams, _pong_handler)]
+
+
+def test_colliding_extension_methods_are_rejected_at_registration() -> None:
+    """SDK-defined: two extensions binding the same method would silently last-write-win;
+    the collision is rejected when the second extension is applied."""
+    with pytest.raises(ValueError) as exc_info:
+        MCPServer("test", extensions=[_MethodExt(), _OtherMethodExt()])
+    assert str(exc_info.value) == snapshot(
+        "Extension 'com.example/other-method' binds method 'com.example/ping', which is already "
+        "registered; extension methods are additive and cannot replace another handler"
+    )
+
+
+_NEEDS_EXT = "com.example/needed"
+
+
+class _RequiresExt(Extension):
+    """A tool that requires the client to have declared `com.example/needed`."""
+
+    identifier = _NEEDS_EXT
+
+    def tools(self):
+        def guarded(ctx: Context) -> str:
+            require_client_extension(ctx.request_context, _NEEDS_EXT)
+            return "ok"
+
+        return [ToolBinding(fn=guarded)]
+
+
+async def test_require_client_extension_passes_when_client_declared_it() -> None:
+    """SDK-defined: `require_client_extension` is a no-op when the client advertised the id."""
+    server = MCPServer("test", extensions=[_RequiresExt()])
+
+    async with Client(server, extensions=[advertise(_NEEDS_EXT)]) as client:
+        result = await client.call_tool("guarded", {})
+
+    assert result == snapshot(CallToolResult(content=[TextContent(text="ok")], structured_content={"result": "ok"}))
+
+
+async def test_require_client_extension_raises_minus_32021_when_client_did_not_declare_it() -> None:
+    """SDK-defined: `require_client_extension` raises the -32021 missing-required-capability
+    error when the client did not advertise the id."""
+    server = MCPServer("test", extensions=[_RequiresExt()])
+
+    async with Client(server) as client:
+        with pytest.raises(MCPError) as exc_info:
+            await client.call_tool("guarded", {})
+
+    assert exc_info.value.code == MISSING_REQUIRED_CLIENT_CAPABILITY
+    assert exc_info.value.error.data == snapshot({"requiredCapabilities": {"extensions": {_NEEDS_EXT: {}}}})