modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
This commit is contained in:
Max Isbey 2026-02-12 15:55:54 +00:00 • committed by GitHub
commit 0a22a9dc33
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 3563 additions and 5481 deletions
Download patch file Download diff file Expand all files Collapse all files

30
tests/shared/test_memory.py
View file

@ -1,30 +0,0 @@
import pytest
from mcp import Client
from mcp.server import Server
from mcp.types import EmptyResult, Resource
@pytest.fixture
def mcp_server() -> Server:
server = Server(name="test_server")
@server.list_resources()
async def handle_list_resources(): # pragma: no cover
return [
Resource(
uri="memory://test",
name="Test Resource",
description="A test resource",
)
]
return server
@pytest.mark.anyio
async def test_memory_server_and_client_connection(mcp_server: Server):
"""Shows how a client and server can communicate over memory streams."""
async with Client(mcp_server) as client:
response = await client.send_ping()
assert isinstance(response, EmptyResult)

177
tests/shared/test_progress_notifications.py
View file

@ -6,7 +6,7 @@ import pytest
from mcp import Client, types
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
@ -35,9 +35,6 @@ async def test_bidirectional_progress_notifications():
capabilities=server.get_capabilities(NotificationOptions(), {}),
),
) as server_session:
global serv_sesh
serv_sesh = server_session
async for message in server_session.incoming_messages:
try:
await server._handle_message(message, server_session, {})
@ -52,79 +49,73 @@ async def test_bidirectional_progress_notifications():
server_progress_token = "server_token_123"
client_progress_token = "client_token_456"
# Create a server with progress capability
server = Server(name="ProgressTestServer")
# Register progress handler
@server.progress_notification()
async def handle_progress(
progress_token: str | int,
progress: float,
total: float | None,
message: str | None,
):
async def handle_progress(ctx: ServerRequestContext, params: types.ProgressNotificationParams) -> None:
server_progress_updates.append(
{
"token": progress_token,
"progress": progress,
"total": total,
"message": message,
"token": params.progress_token,
"progress": params.progress,
"total": params.total,
"message": params.message,
}
)
# Register list tool handler
@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
return [
types.Tool(
name="test_tool",
description="A tool that sends progress notifications <o/",
input_schema={},
)
]
async def handle_list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
return types.ListToolsResult(
tools=[
types.Tool(
name="test_tool",
description="A tool that sends progress notifications <o/",
input_schema={},
)
]
)
# Register tool handler
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[types.TextContent]:
async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
# Make sure we received a progress token
if name == "test_tool":
if arguments and "_meta" in arguments:
progressToken = arguments["_meta"]["progressToken"]
if params.name == "test_tool":
assert params.meta is not None
progress_token = params.meta.get("progress_token")
assert progress_token is not None
assert progress_token == client_progress_token
if not progressToken: # pragma: no cover
raise ValueError("Empty progress token received")
# Send progress notifications using ctx.session
await ctx.session.send_progress_notification(
progress_token=progress_token,
progress=0.25,
total=1.0,
message="Server progress 25%",
)
if progressToken != client_progress_token: # pragma: no cover
raise ValueError("Server sending back incorrect progressToken")
await ctx.session.send_progress_notification(
progress_token=progress_token,
progress=0.5,
total=1.0,
message="Server progress 50%",
)
# Send progress notifications
await serv_sesh.send_progress_notification(
progress_token=progressToken,
progress=0.25,
total=1.0,
message="Server progress 25%",
)
await ctx.session.send_progress_notification(
progress_token=progress_token,
progress=1.0,
total=1.0,
message="Server progress 100%",
)
await serv_sesh.send_progress_notification(
progress_token=progressToken,
progress=0.5,
total=1.0,
message="Server progress 50%",
)
return types.CallToolResult(content=[types.TextContent(type="text", text="Tool executed successfully")])
await serv_sesh.send_progress_notification(
progress_token=progressToken,
progress=1.0,
total=1.0,
message="Server progress 100%",
)
raise ValueError(f"Unknown tool: {params.name}") # pragma: no cover
else: # pragma: no cover
raise ValueError("Progress token not sent.")
return [types.TextContent(type="text", text="Tool executed successfully")]
raise ValueError(f"Unknown tool: {name}") # pragma: no cover
# Create a server with progress capability
server = Server(
name="ProgressTestServer",
on_progress=handle_progress,
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
)
# Client message handler to store progress notifications
async def handle_client_message(
@ -164,7 +155,7 @@ async def test_bidirectional_progress_notifications():
await client_session.list_tools()
# Call test_tool with progress token
await client_session.call_tool("test_tool", {"_meta": {"progressToken": client_progress_token}})
await client_session.call_tool("test_tool", meta={"progress_token": client_progress_token})
# Send progress notifications from client to server
await client_session.send_progress_notification(
@ -217,22 +208,21 @@ async def test_progress_context_manager():
# Track progress updates
server_progress_updates: list[dict[str, Any]] = []
server = Server(name="ProgressContextTestServer")
progress_token = None
# Register progress handler
@server.progress_notification()
async def handle_progress(
progress_token: str | int,
progress: float,
total: float | None,
message: str | None,
):
async def handle_progress(ctx: ServerRequestContext, params: types.ProgressNotificationParams) -> None:
server_progress_updates.append(
{"token": progress_token, "progress": progress, "total": total, "message": message}
{
"token": params.progress_token,
"progress": params.progress,
"total": params.total,
"message": params.message,
}
)
server = Server(name="ProgressContextTestServer", on_progress=handle_progress)
# Run server session to receive progress updates
async def run_server():
# Create a server session
@ -334,30 +324,37 @@ async def test_progress_callback_exception_logging():
raise ValueError("Progress callback failed!")
# Create a server with a tool that sends progress notifications
server = Server(name="TestProgressServer")
@server.call_tool()
async def handle_call_tool(name: str, arguments: Any) -> list[types.TextContent]:
if name == "progress_tool":
async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
if params.name == "progress_tool":
assert ctx.request_id is not None
# Send a progress notification
await server.request_context.session.send_progress_notification(
progress_token=server.request_context.request_id,
await ctx.session.send_progress_notification(
progress_token=ctx.request_id,
progress=50.0,
total=100.0,
message="Halfway done",
)
return [types.TextContent(type="text", text="progress_result")]
raise ValueError(f"Unknown tool: {name}") # pragma: no cover
return types.CallToolResult(content=[types.TextContent(type="text", text="progress_result")])
raise ValueError(f"Unknown tool: {params.name}") # pragma: no cover
@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
return [
types.Tool(
name="progress_tool",
description="A tool that sends progress notifications",
input_schema={},
)
]
async def handle_list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
return types.ListToolsResult(
tools=[
types.Tool(
name="progress_tool",
description="A tool that sends progress notifications",
input_schema={},
)
]
)
server = Server(
name="TestProgressServer",
on_call_tool=handle_call_tool,
on_list_tools=handle_list_tools,
)
# Test with mocked logging
with patch("mcp.shared.session.logging.exception", side_effect=mock_log_exception):

39
tests/shared/test_session.py
View file

@ -1,11 +1,9 @@
from typing import Any
import anyio
import pytest
from mcp import Client, types
from mcp.client.session import ClientSession
from mcp.server.lowlevel.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.shared.exceptions import MCPError
from mcp.shared.memory import create_client_server_memory_streams
from mcp.shared.message import SessionMessage
@ -17,7 +15,6 @@ from mcp.types import (
JSONRPCError,
JSONRPCRequest,
JSONRPCResponse,
TextContent,
)
@ -42,29 +39,25 @@ async def test_request_cancellation():
request_id = None
# Create a server with a slow tool
server = Server(name="TestSessionServer")
# Register the tool handler
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[TextContent]:
async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
nonlocal request_id, ev_tool_called
if name == "slow_tool":
request_id = server.request_context.request_id
if params.name == "slow_tool":
request_id = ctx.request_id
ev_tool_called.set()
await anyio.sleep(10) # Long enough to ensure we can cancel
return [] # pragma: no cover
raise ValueError(f"Unknown tool: {name}") # pragma: no cover
return types.CallToolResult(content=[]) # pragma: no cover
raise ValueError(f"Unknown tool: {params.name}") # pragma: no cover
# Register the tool so it shows up in list_tools
@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
return [
types.Tool(
name="slow_tool",
description="A slow tool that takes 10 seconds to complete",
input_schema={},
)
]
async def handle_list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
raise NotImplementedError
server = Server(
name="TestSessionServer",
on_call_tool=handle_call_tool,
on_list_tools=handle_list_tools,
)
async def make_request(client: Client):
nonlocal ev_cancelled

152
tests/shared/test_sse.py
View file

@ -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:

586
tests/shared/test_streamable_http.py
View file

@ -10,7 +10,9 @@ import multiprocessing
import socket
import time
import traceback
from collections.abc import Generator
from collections.abc import AsyncIterator, Generator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any
from unittest.mock import MagicMock
from urllib.parse import urlparse
@ -28,7 +30,7 @@ from starlette.routing import Mount
from mcp import MCPError, types
from mcp.client.session import ClientSession
from mcp.client.streamable_http import StreamableHTTPTransport, streamable_http_client
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.streamable_http import (
MCP_PROTOCOL_VERSION_HEADER,
MCP_SESSION_ID_HEADER,
@ -50,7 +52,19 @@ from mcp.shared._httpx_utils import (
)
from mcp.shared.message import ClientMessageMetadata, ServerMessageMetadata, SessionMessage
from mcp.shared.session import RequestResponder
from mcp.types import InitializeResult, JSONRPCRequest, TextContent, TextResourceContents, Tool
from mcp.types import (
CallToolRequestParams,
CallToolResult,
InitializeResult,
JSONRPCRequest,
ListToolsResult,
PaginatedRequestParams,
ReadResourceRequestParams,
ReadResourceResult,
TextContent,
TextResourceContents,
Tool,
)
from tests.test_helpers import wait_for_server
# Test constants
@ -124,263 +138,258 @@ class SimpleEventStore(EventStore):
return target_stream_id
# Test server implementation that follows MCP protocol
class ServerTest(Server): # pragma: no cover
def __init__(self):
super().__init__(SERVER_NAME)
self._lock = None # Will be initialized in async context
@dataclass
class ServerState:
lock: anyio.Event = field(default_factory=anyio.Event)
@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 ValueError(f"Unknown resource: {uri}")
@asynccontextmanager
async def _server_lifespan(_server: Server[ServerState]) -> AsyncIterator[ServerState]: # pragma: no cover
yield ServerState()
@self.list_tools()
async def handle_list_tools() -> list[Tool]:
return [
Tool(
name="test_tool",
description="A test tool",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="test_tool_with_standalone_notification",
description="A test tool that sends a notification",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="long_running_with_checkpoints",
description="A long-running tool that sends periodic notifications",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="test_sampling_tool",
description="A tool that triggers server-side sampling",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="wait_for_lock_with_notification",
description="A tool that sends a notification and waits for lock",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="release_lock",
description="A tool that releases the lock",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="tool_with_stream_close",
description="A tool that closes SSE stream mid-operation",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="tool_with_multiple_notifications_and_close",
description="Tool that sends notification1, closes stream, sends notification2, notification3",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="tool_with_multiple_stream_closes",
description="Tool that closes SSE stream multiple times during execution",
input_schema={
"type": "object",
"properties": {
"checkpoints": {"type": "integer", "default": 3},
"sleep_time": {"type": "number", "default": 0.2},
},
async def _handle_read_resource( # pragma: no cover
ctx: ServerRequestContext[ServerState], 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 ValueError(f"Unknown resource: {uri}")
return ReadResourceResult(contents=[TextResourceContents(uri=uri, text=text, mime_type="text/plain")])
async def _handle_list_tools( # pragma: no cover
ctx: ServerRequestContext[ServerState], params: PaginatedRequestParams | None
) -> ListToolsResult:
return ListToolsResult(
tools=[
Tool(
name="test_tool",
description="A test tool",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="test_tool_with_standalone_notification",
description="A test tool that sends a notification",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="long_running_with_checkpoints",
description="A long-running tool that sends periodic notifications",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="test_sampling_tool",
description="A tool that triggers server-side sampling",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="wait_for_lock_with_notification",
description="A tool that sends a notification and waits for lock",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="release_lock",
description="A tool that releases the lock",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="tool_with_stream_close",
description="A tool that closes SSE stream mid-operation",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="tool_with_multiple_notifications_and_close",
description="Tool that sends notification1, closes stream, sends notification2, notification3",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="tool_with_multiple_stream_closes",
description="Tool that closes SSE stream multiple times during execution",
input_schema={
"type": "object",
"properties": {
"checkpoints": {"type": "integer", "default": 3},
"sleep_time": {"type": "number", "default": 0.2},
},
),
Tool(
name="tool_with_standalone_stream_close",
description="Tool that closes standalone GET stream mid-operation",
input_schema={"type": "object", "properties": {}},
),
},
),
Tool(
name="tool_with_standalone_stream_close",
description="Tool that closes standalone GET stream mid-operation",
input_schema={"type": "object", "properties": {}},
),
]
)
async def _handle_call_tool( # pragma: no cover
ctx: ServerRequestContext[ServerState], params: CallToolRequestParams
) -> CallToolResult:
name = params.name
args = params.arguments or {}
# When the tool is called, send a notification to test GET stream
if name == "test_tool_with_standalone_notification":
await ctx.session.send_resource_updated(uri="http://test_resource")
return CallToolResult(content=[TextContent(type="text", text=f"Called {name}")])
elif name == "long_running_with_checkpoints":
await ctx.session.send_log_message(
level="info",
data="Tool started",
logger="tool",
related_request_id=ctx.request_id,
)
await anyio.sleep(0.1)
await ctx.session.send_log_message(
level="info",
data="Tool is almost done",
logger="tool",
related_request_id=ctx.request_id,
)
return CallToolResult(content=[TextContent(type="text", text="Completed!")])
elif name == "test_sampling_tool":
sampling_result = await ctx.session.create_message(
messages=[
types.SamplingMessage(
role="user",
content=types.TextContent(type="text", text="Server needs client sampling"),
)
],
max_tokens=100,
related_request_id=ctx.request_id,
)
if sampling_result.content.type == "text":
response = sampling_result.content.text
else:
response = str(sampling_result.content)
return CallToolResult(
content=[
TextContent(
type="text",
text=f"Response from sampling: {response}",
)
]
)
@self.call_tool()
async def handle_call_tool(name: str, args: dict[str, Any]) -> list[TextContent]:
ctx = self.request_context
elif name == "wait_for_lock_with_notification":
await ctx.session.send_log_message(
level="info",
data="First notification before lock",
logger="lock_tool",
related_request_id=ctx.request_id,
)
# When the tool is called, send a notification to test GET stream
if name == "test_tool_with_standalone_notification":
await ctx.session.send_resource_updated(uri="http://test_resource")
return [TextContent(type="text", text=f"Called {name}")]
await ctx.lifespan_context.lock.wait()
elif name == "long_running_with_checkpoints":
# Send notifications that are part of the response stream
# This simulates a long-running tool that sends logs
await ctx.session.send_log_message(
level="info",
data="Second notification after lock",
logger="lock_tool",
related_request_id=ctx.request_id,
)
await ctx.session.send_log_message(
level="info",
data="Tool started",
logger="tool",
related_request_id=ctx.request_id, # need for stream association
)
return CallToolResult(content=[TextContent(type="text", text="Completed")])
await anyio.sleep(0.1)
elif name == "release_lock":
ctx.lifespan_context.lock.set()
return CallToolResult(content=[TextContent(type="text", text="Lock released")])
await ctx.session.send_log_message(
level="info",
data="Tool is almost done",
logger="tool",
related_request_id=ctx.request_id,
)
elif name == "tool_with_stream_close":
await ctx.session.send_log_message(
level="info",
data="Before close",
logger="stream_close_tool",
related_request_id=ctx.request_id,
)
assert ctx.close_sse_stream is not None
await ctx.close_sse_stream()
await anyio.sleep(0.1)
await ctx.session.send_log_message(
level="info",
data="After close",
logger="stream_close_tool",
related_request_id=ctx.request_id,
)
return CallToolResult(content=[TextContent(type="text", text="Done")])
return [TextContent(type="text", text="Completed!")]
elif name == "tool_with_multiple_notifications_and_close":
await ctx.session.send_log_message(
level="info",
data="notification1",
logger="multi_notif_tool",
related_request_id=ctx.request_id,
)
assert ctx.close_sse_stream is not None
await ctx.close_sse_stream()
await anyio.sleep(0.1)
await ctx.session.send_log_message(
level="info",
data="notification2",
logger="multi_notif_tool",
related_request_id=ctx.request_id,
)
await ctx.session.send_log_message(
level="info",
data="notification3",
logger="multi_notif_tool",
related_request_id=ctx.request_id,
)
return CallToolResult(content=[TextContent(type="text", text="All notifications sent")])
elif name == "test_sampling_tool":
# Test sampling by requesting the client to sample a message
sampling_result = await ctx.session.create_message(
messages=[
types.SamplingMessage(
role="user",
content=types.TextContent(type="text", text="Server needs client sampling"),
)
],
max_tokens=100,
related_request_id=ctx.request_id,
)
elif name == "tool_with_multiple_stream_closes":
num_checkpoints = args.get("checkpoints", 3)
sleep_time = args.get("sleep_time", 0.2)
# Return the sampling result in the tool response
# Since we're not passing tools param, result.content is single content
if sampling_result.content.type == "text":
response = sampling_result.content.text
else:
response = str(sampling_result.content)
return [
TextContent(
type="text",
text=f"Response from sampling: {response}",
)
]
for i in range(num_checkpoints):
await ctx.session.send_log_message(
level="info",
data=f"checkpoint_{i}",
logger="multi_close_tool",
related_request_id=ctx.request_id,
)
elif name == "wait_for_lock_with_notification":
# Initialize lock if not already done
if self._lock is None:
self._lock = anyio.Event()
# First send a notification
await ctx.session.send_log_message(
level="info",
data="First notification before lock",
logger="lock_tool",
related_request_id=ctx.request_id,
)
# Now wait for the lock to be released
await self._lock.wait()
# Send second notification after lock is released
await ctx.session.send_log_message(
level="info",
data="Second notification after lock",
logger="lock_tool",
related_request_id=ctx.request_id,
)
return [TextContent(type="text", text="Completed")]
elif name == "release_lock":
assert self._lock is not None, "Lock must be initialized before releasing"
# Release the lock
self._lock.set()
return [TextContent(type="text", text="Lock released")]
elif name == "tool_with_stream_close":
# Send notification before closing
await ctx.session.send_log_message(
level="info",
data="Before close",
logger="stream_close_tool",
related_request_id=ctx.request_id,
)
# Close SSE stream (triggers client reconnect)
assert ctx.close_sse_stream is not None
if ctx.close_sse_stream:
await ctx.close_sse_stream()
# Continue processing (events stored in event_store)
await anyio.sleep(0.1)
await ctx.session.send_log_message(
level="info",
data="After close",
logger="stream_close_tool",
related_request_id=ctx.request_id,
)
return [TextContent(type="text", text="Done")]
elif name == "tool_with_multiple_notifications_and_close":
# Send notification1
await ctx.session.send_log_message(
level="info",
data="notification1",
logger="multi_notif_tool",
related_request_id=ctx.request_id,
)
# Close SSE stream
assert ctx.close_sse_stream is not None
await ctx.close_sse_stream()
# Send notification2, notification3 (stored in event_store)
await anyio.sleep(0.1)
await ctx.session.send_log_message(
level="info",
data="notification2",
logger="multi_notif_tool",
related_request_id=ctx.request_id,
)
await ctx.session.send_log_message(
level="info",
data="notification3",
logger="multi_notif_tool",
related_request_id=ctx.request_id,
)
return [TextContent(type="text", text="All notifications sent")]
await anyio.sleep(sleep_time)
elif name == "tool_with_multiple_stream_closes":
num_checkpoints = args.get("checkpoints", 3)
sleep_time = args.get("sleep_time", 0.2)
return CallToolResult(content=[TextContent(type="text", text=f"Completed {num_checkpoints} checkpoints")])
for i in range(num_checkpoints):
await ctx.session.send_log_message(
level="info",
data=f"checkpoint_{i}",
logger="multi_close_tool",
related_request_id=ctx.request_id,
)
elif name == "tool_with_standalone_stream_close":
await ctx.session.send_resource_updated(uri="http://notification_1")
await anyio.sleep(0.1)
if ctx.close_sse_stream:
await ctx.close_sse_stream()
if ctx.close_standalone_sse_stream:
await ctx.close_standalone_sse_stream()
await anyio.sleep(sleep_time)
await anyio.sleep(1.5)
await ctx.session.send_resource_updated(uri="http://notification_2")
return [TextContent(type="text", text=f"Completed {num_checkpoints} checkpoints")]
return CallToolResult(content=[TextContent(type="text", text="Standalone stream close test done")])
elif name == "tool_with_standalone_stream_close":
# Test for GET stream reconnection
# 1. Send unsolicited notification via GET stream (no related_request_id)
await ctx.session.send_resource_updated(uri="http://notification_1")
return CallToolResult(content=[TextContent(type="text", text=f"Called {name}")])
# Small delay to ensure notification is flushed before closing
await anyio.sleep(0.1)
# 2. Close the standalone GET stream
if ctx.close_standalone_sse_stream:
await ctx.close_standalone_sse_stream()
# 3. Wait for client to reconnect (uses retry_interval from server, default 1000ms)
await anyio.sleep(1.5)
# 4. Send another notification on the new GET stream connection
await ctx.session.send_resource_updated(uri="http://notification_2")
return [TextContent(type="text", text="Standalone stream close test done")]
return [TextContent(type="text", text=f"Called {name}")]
def _create_server() -> Server[ServerState]: # pragma: no cover
return Server(
SERVER_NAME,
lifespan=_server_lifespan,
on_read_resource=_handle_read_resource,
on_list_tools=_handle_list_tools,
on_call_tool=_handle_call_tool,
)
def create_app(
@ -396,7 +405,7 @@ def create_app(
retry_interval: Retry interval in milliseconds for SSE polling.
"""
# Create server instance
server = ServerTest()
server = _create_server()
# Create the session manager
security_settings = TransportSecuritySettings(
@ -1385,69 +1394,68 @@ async def test_streamablehttp_server_sampling(basic_server: None, basic_server_u
# Context-aware server implementation for testing request context propagation
class ContextAwareServerTest(Server): # pragma: no cover
def __init__(self):
super().__init__("ContextAwareServer")
@self.list_tools()
async def handle_list_tools() -> list[Tool]:
return [
Tool(
name="echo_headers",
description="Echo request headers from context",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="echo_context",
description="Echo request context with custom data",
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="Echo request headers from context",
input_schema={"type": "object", "properties": {}},
),
Tool(
name="echo_context",
description="Echo request context with custom data",
input_schema={
"type": "object",
"properties": {
"request_id": {"type": "string"},
},
),
]
"required": ["request_id"],
},
),
]
)
@self.call_tool()
async def handle_call_tool(name: str, args: dict[str, Any]) -> list[TextContent]:
ctx = self.request_context
if name == "echo_headers":
# Access the request object from context
headers_info = {}
if ctx.request and isinstance(ctx.request, Request):
headers_info = dict(ctx.request.headers)
return [TextContent(type="text", text=json.dumps(headers_info))]
async def _handle_context_call_tool( # pragma: no cover
ctx: ServerRequestContext, params: CallToolRequestParams
) -> CallToolResult:
name = params.name
args = params.arguments or {}
elif name == "echo_context":
# Return full context information
context_data: dict[str, Any] = {
"request_id": args.get("request_id"),
"headers": {},
"method": None,
"path": None,
}
if ctx.request and isinstance(ctx.request, Request):
request = ctx.request
context_data["headers"] = dict(request.headers)
context_data["method"] = request.method
context_data["path"] = request.url.path
return [
TextContent(
type="text",
text=json.dumps(context_data),
)
]
if name == "echo_headers":
headers_info: dict[str, Any] = {}
if ctx.request and isinstance(ctx.request, Request):
headers_info = dict(ctx.request.headers)
return CallToolResult(content=[TextContent(type="text", text=json.dumps(headers_info))])
return [TextContent(type="text", text=f"Unknown tool: {name}")]
elif name == "echo_context":
context_data: dict[str, Any] = {
"request_id": args.get("request_id"),
"headers": {},
"method": None,
"path": None,
}
if ctx.request and isinstance(ctx.request, Request):
request = ctx.request
context_data["headers"] = dict(request.headers)
context_data["method"] = request.method
context_data["path"] = request.url.path
return CallToolResult(content=[TextContent(type="text", text=json.dumps(context_data))])
return CallToolResult(content=[TextContent(type="text", text=f"Unknown tool: {name}")])
# Server runner for context-aware testing
def run_context_aware_server(port: int): # pragma: no cover
"""Run the context-aware test server."""
server = ContextAwareServerTest()
server = Server(
"ContextAwareServer",
on_list_tools=_handle_context_list_tools,
on_call_tool=_handle_context_call_tool,
)
session_manager = StreamableHTTPSessionManager(
app=server,

90
tests/shared/test_ws.py
View file

@ -2,7 +2,6 @@ import multiprocessing
import socket
import time
from collections.abc import AsyncGenerator, Generator
from typing import Any
from urllib.parse import urlparse
import anyio
@ -15,9 +14,21 @@ from starlette.websockets import WebSocket
from mcp import MCPError
from mcp.client.session import ClientSession
from mcp.client.websocket import websocket_client
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.websocket import websocket_server
from mcp.types import EmptyResult, InitializeResult, ReadResourceResult, TextContent, TextResourceContents, Tool
from mcp.types import (
CallToolRequestParams,
CallToolResult,
EmptyResult,
InitializeResult,
ListToolsResult,
PaginatedRequestParams,
ReadResourceRequestParams,
ReadResourceResult,
TextContent,
TextResourceContents,
Tool,
)
from tests.test_helpers import wait_for_server
SERVER_NAME = "test_server_for_WS"
@ -35,42 +46,59 @@ def server_url(server_port: int) -> str:
return f"ws://127.0.0.1:{server_port}"
# Test server implementation
class ServerTest(Server): # pragma: no cover
def __init__(self):
super().__init__(SERVER_NAME)
@self.read_resource()
async def handle_read_resource(uri: str) -> str | bytes:
parsed = urlparse(uri)
if parsed.scheme == "foobar":
return f"Read {parsed.netloc}"
elif 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")
@self.list_tools()
async def handle_list_tools() -> list[Tool]:
return [
Tool(
name="test_tool",
description="A test tool",
input_schema={"type": "object", "properties": {}},
async def handle_read_resource( # pragma: no cover
ctx: ServerRequestContext, params: ReadResourceRequestParams
) -> ReadResourceResult:
parsed = urlparse(str(params.uri))
if parsed.scheme == "foobar":
return ReadResourceResult(
contents=[TextResourceContents(uri=str(params.uri), text=f"Read {parsed.netloc}", mime_type="text/plain")]
)
elif parsed.scheme == "slow":
await anyio.sleep(2.0)
return ReadResourceResult(
contents=[
TextResourceContents(
uri=str(params.uri), text=f"Slow response from {parsed.netloc}", mime_type="text/plain"
)
]
)
raise MCPError(code=404, message="OOPS! no resource with that URI was found")
@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_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": {}},
)
]
)
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
def make_server_app() -> Starlette: # pragma: no cover
"""Create test Starlette app with WebSocket transport"""
server = ServerTest()
server = _create_server()
async def handle_ws(websocket: WebSocket):
async with websocket_server(websocket.scope, websocket.receive, websocket.send) as streams: