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
38 changed files with 3563 additions and 5481 deletions
Download patch file Download diff file Expand all files Collapse all files

60
tests/client/test_client.py
View file

@ -11,7 +11,7 @@ from inline_snapshot import snapshot
from mcp import types
from mcp.client._memory import InMemoryTransport
from mcp.client.client import Client
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.mcpserver import MCPServer
from mcp.types import (
CallToolResult,
@ -41,33 +41,36 @@ pytestmark = pytest.mark.anyio
@pytest.fixture
def simple_server() -> Server:
"""Create a simple MCP server for testing."""
server = Server(name="test_server")
@server.list_resources()
async def handle_list_resources():
return [Resource(uri="memory://test", name="Test Resource", description="A test resource")]
async def handle_list_resources(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> ListResourcesResult:
return ListResourcesResult(
resources=[Resource(uri="memory://test", name="Test Resource", description="A test resource")]
)
@server.subscribe_resource()
async def handle_subscribe_resource(uri: str):
pass
async def handle_subscribe_resource(ctx: ServerRequestContext, params: types.SubscribeRequestParams) -> EmptyResult:
return EmptyResult()
@server.unsubscribe_resource()
async def handle_unsubscribe_resource(uri: str):
pass
async def handle_unsubscribe_resource(
ctx: ServerRequestContext, params: types.UnsubscribeRequestParams
) -> EmptyResult:
return EmptyResult()
@server.set_logging_level()
async def handle_set_logging_level(level: str):
pass
async def handle_set_logging_level(ctx: ServerRequestContext, params: types.SetLevelRequestParams) -> EmptyResult:
return EmptyResult()
@server.completion()
async def handle_completion(
ref: types.PromptReference | types.ResourceTemplateReference,
argument: types.CompletionArgument,
context: types.CompletionContext | None,
) -> types.Completion | None:
return types.Completion(values=[])
async def handle_completion(ctx: ServerRequestContext, params: types.CompleteRequestParams) -> types.CompleteResult:
return types.CompleteResult(completion=types.Completion(values=[]))
return server
return Server(
name="test_server",
on_list_resources=handle_list_resources,
on_subscribe_resource=handle_subscribe_resource,
on_unsubscribe_resource=handle_unsubscribe_resource,
on_set_logging_level=handle_set_logging_level,
on_completion=handle_completion,
)
@pytest.fixture
@ -202,19 +205,14 @@ async def test_client_send_progress_notification():
"""Test sending progress notification."""
received_from_client = None
event = anyio.Event()
server = Server(name="test_server")
@server.progress_notification()
async def handle_progress_notification(
progress_token: str | int,
progress: float = 0.0,
total: float | None = None,
message: str | None = None,
) -> None:
async def handle_progress(ctx: ServerRequestContext, params: types.ProgressNotificationParams) -> None:
nonlocal received_from_client
received_from_client = {"progress_token": progress_token, "progress": progress}
received_from_client = {"progress_token": params.progress_token, "progress": params.progress}
event.set()
server = Server(name="test_server", on_progress=handle_progress)
async with Client(server) as client:
await client.send_progress_notification(progress_token="token123", progress=50.0)
await event.wait()

101
tests/client/test_http_unicode.py
View file

@ -8,7 +8,6 @@ import multiprocessing
import socket
from collections.abc import AsyncGenerator, Generator
from contextlib import asynccontextmanager
from typing import Any
import pytest
from starlette.applications import Starlette
@ -17,7 +16,7 @@ from starlette.routing import Mount
from mcp import types
from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.types import TextContent, Tool
from tests.test_helpers import wait_for_server
@ -47,54 +46,56 @@ def run_unicode_server(port: int) -> None: # pragma: no cover
import uvicorn
# Need to recreate the server setup in this process
server = Server(name="unicode_test_server")
@server.list_tools()
async def list_tools() -> list[Tool]:
"""List tools with Unicode descriptions."""
return [
Tool(
name="echo_unicode",
description="🔤 Echo Unicode text - Hello 👋 World 🌍 - Testing 🧪 Unicode ✨",
input_schema={
"type": "object",
"properties": {
"text": {"type": "string", "description": "Text to echo back"},
async def handle_list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
return types.ListToolsResult(
tools=[
Tool(
name="echo_unicode",
description="🔤 Echo Unicode text - Hello 👋 World 🌍 - Testing 🧪 Unicode ✨",
input_schema={
"type": "object",
"properties": {
"text": {"type": "string", "description": "Text to echo back"},
},
"required": ["text"],
},
"required": ["text"],
},
),
]
),
]
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any] | None) -> list[TextContent]:
"""Handle tool calls with Unicode content."""
if name == "echo_unicode":
text = arguments.get("text", "") if arguments else ""
return [
TextContent(
type="text",
text=f"Echo: {text}",
async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
if params.name == "echo_unicode":
text = params.arguments.get("text", "") if params.arguments else ""
return types.CallToolResult(
content=[
TextContent(
type="text",
text=f"Echo: {text}",
)
]
)
else:
raise ValueError(f"Unknown tool: {params.name}")
async def handle_list_prompts(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListPromptsResult:
return types.ListPromptsResult(
prompts=[
types.Prompt(
name="unicode_prompt",
description="Unicode prompt - Слой хранилища, где располагаются",
arguments=[],
)
]
else:
raise ValueError(f"Unknown tool: {name}")
)
@server.list_prompts()
async def list_prompts() -> list[types.Prompt]:
"""List prompts with Unicode names and descriptions."""
return [
types.Prompt(
name="unicode_prompt",
description="Unicode prompt - Слой хранилища, где располагаются",
arguments=[],
)
]
@server.get_prompt()
async def get_prompt(name: str, arguments: dict[str, Any] | None) -> types.GetPromptResult:
"""Get a prompt with Unicode content."""
if name == "unicode_prompt":
async def handle_get_prompt(
ctx: ServerRequestContext, params: types.GetPromptRequestParams
) -> types.GetPromptResult:
if params.name == "unicode_prompt":
return types.GetPromptResult(
messages=[
types.PromptMessage(
@ -106,7 +107,15 @@ def run_unicode_server(port: int) -> None: # pragma: no cover
)
]
)
raise ValueError(f"Unknown prompt: {name}")
raise ValueError(f"Unknown prompt: {params.name}")
server = Server(
name="unicode_test_server",
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
on_list_prompts=handle_list_prompts,
on_get_prompt=handle_get_prompt,
)
# Create the session manager
session_manager = StreamableHTTPSessionManager(

14
tests/client/test_list_methods_cursor.py
View file

@ -3,9 +3,9 @@ from collections.abc import Callable
import pytest
from mcp import Client, types
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.mcpserver import MCPServer
from mcp.types import ListToolsRequest, ListToolsResult
from mcp.types import ListToolsResult
from .conftest import StreamSpyCollection
@ -105,14 +105,16 @@ async def test_list_tools_with_strict_server_validation(
async def test_list_tools_with_lowlevel_server():
"""Test that list_tools works with a lowlevel Server using params."""
server = Server("test-lowlevel")
@server.list_tools()
async def handle_list_tools(request: ListToolsRequest) -> ListToolsResult:
async def handle_list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> ListToolsResult:
# Echo back what cursor we received in the tool description
cursor = request.params.cursor if request.params else None
cursor = params.cursor if params else None
return ListToolsResult(tools=[types.Tool(name="test_tool", description=f"cursor={cursor}", input_schema={})])
server = Server("test-lowlevel", on_list_tools=handle_list_tools)
async with Client(server) as client:
result = await client.list_tools()
assert result.tools[0].description == "cursor=None"

196
tests/client/test_output_schema_validation.py
View file

@ -1,50 +1,41 @@
import inspect
import logging
from contextlib import contextmanager
from typing import Any
from unittest.mock import patch
import jsonschema
import pytest
from mcp import Client
from mcp.server.lowlevel import Server
from mcp.types import Tool
from mcp.server import Server, ServerRequestContext
from mcp.types import (
CallToolRequestParams,
CallToolResult,
ListToolsResult,
PaginatedRequestParams,
TextContent,
Tool,
)
@contextmanager
def bypass_server_output_validation():
"""Context manager that bypasses server-side output validation.
This simulates a malicious or non-compliant server that doesn't validate
its outputs, allowing us to test client-side validation.
"""
# Save the original validate function
original_validate = jsonschema.validate
def _make_server(
tools: list[Tool],
structured_content: dict[str, Any],
) -> Server:
"""Create a low-level server that returns the given structured_content for any tool call."""
# Create a mock that tracks which module is calling it
def selective_mock(instance: Any = None, schema: Any = None, *args: Any, **kwargs: Any) -> None:
# Check the call stack to see where this is being called from
for frame_info in inspect.stack():
# If called from the server module, skip validation
# TODO: fix this as it's a rather gross workaround and will eventually break
# Normalize path separators for cross-platform compatibility
normalized_path = frame_info.filename.replace("\\", "/")
if "mcp/server/lowlevel/server.py" in normalized_path:
return None
# Otherwise, use the real validation (for client-side)
return original_validate(instance=instance, schema=schema, *args, **kwargs)
async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(tools=tools)
with patch("jsonschema.validate", selective_mock):
yield
async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
return CallToolResult(
content=[TextContent(type="text", text="result")],
structured_content=structured_content,
)
return Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
@pytest.mark.anyio
async def test_tool_structured_output_client_side_validation_basemodel():
"""Test that client validates structured content against schema for BaseModel outputs"""
# Create a malicious low-level server that returns invalid structured content
server = Server("test-server")
# Define the expected schema for our tool
output_schema = {
"type": "object",
"properties": {"name": {"type": "string", "title": "Name"}, "age": {"type": "integer", "title": "Age"}},
@ -52,39 +43,27 @@ async def test_tool_structured_output_client_side_validation_basemodel():
"title": "UserOutput",
}
@server.list_tools()
async def list_tools():
return [
server = _make_server(
tools=[
Tool(
name="get_user",
description="Get user data",
input_schema={"type": "object"},
output_schema=output_schema,
)
]
],
structured_content={"name": "John", "age": "invalid"}, # Invalid: age should be int
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
# Return invalid structured content - age is string instead of integer
# The low-level server will wrap this in CallToolResult
return {"name": "John", "age": "invalid"} # Invalid: age should be int
# Test that client validates the structured content
with bypass_server_output_validation():
async with Client(server) as client:
# The client validates structured content and should raise an error
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("get_user", {})
# Verify it's a validation error
assert "Invalid structured content returned by tool get_user" in str(exc_info.value)
async with Client(server) as client:
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("get_user", {})
assert "Invalid structured content returned by tool get_user" in str(exc_info.value)
@pytest.mark.anyio
async def test_tool_structured_output_client_side_validation_primitive():
"""Test that client validates structured content for primitive outputs"""
server = Server("test-server")
# Primitive types are wrapped in {"result": value}
output_schema = {
"type": "object",
"properties": {"result": {"type": "integer", "title": "Result"}},
@ -92,122 +71,95 @@ async def test_tool_structured_output_client_side_validation_primitive():
"title": "calculate_Output",
}
@server.list_tools()
async def list_tools():
return [
server = _make_server(
tools=[
Tool(
name="calculate",
description="Calculate something",
input_schema={"type": "object"},
output_schema=output_schema,
)
]
],
structured_content={"result": "not_a_number"}, # Invalid: should be int
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
# Return invalid structured content - result is string instead of integer
return {"result": "not_a_number"} # Invalid: should be int
with bypass_server_output_validation():
async with Client(server) as client:
# The client validates structured content and should raise an error
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("calculate", {})
assert "Invalid structured content returned by tool calculate" in str(exc_info.value)
async with Client(server) as client:
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("calculate", {})
assert "Invalid structured content returned by tool calculate" in str(exc_info.value)
@pytest.mark.anyio
async def test_tool_structured_output_client_side_validation_dict_typed():
"""Test that client validates dict[str, T] structured content"""
server = Server("test-server")
# dict[str, int] schema
output_schema = {"type": "object", "additionalProperties": {"type": "integer"}, "title": "get_scores_Output"}
@server.list_tools()
async def list_tools():
return [
server = _make_server(
tools=[
Tool(
name="get_scores",
description="Get scores",
input_schema={"type": "object"},
output_schema=output_schema,
)
]
],
structured_content={"alice": "100", "bob": "85"}, # Invalid: values should be int
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
# Return invalid structured content - values should be integers
return {"alice": "100", "bob": "85"} # Invalid: values should be int
with bypass_server_output_validation():
async with Client(server) as client:
# The client validates structured content and should raise an error
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("get_scores", {})
assert "Invalid structured content returned by tool get_scores" in str(exc_info.value)
async with Client(server) as client:
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("get_scores", {})
assert "Invalid structured content returned by tool get_scores" in str(exc_info.value)
@pytest.mark.anyio
async def test_tool_structured_output_client_side_validation_missing_required():
"""Test that client validates missing required fields"""
server = Server("test-server")
output_schema = {
"type": "object",
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}, "email": {"type": "string"}},
"required": ["name", "age", "email"], # All fields required
"required": ["name", "age", "email"],
"title": "PersonOutput",
}
@server.list_tools()
async def list_tools():
return [
server = _make_server(
tools=[
Tool(
name="get_person",
description="Get person data",
input_schema={"type": "object"},
output_schema=output_schema,
)
]
],
structured_content={"name": "John", "age": 30}, # Missing required 'email'
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
# Return structured content missing required field 'email'
return {"name": "John", "age": 30} # Missing required 'email'
with bypass_server_output_validation():
async with Client(server) as client:
# The client validates structured content and should raise an error
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("get_person", {})
assert "Invalid structured content returned by tool get_person" in str(exc_info.value)
async with Client(server) as client:
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("get_person", {})
assert "Invalid structured content returned by tool get_person" in str(exc_info.value)
@pytest.mark.anyio
async def test_tool_not_listed_warning(caplog: pytest.LogCaptureFixture):
"""Test that client logs warning when tool is not in list_tools but has output_schema"""
server = Server("test-server")
@server.list_tools()
async def list_tools() -> list[Tool]:
# Return empty list - tool is not listed
return []
async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(tools=[])
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
# Server still responds to the tool call with structured content
return {"result": 42}
async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
return CallToolResult(
content=[TextContent(type="text", text="result")],
structured_content={"result": 42},
)
server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
# Set logging level to capture warnings
caplog.set_level(logging.WARNING)
with bypass_server_output_validation():
async with Client(server) as client:
# Call a tool that wasn't listed
result = await client.call_tool("mystery_tool", {})
assert result.structured_content == {"result": 42}
assert result.is_error is False
async with Client(server) as client:
result = await client.call_tool("mystery_tool", {})
assert result.structured_content == {"result": 42}
assert result.is_error is False
# Check that warning was logged
assert "Tool mystery_tool not listed" in caplog.text
assert "Tool mystery_tool not listed" in caplog.text

32
tests/client/transports/test_memory.py
View file

@ -2,31 +2,31 @@
import pytest
from mcp import Client
from mcp import Client, types
from mcp.client._memory import InMemoryTransport
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.mcpserver import MCPServer
from mcp.types import Resource
from mcp.types import ListResourcesResult, Resource
@pytest.fixture
def simple_server() -> Server:
"""Create a simple MCP server for testing."""
server = Server(name="test_server")
# pragma: no cover - handler exists only to register a resource capability.
# Transport tests verify stream creation, not handler invocation.
@server.list_resources()
async def handle_list_resources(): # pragma: no cover
return [
Resource(
uri="memory://test",
name="Test Resource",
description="A test resource",
)
]
async def handle_list_resources(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> ListResourcesResult: # pragma: no cover
return ListResourcesResult(
resources=[
Resource(
uri="memory://test",
name="Test Resource",
description="A test resource",
)
]
)
return server
return Server(name="test_server", on_list_resources=handle_list_resources)
@pytest.fixture

541
tests/experimental/tasks/client/test_tasks.py
View file

@ -1,43 +1,38 @@
"""Tests for the experimental client task methods (session.experimental)."""
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any
import anyio
import pytest
from anyio import Event
from anyio.abc import TaskGroup
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp import Client
from mcp.server import Server, ServerRequestContext
from mcp.shared.experimental.tasks.helpers import task_execution
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from mcp.types import (
CallToolRequest,
CallToolRequestParams,
CallToolResult,
CancelTaskRequest,
CancelTaskRequestParams,
CancelTaskResult,
ClientResult,
CreateTaskResult,
GetTaskPayloadRequest,
GetTaskPayloadRequestParams,
GetTaskPayloadResult,
GetTaskRequest,
GetTaskRequestParams,
GetTaskResult,
ListTasksRequest,
ListTasksResult,
ServerNotification,
ServerRequest,
ListToolsResult,
PaginatedRequestParams,
TaskMetadata,
TextContent,
Tool,
)
pytestmark = pytest.mark.anyio
@dataclass
class AppContext:
@ -48,44 +43,53 @@ class AppContext:
task_done_events: dict[str, Event] = field(default_factory=lambda: {})
@pytest.mark.anyio
async def _handle_list_tools(
ctx: ServerRequestContext[AppContext], params: PaginatedRequestParams | None
) -> ListToolsResult:
raise NotImplementedError
async def _handle_call_tool_with_done_event(
ctx: ServerRequestContext[AppContext], params: CallToolRequestParams, *, result_text: str = "Done"
) -> CallToolResult | CreateTaskResult:
app = ctx.lifespan_context
if ctx.experimental.is_task:
task_metadata = ctx.experimental.task_metadata
assert task_metadata is not None
task = await app.store.create_task(task_metadata)
done_event = Event()
app.task_done_events[task.task_id] = done_event
async def do_work() -> None:
async with task_execution(task.task_id, app.store) as task_ctx:
await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text=result_text)]))
done_event.set()
app.task_group.start_soon(do_work)
return CreateTaskResult(task=task)
raise NotImplementedError
def _make_lifespan(store: InMemoryTaskStore, task_done_events: dict[str, Event]):
@asynccontextmanager
async def app_lifespan(server: Server[AppContext]) -> AsyncIterator[AppContext]:
async with anyio.create_task_group() as tg:
yield AppContext(task_group=tg, store=store, task_done_events=task_done_events)
return app_lifespan
async def test_session_experimental_get_task() -> None:
"""Test session.experimental.get_task() method."""
# Note: We bypass the normal lifespan mechanism
server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment]
store = InMemoryTaskStore()
task_done_events: dict[str, Event] = {}
@server.list_tools()
async def list_tools():
return [Tool(name="test_tool", description="Test", input_schema={"type": "object"})]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
ctx = server.request_context
async def handle_get_task(ctx: ServerRequestContext[AppContext], params: GetTaskRequestParams) -> GetTaskResult:
app = ctx.lifespan_context
if ctx.experimental.is_task:
task_metadata = ctx.experimental.task_metadata
assert task_metadata is not None
task = await app.store.create_task(task_metadata)
done_event = Event()
app.task_done_events[task.task_id] = done_event
async def do_work():
async with task_execution(task.task_id, app.store) as task_ctx:
await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")]))
done_event.set()
app.task_group.start_soon(do_work)
return CreateTaskResult(task=task)
raise NotImplementedError
@server.experimental.get_task()
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
app = server.request_context.lifespan_context
task = await app.store.get_task(request.params.task_id)
assert task is not None, f"Test setup error: task {request.params.task_id} should exist"
task = await app.store.get_task(params.task_id)
assert task is not None, f"Test setup error: task {params.task_id} should exist"
return GetTaskResult(
task_id=task.task_id,
status=task.status,
@ -96,147 +100,116 @@ async def test_session_experimental_get_task() -> None:
poll_interval=task.poll_interval,
)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server: Server[AppContext] = Server(
"test-server",
lifespan=_make_lifespan(store, task_done_events),
on_list_tools=_handle_list_tools,
on_call_tool=_handle_call_tool_with_done_event,
)
server.experimental.enable_tasks(on_get_task=handle_get_task)
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None: ... # pragma: no branch
async def run_server(app_context: AppContext):
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
async with Client(server) as client:
# Create a task
create_result = await client.session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="test_tool",
arguments={},
task=TaskMetadata(ttl=60000),
)
),
) as server_session:
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
CreateTaskResult,
)
task_id = create_result.task.task_id
async with anyio.create_task_group() as tg:
app_context = AppContext(task_group=tg, store=store)
tg.start_soon(run_server, app_context)
# Wait for task to complete
await task_done_events[task_id].wait()
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
await client_session.initialize()
# Use session.experimental to get task status
task_status = await client.session.experimental.get_task(task_id)
# Create a task
create_result = await client_session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="test_tool",
arguments={},
task=TaskMetadata(ttl=60000),
)
),
CreateTaskResult,
)
task_id = create_result.task.task_id
# Wait for task to complete
await app_context.task_done_events[task_id].wait()
# Use session.experimental to get task status
task_status = await client_session.experimental.get_task(task_id)
assert task_status.task_id == task_id
assert task_status.status == "completed"
tg.cancel_scope.cancel()
assert task_status.task_id == task_id
assert task_status.status == "completed"
@pytest.mark.anyio
async def test_session_experimental_get_task_result() -> None:
"""Test session.experimental.get_task_result() method."""
server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment]
store = InMemoryTaskStore()
task_done_events: dict[str, Event] = {}
@server.list_tools()
async def list_tools():
return [Tool(name="test_tool", description="Test", input_schema={"type": "object"})]
async def handle_call_tool(
ctx: ServerRequestContext[AppContext], params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
return await _handle_call_tool_with_done_event(ctx, params, result_text="Task result content")
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
ctx = server.request_context
app = ctx.lifespan_context
if ctx.experimental.is_task:
task_metadata = ctx.experimental.task_metadata
assert task_metadata is not None
task = await app.store.create_task(task_metadata)
done_event = Event()
app.task_done_events[task.task_id] = done_event
async def do_work():
async with task_execution(task.task_id, app.store) as task_ctx:
await task_ctx.complete(
CallToolResult(content=[TextContent(type="text", text="Task result content")])
)
done_event.set()
app.task_group.start_soon(do_work)
return CreateTaskResult(task=task)
raise NotImplementedError
@server.experimental.get_task_result()
async def handle_get_task_result(
request: GetTaskPayloadRequest,
ctx: ServerRequestContext[AppContext], params: GetTaskPayloadRequestParams
) -> GetTaskPayloadResult:
app = server.request_context.lifespan_context
result = await app.store.get_result(request.params.task_id)
assert result is not None, f"Test setup error: result for {request.params.task_id} should exist"
app = ctx.lifespan_context
result = await app.store.get_result(params.task_id)
assert result is not None, f"Test setup error: result for {params.task_id} should exist"
assert isinstance(result, CallToolResult)
return GetTaskPayloadResult(**result.model_dump())
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server: Server[AppContext] = Server(
"test-server",
lifespan=_make_lifespan(store, task_done_events),
on_list_tools=_handle_list_tools,
on_call_tool=handle_call_tool,
)
server.experimental.enable_tasks(on_task_result=handle_get_task_result)
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None: ... # pragma: no branch
async def run_server(app_context: AppContext):
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
async with Client(server) as client:
# Create a task
create_result = await client.session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="test_tool",
arguments={},
task=TaskMetadata(ttl=60000),
)
),
) as server_session:
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
CreateTaskResult,
)
task_id = create_result.task.task_id
async with anyio.create_task_group() as tg:
app_context = AppContext(task_group=tg, store=store)
tg.start_soon(run_server, app_context)
# Wait for task to complete
await task_done_events[task_id].wait()
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
await client_session.initialize()
# Use TaskClient to get task result
task_result = await client.session.experimental.get_task_result(task_id, CallToolResult)
# Create a task
create_result = await client_session.send_request(
assert len(task_result.content) == 1
content = task_result.content[0]
assert isinstance(content, TextContent)
assert content.text == "Task result content"
async def test_session_experimental_list_tasks() -> None:
"""Test TaskClient.list_tasks() method."""
store = InMemoryTaskStore()
task_done_events: dict[str, Event] = {}
async def handle_list_tasks(
ctx: ServerRequestContext[AppContext], params: PaginatedRequestParams | None
) -> ListTasksResult:
app = ctx.lifespan_context
cursor = params.cursor if params else None
tasks_list, next_cursor = await app.store.list_tasks(cursor=cursor)
return ListTasksResult(tasks=tasks_list, next_cursor=next_cursor)
server: Server[AppContext] = Server(
"test-server",
lifespan=_make_lifespan(store, task_done_events),
on_list_tools=_handle_list_tools,
on_call_tool=_handle_call_tool_with_done_event,
)
server.experimental.enable_tasks(on_list_tasks=handle_list_tasks)
async with Client(server) as client:
# Create two tasks
for _ in range(2):
create_result = await client.session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="test_tool",
@ -246,130 +219,22 @@ async def test_session_experimental_get_task_result() -> None:
),
CreateTaskResult,
)
task_id = create_result.task.task_id
await task_done_events[create_result.task.task_id].wait()
# Wait for task to complete
await app_context.task_done_events[task_id].wait()
# Use TaskClient to list tasks
list_result = await client.session.experimental.list_tasks()
# Use TaskClient to get task result
task_result = await client_session.experimental.get_task_result(task_id, CallToolResult)
assert len(task_result.content) == 1
content = task_result.content[0]
assert isinstance(content, TextContent)
assert content.text == "Task result content"
tg.cancel_scope.cancel()
assert len(list_result.tasks) == 2
@pytest.mark.anyio
async def test_session_experimental_list_tasks() -> None:
"""Test TaskClient.list_tasks() method."""
server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment]
store = InMemoryTaskStore()
@server.list_tools()
async def list_tools():
return [Tool(name="test_tool", description="Test", input_schema={"type": "object"})]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
ctx = server.request_context
app = ctx.lifespan_context
if ctx.experimental.is_task:
task_metadata = ctx.experimental.task_metadata
assert task_metadata is not None
task = await app.store.create_task(task_metadata)
done_event = Event()
app.task_done_events[task.task_id] = done_event
async def do_work():
async with task_execution(task.task_id, app.store) as task_ctx:
await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")]))
done_event.set()
app.task_group.start_soon(do_work)
return CreateTaskResult(task=task)
raise NotImplementedError
@server.experimental.list_tasks()
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
app = server.request_context.lifespan_context
tasks_list, next_cursor = await app.store.list_tasks(cursor=request.params.cursor if request.params else None)
return ListTasksResult(tasks=tasks_list, next_cursor=next_cursor)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None: ... # pragma: no branch
async def run_server(app_context: AppContext):
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
),
) as server_session:
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
async with anyio.create_task_group() as tg:
app_context = AppContext(task_group=tg, store=store)
tg.start_soon(run_server, app_context)
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
await client_session.initialize()
# Create two tasks
for _ in range(2):
create_result = await client_session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="test_tool",
arguments={},
task=TaskMetadata(ttl=60000),
)
),
CreateTaskResult,
)
await app_context.task_done_events[create_result.task.task_id].wait()
# Use TaskClient to list tasks
list_result = await client_session.experimental.list_tasks()
assert len(list_result.tasks) == 2
tg.cancel_scope.cancel()
@pytest.mark.anyio
async def test_session_experimental_cancel_task() -> None:
"""Test TaskClient.cancel_task() method."""
server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment]
store = InMemoryTaskStore()
task_done_events: dict[str, Event] = {}
@server.list_tools()
async def list_tools():
return [Tool(name="test_tool", description="Test", input_schema={"type": "object"})]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
ctx = server.request_context
async def handle_call_tool_no_work(
ctx: ServerRequestContext[AppContext], params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
app = ctx.lifespan_context
if ctx.experimental.is_task:
task_metadata = ctx.experimental.task_metadata
@ -377,14 +242,12 @@ async def test_session_experimental_cancel_task() -> None:
task = await app.store.create_task(task_metadata)
# Don't start any work - task stays in "working" status
return CreateTaskResult(task=task)
raise NotImplementedError
@server.experimental.get_task()
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
app = server.request_context.lifespan_context
task = await app.store.get_task(request.params.task_id)
assert task is not None, f"Test setup error: task {request.params.task_id} should exist"
async def handle_get_task(ctx: ServerRequestContext[AppContext], params: GetTaskRequestParams) -> GetTaskResult:
app = ctx.lifespan_context
task = await app.store.get_task(params.task_id)
assert task is not None, f"Test setup error: task {params.task_id} should exist"
return GetTaskResult(
task_id=task.task_id,
status=task.status,
@ -395,14 +258,14 @@ async def test_session_experimental_cancel_task() -> None:
poll_interval=task.poll_interval,
)
@server.experimental.cancel_task()
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
app = server.request_context.lifespan_context
task = await app.store.get_task(request.params.task_id)
assert task is not None, f"Test setup error: task {request.params.task_id} should exist"
await app.store.update_task(request.params.task_id, status="cancelled")
# CancelTaskResult extends Task, so we need to return the updated task info
updated_task = await app.store.get_task(request.params.task_id)
async def handle_cancel_task(
ctx: ServerRequestContext[AppContext], params: CancelTaskRequestParams
) -> CancelTaskResult:
app = ctx.lifespan_context
task = await app.store.get_task(params.task_id)
assert task is not None, f"Test setup error: task {params.task_id} should exist"
await app.store.update_task(params.task_id, status="cancelled")
updated_task = await app.store.get_task(params.task_id)
assert updated_task is not None
return CancelTaskResult(
task_id=updated_task.task_id,
@ -412,63 +275,35 @@ async def test_session_experimental_cancel_task() -> None:
ttl=updated_task.ttl,
)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server: Server[AppContext] = Server(
"test-server",
lifespan=_make_lifespan(store, task_done_events),
on_list_tools=_handle_list_tools,
on_call_tool=handle_call_tool_no_work,
)
server.experimental.enable_tasks(on_get_task=handle_get_task, on_cancel_task=handle_cancel_task)
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None: ... # pragma: no branch
async def run_server(app_context: AppContext):
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
async with Client(server) as client:
# Create a task (but don't complete it)
create_result = await client.session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="test_tool",
arguments={},
task=TaskMetadata(ttl=60000),
)
),
) as server_session:
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
CreateTaskResult,
)
task_id = create_result.task.task_id
async with anyio.create_task_group() as tg:
app_context = AppContext(task_group=tg, store=store)
tg.start_soon(run_server, app_context)
# Verify task is working
status_before = await client.session.experimental.get_task(task_id)
assert status_before.status == "working"
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
await client_session.initialize()
# Cancel the task
await client.session.experimental.cancel_task(task_id)
# Create a task (but don't complete it)
create_result = await client_session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="test_tool",
arguments={},
task=TaskMetadata(ttl=60000),
)
),
CreateTaskResult,
)
task_id = create_result.task.task_id
# Verify task is working
status_before = await client_session.experimental.get_task(task_id)
assert status_before.status == "working"
# Cancel the task
await client_session.experimental.cancel_task(task_id)
# Verify task is cancelled
status_after = await client_session.experimental.get_task(task_id)
assert status_after.status == "cancelled"
tg.cancel_scope.cancel()
# Verify task is cancelled
status_after = await client.session.experimental.get_task(task_id)
assert status_after.status == "cancelled"

329
tests/experimental/tasks/server/test_integration.py
View file

@ -8,46 +8,37 @@ These tests demonstrate the full task lifecycle:
5. Client retrieves result with tasks/result
"""
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any
import anyio
import pytest
from anyio import Event
from anyio.abc import TaskGroup
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp import Client
from mcp.server import Server, ServerRequestContext
from mcp.shared.experimental.tasks.helpers import task_execution
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from mcp.types import (
TASK_REQUIRED,
CallToolRequest,
CallToolRequestParams,
CallToolResult,
ClientResult,
CreateTaskResult,
GetTaskPayloadRequest,
GetTaskPayloadRequestParams,
GetTaskPayloadResult,
GetTaskRequest,
GetTaskRequestParams,
GetTaskResult,
ListTasksRequest,
ListTasksResult,
ServerNotification,
ServerRequest,
ListToolsResult,
PaginatedRequestParams,
TaskMetadata,
TextContent,
Tool,
ToolExecution,
)
pytestmark = pytest.mark.anyio
@dataclass
class AppContext:
@ -55,77 +46,57 @@ class AppContext:
task_group: TaskGroup
store: InMemoryTaskStore
# Events to signal when tasks complete (for testing without sleeps)
task_done_events: dict[str, Event] = field(default_factory=lambda: {})
@pytest.mark.anyio
def _make_lifespan(store: InMemoryTaskStore, task_done_events: dict[str, Event]):
@asynccontextmanager
async def app_lifespan(server: Server[AppContext]) -> AsyncIterator[AppContext]:
async with anyio.create_task_group() as tg:
yield AppContext(task_group=tg, store=store, task_done_events=task_done_events)
return app_lifespan
async def test_task_lifecycle_with_task_execution() -> None:
"""Test the complete task lifecycle using the task_execution pattern.
This demonstrates the recommended way to implement task-augmented tools:
1. Create task in store
2. Spawn work using task_execution() context manager
3. Return CreateTaskResult immediately
4. Work executes in background, auto-fails on exception
"""
# Note: We bypass the normal lifespan mechanism and pass context directly to _handle_message
server: Server[AppContext, Any] = Server("test-tasks") # type: ignore[assignment]
"""Test the complete task lifecycle using the task_execution pattern."""
store = InMemoryTaskStore()
task_done_events: dict[str, Event] = {}
@server.list_tools()
async def list_tools():
return [
Tool(
name="process_data",
description="Process data asynchronously",
input_schema={
"type": "object",
"properties": {"input": {"type": "string"}},
},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
async def handle_list_tools(
ctx: ServerRequestContext[AppContext], params: PaginatedRequestParams | None
) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(
ctx: ServerRequestContext[AppContext], params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
app = ctx.lifespan_context
if name == "process_data" and ctx.experimental.is_task:
# 1. Create task in store
if params.name == "process_data" and ctx.experimental.is_task:
task_metadata = ctx.experimental.task_metadata
assert task_metadata is not None
task = await app.store.create_task(task_metadata)
# 2. Create event to signal completion (for testing)
done_event = Event()
app.task_done_events[task.task_id] = done_event
# 3. Define work function using task_execution for safety
async def do_work():
async def do_work() -> None:
async with task_execution(task.task_id, app.store) as task_ctx:
await task_ctx.update_status("Processing input...")
# Simulate work
input_value = arguments.get("input", "")
input_value = (params.arguments or {}).get("input", "")
result_text = f"Processed: {input_value.upper()}"
await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text=result_text)]))
# Signal completion
done_event.set()
# 4. Spawn work in task group (from lifespan_context)
app.task_group.start_soon(do_work)
# 5. Return CreateTaskResult immediately
return CreateTaskResult(task=task)
raise NotImplementedError
# Register task query handlers (delegate to store)
@server.experimental.get_task()
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
app = server.request_context.lifespan_context
task = await app.store.get_task(request.params.task_id)
assert task is not None, f"Test setup error: task {request.params.task_id} should exist"
async def handle_get_task(ctx: ServerRequestContext[AppContext], params: GetTaskRequestParams) -> GetTaskResult:
app = ctx.lifespan_context
task = await app.store.get_task(params.task_id)
assert task is not None, f"Test setup error: task {params.task_id} should exist"
return GetTaskResult(
task_id=task.task_id,
status=task.status,
@ -136,134 +107,91 @@ async def test_task_lifecycle_with_task_execution() -> None:
poll_interval=task.poll_interval,
)
@server.experimental.get_task_result()
async def handle_get_task_result(
request: GetTaskPayloadRequest,
ctx: ServerRequestContext[AppContext], params: GetTaskPayloadRequestParams
) -> GetTaskPayloadResult:
app = server.request_context.lifespan_context
result = await app.store.get_result(request.params.task_id)
assert result is not None, f"Test setup error: result for {request.params.task_id} should exist"
app = ctx.lifespan_context
result = await app.store.get_result(params.task_id)
assert result is not None, f"Test setup error: result for {params.task_id} should exist"
assert isinstance(result, CallToolResult)
# Return as GetTaskPayloadResult (which accepts extra fields)
return GetTaskPayloadResult(**result.model_dump())
@server.experimental.list_tasks()
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
async def handle_list_tasks(
ctx: ServerRequestContext[AppContext], params: PaginatedRequestParams | None
) -> ListTasksResult:
raise NotImplementedError
# Set up client-server communication
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server: Server[AppContext] = Server(
"test-tasks",
lifespan=_make_lifespan(store, task_done_events),
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
)
server.experimental.enable_tasks(
on_get_task=handle_get_task,
on_task_result=handle_get_task_result,
on_list_tasks=handle_list_tasks,
)
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None: ... # pragma: no cover
async def run_server(app_context: AppContext):
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
async with Client(server) as client:
# Step 1: Send task-augmented tool call
create_result = await client.session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="process_data",
arguments={"input": "hello world"},
task=TaskMetadata(ttl=60000),
),
),
) as server_session:
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
CreateTaskResult,
)
async with anyio.create_task_group() as tg:
# Create app context with task group and store
app_context = AppContext(task_group=tg, store=store)
tg.start_soon(run_server, app_context)
assert isinstance(create_result, CreateTaskResult)
assert create_result.task.status == "working"
task_id = create_result.task.task_id
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
await client_session.initialize()
# Step 2: Wait for task to complete
await task_done_events[task_id].wait()
# === Step 1: Send task-augmented tool call ===
create_result = await client_session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="process_data",
arguments={"input": "hello world"},
task=TaskMetadata(ttl=60000),
),
),
CreateTaskResult,
)
task_status = await client.session.experimental.get_task(task_id)
assert task_status.task_id == task_id
assert task_status.status == "completed"
assert isinstance(create_result, CreateTaskResult)
assert create_result.task.status == "working"
task_id = create_result.task.task_id
# Step 3: Retrieve the actual result
task_result = await client.session.experimental.get_task_result(task_id, CallToolResult)
# === Step 2: Wait for task to complete ===
await app_context.task_done_events[task_id].wait()
task_status = await client_session.send_request(
GetTaskRequest(params=GetTaskRequestParams(task_id=task_id)),
GetTaskResult,
)
assert task_status.task_id == task_id
assert task_status.status == "completed"
# === Step 3: Retrieve the actual result ===
task_result = await client_session.send_request(
GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(task_id=task_id)),
CallToolResult,
)
assert len(task_result.content) == 1
content = task_result.content[0]
assert isinstance(content, TextContent)
assert content.text == "Processed: HELLO WORLD"
tg.cancel_scope.cancel()
assert len(task_result.content) == 1
content = task_result.content[0]
assert isinstance(content, TextContent)
assert content.text == "Processed: HELLO WORLD"
@pytest.mark.anyio
async def test_task_auto_fails_on_exception() -> None:
"""Test that task_execution automatically fails the task on unhandled exception."""
# Note: We bypass the normal lifespan mechanism and pass context directly to _handle_message
server: Server[AppContext, Any] = Server("test-tasks-failure") # type: ignore[assignment]
store = InMemoryTaskStore()
task_done_events: dict[str, Event] = {}
@server.list_tools()
async def list_tools():
return [
Tool(
name="failing_task",
description="A task that fails",
input_schema={"type": "object", "properties": {}},
)
]
async def handle_list_tools(
ctx: ServerRequestContext[AppContext], params: PaginatedRequestParams | None
) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(
ctx: ServerRequestContext[AppContext], params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
app = ctx.lifespan_context
if name == "failing_task" and ctx.experimental.is_task:
if params.name == "failing_task" and ctx.experimental.is_task:
task_metadata = ctx.experimental.task_metadata
assert task_metadata is not None
task = await app.store.create_task(task_metadata)
# Create event to signal completion (for testing)
done_event = Event()
app.task_done_events[task.task_id] = done_event
async def do_failing_work():
async def do_failing_work() -> None:
async with task_execution(task.task_id, app.store) as task_ctx:
await task_ctx.update_status("About to fail...")
raise RuntimeError("Something went wrong!")
# Note: complete() is never called, but task_execution
# will automatically call fail() due to the exception
# This line is reached because task_execution suppresses the exception
done_event.set()
@ -272,11 +200,10 @@ async def test_task_auto_fails_on_exception() -> None:
raise NotImplementedError
@server.experimental.get_task()
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
app = server.request_context.lifespan_context
task = await app.store.get_task(request.params.task_id)
assert task is not None, f"Test setup error: task {request.params.task_id} should exist"
async def handle_get_task(ctx: ServerRequestContext[AppContext], params: GetTaskRequestParams) -> GetTaskResult:
app = ctx.lifespan_context
task = await app.store.get_task(params.task_id)
assert task is not None, f"Test setup error: task {params.task_id} should exist"
return GetTaskResult(
task_id=task.task_id,
status=task.status,
@ -287,64 +214,34 @@ async def test_task_auto_fails_on_exception() -> None:
poll_interval=task.poll_interval,
)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server: Server[AppContext] = Server(
"test-tasks-failure",
lifespan=_make_lifespan(store, task_done_events),
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
)
server.experimental.enable_tasks(on_get_task=handle_get_task)
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None: ... # pragma: no cover
async def run_server(app_context: AppContext):
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
async with Client(server) as client:
# Send task request
create_result = await client.session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="failing_task",
arguments={},
task=TaskMetadata(ttl=60000),
),
),
) as server_session:
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
CreateTaskResult,
)
async with anyio.create_task_group() as tg:
app_context = AppContext(task_group=tg, store=store)
tg.start_soon(run_server, app_context)
task_id = create_result.task.task_id
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
await client_session.initialize()
# Wait for task to complete (even though it fails)
await task_done_events[task_id].wait()
# Send task request
create_result = await client_session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="failing_task",
arguments={},
task=TaskMetadata(ttl=60000),
),
),
CreateTaskResult,
)
# Check that task was auto-failed
task_status = await client.session.experimental.get_task(task_id)
task_id = create_result.task.task_id
# Wait for task to complete (even though it fails)
await app_context.task_done_events[task_id].wait()
# Check that task was auto-failed
task_status = await client_session.send_request(
GetTaskRequest(params=GetTaskRequestParams(task_id=task_id)), GetTaskResult
)
assert task_status.status == "failed"
assert task_status.status_message == "Something went wrong!"
tg.cancel_scope.cancel()
assert task_status.status == "failed"
assert task_status.status_message == "Something went wrong!"

449
tests/experimental/tasks/server/test_run_task_flow.py
View file

@ -8,159 +8,102 @@ This tests the recommended user flow:
These are integration tests that verify the complete flow works end-to-end.
"""
from typing import Any
from unittest.mock import Mock
import anyio
import pytest
from anyio import Event
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp import Client
from mcp.server import Server, ServerRequestContext
from mcp.server.experimental.request_context import Experimental
from mcp.server.experimental.task_context import ServerTaskContext
from mcp.server.experimental.task_support import TaskSupport
from mcp.server.lowlevel import NotificationOptions
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
from mcp.shared.message import SessionMessage
from mcp.types import (
TASK_REQUIRED,
CallToolRequestParams,
CallToolResult,
CancelTaskRequest,
CancelTaskResult,
CreateTaskResult,
GetTaskPayloadRequest,
GetTaskPayloadResult,
GetTaskRequest,
GetTaskRequestParams,
GetTaskResult,
ListTasksRequest,
ListTasksResult,
ListToolsResult,
PaginatedRequestParams,
TextContent,
Tool,
ToolExecution,
)
pytestmark = pytest.mark.anyio
async def _handle_list_tools_simple_task(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListToolsResult:
raise NotImplementedError
@pytest.mark.anyio
async def test_run_task_basic_flow() -> None:
"""Test the basic run_task flow without elicitation.
1. enable_tasks() sets up handlers
2. Client calls tool with task field
3. run_task() spawns work, returns CreateTaskResult
4. Work completes in background
5. Client polls and sees completed status
"""
server = Server("test-run-task")
# One-line setup
server.experimental.enable_tasks()
# Track when work completes and capture received meta
"""Test the basic run_task flow without elicitation."""
work_completed = Event()
received_meta: list[str | None] = [None]
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="simple_task",
description="A simple task",
input_schema={"type": "object", "properties": {"input": {"type": "string"}}},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(
ctx: ServerRequestContext, params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
ctx.experimental.validate_task_mode(TASK_REQUIRED)
# Capture the meta from the request (if present)
if ctx.meta is not None: # pragma: no branch
received_meta[0] = ctx.meta.get("custom_field")
async def work(task: ServerTaskContext) -> CallToolResult:
await task.update_status("Working...")
input_val = arguments.get("input", "default")
input_val = (params.arguments or {}).get("input", "default")
result = CallToolResult(content=[TextContent(type="text", text=f"Processed: {input_val}")])
work_completed.set()
return result
return await ctx.experimental.run_task(work)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server = Server(
"test-run-task",
on_list_tools=_handle_list_tools_simple_task,
on_call_tool=handle_call_tool,
)
server.experimental.enable_tasks()
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
async with Client(server) as client:
result = await client.session.experimental.call_tool_as_task(
"simple_task",
{"input": "hello"},
meta={"custom_field": "test_value"},
)
async def run_client() -> None:
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
# Initialize
await client_session.initialize()
task_id = result.task.task_id
assert result.task.status == "working"
# Call tool as task (with meta to test that code path)
result = await client_session.experimental.call_tool_as_task(
"simple_task",
{"input": "hello"},
meta={"custom_field": "test_value"},
)
with anyio.fail_after(5):
await work_completed.wait()
# Should get CreateTaskResult
task_id = result.task.task_id
assert result.task.status == "working"
with anyio.fail_after(5):
while True:
task_status = await client.session.experimental.get_task(task_id)
if task_status.status == "completed": # pragma: no branch
break
# Wait for work to complete
with anyio.fail_after(5):
await work_completed.wait()
# Poll until task status is completed
with anyio.fail_after(5):
while True:
task_status = await client_session.experimental.get_task(task_id)
if task_status.status == "completed": # pragma: no branch
break
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
# Verify the meta was passed through correctly
assert received_meta[0] == "test_value"
@pytest.mark.anyio
async def test_run_task_auto_fails_on_exception() -> None:
"""Test that run_task automatically fails the task when work raises."""
server = Server("test-run-task-fail")
server.experimental.enable_tasks()
work_failed = Event()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="failing_task",
description="A task that fails",
input_schema={"type": "object"},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(
ctx: ServerRequestContext, params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
@ -169,42 +112,29 @@ async def test_run_task_auto_fails_on_exception() -> None:
return await ctx.experimental.run_task(work)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server = Server(
"test-run-task-fail",
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
)
server.experimental.enable_tasks()
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(),
)
async with Client(server) as client:
result = await client.session.experimental.call_tool_as_task("failing_task", {})
task_id = result.task.task_id
async def run_client() -> None:
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
await client_session.initialize()
with anyio.fail_after(5):
await work_failed.wait()
result = await client_session.experimental.call_tool_as_task("failing_task", {})
task_id = result.task.task_id
with anyio.fail_after(5):
while True:
task_status = await client.session.experimental.get_task(task_id)
if task_status.status == "failed": # pragma: no branch
break
# Wait for work to fail
with anyio.fail_after(5):
await work_failed.wait()
# Poll until task status is failed
with anyio.fail_after(5):
while True:
task_status = await client_session.experimental.get_task(task_id)
if task_status.status == "failed": # pragma: no branch
break
assert "Something went wrong" in (task_status.status_message or "")
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
assert "Something went wrong" in (task_status.status_message or "")
@pytest.mark.anyio
async def test_enable_tasks_auto_registers_handlers() -> None:
"""Test that enable_tasks() auto-registers get_task, list_tasks, cancel_task handlers."""
server = Server("test-enable-tasks")
@ -221,63 +151,41 @@ async def test_enable_tasks_auto_registers_handlers() -> None:
assert caps_after.tasks is not None
assert caps_after.tasks.list is not None
assert caps_after.tasks.cancel is not None
# Verify nested call capability is present
assert caps_after.tasks.requests is not None
assert caps_after.tasks.requests.tools is not None
assert caps_after.tasks.requests.tools.call is not None
@pytest.mark.anyio
async def test_enable_tasks_with_custom_store_and_queue() -> None:
"""Test that enable_tasks() uses provided store and queue instead of defaults."""
server = Server("test-custom-store-queue")
# Create custom store and queue
custom_store = InMemoryTaskStore()
custom_queue = InMemoryTaskMessageQueue()
# Enable tasks with custom implementations
task_support = server.experimental.enable_tasks(store=custom_store, queue=custom_queue)
# Verify our custom implementations are used
assert task_support.store is custom_store
assert task_support.queue is custom_queue
@pytest.mark.anyio
async def test_enable_tasks_skips_default_handlers_when_custom_registered() -> None:
"""Test that enable_tasks() doesn't override already-registered handlers."""
server = Server("test-custom-handlers")
# Register custom handlers BEFORE enable_tasks (never called, just for registration)
@server.experimental.get_task()
async def custom_get_task(req: GetTaskRequest) -> GetTaskResult:
# Register custom handlers via enable_tasks kwargs
async def custom_get_task(ctx: ServerRequestContext, params: GetTaskRequestParams) -> GetTaskResult:
raise NotImplementedError
@server.experimental.get_task_result()
async def custom_get_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
raise NotImplementedError
server.experimental.enable_tasks(on_get_task=custom_get_task)
@server.experimental.list_tasks()
async def custom_list_tasks(req: ListTasksRequest) -> ListTasksResult:
raise NotImplementedError
@server.experimental.cancel_task()
async def custom_cancel_task(req: CancelTaskRequest) -> CancelTaskResult:
raise NotImplementedError
# Now enable tasks - should NOT override our custom handlers
server.experimental.enable_tasks()
# Verify our custom handlers are still registered (not replaced by defaults)
# The handlers dict should contain our custom handlers
assert GetTaskRequest in server.request_handlers
assert GetTaskPayloadRequest in server.request_handlers
assert ListTasksRequest in server.request_handlers
assert CancelTaskRequest in server.request_handlers
# Verify handler is registered
assert server._has_handler("tasks/get")
assert server._has_handler("tasks/list")
assert server._has_handler("tasks/cancel")
assert server._has_handler("tasks/result")
@pytest.mark.anyio
async def test_run_task_without_enable_tasks_raises() -> None:
"""Test that run_task raises when enable_tasks() wasn't called."""
experimental = Experimental(
@ -294,7 +202,6 @@ async def test_run_task_without_enable_tasks_raises() -> None:
await experimental.run_task(work)
@pytest.mark.anyio
async def test_task_support_task_group_before_run_raises() -> None:
"""Test that accessing task_group before run() raises RuntimeError."""
task_support = TaskSupport.in_memory()
@ -303,7 +210,6 @@ async def test_task_support_task_group_before_run_raises() -> None:
_ = task_support.task_group
@pytest.mark.anyio
async def test_run_task_without_session_raises() -> None:
"""Test that run_task raises when session is not available."""
task_support = TaskSupport.in_memory()
@ -322,7 +228,6 @@ async def test_run_task_without_session_raises() -> None:
await experimental.run_task(work)
@pytest.mark.anyio
async def test_run_task_without_task_metadata_raises() -> None:
"""Test that run_task raises when request is not task-augmented."""
task_support = TaskSupport.in_memory()
@ -342,29 +247,17 @@ async def test_run_task_without_task_metadata_raises() -> None:
await experimental.run_task(work)
@pytest.mark.anyio
async def test_run_task_with_model_immediate_response() -> None:
"""Test that run_task includes model_immediate_response in CreateTaskResult._meta."""
server = Server("test-run-task-immediate")
server.experimental.enable_tasks()
work_completed = Event()
immediate_response_text = "Processing your request..."
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="task_with_immediate",
description="A task with immediate response",
input_schema={"type": "object"},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(
ctx: ServerRequestContext, params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
@ -373,164 +266,102 @@ async def test_run_task_with_model_immediate_response() -> None:
return await ctx.experimental.run_task(work, model_immediate_response=immediate_response_text)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(),
)
async def run_client() -> None:
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
await client_session.initialize()
result = await client_session.experimental.call_tool_as_task("task_with_immediate", {})
# Verify the immediate response is in _meta
assert result.meta is not None
assert "io.modelcontextprotocol/model-immediate-response" in result.meta
assert result.meta["io.modelcontextprotocol/model-immediate-response"] == immediate_response_text
with anyio.fail_after(5):
await work_completed.wait()
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
@pytest.mark.anyio
async def test_run_task_doesnt_complete_if_already_terminal() -> None:
"""Test that run_task doesn't auto-complete if work manually completed the task."""
server = Server("test-already-complete")
server = Server(
"test-run-task-immediate",
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
)
server.experimental.enable_tasks()
async with Client(server) as client:
result = await client.session.experimental.call_tool_as_task("task_with_immediate", {})
assert result.meta is not None
assert "io.modelcontextprotocol/model-immediate-response" in result.meta
assert result.meta["io.modelcontextprotocol/model-immediate-response"] == immediate_response_text
with anyio.fail_after(5):
await work_completed.wait()
async def test_run_task_doesnt_complete_if_already_terminal() -> None:
"""Test that run_task doesn't auto-complete if work manually completed the task."""
work_completed = Event()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="manual_complete_task",
description="A task that manually completes",
input_schema={"type": "object"},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(
ctx: ServerRequestContext, params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
# Manually complete the task before returning
manual_result = CallToolResult(content=[TextContent(type="text", text="Manually completed")])
await task.complete(manual_result, notify=False)
work_completed.set()
# Return a different result - but it should be ignored since task is already terminal
return CallToolResult(content=[TextContent(type="text", text="This should be ignored")])
return await ctx.experimental.run_task(work)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(),
)
async def run_client() -> None:
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
await client_session.initialize()
result = await client_session.experimental.call_tool_as_task("manual_complete_task", {})
task_id = result.task.task_id
with anyio.fail_after(5):
await work_completed.wait()
# Poll until task status is completed
with anyio.fail_after(5):
while True:
status = await client_session.experimental.get_task(task_id)
if status.status == "completed": # pragma: no branch
break
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
@pytest.mark.anyio
async def test_run_task_doesnt_fail_if_already_terminal() -> None:
"""Test that run_task doesn't auto-fail if work manually failed/cancelled the task."""
server = Server("test-already-failed")
server = Server(
"test-already-complete",
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
)
server.experimental.enable_tasks()
async with Client(server) as client:
result = await client.session.experimental.call_tool_as_task("manual_complete_task", {})
task_id = result.task.task_id
with anyio.fail_after(5):
await work_completed.wait()
with anyio.fail_after(5):
while True:
status = await client.session.experimental.get_task(task_id)
if status.status == "completed": # pragma: no branch
break
async def test_run_task_doesnt_fail_if_already_terminal() -> None:
"""Test that run_task doesn't auto-fail if work manually failed/cancelled the task."""
work_completed = Event()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="manual_cancel_task",
description="A task that manually cancels then raises",
input_schema={"type": "object"},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(
ctx: ServerRequestContext, params: CallToolRequestParams
) -> CallToolResult | CreateTaskResult:
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
# Manually fail the task first
await task.fail("Manually failed", notify=False)
work_completed.set()
# Then raise - but the auto-fail should be skipped since task is already terminal
raise RuntimeError("This error should not change status")
return await ctx.experimental.run_task(work)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server = Server(
"test-already-failed",
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
)
server.experimental.enable_tasks()
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(),
)
async with Client(server) as client:
result = await client.session.experimental.call_tool_as_task("manual_cancel_task", {})
task_id = result.task.task_id
async def run_client() -> None:
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
await client_session.initialize()
with anyio.fail_after(5):
await work_completed.wait()
result = await client_session.experimental.call_tool_as_task("manual_cancel_task", {})
task_id = result.task.task_id
with anyio.fail_after(5):
while True:
status = await client.session.experimental.get_task(task_id)
if status.status == "failed": # pragma: no branch
break
with anyio.fail_after(5):
await work_completed.wait()
# Poll until task status is failed
with anyio.fail_after(5):
while True:
status = await client_session.experimental.get_task(task_id)
if status.status == "failed": # pragma: no branch
break
# Task should still be failed (from manual fail, not auto-fail from exception)
assert status.status_message == "Manually failed" # Not "This error should not change status"
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
assert status.status_message == "Manually failed"

462
tests/experimental/tasks/server/test_server.py
View file

@ -6,8 +6,9 @@ from typing import Any
import anyio
import pytest
from mcp import Client
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
@ -23,7 +24,6 @@ from mcp.types import (
CallToolRequest,
CallToolRequestParams,
CallToolResult,
CancelTaskRequest,
CancelTaskRequestParams,
CancelTaskResult,
ClientResult,
@ -31,21 +31,18 @@ from mcp.types import (
GetTaskPayloadRequest,
GetTaskPayloadRequestParams,
GetTaskPayloadResult,
GetTaskRequest,
GetTaskRequestParams,
GetTaskResult,
JSONRPCError,
JSONRPCNotification,
JSONRPCResponse,
ListTasksRequest,
ListTasksResult,
ListToolsRequest,
ListToolsResult,
PaginatedRequestParams,
SamplingMessage,
ServerCapabilities,
ServerNotification,
ServerRequest,
ServerResult,
Task,
TaskMetadata,
TextContent,
@ -53,57 +50,37 @@ from mcp.types import (
ToolExecution,
)
pytestmark = pytest.mark.anyio
@pytest.mark.anyio
async def test_list_tasks_handler() -> None:
"""Test that experimental list_tasks handler works."""
server = Server("test")
"""Test that experimental list_tasks handler works via Client."""
now = datetime.now(timezone.utc)
test_tasks = [
Task(
task_id="task-1",
status="working",
created_at=now,
last_updated_at=now,
ttl=60000,
poll_interval=1000,
),
Task(
task_id="task-2",
status="completed",
created_at=now,
last_updated_at=now,
ttl=60000,
poll_interval=1000,
),
Task(task_id="task-1", status="working", created_at=now, last_updated_at=now, ttl=60000, poll_interval=1000),
Task(task_id="task-2", status="completed", created_at=now, last_updated_at=now, ttl=60000, poll_interval=1000),
]
@server.experimental.list_tasks()
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
async def handle_list_tasks(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListTasksResult:
return ListTasksResult(tasks=test_tasks)
handler = server.request_handlers[ListTasksRequest]
request = ListTasksRequest(method="tasks/list")
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, ListTasksResult)
assert len(result.tasks) == 2
assert result.tasks[0].task_id == "task-1"
assert result.tasks[1].task_id == "task-2"
@pytest.mark.anyio
async def test_get_task_handler() -> None:
"""Test that experimental get_task handler works."""
server = Server("test")
server.experimental.enable_tasks(on_list_tasks=handle_list_tasks)
@server.experimental.get_task()
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
async with Client(server) as client:
result = await client.session.experimental.list_tasks()
assert len(result.tasks) == 2
assert result.tasks[0].task_id == "task-1"
assert result.tasks[1].task_id == "task-2"
async def test_get_task_handler() -> None:
"""Test that experimental get_task handler works via Client."""
async def handle_get_task(ctx: ServerRequestContext, params: GetTaskRequestParams) -> GetTaskResult:
now = datetime.now(timezone.utc)
return GetTaskResult(
task_id=request.params.task_id,
task_id=params.task_id,
status="working",
created_at=now,
last_updated_at=now,
@ -111,85 +88,69 @@ async def test_get_task_handler() -> None:
poll_interval=1000,
)
handler = server.request_handlers[GetTaskRequest]
request = GetTaskRequest(
method="tasks/get",
params=GetTaskRequestParams(task_id="test-task-123"),
)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, GetTaskResult)
assert result.task_id == "test-task-123"
assert result.status == "working"
@pytest.mark.anyio
async def test_get_task_result_handler() -> None:
"""Test that experimental get_task_result handler works."""
server = Server("test")
server.experimental.enable_tasks(on_get_task=handle_get_task)
@server.experimental.get_task_result()
async def handle_get_task_result(request: GetTaskPayloadRequest) -> GetTaskPayloadResult:
async with Client(server) as client:
result = await client.session.experimental.get_task("test-task-123")
assert result.task_id == "test-task-123"
assert result.status == "working"
async def test_get_task_result_handler() -> None:
"""Test that experimental get_task_result handler works via Client."""
async def handle_get_task_result(
ctx: ServerRequestContext, params: GetTaskPayloadRequestParams
) -> GetTaskPayloadResult:
return GetTaskPayloadResult()
handler = server.request_handlers[GetTaskPayloadRequest]
request = GetTaskPayloadRequest(
method="tasks/result",
params=GetTaskPayloadRequestParams(task_id="test-task-123"),
)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, GetTaskPayloadResult)
@pytest.mark.anyio
async def test_cancel_task_handler() -> None:
"""Test that experimental cancel_task handler works."""
server = Server("test")
server.experimental.enable_tasks(on_task_result=handle_get_task_result)
@server.experimental.cancel_task()
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
async with Client(server) as client:
result = await client.session.send_request(
GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(task_id="test-task-123")),
GetTaskPayloadResult,
)
assert isinstance(result, GetTaskPayloadResult)
async def test_cancel_task_handler() -> None:
"""Test that experimental cancel_task handler works via Client."""
async def handle_cancel_task(ctx: ServerRequestContext, params: CancelTaskRequestParams) -> CancelTaskResult:
now = datetime.now(timezone.utc)
return CancelTaskResult(
task_id=request.params.task_id,
task_id=params.task_id,
status="cancelled",
created_at=now,
last_updated_at=now,
ttl=60000,
)
handler = server.request_handlers[CancelTaskRequest]
request = CancelTaskRequest(
method="tasks/cancel",
params=CancelTaskRequestParams(task_id="test-task-123"),
)
result = await handler(request)
server = Server("test")
server.experimental.enable_tasks(on_cancel_task=handle_cancel_task)
assert isinstance(result, ServerResult)
assert isinstance(result, CancelTaskResult)
assert result.task_id == "test-task-123"
assert result.status == "cancelled"
async with Client(server) as client:
result = await client.session.experimental.cancel_task("test-task-123")
assert result.task_id == "test-task-123"
assert result.status == "cancelled"
@pytest.mark.anyio
async def test_server_capabilities_include_tasks() -> None:
"""Test that server capabilities include tasks when handlers are registered."""
server = Server("test")
@server.experimental.list_tasks()
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
async def noop_list_tasks(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListTasksResult:
raise NotImplementedError
@server.experimental.cancel_task()
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
async def noop_cancel_task(ctx: ServerRequestContext, params: CancelTaskRequestParams) -> CancelTaskResult:
raise NotImplementedError
capabilities = server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
)
server.experimental.enable_tasks(on_list_tasks=noop_list_tasks, on_cancel_task=noop_cancel_task)
capabilities = server.get_capabilities(notification_options=NotificationOptions(), experimental_capabilities={})
assert capabilities.tasks is not None
assert capabilities.tasks.list is not None
@ -198,259 +159,164 @@ async def test_server_capabilities_include_tasks() -> None:
assert capabilities.tasks.requests.tools is not None
@pytest.mark.anyio
async def test_server_capabilities_partial_tasks() -> None:
@pytest.mark.skip(
reason="TODO(maxisbey): enable_tasks registers default handlers for all task methods, "
"so partial capabilities aren't possible yet. Low-level API should support "
"selectively enabling/disabling task capabilities."
)
async def test_server_capabilities_partial_tasks() -> None: # pragma: no cover
"""Test capabilities with only some task handlers registered."""
server = Server("test")
@server.experimental.list_tasks()
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
async def noop_list_tasks(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListTasksResult:
raise NotImplementedError
# Only list_tasks registered, not cancel_task
server.experimental.enable_tasks(on_list_tasks=noop_list_tasks)
capabilities = server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
)
capabilities = server.get_capabilities(notification_options=NotificationOptions(), experimental_capabilities={})
assert capabilities.tasks is not None
assert capabilities.tasks.list is not None
assert capabilities.tasks.cancel is None # Not registered
@pytest.mark.anyio
async def test_tool_with_task_execution_metadata() -> None:
"""Test that tools can declare task execution mode."""
server = Server("test")
@server.list_tools()
async def list_tools():
return [
Tool(
name="quick_tool",
description="Fast tool",
input_schema={"type": "object", "properties": {}},
execution=ToolExecution(task_support=TASK_FORBIDDEN),
),
Tool(
name="long_tool",
description="Long running tool",
input_schema={"type": "object", "properties": {}},
execution=ToolExecution(task_support=TASK_REQUIRED),
),
Tool(
name="flexible_tool",
description="Can be either",
input_schema={"type": "object", "properties": {}},
execution=ToolExecution(task_support=TASK_OPTIONAL),
),
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(
tools=[
Tool(
name="quick_tool",
description="Fast tool",
input_schema={"type": "object", "properties": {}},
execution=ToolExecution(task_support=TASK_FORBIDDEN),
),
Tool(
name="long_tool",
description="Long running tool",
input_schema={"type": "object", "properties": {}},
execution=ToolExecution(task_support=TASK_REQUIRED),
),
Tool(
name="flexible_tool",
description="Can be either",
input_schema={"type": "object", "properties": {}},
execution=ToolExecution(task_support=TASK_OPTIONAL),
),
]
)
tools_handler = server.request_handlers[ListToolsRequest]
request = ListToolsRequest(method="tools/list")
result = await tools_handler(request)
server = Server("test", on_list_tools=handle_list_tools)
assert isinstance(result, ServerResult)
assert isinstance(result, ListToolsResult)
tools = result.tools
async with Client(server) as client:
result = await client.list_tools()
tools = result.tools
assert tools[0].execution is not None
assert tools[0].execution.task_support == TASK_FORBIDDEN
assert tools[1].execution is not None
assert tools[1].execution.task_support == TASK_REQUIRED
assert tools[2].execution is not None
assert tools[2].execution.task_support == TASK_OPTIONAL
assert tools[0].execution is not None
assert tools[0].execution.task_support == TASK_FORBIDDEN
assert tools[1].execution is not None
assert tools[1].execution.task_support == TASK_REQUIRED
assert tools[2].execution is not None
assert tools[2].execution.task_support == TASK_OPTIONAL
@pytest.mark.anyio
async def test_task_metadata_in_call_tool_request() -> None:
"""Test that task metadata is accessible via RequestContext when calling a tool."""
server = Server("test")
"""Test that task metadata is accessible via ctx when calling a tool."""
captured_task_metadata: TaskMetadata | None = None
@server.list_tools()
async def list_tools():
return [
Tool(
name="long_task",
description="A long running task",
input_schema={"type": "object", "properties": {}},
execution=ToolExecution(task_support="optional"),
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
nonlocal captured_task_metadata
ctx = server.request_context
captured_task_metadata = ctx.experimental.task_metadata
return [TextContent(type="text", text="done")]
return CallToolResult(content=[TextContent(type="text", text="done")])
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server = Server("test", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None: ... # pragma: no branch
async def run_server():
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
async with Client(server) as client:
# Call tool with task metadata
await client.session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="long_task",
arguments={},
task=TaskMetadata(ttl=60000),
),
),
) as server_session:
async with anyio.create_task_group() as tg:
async def handle_messages():
async for message in server_session.incoming_messages: # pragma: no branch
await server._handle_message(message, server_session, {}, False)
tg.start_soon(handle_messages)
await anyio.sleep_forever()
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
await client_session.initialize()
# Call tool with task metadata
await client_session.send_request(
CallToolRequest(
params=CallToolRequestParams(
name="long_task",
arguments={},
task=TaskMetadata(ttl=60000),
),
),
CallToolResult,
)
tg.cancel_scope.cancel()
CallToolResult,
)
assert captured_task_metadata is not None
assert captured_task_metadata.ttl == 60000
@pytest.mark.anyio
async def test_task_metadata_is_task_property() -> None:
"""Test that RequestContext.experimental.is_task works correctly."""
server = Server("test")
"""Test that ctx.experimental.is_task works correctly."""
is_task_values: list[bool] = []
@server.list_tools()
async def list_tools():
return [
Tool(
name="test_tool",
description="Test tool",
input_schema={"type": "object", "properties": {}},
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
ctx = server.request_context
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
is_task_values.append(ctx.experimental.is_task)
return [TextContent(type="text", text="done")]
return CallToolResult(content=[TextContent(type="text", text="done")])
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
server = Server("test", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None: ... # pragma: no branch
async with Client(server) as client:
# Call without task metadata
await client.session.send_request(
CallToolRequest(params=CallToolRequestParams(name="test_tool", arguments={})),
CallToolResult,
)
async def run_server():
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
# Call with task metadata
await client.session.send_request(
CallToolRequest(
params=CallToolRequestParams(name="test_tool", arguments={}, task=TaskMetadata(ttl=60000)),
),
) as server_session:
async with anyio.create_task_group() as tg:
async def handle_messages():
async for message in server_session.incoming_messages: # pragma: no branch
await server._handle_message(message, server_session, {}, False)
tg.start_soon(handle_messages)
await anyio.sleep_forever()
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
await client_session.initialize()
# Call without task metadata
await client_session.send_request(
CallToolRequest(params=CallToolRequestParams(name="test_tool", arguments={})),
CallToolResult,
)
# Call with task metadata
await client_session.send_request(
CallToolRequest(
params=CallToolRequestParams(name="test_tool", arguments={}, task=TaskMetadata(ttl=60000)),
),
CallToolResult,
)
tg.cancel_scope.cancel()
CallToolResult,
)
assert len(is_task_values) == 2
assert is_task_values[0] is False # First call without task
assert is_task_values[1] is True # Second call with task
@pytest.mark.anyio
async def test_update_capabilities_no_handlers() -> None:
"""Test that update_capabilities returns early when no task handlers are registered."""
server = Server("test-no-handlers")
# Access experimental to initialize it, but don't register any task handlers
_ = server.experimental
caps = server.get_capabilities(NotificationOptions(), {})
# Without any task handlers registered, tasks capability should be None
assert caps.tasks is None
@pytest.mark.anyio
async def test_default_task_handlers_via_enable_tasks() -> None:
"""Test that enable_tasks() auto-registers working default handlers.
async def test_update_capabilities_partial_handlers() -> None:
"""Test that update_capabilities skips list/cancel when only tasks/get is registered."""
server = Server("test-partial")
# Access .experimental to create the ExperimentalHandlers instance
exp = server.experimental
# Second access returns the same cached instance
assert server.experimental is exp
This exercises the default handlers in lowlevel/experimental.py:
- _default_get_task (task not found)
- _default_get_task_result
- _default_list_tasks
- _default_cancel_task
"""
async def noop_get(ctx: ServerRequestContext, params: GetTaskRequestParams) -> GetTaskResult:
raise NotImplementedError
server._add_request_handler("tasks/get", noop_get)
caps = server.get_capabilities(NotificationOptions(), {})
assert caps.tasks is not None
assert caps.tasks.list is None
assert caps.tasks.cancel is None
async def test_default_task_handlers_via_enable_tasks() -> None:
"""Test that enable_tasks() auto-registers working default handlers."""
server = Server("test-default-handlers")
# Enable tasks with default handlers (no custom handlers registered)
task_support = server.experimental.enable_tasks()
store = task_support.store
@ -493,24 +359,18 @@ async def test_default_task_handlers_via_enable_tasks() -> None:
task = await store.create_task(TaskMetadata(ttl=60000))
# Test list_tasks (default handler)
list_result = await client_session.send_request(ListTasksRequest(), ListTasksResult)
list_result = await client_session.experimental.list_tasks()
assert len(list_result.tasks) == 1
assert list_result.tasks[0].task_id == task.task_id
# Test get_task (default handler - found)
get_result = await client_session.send_request(
GetTaskRequest(params=GetTaskRequestParams(task_id=task.task_id)),
GetTaskResult,
)
get_result = await client_session.experimental.get_task(task.task_id)
assert get_result.task_id == task.task_id
assert get_result.status == "working"
# Test get_task (default handler - not found path)
with pytest.raises(MCPError, match="not found"):
await client_session.send_request(
GetTaskRequest(params=GetTaskRequestParams(task_id="nonexistent-task")),
GetTaskResult,
)
await client_session.experimental.get_task("nonexistent-task")
# Create a completed task to test get_task_result
completed_task = await store.create_task(TaskMetadata(ttl=60000))
@ -529,9 +389,7 @@ async def test_default_task_handlers_via_enable_tasks() -> None:
assert "io.modelcontextprotocol/related-task" in payload_result.meta
# Test cancel_task (default handler)
cancel_result = await client_session.send_request(
CancelTaskRequest(params=CancelTaskRequestParams(task_id=task.task_id)), CancelTaskResult
)
cancel_result = await client_session.experimental.cancel_task(task.task_id)
assert cancel_result.task_id == task.task_id
assert cancel_result.status == "cancelled"

149
tests/experimental/tasks/test_elicitation_scenarios.py
View file

@ -17,7 +17,7 @@ from anyio import Event
from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.experimental.task_context import ServerTaskContext
from mcp.server.lowlevel import NotificationOptions
from mcp.shared._context import RequestContext
@ -26,6 +26,7 @@ from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.message import SessionMessage
from mcp.types import (
TASK_REQUIRED,
CallToolRequestParams,
CallToolResult,
CreateMessageRequestParams,
CreateMessageResult,
@ -35,11 +36,12 @@ from mcp.types import (
ErrorData,
GetTaskPayloadResult,
GetTaskResult,
ListToolsResult,
PaginatedRequestParams,
SamplingMessage,
TaskMetadata,
TextContent,
Tool,
ToolExecution,
)
@ -181,24 +183,21 @@ async def test_scenario1_normal_tool_normal_elicitation() -> None:
Server calls session.elicit() directly, client responds immediately.
"""
server = Server("test-scenario1")
elicit_received = Event()
tool_result: list[str] = []
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="confirm_action",
description="Confirm an action",
input_schema={"type": "object"},
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult:
ctx = server.request_context
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(
tools=[
Tool(
name="confirm_action",
description="Confirm an action",
input_schema={"type": "object"},
)
]
)
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
# Normal elicitation - expects immediate response
result = await ctx.session.elicit(
message="Please confirm the action",
@ -209,6 +208,8 @@ async def test_scenario1_normal_tool_normal_elicitation() -> None:
tool_result.append("confirmed" if confirmed else "cancelled")
return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")])
server = Server("test-scenario1", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
# Elicitation callback for client
async def elicitation_callback(
context: RequestContext[ClientSession],
@ -262,27 +263,24 @@ async def test_scenario2_normal_tool_task_augmented_elicitation() -> None:
Server calls session.experimental.elicit_as_task(), client creates a task
for the elicitation and returns CreateTaskResult. Server polls client.
"""
server = Server("test-scenario2")
elicit_received = Event()
tool_result: list[str] = []
# Client-side task store for handling task-augmented elicitation
client_task_store = InMemoryTaskStore()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="confirm_action",
description="Confirm an action",
input_schema={"type": "object"},
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult:
ctx = server.request_context
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(
tools=[
Tool(
name="confirm_action",
description="Confirm an action",
input_schema={"type": "object"},
)
]
)
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
# Task-augmented elicitation - server polls client
result = await ctx.session.experimental.elicit_as_task(
message="Please confirm the action",
@ -294,6 +292,7 @@ async def test_scenario2_normal_tool_task_augmented_elicitation() -> None:
tool_result.append("confirmed" if confirmed else "cancelled")
return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")])
server = Server("test-scenario2", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
task_handlers = create_client_task_handlers(client_task_store, elicit_received)
# Set up streams
@ -342,26 +341,13 @@ async def test_scenario3_task_augmented_tool_normal_elicitation() -> None:
Client calls tool as task. Inside the task, server uses task.elicit()
which queues the request and delivers via tasks/result.
"""
server = Server("test-scenario3")
server.experimental.enable_tasks()
elicit_received = Event()
work_completed = Event()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="confirm_action",
description="Confirm an action",
input_schema={"type": "object"},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CreateTaskResult:
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
@ -377,6 +363,9 @@ async def test_scenario3_task_augmented_tool_normal_elicitation() -> None:
return await ctx.experimental.run_task(work)
server = Server("test-scenario3", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
server.experimental.enable_tasks()
# Elicitation callback for client
async def elicitation_callback(
context: RequestContext[ClientSession],
@ -452,29 +441,16 @@ async def test_scenario4_task_augmented_tool_task_augmented_elicitation() -> Non
5. Server gets the ElicitResult and completes the tool task
6. Client's tasks/result returns with the CallToolResult
"""
server = Server("test-scenario4")
server.experimental.enable_tasks()
elicit_received = Event()
work_completed = Event()
# Client-side task store for handling task-augmented elicitation
client_task_store = InMemoryTaskStore()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="confirm_action",
description="Confirm an action",
input_schema={"type": "object"},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CreateTaskResult:
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
@ -491,6 +467,8 @@ async def test_scenario4_task_augmented_tool_task_augmented_elicitation() -> Non
return await ctx.experimental.run_task(work)
server = Server("test-scenario4", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
server.experimental.enable_tasks()
task_handlers = create_client_task_handlers(client_task_store, elicit_received)
# Set up streams
@ -553,27 +531,24 @@ async def test_scenario2_sampling_normal_tool_task_augmented_sampling() -> None:
Server calls session.experimental.create_message_as_task(), client creates
a task for the sampling and returns CreateTaskResult. Server polls client.
"""
server = Server("test-scenario2-sampling")
sampling_received = Event()
tool_result: list[str] = []
# Client-side task store for handling task-augmented sampling
client_task_store = InMemoryTaskStore()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="generate_text",
description="Generate text using sampling",
input_schema={"type": "object"},
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult:
ctx = server.request_context
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(
tools=[
Tool(
name="generate_text",
description="Generate text using sampling",
input_schema={"type": "object"},
)
]
)
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
# Task-augmented sampling - server polls client
result = await ctx.session.experimental.create_message_as_task(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
@ -587,6 +562,7 @@ async def test_scenario2_sampling_normal_tool_task_augmented_sampling() -> None:
tool_result.append(response_text)
return CallToolResult(content=[TextContent(type="text", text=response_text)])
server = Server("test-scenario2-sampling", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
task_handlers = create_sampling_task_handlers(client_task_store, sampling_received)
# Set up streams
@ -636,29 +612,16 @@ async def test_scenario4_sampling_task_augmented_tool_task_augmented_sampling()
which sends task-augmented sampling. Client creates its own task for the sampling,
and server polls the client.
"""
server = Server("test-scenario4-sampling")
server.experimental.enable_tasks()
sampling_received = Event()
work_completed = Event()
# Client-side task store for handling task-augmented sampling
client_task_store = InMemoryTaskStore()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="generate_text",
description="Generate text using sampling",
input_schema={"type": "object"},
execution=ToolExecution(task_support=TASK_REQUIRED),
)
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
raise NotImplementedError
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult:
ctx = server.request_context
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CreateTaskResult:
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
@ -677,6 +640,8 @@ async def test_scenario4_sampling_task_augmented_tool_task_augmented_sampling()
return await ctx.experimental.run_task(work)
server = Server("test-scenario4-sampling", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
server.experimental.enable_tasks()
task_handlers = create_sampling_task_handlers(client_task_store, sampling_received)
# Set up streams

78
tests/experimental/tasks/test_spec_compliance.py
View file

@ -10,17 +10,17 @@ from datetime import datetime, timezone
import pytest
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.lowlevel import NotificationOptions
from mcp.shared.experimental.tasks.helpers import MODEL_IMMEDIATE_RESPONSE_KEY
from mcp.types import (
CancelTaskRequest,
CancelTaskRequestParams,
CancelTaskResult,
CreateTaskResult,
GetTaskRequest,
GetTaskRequestParams,
GetTaskResult,
ListTasksRequest,
ListTasksResult,
PaginatedRequestParams,
ServerCapabilities,
Task,
)
@ -44,13 +44,22 @@ def test_server_without_task_handlers_has_no_tasks_capability() -> None:
assert caps.tasks is None
async def _noop_get_task(ctx: ServerRequestContext, params: GetTaskRequestParams) -> GetTaskResult:
raise NotImplementedError
async def _noop_list_tasks(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListTasksResult:
raise NotImplementedError
async def _noop_cancel_task(ctx: ServerRequestContext, params: CancelTaskRequestParams) -> CancelTaskResult:
raise NotImplementedError
def test_server_with_list_tasks_handler_declares_list_capability() -> None:
"""Server with list_tasks handler declares tasks.list capability."""
server: Server = Server("test")
@server.experimental.list_tasks()
async def handle_list(req: ListTasksRequest) -> ListTasksResult:
raise NotImplementedError
server.experimental.enable_tasks(on_list_tasks=_noop_list_tasks)
caps = _get_capabilities(server)
assert caps.tasks is not None
@ -60,10 +69,7 @@ def test_server_with_list_tasks_handler_declares_list_capability() -> None:
def test_server_with_cancel_task_handler_declares_cancel_capability() -> None:
"""Server with cancel_task handler declares tasks.cancel capability."""
server: Server = Server("test")
@server.experimental.cancel_task()
async def handle_cancel(req: CancelTaskRequest) -> CancelTaskResult:
raise NotImplementedError
server.experimental.enable_tasks(on_cancel_task=_noop_cancel_task)
caps = _get_capabilities(server)
assert caps.tasks is not None
@ -75,10 +81,7 @@ def test_server_with_get_task_handler_declares_requests_tools_call_capability()
(get_task is required for task-augmented tools/call support)
"""
server: Server = Server("test")
@server.experimental.get_task()
async def handle_get(req: GetTaskRequest) -> GetTaskResult:
raise NotImplementedError
server.experimental.enable_tasks(on_get_task=_noop_get_task)
caps = _get_capabilities(server)
assert caps.tasks is not None
@ -86,28 +89,30 @@ def test_server_with_get_task_handler_declares_requests_tools_call_capability()
assert caps.tasks.requests.tools is not None
def test_server_without_list_handler_has_no_list_capability() -> None:
@pytest.mark.skip(
reason="TODO(maxisbey): enable_tasks registers default handlers for all task methods, "
"so partial capabilities aren't possible yet. Low-level API should support "
"selectively enabling/disabling task capabilities."
)
def test_server_without_list_handler_has_no_list_capability() -> None: # pragma: no cover
"""Server without list_tasks handler has no tasks.list capability."""
server: Server = Server("test")
# Register only get_task (not list_tasks)
@server.experimental.get_task()
async def handle_get(req: GetTaskRequest) -> GetTaskResult:
raise NotImplementedError
server.experimental.enable_tasks(on_get_task=_noop_get_task)
caps = _get_capabilities(server)
assert caps.tasks is not None
assert caps.tasks.list is None
def test_server_without_cancel_handler_has_no_cancel_capability() -> None:
@pytest.mark.skip(
reason="TODO(maxisbey): enable_tasks registers default handlers for all task methods, "
"so partial capabilities aren't possible yet. Low-level API should support "
"selectively enabling/disabling task capabilities."
)
def test_server_without_cancel_handler_has_no_cancel_capability() -> None: # pragma: no cover
"""Server without cancel_task handler has no tasks.cancel capability."""
server: Server = Server("test")
# Register only get_task (not cancel_task)
@server.experimental.get_task()
async def handle_get(req: GetTaskRequest) -> GetTaskResult:
raise NotImplementedError
server.experimental.enable_tasks(on_get_task=_noop_get_task)
caps = _get_capabilities(server)
assert caps.tasks is not None
@ -117,18 +122,11 @@ def test_server_without_cancel_handler_has_no_cancel_capability() -> None:
def test_server_with_all_task_handlers_has_full_capability() -> None:
"""Server with all task handlers declares complete tasks capability."""
server: Server = Server("test")
@server.experimental.list_tasks()
async def handle_list(req: ListTasksRequest) -> ListTasksResult:
raise NotImplementedError
@server.experimental.cancel_task()
async def handle_cancel(req: CancelTaskRequest) -> CancelTaskResult:
raise NotImplementedError
@server.experimental.get_task()
async def handle_get(req: GetTaskRequest) -> GetTaskResult:
raise NotImplementedError
server.experimental.enable_tasks(
on_list_tasks=_noop_list_tasks,
on_cancel_task=_noop_cancel_task,
on_get_task=_noop_get_task,
)
caps = _get_capabilities(server)
assert caps.tasks is not None

31
tests/issues/test_129_resource_templates.py
View file

@ -1,15 +1,13 @@
import pytest
from mcp import types
from mcp import Client
from mcp.server.mcpserver import MCPServer
@pytest.mark.anyio
async def test_resource_templates():
# Create an MCP server
mcp = MCPServer("Demo")
# Add a dynamic greeting resource
@mcp.resource("greeting://{name}")
def get_greeting(name: str) -> str: # pragma: no cover
"""Get a personalized greeting"""
@ -20,23 +18,16 @@ async def test_resource_templates():
"""Dynamic user data"""
return f"Profile data for user {user_id}"
# Get the list of resource templates using the underlying server
# Note: list_resource_templates() returns a decorator that wraps the handler
# The handler returns a ServerResult with a ListResourceTemplatesResult inside
result = await mcp._lowlevel_server.request_handlers[types.ListResourceTemplatesRequest](
types.ListResourceTemplatesRequest(params=None)
)
assert isinstance(result, types.ListResourceTemplatesResult)
templates = result.resource_templates
async with Client(mcp) as client:
result = await client.list_resource_templates()
templates = result.resource_templates
# Verify we get both templates back
assert len(templates) == 2
assert len(templates) == 2
# Verify template details
greeting_template = next(t for t in templates if t.name == "get_greeting")
assert greeting_template.uri_template == "greeting://{name}"
assert greeting_template.description == "Get a personalized greeting"
greeting_template = next(t for t in templates if t.name == "get_greeting")
assert greeting_template.uri_template == "greeting://{name}"
assert greeting_template.description == "Get a personalized greeting"
profile_template = next(t for t in templates if t.name == "get_user_profile")
assert profile_template.uri_template == "users://{user_id}/profile"
assert profile_template.description == "Dynamic user data"
profile_template = next(t for t in templates if t.name == "get_user_profile")
assert profile_template.uri_template == "users://{user_id}/profile"
assert profile_template.description == "Dynamic user data"

39
tests/issues/test_152_resource_mime_type.py
View file

@ -3,9 +3,16 @@ import base64
import pytest
from mcp import Client, types
from mcp.server.lowlevel import Server
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.server import Server, ServerRequestContext
from mcp.server.mcpserver import MCPServer
from mcp.types import (
BlobResourceContents,
ListResourcesResult,
PaginatedRequestParams,
ReadResourceRequestParams,
ReadResourceResult,
TextResourceContents,
)
pytestmark = pytest.mark.anyio
@ -58,7 +65,6 @@ async def test_mcpserver_resource_mime_type():
async def test_lowlevel_resource_mime_type():
"""Test that mime_type parameter is respected for resources."""
server = Server("test")
# Create a small test image as bytes
image_bytes = b"fake_image_data"
@ -74,17 +80,24 @@ async def test_lowlevel_resource_mime_type():
),
]
@server.list_resources()
async def handle_list_resources():
return test_resources
async def handle_list_resources(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListResourcesResult:
return ListResourcesResult(resources=test_resources)
@server.read_resource()
async def handle_read_resource(uri: str):
if str(uri) == "test://image":
return [ReadResourceContents(content=base64_string, mime_type="image/png")]
elif str(uri) == "test://image_bytes":
return [ReadResourceContents(content=bytes(image_bytes), mime_type="image/png")]
raise Exception(f"Resource not found: {uri}") # pragma: no cover
resource_contents: dict[str, list[TextResourceContents | BlobResourceContents]] = {
"test://image": [TextResourceContents(uri="test://image", text=base64_string, mime_type="image/png")],
"test://image_bytes": [
BlobResourceContents(
uri="test://image_bytes", blob=base64.b64encode(image_bytes).decode("utf-8"), mime_type="image/png"
)
],
}
async def handle_read_resource(ctx: ServerRequestContext, params: ReadResourceRequestParams) -> ReadResourceResult:
return ReadResourceResult(contents=resource_contents[str(params.uri)])
server = Server("test", on_list_resources=handle_list_resources, on_read_resource=handle_read_resource)
# Test that resources are listed with correct mime type
async with Client(server) as client:

67
tests/issues/test_1574_resource_uri_validation.py
View file

@ -13,8 +13,14 @@ These tests verify the fix works end-to-end through the JSON-RPC protocol.
import pytest
from mcp import Client, types
from mcp.server.lowlevel import Server
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.server import Server, ServerRequestContext
from mcp.types import (
ListResourcesResult,
PaginatedRequestParams,
ReadResourceRequestParams,
ReadResourceResult,
TextResourceContents,
)
pytestmark = pytest.mark.anyio
@ -26,24 +32,24 @@ async def test_relative_uri_roundtrip():
the server would fail to serialize resources with relative URIs,
or the URI would be transformed during the roundtrip.
"""
server = Server("test")
@server.list_resources()
async def list_resources():
return [
types.Resource(name="user", uri="users/me"),
types.Resource(name="config", uri="./config"),
types.Resource(name="parent", uri="../parent/resource"),
]
async def handle_list_resources(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListResourcesResult:
return ListResourcesResult(
resources=[
types.Resource(name="user", uri="users/me"),
types.Resource(name="config", uri="./config"),
types.Resource(name="parent", uri="../parent/resource"),
]
)
@server.read_resource()
async def read_resource(uri: str):
return [
ReadResourceContents(
content=f"data for {uri}",
mime_type="text/plain",
)
]
async def handle_read_resource(ctx: ServerRequestContext, params: ReadResourceRequestParams) -> ReadResourceResult:
return ReadResourceResult(
contents=[TextResourceContents(uri=str(params.uri), text=f"data for {params.uri}", mime_type="text/plain")]
)
server = Server("test", on_list_resources=handle_list_resources, on_read_resource=handle_read_resource)
async with Client(server) as client:
# List should return the exact URIs we specified
@ -67,18 +73,23 @@ async def test_custom_scheme_uri_roundtrip():
Some MCP servers use custom schemes like "custom://resource".
These should work end-to-end.
"""
server = Server("test")
@server.list_resources()
async def list_resources():
return [
types.Resource(name="custom", uri="custom://my-resource"),
types.Resource(name="file", uri="file:///path/to/file"),
]
async def handle_list_resources(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListResourcesResult:
return ListResourcesResult(
resources=[
types.Resource(name="custom", uri="custom://my-resource"),
types.Resource(name="file", uri="file:///path/to/file"),
]
)
@server.read_resource()
async def read_resource(uri: str):
return [ReadResourceContents(content="data", mime_type="text/plain")]
async def handle_read_resource(ctx: ServerRequestContext, params: ReadResourceRequestParams) -> ReadResourceResult:
return ReadResourceResult(
contents=[TextResourceContents(uri=str(params.uri), text="data", mime_type="text/plain")]
)
server = Server("test", on_list_resources=handle_list_resources, on_read_resource=handle_read_resource)
async with Client(server) as client:
resources = await client.list_resources()

91
tests/issues/test_342_base64_encoding.py
View file

@ -1,83 +1,52 @@
"""Test for base64 encoding issue in MCP server.
This test demonstrates the issue in server.py where the server uses
urlsafe_b64encode but the BlobResourceContents validator expects standard
base64 encoding.
The test should FAIL before fixing server.py to use b64encode instead of
urlsafe_b64encode.
After the fix, the test should PASS.
This test verifies that binary resource data is encoded with standard base64
(not urlsafe_b64encode), so BlobResourceContents validation succeeds.
"""
import base64
from typing import cast
import pytest
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.server.lowlevel.server import Server
from mcp.types import (
BlobResourceContents,
ReadResourceRequest,
ReadResourceRequestParams,
ReadResourceResult,
ServerResult,
)
from mcp import Client
from mcp.server.mcpserver import MCPServer
from mcp.types import BlobResourceContents
pytestmark = pytest.mark.anyio
@pytest.mark.anyio
async def test_server_base64_encoding_issue():
"""Tests that server response can be validated by BlobResourceContents.
async def test_server_base64_encoding():
"""Tests that binary resource data round-trips correctly through base64 encoding.
This test will:
1. Set up a server that returns binary data
2. Extract the base64-encoded blob from the server's response
3. Verify the encoded data can be properly validated by BlobResourceContents
BEFORE FIX: The test will fail because server uses urlsafe_b64encode
AFTER FIX: The test will pass because server uses standard b64encode
The test uses binary data that produces different results with urlsafe vs standard
base64, ensuring the server uses standard encoding.
"""
server = Server("test")
mcp = MCPServer("test")
# Create binary data that will definitely result in + and / characters
# when encoded with standard base64
binary_data = bytes(list(range(255)) * 4)
# Register a resource handler that returns our test data
@server.read_resource()
async def read_resource(uri: str) -> list[ReadResourceContents]:
return [ReadResourceContents(content=binary_data, mime_type="application/octet-stream")]
# Get the handler directly from the server
handler = server.request_handlers[ReadResourceRequest]
# Create a request
request = ReadResourceRequest(
params=ReadResourceRequestParams(uri="test://resource"),
)
# Call the handler to get the response
result: ServerResult = await handler(request)
# After (fixed code):
read_result: ReadResourceResult = cast(ReadResourceResult, result)
blob_content = read_result.contents[0]
# First verify our test data actually produces different encodings
# Sanity check: our test data produces different encodings
urlsafe_b64 = base64.urlsafe_b64encode(binary_data).decode()
standard_b64 = base64.b64encode(binary_data).decode()
assert urlsafe_b64 != standard_b64, "Test data doesn't demonstrate"
" encoding difference"
assert urlsafe_b64 != standard_b64, "Test data doesn't demonstrate encoding difference"
# Now validate the server's output with BlobResourceContents.model_validate
# Before the fix: This should fail with "Invalid base64" because server
# uses urlsafe_b64encode
# After the fix: This should pass because server will use standard b64encode
model_dict = blob_content.model_dump()
@mcp.resource("test://binary", mime_type="application/octet-stream")
def get_binary() -> bytes:
"""Return binary test data."""
return binary_data
# Direct validation - this will fail before fix, pass after fix
blob_model = BlobResourceContents.model_validate(model_dict)
async with Client(mcp) as client:
result = await client.read_resource("test://binary")
assert len(result.contents) == 1
# Verify we can decode the data back correctly
decoded = base64.b64decode(blob_model.blob)
assert decoded == binary_data
blob_content = result.contents[0]
assert isinstance(blob_content, BlobResourceContents)
# Verify standard base64 was used (not urlsafe)
assert blob_content.blob == standard_b64
# Verify we can decode the data back correctly
decoded = base64.b64decode(blob_content.blob)
assert decoded == binary_data

52
tests/issues/test_88_random_error.py
View file

@ -1,8 +1,6 @@
"""Test to reproduce issue #88: Random error thrown on response."""
from collections.abc import Sequence
from pathlib import Path
from typing import Any
import anyio
import pytest
@ -11,10 +9,10 @@ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStre
from mcp import types
from mcp.client.session import ClientSession
from mcp.server.lowlevel import Server
from mcp.server import Server, ServerRequestContext
from mcp.shared.exceptions import MCPError
from mcp.shared.message import SessionMessage
from mcp.types import ContentBlock, TextContent
from mcp.types import CallToolRequestParams, CallToolResult, ListToolsResult, PaginatedRequestParams, TextContent
@pytest.mark.anyio
@ -32,36 +30,38 @@ async def test_notification_validation_error(tmp_path: Path):
- Slow operations use minimal timeout (10ms) for quick test execution
"""
server = Server(name="test")
request_count = 0
slow_request_lock = anyio.Event()
@server.list_tools()
async def list_tools() -> list[types.Tool]:
return [
types.Tool(
name="slow",
description="A slow tool",
input_schema={"type": "object"},
),
types.Tool(
name="fast",
description="A fast tool",
input_schema={"type": "object"},
),
]
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(
tools=[
types.Tool(
name="slow",
description="A slow tool",
input_schema={"type": "object"},
),
types.Tool(
name="fast",
description="A fast tool",
input_schema={"type": "object"},
),
]
)
@server.call_tool()
async def slow_tool(name: str, arguments: dict[str, Any]) -> Sequence[ContentBlock]:
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
nonlocal request_count
request_count += 1
assert params.name in ("slow", "fast"), f"Unknown tool: {params.name}"
if name == "slow":
if params.name == "slow":
await slow_request_lock.wait() # it should timeout here
return [TextContent(type="text", text=f"slow {request_count}")]
elif name == "fast":
return [TextContent(type="text", text=f"fast {request_count}")]
return [TextContent(type="text", text=f"unknown {request_count}")] # pragma: no cover
text = f"slow {request_count}"
else:
text = f"fast {request_count}"
return CallToolResult(content=[TextContent(type="text", text=text)])
server = Server(name="test", on_list_tools=handle_list_tools, on_call_tool=handle_call_tool)
async def server_handler(
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],

292
tests/server/lowlevel/test_func_inspection.py
View file

@ -1,292 +0,0 @@
"""Unit tests for func_inspection module.
Tests the create_call_wrapper function which determines how to call handler functions
with different parameter signatures and type hints.
"""
from typing import Any, Generic, TypeVar
import pytest
from mcp.server.lowlevel.func_inspection import create_call_wrapper
from mcp.types import ListPromptsRequest, ListResourcesRequest, ListToolsRequest, PaginatedRequestParams
T = TypeVar("T")
@pytest.mark.anyio
async def test_no_params_returns_deprecated_wrapper() -> None:
"""Test: def foo() - should call without request."""
called_without_request = False
async def handler() -> list[str]:
nonlocal called_without_request
called_without_request = True
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should call handler without passing request
request = ListPromptsRequest(method="prompts/list", params=None)
result = await wrapper(request)
assert called_without_request is True
assert result == ["test"]
@pytest.mark.anyio
async def test_param_with_default_returns_deprecated_wrapper() -> None:
"""Test: def foo(thing: int = 1) - should call without request."""
called_without_request = False
async def handler(thing: int = 1) -> list[str]:
nonlocal called_without_request
called_without_request = True
return [f"test-{thing}"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should call handler without passing request (uses default value)
request = ListPromptsRequest(method="prompts/list", params=None)
result = await wrapper(request)
assert called_without_request is True
assert result == ["test-1"]
@pytest.mark.anyio
async def test_typed_request_param_passes_request() -> None:
"""Test: def foo(req: ListPromptsRequest) - should pass request through."""
received_request = None
async def handler(req: ListPromptsRequest) -> list[str]:
nonlocal received_request
received_request = req
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should pass request to handler
request = ListPromptsRequest(method="prompts/list", params=PaginatedRequestParams(cursor="test-cursor"))
await wrapper(request)
assert received_request is not None
assert received_request is request
params = getattr(received_request, "params", None)
assert params is not None
assert params.cursor == "test-cursor"
@pytest.mark.anyio
async def test_typed_request_with_default_param_passes_request() -> None:
"""Test: def foo(req: ListPromptsRequest, thing: int = 1) - should pass request through."""
received_request = None
received_thing = None
async def handler(req: ListPromptsRequest, thing: int = 1) -> list[str]:
nonlocal received_request, received_thing
received_request = req
received_thing = thing
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should pass request to handler
request = ListPromptsRequest(method="prompts/list", params=None)
await wrapper(request)
assert received_request is request
assert received_thing == 1 # default value
@pytest.mark.anyio
async def test_optional_typed_request_with_default_none_is_deprecated() -> None:
"""Test: def foo(thing: int = 1, req: ListPromptsRequest | None = None) - old style."""
called_without_request = False
async def handler(thing: int = 1, req: ListPromptsRequest | None = None) -> list[str]:
nonlocal called_without_request
called_without_request = True
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should call handler without passing request
request = ListPromptsRequest(method="prompts/list", params=None)
result = await wrapper(request)
assert called_without_request is True
assert result == ["test"]
@pytest.mark.anyio
async def test_untyped_request_param_is_deprecated() -> None:
"""Test: def foo(req) - should call without request."""
called = False
async def handler(req): # type: ignore[no-untyped-def] # pyright: ignore[reportMissingParameterType] # pragma: no cover
nonlocal called
called = True
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest) # pyright: ignore[reportUnknownArgumentType]
# Wrapper should call handler without passing request, which will fail because req is required
request = ListPromptsRequest(method="prompts/list", params=None)
# This will raise TypeError because handler expects 'req' but wrapper doesn't provide it
with pytest.raises(TypeError, match="missing 1 required positional argument"):
await wrapper(request)
@pytest.mark.anyio
async def test_any_typed_request_param_is_deprecated() -> None:
"""Test: def foo(req: Any) - should call without request."""
async def handler(req: Any) -> list[str]: # pragma: no cover
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should call handler without passing request, which will fail because req is required
request = ListPromptsRequest(method="prompts/list", params=None)
# This will raise TypeError because handler expects 'req' but wrapper doesn't provide it
with pytest.raises(TypeError, match="missing 1 required positional argument"):
await wrapper(request)
@pytest.mark.anyio
async def test_generic_typed_request_param_is_deprecated() -> None:
"""Test: def foo(req: Generic[T]) - should call without request."""
async def handler(req: Generic[T]) -> list[str]: # pyright: ignore[reportGeneralTypeIssues] # pragma: no cover
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should call handler without passing request, which will fail because req is required
request = ListPromptsRequest(method="prompts/list", params=None)
# This will raise TypeError because handler expects 'req' but wrapper doesn't provide it
with pytest.raises(TypeError, match="missing 1 required positional argument"):
await wrapper(request)
@pytest.mark.anyio
async def test_wrong_typed_request_param_is_deprecated() -> None:
"""Test: def foo(req: str) - should call without request."""
async def handler(req: str) -> list[str]: # pragma: no cover
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should call handler without passing request, which will fail because req is required
request = ListPromptsRequest(method="prompts/list", params=None)
# This will raise TypeError because handler expects 'req' but wrapper doesn't provide it
with pytest.raises(TypeError, match="missing 1 required positional argument"):
await wrapper(request)
@pytest.mark.anyio
async def test_required_param_before_typed_request_attempts_to_pass() -> None:
"""Test: def foo(thing: int, req: ListPromptsRequest) - attempts to pass request (will fail at runtime)."""
received_request = None
async def handler(thing: int, req: ListPromptsRequest) -> list[str]: # pragma: no cover
nonlocal received_request
received_request = req
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper will attempt to pass request, but it will fail at runtime
# because 'thing' is required and has no default
request = ListPromptsRequest(method="prompts/list", params=None)
# This will raise TypeError because 'thing' is missing
with pytest.raises(TypeError, match="missing 1 required positional argument: 'thing'"):
await wrapper(request)
@pytest.mark.anyio
async def test_positional_only_param_with_correct_type() -> None:
"""Test: def foo(req: ListPromptsRequest, /) - should pass request through."""
received_request = None
async def handler(req: ListPromptsRequest, /) -> list[str]:
nonlocal received_request
received_request = req
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should pass request to handler
request = ListPromptsRequest(method="prompts/list", params=None)
await wrapper(request)
assert received_request is request
@pytest.mark.anyio
async def test_keyword_only_param_with_correct_type() -> None:
"""Test: def foo(*, req: ListPromptsRequest) - should pass request through."""
received_request = None
async def handler(*, req: ListPromptsRequest) -> list[str]:
nonlocal received_request
received_request = req
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Wrapper should pass request to handler with keyword argument
request = ListPromptsRequest(method="prompts/list", params=None)
await wrapper(request)
assert received_request is request
@pytest.mark.anyio
async def test_different_request_types() -> None:
"""Test that wrapper works with different request types."""
# Test with ListResourcesRequest
received_request = None
async def handler(req: ListResourcesRequest) -> list[str]:
nonlocal received_request
received_request = req
return ["test"]
wrapper = create_call_wrapper(handler, ListResourcesRequest)
request = ListResourcesRequest(method="resources/list", params=None)
await wrapper(request)
assert received_request is request
# Test with ListToolsRequest
received_request = None
async def handler2(req: ListToolsRequest) -> list[str]:
nonlocal received_request
received_request = req
return ["test"]
wrapper2 = create_call_wrapper(handler2, ListToolsRequest)
request2 = ListToolsRequest(method="tools/list", params=None)
await wrapper2(request2)
assert received_request is request2
@pytest.mark.anyio
async def test_mixed_params_with_typed_request() -> None:
"""Test: def foo(a: str, req: ListPromptsRequest, b: int = 5) - attempts to pass request."""
async def handler(a: str, req: ListPromptsRequest, b: int = 5) -> list[str]: # pragma: no cover
return ["test"]
wrapper = create_call_wrapper(handler, ListPromptsRequest)
# Will fail at runtime due to missing 'a'
request = ListPromptsRequest(method="prompts/list", params=None)
with pytest.raises(TypeError, match="missing 1 required positional argument: 'a'"):
await wrapper(request)

143
tests/server/lowlevel/test_server_listing.py
View file

@ -1,20 +1,16 @@
"""Basic tests for list_prompts, list_resources, and list_tools decorators without pagination."""
import warnings
"""Basic tests for list_prompts, list_resources, and list_tools handlers without pagination."""
import pytest
from mcp.server import Server
from mcp import Client
from mcp.server import Server, ServerRequestContext
from mcp.types import (
ListPromptsRequest,
ListPromptsResult,
ListResourcesRequest,
ListResourcesResult,
ListToolsRequest,
ListToolsResult,
PaginatedRequestParams,
Prompt,
Resource,
ServerResult,
Tool,
)
@ -22,60 +18,44 @@ from mcp.types import (
@pytest.mark.anyio
async def test_list_prompts_basic() -> None:
"""Test basic prompt listing without pagination."""
server = Server("test")
test_prompts = [
Prompt(name="prompt1", description="First prompt"),
Prompt(name="prompt2", description="Second prompt"),
]
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
async def handle_list_prompts(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListPromptsResult:
return ListPromptsResult(prompts=test_prompts)
@server.list_prompts()
async def handle_list_prompts() -> list[Prompt]:
return test_prompts
handler = server.request_handlers[ListPromptsRequest]
request = ListPromptsRequest(method="prompts/list", params=None)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, ListPromptsResult)
assert result.prompts == test_prompts
server = Server("test", on_list_prompts=handle_list_prompts)
async with Client(server) as client:
result = await client.list_prompts()
assert result.prompts == test_prompts
@pytest.mark.anyio
async def test_list_resources_basic() -> None:
"""Test basic resource listing without pagination."""
server = Server("test")
test_resources = [
Resource(uri="file:///test1.txt", name="Test 1"),
Resource(uri="file:///test2.txt", name="Test 2"),
]
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
async def handle_list_resources(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListResourcesResult:
return ListResourcesResult(resources=test_resources)
@server.list_resources()
async def handle_list_resources() -> list[Resource]:
return test_resources
handler = server.request_handlers[ListResourcesRequest]
request = ListResourcesRequest(method="resources/list", params=None)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, ListResourcesResult)
assert result.resources == test_resources
server = Server("test", on_list_resources=handle_list_resources)
async with Client(server) as client:
result = await client.list_resources()
assert result.resources == test_resources
@pytest.mark.anyio
async def test_list_tools_basic() -> None:
"""Test basic tool listing without pagination."""
server = Server("test")
test_tools = [
Tool(
name="tool1",
@ -102,80 +82,53 @@ async def test_list_tools_basic() -> None:
),
]
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(tools=test_tools)
@server.list_tools()
async def handle_list_tools() -> list[Tool]:
return test_tools
handler = server.request_handlers[ListToolsRequest]
request = ListToolsRequest(method="tools/list", params=None)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, ListToolsResult)
assert result.tools == test_tools
server = Server("test", on_list_tools=handle_list_tools)
async with Client(server) as client:
result = await client.list_tools()
assert result.tools == test_tools
@pytest.mark.anyio
async def test_list_prompts_empty() -> None:
"""Test listing with empty results."""
server = Server("test")
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
async def handle_list_prompts(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListPromptsResult:
return ListPromptsResult(prompts=[])
@server.list_prompts()
async def handle_list_prompts() -> list[Prompt]:
return []
handler = server.request_handlers[ListPromptsRequest]
request = ListPromptsRequest(method="prompts/list", params=None)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, ListPromptsResult)
assert result.prompts == []
server = Server("test", on_list_prompts=handle_list_prompts)
async with Client(server) as client:
result = await client.list_prompts()
assert result.prompts == []
@pytest.mark.anyio
async def test_list_resources_empty() -> None:
"""Test listing with empty results."""
server = Server("test")
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
async def handle_list_resources(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListResourcesResult:
return ListResourcesResult(resources=[])
@server.list_resources()
async def handle_list_resources() -> list[Resource]:
return []
handler = server.request_handlers[ListResourcesRequest]
request = ListResourcesRequest(method="resources/list", params=None)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, ListResourcesResult)
assert result.resources == []
server = Server("test", on_list_resources=handle_list_resources)
async with Client(server) as client:
result = await client.list_resources()
assert result.resources == []
@pytest.mark.anyio
async def test_list_tools_empty() -> None:
"""Test listing with empty results."""
server = Server("test")
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(tools=[])
@server.list_tools()
async def handle_list_tools() -> list[Tool]:
return []
handler = server.request_handlers[ListToolsRequest]
request = ListToolsRequest(method="tools/list", params=None)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result, ListToolsResult)
assert result.tools == []
server = Server("test", on_list_tools=handle_list_tools)
async with Client(server) as client:
result = await client.list_tools()
assert result.tools == []

124
tests/server/lowlevel/test_server_pagination.py
View file

@ -1,111 +1,83 @@
import pytest
from mcp.server import Server
from mcp import Client
from mcp.server import Server, ServerRequestContext
from mcp.types import (
ListPromptsRequest,
ListPromptsResult,
ListResourcesRequest,
ListResourcesResult,
ListToolsRequest,
ListToolsResult,
PaginatedRequestParams,
ServerResult,
)
@pytest.mark.anyio
async def test_list_prompts_pagination() -> None:
server = Server("test")
test_cursor = "test-cursor-123"
received_params: PaginatedRequestParams | None = None
# Track what request was received
received_request: ListPromptsRequest | None = None
@server.list_prompts()
async def handle_list_prompts(request: ListPromptsRequest) -> ListPromptsResult:
nonlocal received_request
received_request = request
async def handle_list_prompts(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListPromptsResult:
nonlocal received_params
received_params = params
return ListPromptsResult(prompts=[], next_cursor="next")
handler = server.request_handlers[ListPromptsRequest]
server = Server("test", on_list_prompts=handle_list_prompts)
async with Client(server) as client:
# No cursor provided
await client.list_prompts()
assert received_params is not None
assert received_params.cursor is None
# Test: No cursor provided -> handler receives request with None params
request = ListPromptsRequest(method="prompts/list", params=None)
result = await handler(request)
assert received_request is not None
assert received_request.params is None
assert isinstance(result, ServerResult)
# Test: Cursor provided -> handler receives request with cursor in params
request_with_cursor = ListPromptsRequest(method="prompts/list", params=PaginatedRequestParams(cursor=test_cursor))
result2 = await handler(request_with_cursor)
assert received_request is not None
assert received_request.params is not None
assert received_request.params.cursor == test_cursor
assert isinstance(result2, ServerResult)
# Cursor provided
await client.list_prompts(cursor=test_cursor)
assert received_params is not None
assert received_params.cursor == test_cursor
@pytest.mark.anyio
async def test_list_resources_pagination() -> None:
server = Server("test")
test_cursor = "resource-cursor-456"
received_params: PaginatedRequestParams | None = None
# Track what request was received
received_request: ListResourcesRequest | None = None
@server.list_resources()
async def handle_list_resources(request: ListResourcesRequest) -> ListResourcesResult:
nonlocal received_request
received_request = request
async def handle_list_resources(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListResourcesResult:
nonlocal received_params
received_params = params
return ListResourcesResult(resources=[], next_cursor="next")
handler = server.request_handlers[ListResourcesRequest]
server = Server("test", on_list_resources=handle_list_resources)
async with Client(server) as client:
# No cursor provided
await client.list_resources()
assert received_params is not None
assert received_params.cursor is None
# Test: No cursor provided -> handler receives request with None params
request = ListResourcesRequest(method="resources/list", params=None)
result = await handler(request)
assert received_request is not None
assert received_request.params is None
assert isinstance(result, ServerResult)
# Test: Cursor provided -> handler receives request with cursor in params
request_with_cursor = ListResourcesRequest(
method="resources/list", params=PaginatedRequestParams(cursor=test_cursor)
)
result2 = await handler(request_with_cursor)
assert received_request is not None
assert received_request.params is not None
assert received_request.params.cursor == test_cursor
assert isinstance(result2, ServerResult)
# Cursor provided
await client.list_resources(cursor=test_cursor)
assert received_params is not None
assert received_params.cursor == test_cursor
@pytest.mark.anyio
async def test_list_tools_pagination() -> None:
server = Server("test")
test_cursor = "tools-cursor-789"
received_params: PaginatedRequestParams | None = None
# Track what request was received
received_request: ListToolsRequest | None = None
@server.list_tools()
async def handle_list_tools(request: ListToolsRequest) -> ListToolsResult:
nonlocal received_request
received_request = request
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
nonlocal received_params
received_params = params
return ListToolsResult(tools=[], next_cursor="next")
handler = server.request_handlers[ListToolsRequest]
server = Server("test", on_list_tools=handle_list_tools)
async with Client(server) as client:
# No cursor provided
await client.list_tools()
assert received_params is not None
assert received_params.cursor is None
# Test: No cursor provided -> handler receives request with None params
request = ListToolsRequest(method="tools/list", params=None)
result = await handler(request)
assert received_request is not None
assert received_request.params is None
assert isinstance(result, ServerResult)
# Test: Cursor provided -> handler receives request with cursor in params
request_with_cursor = ListToolsRequest(method="tools/list", params=PaginatedRequestParams(cursor=test_cursor))
result2 = await handler(request_with_cursor)
assert received_request is not None
assert received_request.params is not None
assert received_request.params.cursor == test_cursor
assert isinstance(result2, ServerResult)
# Cursor provided
await client.list_tools(cursor=test_cursor)
assert received_params is not None
assert received_params.cursor == test_cursor

3
tests/server/mcpserver/auth/test_auth_integration.py
View file

@ -21,7 +21,8 @@ from mcp.server.auth.provider import (
RefreshToken,
construct_redirect_uri,
)
from mcp.server.auth.routes import ClientRegistrationOptions, RevocationOptions, create_auth_routes
from mcp.server.auth.routes import create_auth_routes
from mcp.server.auth.settings import ClientRegistrationOptions, RevocationOptions
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken

Some files were not shown because too many files have changed in this diff Show more