refactor: replace lowlevel Server decorators with on_* constructor kwargs (#1985)

tests/shared/test_sse.py

@@ -22,15 +22,20 @@ import mcp.client.sse
 from mcp import types
 from mcp.client.session import ClientSession
 from mcp.client.sse import _extract_session_id_from_endpoint, sse_client
-from mcp.server import Server
+from mcp.server import Server, ServerRequestContext
 from mcp.server.sse import SseServerTransport
 from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared.exceptions import MCPError
 from mcp.types import (
+    CallToolRequestParams,
+    CallToolResult,
     EmptyResult,
     Implementation,
     InitializeResult,
     JSONRPCResponse,
+    ListToolsResult,
+    PaginatedRequestParams,
+    ReadResourceRequestParams,
     ReadResourceResult,
     ServerCapabilities,
     TextContent,
@@ -54,36 +59,48 @@ def server_url(server_port: int) -> str:
     return f"http://127.0.0.1:{server_port}"
 
 
-# Test server implementation
-class ServerTest(Server):  # pragma: no cover
-    def __init__(self):
-        super().__init__(SERVER_NAME)
+async def _handle_read_resource(  # pragma: no cover
+    ctx: ServerRequestContext, params: ReadResourceRequestParams
+) -> ReadResourceResult:
+    uri = str(params.uri)
+    parsed = urlparse(uri)
+    if parsed.scheme == "foobar":
+        text = f"Read {parsed.netloc}"
+    elif parsed.scheme == "slow":
+        await anyio.sleep(2.0)
+        text = f"Slow response from {parsed.netloc}"
+    else:
+        raise MCPError(code=404, message="OOPS! no resource with that URI was found")
+    return ReadResourceResult(contents=[TextResourceContents(uri=uri, text=text, mime_type="text/plain")])
 
-        @self.read_resource()
-        async def handle_read_resource(uri: str) -> str | bytes:
-            parsed = urlparse(uri)
-            if parsed.scheme == "foobar":
-                return f"Read {parsed.netloc}"
-            if parsed.scheme == "slow":
-                # Simulate a slow resource
-                await anyio.sleep(2.0)
-                return f"Slow response from {parsed.netloc}"
 
-            raise MCPError(code=404, message="OOPS! no resource with that URI was found")
+async def _handle_list_tools(  # pragma: no cover
+    ctx: ServerRequestContext, params: PaginatedRequestParams | None
+) -> ListToolsResult:
+    return ListToolsResult(
+        tools=[
+            Tool(
+                name="test_tool",
+                description="A test tool",
+                input_schema={"type": "object", "properties": {}},
+            )
+        ]
+    )
 
-        @self.list_tools()
-        async def handle_list_tools() -> list[Tool]:
-            return [
-                Tool(
-                    name="test_tool",
-                    description="A test tool",
-                    input_schema={"type": "object", "properties": {}},
-                )
-            ]
 
-        @self.call_tool()
-        async def handle_call_tool(name: str, args: dict[str, Any]) -> list[TextContent]:
-            return [TextContent(type="text", text=f"Called {name}")]
+async def _handle_call_tool(  # pragma: no cover
+    ctx: ServerRequestContext, params: CallToolRequestParams
+) -> CallToolResult:
+    return CallToolResult(content=[TextContent(type="text", text=f"Called {params.name}")])
+
+
+def _create_server() -> Server:  # pragma: no cover
+    return Server(
+        SERVER_NAME,
+        on_read_resource=_handle_read_resource,
+        on_list_tools=_handle_list_tools,
+        on_call_tool=_handle_call_tool,
+    )
 
 
 # Test fixtures
@@ -94,7 +111,7 @@ def make_server_app() -> Starlette:  # pragma: no cover
         allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"]
     )
     sse = SseServerTransport("/messages/", security_settings=security_settings)
-    server = ServerTest()
+    server = _create_server()
 
     async def handle_sse(request: Request) -> Response:
         async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
@@ -336,47 +353,46 @@ async def test_sse_client_basic_connection_mounted_app(mounted_server: None, ser
             assert isinstance(ping_result, EmptyResult)
 
 
-# Test server with request context that returns headers in the response
-class RequestContextServer(Server[object, Request]):  # pragma: no cover
-    def __init__(self):
-        super().__init__("request_context_server")
+async def _handle_context_call_tool(  # pragma: no cover
+    ctx: ServerRequestContext, params: CallToolRequestParams
+) -> CallToolResult:
+    headers_info: dict[str, Any] = {}
+    if ctx.request:
+        headers_info = dict(ctx.request.headers)
 
-        @self.call_tool()
-        async def handle_call_tool(name: str, args: dict[str, Any]) -> list[TextContent]:
-            headers_info = {}
-            context = self.request_context
-            if context.request:
-                headers_info = dict(context.request.headers)
+    if params.name == "echo_headers":
+        return CallToolResult(content=[TextContent(type="text", text=json.dumps(headers_info))])
+    elif params.name == "echo_context":
+        context_data = {
+            "request_id": (params.arguments or {}).get("request_id"),
+            "headers": headers_info,
+        }
+        return CallToolResult(content=[TextContent(type="text", text=json.dumps(context_data))])
 
-            if name == "echo_headers":
-                return [TextContent(type="text", text=json.dumps(headers_info))]
-            elif name == "echo_context":
-                context_data = {
-                    "request_id": args.get("request_id"),
-                    "headers": headers_info,
-                }
-                return [TextContent(type="text", text=json.dumps(context_data))]
+    return CallToolResult(content=[TextContent(type="text", text=f"Called {params.name}")])
 
-            return [TextContent(type="text", text=f"Called {name}")]
 
-        @self.list_tools()
-        async def handle_list_tools() -> list[Tool]:
-            return [
-                Tool(
-                    name="echo_headers",
-                    description="Echoes request headers",
-                    input_schema={"type": "object", "properties": {}},
-                ),
-                Tool(
-                    name="echo_context",
-                    description="Echoes request context",
-                    input_schema={
-                        "type": "object",
-                        "properties": {"request_id": {"type": "string"}},
-                        "required": ["request_id"],
-                    },
-                ),
-            ]
+async def _handle_context_list_tools(  # pragma: no cover
+    ctx: ServerRequestContext, params: PaginatedRequestParams | None
+) -> ListToolsResult:
+    return ListToolsResult(
+        tools=[
+            Tool(
+                name="echo_headers",
+                description="Echoes request headers",
+                input_schema={"type": "object", "properties": {}},
+            ),
+            Tool(
+                name="echo_context",
+                description="Echoes request context",
+                input_schema={
+                    "type": "object",
+                    "properties": {"request_id": {"type": "string"}},
+                    "required": ["request_id"],
+                },
+            ),
+        ]
+    )
 
 
 def run_context_server(server_port: int) -> None:  # pragma: no cover
@@ -386,7 +402,11 @@ def run_context_server(server_port: int) -> None:  # pragma: no cover
         allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"]
     )
     sse = SseServerTransport("/messages/", security_settings=security_settings)
-    context_server = RequestContextServer()
+    context_server = Server(
+        "request_context_server",
+        on_call_tool=_handle_context_call_tool,
+        on_list_tools=_handle_context_list_tools,
+    )
 
     async def handle_sse(request: Request) -> Response:
         async with sse.connect_sse(request.scope, request.receive, request._send) as streams: