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

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

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

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