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

SEP-1686: Tasks (#1645)

Browse source
This commit is contained in:
Max Isbey 2025-11-28 18:51:58 +00:00 • committed by GitHub
commit c92bb2f7ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 14230 additions and 120 deletions
Download patch file Download diff file Expand all files Collapse all files

0
tests/experimental/tasks/server/__init__.py Normal file
View file

183
tests/experimental/tasks/server/test_context.py Normal file
View file

@ -0,0 +1,183 @@
"""Tests for TaskContext and helper functions."""
import pytest
from mcp.shared.experimental.tasks.context import TaskContext
from mcp.shared.experimental.tasks.helpers import create_task_state, task_execution
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.types import CallToolResult, TaskMetadata, TextContent
@pytest.mark.anyio
async def test_task_context_properties() -> None:
"""Test TaskContext basic properties."""
store = InMemoryTaskStore()
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
ctx = TaskContext(task, store)
assert ctx.task_id == task.taskId
assert ctx.task.taskId == task.taskId
assert ctx.task.status == "working"
assert ctx.is_cancelled is False
store.cleanup()
@pytest.mark.anyio
async def test_task_context_update_status() -> None:
"""Test TaskContext.update_status."""
store = InMemoryTaskStore()
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
ctx = TaskContext(task, store)
await ctx.update_status("Processing step 1...")
# Check status message was updated
updated = await store.get_task(task.taskId)
assert updated is not None
assert updated.statusMessage == "Processing step 1..."
store.cleanup()
@pytest.mark.anyio
async def test_task_context_complete() -> None:
"""Test TaskContext.complete."""
store = InMemoryTaskStore()
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
ctx = TaskContext(task, store)
result = CallToolResult(content=[TextContent(type="text", text="Done!")])
await ctx.complete(result)
# Check task status
updated = await store.get_task(task.taskId)
assert updated is not None
assert updated.status == "completed"
# Check result is stored
stored_result = await store.get_result(task.taskId)
assert stored_result is not None
store.cleanup()
@pytest.mark.anyio
async def test_task_context_fail() -> None:
"""Test TaskContext.fail."""
store = InMemoryTaskStore()
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
ctx = TaskContext(task, store)
await ctx.fail("Something went wrong!")
# Check task status
updated = await store.get_task(task.taskId)
assert updated is not None
assert updated.status == "failed"
assert updated.statusMessage == "Something went wrong!"
store.cleanup()
@pytest.mark.anyio
async def test_task_context_cancellation() -> None:
"""Test TaskContext cancellation request."""
store = InMemoryTaskStore()
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
ctx = TaskContext(task, store)
assert ctx.is_cancelled is False
ctx.request_cancellation()
assert ctx.is_cancelled is True
store.cleanup()
def test_create_task_state_generates_id() -> None:
"""create_task_state generates a unique task ID when none provided."""
task1 = create_task_state(TaskMetadata(ttl=60000))
task2 = create_task_state(TaskMetadata(ttl=60000))
assert task1.taskId != task2.taskId
def test_create_task_state_uses_provided_id() -> None:
"""create_task_state uses the provided task ID."""
task = create_task_state(TaskMetadata(ttl=60000), task_id="my-task-123")
assert task.taskId == "my-task-123"
def test_create_task_state_null_ttl() -> None:
"""create_task_state handles null TTL."""
task = create_task_state(TaskMetadata(ttl=None))
assert task.ttl is None
def test_create_task_state_has_created_at() -> None:
"""create_task_state sets createdAt timestamp."""
task = create_task_state(TaskMetadata(ttl=60000))
assert task.createdAt is not None
@pytest.mark.anyio
async def test_task_execution_provides_context() -> None:
"""task_execution provides a TaskContext for the task."""
store = InMemoryTaskStore()
await store.create_task(TaskMetadata(ttl=60000), task_id="exec-test-1")
async with task_execution("exec-test-1", store) as ctx:
assert ctx.task_id == "exec-test-1"
assert ctx.task.status == "working"
store.cleanup()
@pytest.mark.anyio
async def test_task_execution_auto_fails_on_exception() -> None:
"""task_execution automatically fails task on unhandled exception."""
store = InMemoryTaskStore()
await store.create_task(TaskMetadata(ttl=60000), task_id="exec-fail-1")
async with task_execution("exec-fail-1", store):
raise RuntimeError("Oops!")
# Task should be failed
failed_task = await store.get_task("exec-fail-1")
assert failed_task is not None
assert failed_task.status == "failed"
assert "Oops!" in (failed_task.statusMessage or "")
store.cleanup()
@pytest.mark.anyio
async def test_task_execution_doesnt_fail_if_already_terminal() -> None:
"""task_execution doesn't re-fail if task already terminal."""
store = InMemoryTaskStore()
await store.create_task(TaskMetadata(ttl=60000), task_id="exec-term-1")
async with task_execution("exec-term-1", store) as ctx:
# Complete the task first
await ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")]))
# Then raise - shouldn't change status
raise RuntimeError("This shouldn't matter")
# Task should remain completed
final_task = await store.get_task("exec-term-1")
assert final_task is not None
assert final_task.status == "completed"
store.cleanup()
@pytest.mark.anyio
async def test_task_execution_not_found() -> None:
"""task_execution raises ValueError for non-existent task."""
store = InMemoryTaskStore()
with pytest.raises(ValueError, match="not found"):
async with task_execution("nonexistent", store):
...

357
tests/experimental/tasks/server/test_integration.py Normal file
View file

@ -0,0 +1,357 @@
"""End-to-end integration tests for tasks functionality.
These tests demonstrate the full task lifecycle:
1. Client sends task-augmented request (tools/call with task metadata)
2. Server creates task and returns CreateTaskResult immediately
3. Background work executes (using task_execution context manager)
4. Client polls with tasks/get
5. Client retrieves result with tasks/result
"""
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.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,
ClientRequest,
ClientResult,
CreateTaskResult,
GetTaskPayloadRequest,
GetTaskPayloadRequestParams,
GetTaskPayloadResult,
GetTaskRequest,
GetTaskRequestParams,
GetTaskResult,
ListTasksRequest,
ListTasksResult,
ServerNotification,
ServerRequest,
TaskMetadata,
TextContent,
Tool,
ToolExecution,
)
@dataclass
class AppContext:
"""Application context passed via lifespan_context."""
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
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]
store = InMemoryTaskStore()
@server.list_tools()
async def list_tools():
return [
Tool(
name="process_data",
description="Process data asynchronously",
inputSchema={
"type": "object",
"properties": {"input": {"type": "string"}},
},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@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 name == "process_data" and ctx.experimental.is_task:
# 1. Create task in store
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.taskId] = done_event
# 3. Define work function using task_execution for safety
async def do_work():
async with task_execution(task.taskId, app.store) as task_ctx:
await task_ctx.update_status("Processing input...")
# Simulate work
input_value = arguments.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.taskId)
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
return GetTaskResult(
taskId=task.taskId,
status=task.status,
statusMessage=task.statusMessage,
createdAt=task.createdAt,
lastUpdatedAt=task.lastUpdatedAt,
ttl=task.ttl,
pollInterval=task.pollInterval,
)
@server.experimental.get_task_result()
async def handle_get_task_result(
request: GetTaskPayloadRequest,
) -> GetTaskPayloadResult:
app = server.request_context.lifespan_context
result = await app.store.get_result(request.params.taskId)
assert result is not None, f"Test setup error: result for {request.params.taskId} 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:
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)
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={},
),
),
) 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:
# Create app context with task group and store
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()
# === Step 1: Send task-augmented tool call ===
create_result = await client_session.send_request(
ClientRequest(
CallToolRequest(
params=CallToolRequestParams(
name="process_data",
arguments={"input": "hello world"},
task=TaskMetadata(ttl=60000),
),
)
),
CreateTaskResult,
)
assert isinstance(create_result, CreateTaskResult)
assert create_result.task.status == "working"
task_id = create_result.task.taskId
# === Step 2: Wait for task to complete ===
await app_context.task_done_events[task_id].wait()
task_status = await client_session.send_request(
ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId=task_id))),
GetTaskResult,
)
assert task_status.taskId == task_id
assert task_status.status == "completed"
# === Step 3: Retrieve the actual result ===
task_result = await client_session.send_request(
ClientRequest(GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=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()
@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()
@server.list_tools()
async def list_tools():
return [
Tool(
name="failing_task",
description="A task that fails",
inputSchema={"type": "object", "properties": {}},
)
]
@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 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.taskId] = done_event
async def do_failing_work():
async with task_execution(task.taskId, 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()
app.task_group.start_soon(do_failing_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.taskId)
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
return GetTaskResult(
taskId=task.taskId,
status=task.status,
statusMessage=task.statusMessage,
createdAt=task.createdAt,
lastUpdatedAt=task.lastUpdatedAt,
ttl=task.ttl,
pollInterval=task.pollInterval,
)
# 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 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={},
),
),
) 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()
# Send task request
create_result = await client_session.send_request(
ClientRequest(
CallToolRequest(
params=CallToolRequestParams(
name="failing_task",
arguments={},
task=TaskMetadata(ttl=60000),
),
)
),
CreateTaskResult,
)
task_id = create_result.task.taskId
# 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(
ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId=task_id))),
GetTaskResult,
)
assert task_status.status == "failed"
assert task_status.statusMessage == "Something went wrong!"
tg.cancel_scope.cancel()

538
tests/experimental/tasks/server/test_run_task_flow.py Normal file
View file

@ -0,0 +1,538 @@
"""
Tests for the simplified task API: enable_tasks() + run_task()
This tests the recommended user flow:
1. server.experimental.enable_tasks() - one-line setup
2. ctx.experimental.run_task(work) - spawns work, returns CreateTaskResult
3. work function uses ServerTaskContext for elicit/create_message
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.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,
CallToolResult,
CancelTaskRequest,
CancelTaskResult,
CreateTaskResult,
GetTaskPayloadRequest,
GetTaskPayloadResult,
GetTaskRequest,
GetTaskResult,
ListTasksRequest,
ListTasksResult,
TextContent,
Tool,
ToolExecution,
)
@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
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",
inputSchema={"type": "object", "properties": {"input": {"type": "string"}}},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
ctx.experimental.validate_task_mode(TASK_REQUIRED)
# Capture the meta from the request (if present)
if ctx.meta is not None and ctx.meta.model_extra: # pragma: no branch
received_meta[0] = ctx.meta.model_extra.get("custom_field")
async def work(task: ServerTaskContext) -> CallToolResult:
await task.update_status("Working...")
input_val = arguments.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)
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 def run_client() -> None:
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
# Initialize
await client_session.initialize()
# 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"},
)
# Should get CreateTaskResult
task_id = result.task.taskId
assert result.task.status == "working"
# 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",
inputSchema={"type": "object"},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
work_failed.set()
raise RuntimeError("Something went wrong!")
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("failing_task", {})
task_id = result.task.taskId
# 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.statusMessage or "")
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
@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")
# Before enable_tasks, no task capabilities
caps_before = server.get_capabilities(NotificationOptions(), {})
assert caps_before.tasks is None
# Enable tasks
server.experimental.enable_tasks()
# After enable_tasks, should have task capabilities
caps_after = server.get_capabilities(NotificationOptions(), {})
assert caps_after.tasks is not None
assert caps_after.tasks.list is not None
assert caps_after.tasks.cancel 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:
raise NotImplementedError
@server.experimental.get_task_result()
async def custom_get_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
raise NotImplementedError
@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
@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(
task_metadata=None,
_client_capabilities=None,
_session=None,
_task_support=None, # Not enabled
)
async def work(task: ServerTaskContext) -> CallToolResult:
raise NotImplementedError
with pytest.raises(RuntimeError, match="Task support not enabled"):
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()
with pytest.raises(RuntimeError, match="TaskSupport not running"):
_ = 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()
experimental = Experimental(
task_metadata=None,
_client_capabilities=None,
_session=None, # No session
_task_support=task_support,
)
async def work(task: ServerTaskContext) -> CallToolResult:
raise NotImplementedError
with pytest.raises(RuntimeError, match="Session not available"):
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()
mock_session = Mock()
experimental = Experimental(
task_metadata=None, # Not a task-augmented request
_client_capabilities=None,
_session=mock_session,
_task_support=task_support,
)
async def work(task: ServerTaskContext) -> CallToolResult:
raise NotImplementedError
with pytest.raises(RuntimeError, match="Request is not task-augmented"):
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",
inputSchema={"type": "object"},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
work_completed.set()
return CallToolResult(content=[TextContent(type="text", text="Done")])
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.experimental.enable_tasks()
work_completed = Event()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="manual_complete_task",
description="A task that manually completes",
inputSchema={"type": "object"},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
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.taskId
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.experimental.enable_tasks()
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",
inputSchema={"type": "object"},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
ctx = server.request_context
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)
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_cancel_task", {})
task_id = result.task.taskId
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.statusMessage == "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)

965
tests/experimental/tasks/server/test_server.py Normal file
View file

@ -0,0 +1,965 @@
"""Tests for server-side task support (handlers, capabilities, integration)."""
from datetime import datetime, timezone
from typing import Any
import anyio
import pytest
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.shared.exceptions import McpError
from mcp.shared.message import ServerMessageMetadata, SessionMessage
from mcp.shared.response_router import ResponseRouter
from mcp.shared.session import RequestResponder
from mcp.types import (
INVALID_REQUEST,
TASK_FORBIDDEN,
TASK_OPTIONAL,
TASK_REQUIRED,
CallToolRequest,
CallToolRequestParams,
CallToolResult,
CancelTaskRequest,
CancelTaskRequestParams,
CancelTaskResult,
ClientRequest,
ClientResult,
ErrorData,
GetTaskPayloadRequest,
GetTaskPayloadRequestParams,
GetTaskPayloadResult,
GetTaskRequest,
GetTaskRequestParams,
GetTaskResult,
JSONRPCError,
JSONRPCMessage,
JSONRPCNotification,
JSONRPCResponse,
ListTasksRequest,
ListTasksResult,
ListToolsRequest,
ListToolsResult,
SamplingMessage,
ServerCapabilities,
ServerNotification,
ServerRequest,
ServerResult,
Task,
TaskMetadata,
TextContent,
Tool,
ToolExecution,
)
@pytest.mark.anyio
async def test_list_tasks_handler() -> None:
"""Test that experimental list_tasks handler works."""
server = Server("test")
now = datetime.now(timezone.utc)
test_tasks = [
Task(
taskId="task-1",
status="working",
createdAt=now,
lastUpdatedAt=now,
ttl=60000,
pollInterval=1000,
),
Task(
taskId="task-2",
status="completed",
createdAt=now,
lastUpdatedAt=now,
ttl=60000,
pollInterval=1000,
),
]
@server.experimental.list_tasks()
async def handle_list_tasks(request: ListTasksRequest) -> 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.root, ListTasksResult)
assert len(result.root.tasks) == 2
assert result.root.tasks[0].taskId == "task-1"
assert result.root.tasks[1].taskId == "task-2"
@pytest.mark.anyio
async def test_get_task_handler() -> None:
"""Test that experimental get_task handler works."""
server = Server("test")
@server.experimental.get_task()
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
now = datetime.now(timezone.utc)
return GetTaskResult(
taskId=request.params.taskId,
status="working",
createdAt=now,
lastUpdatedAt=now,
ttl=60000,
pollInterval=1000,
)
handler = server.request_handlers[GetTaskRequest]
request = GetTaskRequest(
method="tasks/get",
params=GetTaskRequestParams(taskId="test-task-123"),
)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result.root, GetTaskResult)
assert result.root.taskId == "test-task-123"
assert result.root.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.get_task_result()
async def handle_get_task_result(request: GetTaskPayloadRequest) -> GetTaskPayloadResult:
return GetTaskPayloadResult()
handler = server.request_handlers[GetTaskPayloadRequest]
request = GetTaskPayloadRequest(
method="tasks/result",
params=GetTaskPayloadRequestParams(taskId="test-task-123"),
)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result.root, GetTaskPayloadResult)
@pytest.mark.anyio
async def test_cancel_task_handler() -> None:
"""Test that experimental cancel_task handler works."""
server = Server("test")
@server.experimental.cancel_task()
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
now = datetime.now(timezone.utc)
return CancelTaskResult(
taskId=request.params.taskId,
status="cancelled",
createdAt=now,
lastUpdatedAt=now,
ttl=60000,
)
handler = server.request_handlers[CancelTaskRequest]
request = CancelTaskRequest(
method="tasks/cancel",
params=CancelTaskRequestParams(taskId="test-task-123"),
)
result = await handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result.root, CancelTaskResult)
assert result.root.taskId == "test-task-123"
assert result.root.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:
raise NotImplementedError
@server.experimental.cancel_task()
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
raise NotImplementedError
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 not None
assert capabilities.tasks.requests is not None
assert capabilities.tasks.requests.tools is not None
@pytest.mark.anyio
async def test_server_capabilities_partial_tasks() -> None:
"""Test capabilities with only some task handlers registered."""
server = Server("test")
@server.experimental.list_tasks()
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
raise NotImplementedError
# Only list_tasks registered, not cancel_task
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",
inputSchema={"type": "object", "properties": {}},
execution=ToolExecution(taskSupport=TASK_FORBIDDEN),
),
Tool(
name="long_tool",
description="Long running tool",
inputSchema={"type": "object", "properties": {}},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
),
Tool(
name="flexible_tool",
description="Can be either",
inputSchema={"type": "object", "properties": {}},
execution=ToolExecution(taskSupport=TASK_OPTIONAL),
),
]
tools_handler = server.request_handlers[ListToolsRequest]
request = ListToolsRequest(method="tools/list")
result = await tools_handler(request)
assert isinstance(result, ServerResult)
assert isinstance(result.root, ListToolsResult)
tools = result.root.tools
assert tools[0].execution is not None
assert tools[0].execution.taskSupport == TASK_FORBIDDEN
assert tools[1].execution is not None
assert tools[1].execution.taskSupport == TASK_REQUIRED
assert tools[2].execution is not None
assert tools[2].execution.taskSupport == 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")
captured_task_metadata: TaskMetadata | None = None
@server.list_tools()
async def list_tools():
return [
Tool(
name="long_task",
description="A long running task",
inputSchema={"type": "object", "properties": {}},
execution=ToolExecution(taskSupport="optional"),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
nonlocal captured_task_metadata
ctx = server.request_context
captured_task_metadata = ctx.experimental.task_metadata
return [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)
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={},
),
),
) as server_session:
async with anyio.create_task_group() as tg:
async def handle_messages():
async for message in server_session.incoming_messages:
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(
ClientRequest(
CallToolRequest(
params=CallToolRequestParams(
name="long_task",
arguments={},
task=TaskMetadata(ttl=60000),
),
)
),
CallToolResult,
)
tg.cancel_scope.cancel()
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")
is_task_values: list[bool] = []
@server.list_tools()
async def list_tools():
return [
Tool(
name="test_tool",
description="Test tool",
inputSchema={"type": "object", "properties": {}},
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
ctx = server.request_context
is_task_values.append(ctx.experimental.is_task)
return [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)
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={},
),
),
) as server_session:
async with anyio.create_task_group() as tg:
async def handle_messages():
async for message in server_session.incoming_messages:
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(
ClientRequest(
CallToolRequest(
params=CallToolRequestParams(name="test_tool", arguments={}),
)
),
CallToolResult,
)
# Call with task metadata
await client_session.send_request(
ClientRequest(
CallToolRequest(
params=CallToolRequestParams(
name="test_tool",
arguments={},
task=TaskMetadata(ttl=60000),
),
)
),
CallToolResult,
)
tg.cancel_scope.cancel()
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.
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
"""
server = Server("test-default-handlers")
# Enable tasks with default handlers (no custom handlers registered)
task_support = server.experimental.enable_tasks()
store = task_support.store
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() -> None:
async with task_support.run():
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:
task_support.configure_session(server_session)
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, {}, False)
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()
# Create a task directly in the store for testing
task = await store.create_task(TaskMetadata(ttl=60000))
# Test list_tasks (default handler)
list_result = await client_session.send_request(
ClientRequest(ListTasksRequest()),
ListTasksResult,
)
assert len(list_result.tasks) == 1
assert list_result.tasks[0].taskId == task.taskId
# Test get_task (default handler - found)
get_result = await client_session.send_request(
ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId=task.taskId))),
GetTaskResult,
)
assert get_result.taskId == task.taskId
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(
ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId="nonexistent-task"))),
GetTaskResult,
)
# Create a completed task to test get_task_result
completed_task = await store.create_task(TaskMetadata(ttl=60000))
await store.store_result(
completed_task.taskId, CallToolResult(content=[TextContent(type="text", text="Test result")])
)
await store.update_task(completed_task.taskId, status="completed")
# Test get_task_result (default handler)
payload_result = await client_session.send_request(
ClientRequest(GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=completed_task.taskId))),
GetTaskPayloadResult,
)
# The result should have the related-task metadata
assert payload_result.meta is not None
assert "io.modelcontextprotocol/related-task" in payload_result.meta
# Test cancel_task (default handler)
cancel_result = await client_session.send_request(
ClientRequest(CancelTaskRequest(params=CancelTaskRequestParams(taskId=task.taskId))),
CancelTaskResult,
)
assert cancel_result.taskId == task.taskId
assert cancel_result.status == "cancelled"
tg.cancel_scope.cancel()
@pytest.mark.anyio
async def test_build_elicit_form_request() -> None:
"""Test that _build_elicit_form_request builds a proper elicitation request."""
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)
try:
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=ServerCapabilities(),
),
) as server_session:
# Test without task_id
request = server_session._build_elicit_form_request(
message="Test message",
requestedSchema={"type": "object", "properties": {"answer": {"type": "string"}}},
)
assert request.method == "elicitation/create"
assert request.params is not None
assert request.params["message"] == "Test message"
# Test with related_task_id (adds related-task metadata)
request_with_task = server_session._build_elicit_form_request(
message="Task message",
requestedSchema={"type": "object"},
related_task_id="test-task-123",
)
assert request_with_task.method == "elicitation/create"
assert request_with_task.params is not None
assert "_meta" in request_with_task.params
assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"]
assert (
request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"] == "test-task-123"
)
finally: # pragma: no cover
await server_to_client_send.aclose()
await server_to_client_receive.aclose()
await client_to_server_send.aclose()
await client_to_server_receive.aclose()
@pytest.mark.anyio
async def test_build_elicit_url_request() -> None:
"""Test that _build_elicit_url_request builds a proper URL mode elicitation request."""
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)
try:
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=ServerCapabilities(),
),
) as server_session:
# Test without related_task_id
request = server_session._build_elicit_url_request(
message="Please authorize with GitHub",
url="https://github.com/login/oauth/authorize",
elicitation_id="oauth-123",
)
assert request.method == "elicitation/create"
assert request.params is not None
assert request.params["message"] == "Please authorize with GitHub"
assert request.params["url"] == "https://github.com/login/oauth/authorize"
assert request.params["elicitationId"] == "oauth-123"
assert request.params["mode"] == "url"
# Test with related_task_id (adds related-task metadata)
request_with_task = server_session._build_elicit_url_request(
message="OAuth required",
url="https://example.com/oauth",
elicitation_id="oauth-456",
related_task_id="test-task-789",
)
assert request_with_task.method == "elicitation/create"
assert request_with_task.params is not None
assert "_meta" in request_with_task.params
assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"]
assert (
request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"] == "test-task-789"
)
finally: # pragma: no cover
await server_to_client_send.aclose()
await server_to_client_receive.aclose()
await client_to_server_send.aclose()
await client_to_server_receive.aclose()
@pytest.mark.anyio
async def test_build_create_message_request() -> None:
"""Test that _build_create_message_request builds a proper sampling request."""
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)
try:
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=ServerCapabilities(),
),
) as server_session:
messages = [
SamplingMessage(role="user", content=TextContent(type="text", text="Hello")),
]
# Test without task_id
request = server_session._build_create_message_request(
messages=messages,
max_tokens=100,
system_prompt="You are helpful",
)
assert request.method == "sampling/createMessage"
assert request.params is not None
assert request.params["maxTokens"] == 100
# Test with related_task_id (adds related-task metadata)
request_with_task = server_session._build_create_message_request(
messages=messages,
max_tokens=50,
related_task_id="sampling-task-456",
)
assert request_with_task.method == "sampling/createMessage"
assert request_with_task.params is not None
assert "_meta" in request_with_task.params
assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"]
assert (
request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"]
== "sampling-task-456"
)
finally: # pragma: no cover
await server_to_client_send.aclose()
await server_to_client_receive.aclose()
await client_to_server_send.aclose()
await client_to_server_receive.aclose()
@pytest.mark.anyio
async def test_send_message() -> None:
"""Test that send_message sends a raw session message."""
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)
try:
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=ServerCapabilities(),
),
) as server_session:
# Create a test message
notification = JSONRPCNotification(jsonrpc="2.0", method="test/notification")
message = SessionMessage(
message=JSONRPCMessage(notification),
metadata=ServerMessageMetadata(related_request_id="test-req-1"),
)
# Send the message
await server_session.send_message(message)
# Verify it was sent to the stream
received = await server_to_client_receive.receive()
assert isinstance(received.message.root, JSONRPCNotification)
assert received.message.root.method == "test/notification"
finally: # pragma: no cover
await server_to_client_send.aclose()
await server_to_client_receive.aclose()
await client_to_server_send.aclose()
await client_to_server_receive.aclose()
@pytest.mark.anyio
async def test_response_routing_success() -> None:
"""Test that response routing works for success responses."""
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)
# Track routed responses with event for synchronization
routed_responses: list[dict[str, Any]] = []
response_received = anyio.Event()
class TestRouter(ResponseRouter):
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
routed_responses.append({"id": request_id, "response": response})
response_received.set()
return True # Handled
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
raise NotImplementedError
try:
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=ServerCapabilities(),
),
) as server_session:
router = TestRouter()
server_session.add_response_router(router)
# Simulate receiving a response from client
response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"})
message = SessionMessage(message=JSONRPCMessage(response))
# Send from "client" side
await client_to_server_send.send(message)
# Wait for response to be routed
with anyio.fail_after(5):
await response_received.wait()
# Verify response was routed
assert len(routed_responses) == 1
assert routed_responses[0]["id"] == "test-req-1"
assert routed_responses[0]["response"]["status"] == "ok"
finally: # pragma: no cover
await server_to_client_send.aclose()
await server_to_client_receive.aclose()
await client_to_server_send.aclose()
await client_to_server_receive.aclose()
@pytest.mark.anyio
async def test_response_routing_error() -> None:
"""Test that error routing works for error responses."""
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)
# Track routed errors with event for synchronization
routed_errors: list[dict[str, Any]] = []
error_received = anyio.Event()
class TestRouter(ResponseRouter):
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
raise NotImplementedError
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
routed_errors.append({"id": request_id, "error": error})
error_received.set()
return True # Handled
try:
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=ServerCapabilities(),
),
) as server_session:
router = TestRouter()
server_session.add_response_router(router)
# Simulate receiving an error response from client
error_data = ErrorData(code=INVALID_REQUEST, message="Test error")
error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data)
message = SessionMessage(message=JSONRPCMessage(error_response))
# Send from "client" side
await client_to_server_send.send(message)
# Wait for error to be routed
with anyio.fail_after(5):
await error_received.wait()
# Verify error was routed
assert len(routed_errors) == 1
assert routed_errors[0]["id"] == "test-req-2"
assert routed_errors[0]["error"].message == "Test error"
finally: # pragma: no cover
await server_to_client_send.aclose()
await server_to_client_receive.aclose()
await client_to_server_send.aclose()
await client_to_server_receive.aclose()
@pytest.mark.anyio
async def test_response_routing_skips_non_matching_routers() -> None:
"""Test that routing continues to next router when first doesn't match."""
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)
# Track which routers were called
router_calls: list[str] = []
response_received = anyio.Event()
class NonMatchingRouter(ResponseRouter):
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
router_calls.append("non_matching_response")
return False # Doesn't handle it
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
raise NotImplementedError
class MatchingRouter(ResponseRouter):
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
router_calls.append("matching_response")
response_received.set()
return True # Handles it
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
raise NotImplementedError
try:
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=ServerCapabilities(),
),
) as server_session:
# Add non-matching router first, then matching router
server_session.add_response_router(NonMatchingRouter())
server_session.add_response_router(MatchingRouter())
# Send a response - should skip first router and be handled by second
response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"})
message = SessionMessage(message=JSONRPCMessage(response))
await client_to_server_send.send(message)
with anyio.fail_after(5):
await response_received.wait()
# Verify both routers were called (first returned False, second returned True)
assert router_calls == ["non_matching_response", "matching_response"]
finally: # pragma: no cover
await server_to_client_send.aclose()
await server_to_client_receive.aclose()
await client_to_server_send.aclose()
await client_to_server_receive.aclose()
@pytest.mark.anyio
async def test_error_routing_skips_non_matching_routers() -> None:
"""Test that error routing continues to next router when first doesn't match."""
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)
# Track which routers were called
router_calls: list[str] = []
error_received = anyio.Event()
class NonMatchingRouter(ResponseRouter):
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
raise NotImplementedError
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
router_calls.append("non_matching_error")
return False # Doesn't handle it
class MatchingRouter(ResponseRouter):
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
raise NotImplementedError
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
router_calls.append("matching_error")
error_received.set()
return True # Handles it
try:
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=ServerCapabilities(),
),
) as server_session:
# Add non-matching router first, then matching router
server_session.add_response_router(NonMatchingRouter())
server_session.add_response_router(MatchingRouter())
# Send an error - should skip first router and be handled by second
error_data = ErrorData(code=INVALID_REQUEST, message="Test error")
error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data)
message = SessionMessage(message=JSONRPCMessage(error_response))
await client_to_server_send.send(message)
with anyio.fail_after(5):
await error_received.wait()
# Verify both routers were called (first returned False, second returned True)
assert router_calls == ["non_matching_error", "matching_error"]
finally: # pragma: no cover
await server_to_client_send.aclose()
await server_to_client_receive.aclose()
await client_to_server_send.aclose()
await client_to_server_receive.aclose()

709
tests/experimental/tasks/server/test_server_task_context.py Normal file
View file

@ -0,0 +1,709 @@
"""Tests for ServerTaskContext."""
import asyncio
from unittest.mock import AsyncMock, Mock
import anyio
import pytest
from mcp.server.experimental.task_context import ServerTaskContext
from mcp.server.experimental.task_result_handler import TaskResultHandler
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
from mcp.types import (
CallToolResult,
ClientCapabilities,
ClientTasksCapability,
ClientTasksRequestsCapability,
Implementation,
InitializeRequestParams,
JSONRPCRequest,
SamplingMessage,
TaskMetadata,
TasksCreateElicitationCapability,
TasksCreateMessageCapability,
TasksElicitationCapability,
TasksSamplingCapability,
TextContent,
)
@pytest.mark.anyio
async def test_server_task_context_properties() -> None:
"""Test ServerTaskContext property accessors."""
store = InMemoryTaskStore()
mock_session = Mock()
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-123")
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
)
assert ctx.task_id == "test-123"
assert ctx.task.taskId == "test-123"
assert ctx.is_cancelled is False
store.cleanup()
@pytest.mark.anyio
async def test_server_task_context_request_cancellation() -> None:
"""Test ServerTaskContext.request_cancellation()."""
store = InMemoryTaskStore()
mock_session = Mock()
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
)
assert ctx.is_cancelled is False
ctx.request_cancellation()
assert ctx.is_cancelled is True
store.cleanup()
@pytest.mark.anyio
async def test_server_task_context_update_status_with_notify() -> None:
"""Test update_status sends notification when notify=True."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.send_notification = AsyncMock()
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
)
await ctx.update_status("Working...", notify=True)
mock_session.send_notification.assert_called_once()
store.cleanup()
@pytest.mark.anyio
async def test_server_task_context_update_status_without_notify() -> None:
"""Test update_status skips notification when notify=False."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.send_notification = AsyncMock()
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
)
await ctx.update_status("Working...", notify=False)
mock_session.send_notification.assert_not_called()
store.cleanup()
@pytest.mark.anyio
async def test_server_task_context_complete_with_notify() -> None:
"""Test complete sends notification when notify=True."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.send_notification = AsyncMock()
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
)
result = CallToolResult(content=[TextContent(type="text", text="Done")])
await ctx.complete(result, notify=True)
mock_session.send_notification.assert_called_once()
store.cleanup()
@pytest.mark.anyio
async def test_server_task_context_fail_with_notify() -> None:
"""Test fail sends notification when notify=True."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.send_notification = AsyncMock()
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
)
await ctx.fail("Something went wrong", notify=True)
mock_session.send_notification.assert_called_once()
store.cleanup()
@pytest.mark.anyio
async def test_elicit_raises_when_client_lacks_capability() -> None:
"""Test that elicit() raises McpError when client doesn't support elicitation."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=False)
queue = InMemoryTaskMessageQueue()
handler = TaskResultHandler(store, queue)
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=handler,
)
with pytest.raises(McpError) as exc_info:
await ctx.elicit(message="Test?", requestedSchema={"type": "object"})
assert "elicitation capability" in exc_info.value.error.message
mock_session.check_client_capability.assert_called_once()
store.cleanup()
@pytest.mark.anyio
async def test_create_message_raises_when_client_lacks_capability() -> None:
"""Test that create_message() raises McpError when client doesn't support sampling."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=False)
queue = InMemoryTaskMessageQueue()
handler = TaskResultHandler(store, queue)
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=handler,
)
with pytest.raises(McpError) as exc_info:
await ctx.create_message(messages=[], max_tokens=100)
assert "sampling capability" in exc_info.value.error.message
mock_session.check_client_capability.assert_called_once()
store.cleanup()
@pytest.mark.anyio
async def test_elicit_raises_without_handler() -> None:
"""Test that elicit() raises when handler is not provided."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=True)
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=None,
)
with pytest.raises(RuntimeError, match="handler is required"):
await ctx.elicit(message="Test?", requestedSchema={"type": "object"})
store.cleanup()
@pytest.mark.anyio
async def test_elicit_url_raises_without_handler() -> None:
"""Test that elicit_url() raises when handler is not provided."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=True)
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=None,
)
with pytest.raises(RuntimeError, match="handler is required for elicit_url"):
await ctx.elicit_url(
message="Please authorize",
url="https://example.com/oauth",
elicitation_id="oauth-123",
)
store.cleanup()
@pytest.mark.anyio
async def test_create_message_raises_without_handler() -> None:
"""Test that create_message() raises when handler is not provided."""
store = InMemoryTaskStore()
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=True)
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=None,
)
with pytest.raises(RuntimeError, match="handler is required"):
await ctx.create_message(messages=[], max_tokens=100)
store.cleanup()
@pytest.mark.anyio
async def test_elicit_queues_request_and_waits_for_response() -> None:
"""Test that elicit() queues request and waits for response."""
store = InMemoryTaskStore()
queue = InMemoryTaskMessageQueue()
handler = TaskResultHandler(store, queue)
task = await store.create_task(TaskMetadata(ttl=60000))
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=True)
mock_session._build_elicit_form_request = Mock(
return_value=JSONRPCRequest(
jsonrpc="2.0",
id="test-req-1",
method="elicitation/create",
params={"message": "Test?", "_meta": {}},
)
)
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=handler,
)
elicit_result = None
async def run_elicit() -> None:
nonlocal elicit_result
elicit_result = await ctx.elicit(
message="Test?",
requestedSchema={"type": "object"},
)
async with anyio.create_task_group() as tg:
tg.start_soon(run_elicit)
# Wait for request to be queued
await queue.wait_for_message(task.taskId)
# Verify task is in input_required status
updated_task = await store.get_task(task.taskId)
assert updated_task is not None
assert updated_task.status == "input_required"
# Dequeue and simulate response
msg = await queue.dequeue(task.taskId)
assert msg is not None
assert msg.resolver is not None
# Resolve with mock elicitation response
msg.resolver.set_result({"action": "accept", "content": {"name": "Alice"}})
# Verify result
assert elicit_result is not None
assert elicit_result.action == "accept"
assert elicit_result.content == {"name": "Alice"}
# Verify task is back to working
final_task = await store.get_task(task.taskId)
assert final_task is not None
assert final_task.status == "working"
store.cleanup()
@pytest.mark.anyio
async def test_elicit_url_queues_request_and_waits_for_response() -> None:
"""Test that elicit_url() queues request and waits for response."""
store = InMemoryTaskStore()
queue = InMemoryTaskMessageQueue()
handler = TaskResultHandler(store, queue)
task = await store.create_task(TaskMetadata(ttl=60000))
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=True)
mock_session._build_elicit_url_request = Mock(
return_value=JSONRPCRequest(
jsonrpc="2.0",
id="test-url-req-1",
method="elicitation/create",
params={"message": "Authorize", "url": "https://example.com", "elicitationId": "123", "mode": "url"},
)
)
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=handler,
)
elicit_result = None
async def run_elicit_url() -> None:
nonlocal elicit_result
elicit_result = await ctx.elicit_url(
message="Authorize",
url="https://example.com/oauth",
elicitation_id="oauth-123",
)
async with anyio.create_task_group() as tg:
tg.start_soon(run_elicit_url)
# Wait for request to be queued
await queue.wait_for_message(task.taskId)
# Verify task is in input_required status
updated_task = await store.get_task(task.taskId)
assert updated_task is not None
assert updated_task.status == "input_required"
# Dequeue and simulate response
msg = await queue.dequeue(task.taskId)
assert msg is not None
assert msg.resolver is not None
# Resolve with mock elicitation response (URL mode just returns action)
msg.resolver.set_result({"action": "accept"})
# Verify result
assert elicit_result is not None
assert elicit_result.action == "accept"
# Verify task is back to working
final_task = await store.get_task(task.taskId)
assert final_task is not None
assert final_task.status == "working"
store.cleanup()
@pytest.mark.anyio
async def test_create_message_queues_request_and_waits_for_response() -> None:
"""Test that create_message() queues request and waits for response."""
store = InMemoryTaskStore()
queue = InMemoryTaskMessageQueue()
handler = TaskResultHandler(store, queue)
task = await store.create_task(TaskMetadata(ttl=60000))
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=True)
mock_session._build_create_message_request = Mock(
return_value=JSONRPCRequest(
jsonrpc="2.0",
id="test-req-2",
method="sampling/createMessage",
params={"messages": [], "maxTokens": 100, "_meta": {}},
)
)
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=handler,
)
sampling_result = None
async def run_sampling() -> None:
nonlocal sampling_result
sampling_result = await ctx.create_message(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
max_tokens=100,
)
async with anyio.create_task_group() as tg:
tg.start_soon(run_sampling)
# Wait for request to be queued
await queue.wait_for_message(task.taskId)
# Verify task is in input_required status
updated_task = await store.get_task(task.taskId)
assert updated_task is not None
assert updated_task.status == "input_required"
# Dequeue and simulate response
msg = await queue.dequeue(task.taskId)
assert msg is not None
assert msg.resolver is not None
# Resolve with mock sampling response
msg.resolver.set_result(
{
"role": "assistant",
"content": {"type": "text", "text": "Hello back!"},
"model": "test-model",
"stopReason": "endTurn",
}
)
# Verify result
assert sampling_result is not None
assert sampling_result.role == "assistant"
assert sampling_result.model == "test-model"
# Verify task is back to working
final_task = await store.get_task(task.taskId)
assert final_task is not None
assert final_task.status == "working"
store.cleanup()
@pytest.mark.anyio
async def test_elicit_restores_status_on_cancellation() -> None:
"""Test that elicit() restores task status to working when cancelled."""
store = InMemoryTaskStore()
queue = InMemoryTaskMessageQueue()
handler = TaskResultHandler(store, queue)
task = await store.create_task(TaskMetadata(ttl=60000))
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=True)
mock_session._build_elicit_form_request = Mock(
return_value=JSONRPCRequest(
jsonrpc="2.0",
id="test-req-cancel",
method="elicitation/create",
params={"message": "Test?", "_meta": {}},
)
)
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=handler,
)
cancelled_error_raised = False
async with anyio.create_task_group() as tg:
async def do_elicit() -> None:
nonlocal cancelled_error_raised
try:
await ctx.elicit(
message="Test?",
requestedSchema={"type": "object"},
)
except anyio.get_cancelled_exc_class():
cancelled_error_raised = True
# Don't re-raise - let the test continue
tg.start_soon(do_elicit)
# Wait for request to be queued
await queue.wait_for_message(task.taskId)
# Verify task is in input_required status
updated_task = await store.get_task(task.taskId)
assert updated_task is not None
assert updated_task.status == "input_required"
# Get the queued message and set cancellation exception on its resolver
msg = await queue.dequeue(task.taskId)
assert msg is not None
assert msg.resolver is not None
# Trigger cancellation by setting exception (use asyncio.CancelledError directly)
msg.resolver.set_exception(asyncio.CancelledError())
# Verify task is back to working after cancellation
final_task = await store.get_task(task.taskId)
assert final_task is not None
assert final_task.status == "working"
assert cancelled_error_raised
store.cleanup()
@pytest.mark.anyio
async def test_create_message_restores_status_on_cancellation() -> None:
"""Test that create_message() restores task status to working when cancelled."""
store = InMemoryTaskStore()
queue = InMemoryTaskMessageQueue()
handler = TaskResultHandler(store, queue)
task = await store.create_task(TaskMetadata(ttl=60000))
mock_session = Mock()
mock_session.check_client_capability = Mock(return_value=True)
mock_session._build_create_message_request = Mock(
return_value=JSONRPCRequest(
jsonrpc="2.0",
id="test-req-cancel-2",
method="sampling/createMessage",
params={"messages": [], "maxTokens": 100, "_meta": {}},
)
)
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=handler,
)
cancelled_error_raised = False
async with anyio.create_task_group() as tg:
async def do_sampling() -> None:
nonlocal cancelled_error_raised
try:
await ctx.create_message(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
max_tokens=100,
)
except anyio.get_cancelled_exc_class():
cancelled_error_raised = True
# Don't re-raise
tg.start_soon(do_sampling)
# Wait for request to be queued
await queue.wait_for_message(task.taskId)
# Verify task is in input_required status
updated_task = await store.get_task(task.taskId)
assert updated_task is not None
assert updated_task.status == "input_required"
# Get the queued message and set cancellation exception on its resolver
msg = await queue.dequeue(task.taskId)
assert msg is not None
assert msg.resolver is not None
# Trigger cancellation by setting exception (use asyncio.CancelledError directly)
msg.resolver.set_exception(asyncio.CancelledError())
# Verify task is back to working after cancellation
final_task = await store.get_task(task.taskId)
assert final_task is not None
assert final_task.status == "working"
assert cancelled_error_raised
store.cleanup()
@pytest.mark.anyio
async def test_elicit_as_task_raises_without_handler() -> None:
"""Test that elicit_as_task() raises when handler is not provided."""
store = InMemoryTaskStore()
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
# Create mock session with proper client capabilities
mock_session = Mock()
mock_session.client_params = InitializeRequestParams(
protocolVersion="2025-01-01",
capabilities=ClientCapabilities(
tasks=ClientTasksCapability(
requests=ClientTasksRequestsCapability(
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
)
)
),
clientInfo=Implementation(name="test", version="1.0"),
)
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=None,
)
with pytest.raises(RuntimeError, match="handler is required for elicit_as_task"):
await ctx.elicit_as_task(message="Test?", requestedSchema={"type": "object"})
store.cleanup()
@pytest.mark.anyio
async def test_create_message_as_task_raises_without_handler() -> None:
"""Test that create_message_as_task() raises when handler is not provided."""
store = InMemoryTaskStore()
queue = InMemoryTaskMessageQueue()
task = await store.create_task(TaskMetadata(ttl=60000))
# Create mock session with proper client capabilities
mock_session = Mock()
mock_session.client_params = InitializeRequestParams(
protocolVersion="2025-01-01",
capabilities=ClientCapabilities(
tasks=ClientTasksCapability(
requests=ClientTasksRequestsCapability(
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
)
)
),
clientInfo=Implementation(name="test", version="1.0"),
)
ctx = ServerTaskContext(
task=task,
store=store,
session=mock_session,
queue=queue,
handler=None,
)
with pytest.raises(RuntimeError, match="handler is required for create_message_as_task"):
await ctx.create_message_as_task(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
max_tokens=100,
)
store.cleanup()

406
tests/experimental/tasks/server/test_store.py Normal file
View file

@ -0,0 +1,406 @@
"""Tests for InMemoryTaskStore."""
from collections.abc import AsyncIterator
from datetime import datetime, timedelta, timezone
import pytest
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.helpers import cancel_task
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.types import INVALID_PARAMS, CallToolResult, TaskMetadata, TextContent
@pytest.fixture
async def store() -> AsyncIterator[InMemoryTaskStore]:
"""Provide a clean InMemoryTaskStore for each test with automatic cleanup."""
store = InMemoryTaskStore()
yield store
store.cleanup()
@pytest.mark.anyio
async def test_create_and_get(store: InMemoryTaskStore) -> None:
"""Test InMemoryTaskStore create and get operations."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
assert task.taskId is not None
assert task.status == "working"
assert task.ttl == 60000
retrieved = await store.get_task(task.taskId)
assert retrieved is not None
assert retrieved.taskId == task.taskId
assert retrieved.status == "working"
@pytest.mark.anyio
async def test_create_with_custom_id(store: InMemoryTaskStore) -> None:
"""Test InMemoryTaskStore create with custom task ID."""
task = await store.create_task(
metadata=TaskMetadata(ttl=60000),
task_id="my-custom-id",
)
assert task.taskId == "my-custom-id"
assert task.status == "working"
retrieved = await store.get_task("my-custom-id")
assert retrieved is not None
assert retrieved.taskId == "my-custom-id"
@pytest.mark.anyio
async def test_create_duplicate_id_raises(store: InMemoryTaskStore) -> None:
"""Test that creating a task with duplicate ID raises."""
await store.create_task(metadata=TaskMetadata(ttl=60000), task_id="duplicate")
with pytest.raises(ValueError, match="already exists"):
await store.create_task(metadata=TaskMetadata(ttl=60000), task_id="duplicate")
@pytest.mark.anyio
async def test_get_nonexistent_returns_none(store: InMemoryTaskStore) -> None:
"""Test that getting a nonexistent task returns None."""
retrieved = await store.get_task("nonexistent")
assert retrieved is None
@pytest.mark.anyio
async def test_update_status(store: InMemoryTaskStore) -> None:
"""Test InMemoryTaskStore status updates."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
updated = await store.update_task(task.taskId, status="completed", status_message="All done!")
assert updated.status == "completed"
assert updated.statusMessage == "All done!"
retrieved = await store.get_task(task.taskId)
assert retrieved is not None
assert retrieved.status == "completed"
assert retrieved.statusMessage == "All done!"
@pytest.mark.anyio
async def test_update_nonexistent_raises(store: InMemoryTaskStore) -> None:
"""Test that updating a nonexistent task raises."""
with pytest.raises(ValueError, match="not found"):
await store.update_task("nonexistent", status="completed")
@pytest.mark.anyio
async def test_store_and_get_result(store: InMemoryTaskStore) -> None:
"""Test InMemoryTaskStore result storage and retrieval."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
# Store result
result = CallToolResult(content=[TextContent(type="text", text="Result data")])
await store.store_result(task.taskId, result)
# Retrieve result
retrieved_result = await store.get_result(task.taskId)
assert retrieved_result == result
@pytest.mark.anyio
async def test_get_result_nonexistent_returns_none(store: InMemoryTaskStore) -> None:
"""Test that getting result for nonexistent task returns None."""
result = await store.get_result("nonexistent")
assert result is None
@pytest.mark.anyio
async def test_get_result_no_result_returns_none(store: InMemoryTaskStore) -> None:
"""Test that getting result when none stored returns None."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
result = await store.get_result(task.taskId)
assert result is None
@pytest.mark.anyio
async def test_list_tasks(store: InMemoryTaskStore) -> None:
"""Test InMemoryTaskStore list operation."""
# Create multiple tasks
for _ in range(3):
await store.create_task(metadata=TaskMetadata(ttl=60000))
tasks, next_cursor = await store.list_tasks()
assert len(tasks) == 3
assert next_cursor is None # Less than page size
@pytest.mark.anyio
async def test_list_tasks_pagination() -> None:
"""Test InMemoryTaskStore pagination."""
# Needs custom page_size, can't use fixture
store = InMemoryTaskStore(page_size=2)
# Create 5 tasks
for _ in range(5):
await store.create_task(metadata=TaskMetadata(ttl=60000))
# First page
tasks, next_cursor = await store.list_tasks()
assert len(tasks) == 2
assert next_cursor is not None
# Second page
tasks, next_cursor = await store.list_tasks(cursor=next_cursor)
assert len(tasks) == 2
assert next_cursor is not None
# Third page (last)
tasks, next_cursor = await store.list_tasks(cursor=next_cursor)
assert len(tasks) == 1
assert next_cursor is None
store.cleanup()
@pytest.mark.anyio
async def test_list_tasks_invalid_cursor(store: InMemoryTaskStore) -> None:
"""Test that invalid cursor raises."""
await store.create_task(metadata=TaskMetadata(ttl=60000))
with pytest.raises(ValueError, match="Invalid cursor"):
await store.list_tasks(cursor="invalid-cursor")
@pytest.mark.anyio
async def test_delete_task(store: InMemoryTaskStore) -> None:
"""Test InMemoryTaskStore delete operation."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
deleted = await store.delete_task(task.taskId)
assert deleted is True
retrieved = await store.get_task(task.taskId)
assert retrieved is None
# Delete non-existent
deleted = await store.delete_task(task.taskId)
assert deleted is False
@pytest.mark.anyio
async def test_get_all_tasks_helper(store: InMemoryTaskStore) -> None:
"""Test the get_all_tasks debugging helper."""
await store.create_task(metadata=TaskMetadata(ttl=60000))
await store.create_task(metadata=TaskMetadata(ttl=60000))
all_tasks = store.get_all_tasks()
assert len(all_tasks) == 2
@pytest.mark.anyio
async def test_store_result_nonexistent_raises(store: InMemoryTaskStore) -> None:
"""Test that storing result for nonexistent task raises ValueError."""
result = CallToolResult(content=[TextContent(type="text", text="Result")])
with pytest.raises(ValueError, match="not found"):
await store.store_result("nonexistent-id", result)
@pytest.mark.anyio
async def test_create_task_with_null_ttl(store: InMemoryTaskStore) -> None:
"""Test creating task with null TTL (never expires)."""
task = await store.create_task(metadata=TaskMetadata(ttl=None))
assert task.ttl is None
# Task should persist (not expire)
retrieved = await store.get_task(task.taskId)
assert retrieved is not None
@pytest.mark.anyio
async def test_task_expiration_cleanup(store: InMemoryTaskStore) -> None:
"""Test that expired tasks are cleaned up lazily."""
# Create a task with very short TTL
task = await store.create_task(metadata=TaskMetadata(ttl=1)) # 1ms TTL
# Manually force the expiry to be in the past
stored = store._tasks.get(task.taskId)
assert stored is not None
stored.expires_at = datetime.now(timezone.utc) - timedelta(seconds=10)
# Task should still exist in internal dict but be expired
assert task.taskId in store._tasks
# Any access operation should clean up expired tasks
# list_tasks triggers cleanup
tasks, _ = await store.list_tasks()
# Expired task should be cleaned up
assert task.taskId not in store._tasks
assert len(tasks) == 0
@pytest.mark.anyio
async def test_task_with_null_ttl_never_expires(store: InMemoryTaskStore) -> None:
"""Test that tasks with null TTL never expire during cleanup."""
# Create task with null TTL
task = await store.create_task(metadata=TaskMetadata(ttl=None))
# Verify internal storage has no expiry
stored = store._tasks.get(task.taskId)
assert stored is not None
assert stored.expires_at is None
# Access operations should NOT remove this task
await store.list_tasks()
await store.get_task(task.taskId)
# Task should still exist
assert task.taskId in store._tasks
retrieved = await store.get_task(task.taskId)
assert retrieved is not None
@pytest.mark.anyio
async def test_terminal_task_ttl_reset(store: InMemoryTaskStore) -> None:
"""Test that TTL is reset when task enters terminal state."""
# Create task with short TTL
task = await store.create_task(metadata=TaskMetadata(ttl=60000)) # 60s
# Get the initial expiry
stored = store._tasks.get(task.taskId)
assert stored is not None
initial_expiry = stored.expires_at
assert initial_expiry is not None
# Update to terminal state (completed)
await store.update_task(task.taskId, status="completed")
# Expiry should be reset to a new time (from now + TTL)
new_expiry = stored.expires_at
assert new_expiry is not None
assert new_expiry >= initial_expiry
@pytest.mark.anyio
async def test_terminal_status_transition_rejected(store: InMemoryTaskStore) -> None:
"""Test that transitions from terminal states are rejected.
Per spec: Terminal states (completed, failed, cancelled) MUST NOT
transition to any other status.
"""
# Test each terminal status
for terminal_status in ("completed", "failed", "cancelled"):
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
# Move to terminal state
await store.update_task(task.taskId, status=terminal_status)
# Attempting to transition to any other status should raise
with pytest.raises(ValueError, match="Cannot transition from terminal status"):
await store.update_task(task.taskId, status="working")
# Also test transitioning to another terminal state
other_terminal = "failed" if terminal_status != "failed" else "completed"
with pytest.raises(ValueError, match="Cannot transition from terminal status"):
await store.update_task(task.taskId, status=other_terminal)
@pytest.mark.anyio
async def test_terminal_status_allows_same_status(store: InMemoryTaskStore) -> None:
"""Test that setting the same terminal status doesn't raise.
This is not a transition, so it should be allowed (no-op).
"""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
await store.update_task(task.taskId, status="completed")
# Setting the same status should not raise
updated = await store.update_task(task.taskId, status="completed")
assert updated.status == "completed"
# Updating just the message should also work
updated = await store.update_task(task.taskId, status_message="Updated message")
assert updated.statusMessage == "Updated message"
@pytest.mark.anyio
async def test_wait_for_update_nonexistent_raises(store: InMemoryTaskStore) -> None:
"""Test that wait_for_update raises for nonexistent task."""
with pytest.raises(ValueError, match="not found"):
await store.wait_for_update("nonexistent-task-id")
@pytest.mark.anyio
async def test_cancel_task_succeeds_for_working_task(store: InMemoryTaskStore) -> None:
"""Test cancel_task helper succeeds for a working task."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
assert task.status == "working"
result = await cancel_task(store, task.taskId)
assert result.taskId == task.taskId
assert result.status == "cancelled"
# Verify store is updated
retrieved = await store.get_task(task.taskId)
assert retrieved is not None
assert retrieved.status == "cancelled"
@pytest.mark.anyio
async def test_cancel_task_rejects_nonexistent_task(store: InMemoryTaskStore) -> None:
"""Test cancel_task raises McpError with INVALID_PARAMS for nonexistent task."""
with pytest.raises(McpError) as exc_info:
await cancel_task(store, "nonexistent-task-id")
assert exc_info.value.error.code == INVALID_PARAMS
assert "not found" in exc_info.value.error.message
@pytest.mark.anyio
async def test_cancel_task_rejects_completed_task(store: InMemoryTaskStore) -> None:
"""Test cancel_task raises McpError with INVALID_PARAMS for completed task."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
await store.update_task(task.taskId, status="completed")
with pytest.raises(McpError) as exc_info:
await cancel_task(store, task.taskId)
assert exc_info.value.error.code == INVALID_PARAMS
assert "terminal state 'completed'" in exc_info.value.error.message
@pytest.mark.anyio
async def test_cancel_task_rejects_failed_task(store: InMemoryTaskStore) -> None:
"""Test cancel_task raises McpError with INVALID_PARAMS for failed task."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
await store.update_task(task.taskId, status="failed")
with pytest.raises(McpError) as exc_info:
await cancel_task(store, task.taskId)
assert exc_info.value.error.code == INVALID_PARAMS
assert "terminal state 'failed'" in exc_info.value.error.message
@pytest.mark.anyio
async def test_cancel_task_rejects_already_cancelled_task(store: InMemoryTaskStore) -> None:
"""Test cancel_task raises McpError with INVALID_PARAMS for already cancelled task."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
await store.update_task(task.taskId, status="cancelled")
with pytest.raises(McpError) as exc_info:
await cancel_task(store, task.taskId)
assert exc_info.value.error.code == INVALID_PARAMS
assert "terminal state 'cancelled'" in exc_info.value.error.message
@pytest.mark.anyio
async def test_cancel_task_succeeds_for_input_required_task(store: InMemoryTaskStore) -> None:
"""Test cancel_task helper succeeds for a task in input_required status."""
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
await store.update_task(task.taskId, status="input_required")
result = await cancel_task(store, task.taskId)
assert result.taskId == task.taskId
assert result.status == "cancelled"

354
tests/experimental/tasks/server/test_task_result_handler.py Normal file
View file

@ -0,0 +1,354 @@
"""Tests for TaskResultHandler."""
from collections.abc import AsyncIterator
from typing import Any
from unittest.mock import AsyncMock, Mock
import anyio
import pytest
from mcp.server.experimental.task_result_handler import TaskResultHandler
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, QueuedMessage
from mcp.shared.experimental.tasks.resolver import Resolver
from mcp.shared.message import SessionMessage
from mcp.types import (
INVALID_REQUEST,
CallToolResult,
ErrorData,
GetTaskPayloadRequest,
GetTaskPayloadRequestParams,
GetTaskPayloadResult,
JSONRPCRequest,
TaskMetadata,
TextContent,
)
@pytest.fixture
async def store() -> AsyncIterator[InMemoryTaskStore]:
"""Provide a clean store for each test."""
s = InMemoryTaskStore()
yield s
s.cleanup()
@pytest.fixture
def queue() -> InMemoryTaskMessageQueue:
"""Provide a clean queue for each test."""
return InMemoryTaskMessageQueue()
@pytest.fixture
def handler(store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue) -> TaskResultHandler:
"""Provide a handler for each test."""
return TaskResultHandler(store, queue)
@pytest.mark.anyio
async def test_handle_returns_result_for_completed_task(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that handle() returns the stored result for a completed task."""
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
result = CallToolResult(content=[TextContent(type="text", text="Done!")])
await store.store_result(task.taskId, result)
await store.update_task(task.taskId, status="completed")
mock_session = Mock()
mock_session.send_message = AsyncMock()
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
response = await handler.handle(request, mock_session, "req-1")
assert response is not None
assert response.meta is not None
assert "io.modelcontextprotocol/related-task" in response.meta
@pytest.mark.anyio
async def test_handle_raises_for_nonexistent_task(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that handle() raises McpError for nonexistent task."""
mock_session = Mock()
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId="nonexistent"))
with pytest.raises(McpError) as exc_info:
await handler.handle(request, mock_session, "req-1")
assert "not found" in exc_info.value.error.message
@pytest.mark.anyio
async def test_handle_returns_empty_result_when_no_result_stored(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that handle() returns minimal result when task completed without stored result."""
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
await store.update_task(task.taskId, status="completed")
mock_session = Mock()
mock_session.send_message = AsyncMock()
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
response = await handler.handle(request, mock_session, "req-1")
assert response is not None
assert response.meta is not None
assert "io.modelcontextprotocol/related-task" in response.meta
@pytest.mark.anyio
async def test_handle_delivers_queued_messages(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that handle() delivers queued messages before returning."""
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
queued_msg = QueuedMessage(
type="notification",
message=JSONRPCRequest(
jsonrpc="2.0",
id="notif-1",
method="test/notification",
params={},
),
)
await queue.enqueue(task.taskId, queued_msg)
await store.update_task(task.taskId, status="completed")
sent_messages: list[SessionMessage] = []
async def track_send(msg: SessionMessage) -> None:
sent_messages.append(msg)
mock_session = Mock()
mock_session.send_message = track_send
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
await handler.handle(request, mock_session, "req-1")
assert len(sent_messages) == 1
@pytest.mark.anyio
async def test_handle_waits_for_task_completion(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that handle() waits for task to complete before returning."""
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
mock_session = Mock()
mock_session.send_message = AsyncMock()
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
result_holder: list[GetTaskPayloadResult | None] = [None]
async def run_handle() -> None:
result_holder[0] = await handler.handle(request, mock_session, "req-1")
async with anyio.create_task_group() as tg:
tg.start_soon(run_handle)
# Wait for handler to start waiting (event gets created when wait starts)
while task.taskId not in store._update_events:
await anyio.sleep(0)
await store.store_result(task.taskId, CallToolResult(content=[TextContent(type="text", text="Done")]))
await store.update_task(task.taskId, status="completed")
assert result_holder[0] is not None
@pytest.mark.anyio
async def test_route_response_resolves_pending_request(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that route_response() resolves a pending request."""
resolver: Resolver[dict[str, Any]] = Resolver()
handler._pending_requests["req-123"] = resolver
result = handler.route_response("req-123", {"status": "ok"})
assert result is True
assert resolver.done()
assert await resolver.wait() == {"status": "ok"}
@pytest.mark.anyio
async def test_route_response_returns_false_for_unknown_request(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that route_response() returns False for unknown request ID."""
result = handler.route_response("unknown-req", {"status": "ok"})
assert result is False
@pytest.mark.anyio
async def test_route_response_returns_false_for_already_done_resolver(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that route_response() returns False if resolver already completed."""
resolver: Resolver[dict[str, Any]] = Resolver()
resolver.set_result({"already": "done"})
handler._pending_requests["req-123"] = resolver
result = handler.route_response("req-123", {"new": "data"})
assert result is False
@pytest.mark.anyio
async def test_route_error_resolves_pending_request_with_exception(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that route_error() sets exception on pending request."""
resolver: Resolver[dict[str, Any]] = Resolver()
handler._pending_requests["req-123"] = resolver
error = ErrorData(code=INVALID_REQUEST, message="Something went wrong")
result = handler.route_error("req-123", error)
assert result is True
assert resolver.done()
with pytest.raises(McpError) as exc_info:
await resolver.wait()
assert exc_info.value.error.message == "Something went wrong"
@pytest.mark.anyio
async def test_route_error_returns_false_for_unknown_request(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that route_error() returns False for unknown request ID."""
error = ErrorData(code=INVALID_REQUEST, message="Error")
result = handler.route_error("unknown-req", error)
assert result is False
@pytest.mark.anyio
async def test_deliver_registers_resolver_for_request_messages(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that _deliver_queued_messages registers resolvers for request messages."""
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
resolver: Resolver[dict[str, Any]] = Resolver()
queued_msg = QueuedMessage(
type="request",
message=JSONRPCRequest(
jsonrpc="2.0",
id="inner-req-1",
method="elicitation/create",
params={},
),
resolver=resolver,
original_request_id="inner-req-1",
)
await queue.enqueue(task.taskId, queued_msg)
mock_session = Mock()
mock_session.send_message = AsyncMock()
await handler._deliver_queued_messages(task.taskId, mock_session, "outer-req-1")
assert "inner-req-1" in handler._pending_requests
assert handler._pending_requests["inner-req-1"] is resolver
@pytest.mark.anyio
async def test_deliver_skips_resolver_registration_when_no_original_id(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that _deliver_queued_messages skips resolver registration when original_request_id is None."""
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
resolver: Resolver[dict[str, Any]] = Resolver()
queued_msg = QueuedMessage(
type="request",
message=JSONRPCRequest(
jsonrpc="2.0",
id="inner-req-1",
method="elicitation/create",
params={},
),
resolver=resolver,
original_request_id=None, # No original request ID
)
await queue.enqueue(task.taskId, queued_msg)
mock_session = Mock()
mock_session.send_message = AsyncMock()
await handler._deliver_queued_messages(task.taskId, mock_session, "outer-req-1")
# Resolver should NOT be registered since original_request_id is None
assert len(handler._pending_requests) == 0
# But the message should still be sent
mock_session.send_message.assert_called_once()
@pytest.mark.anyio
async def test_wait_for_task_update_handles_store_exception(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that _wait_for_task_update handles store exception gracefully."""
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
# Make wait_for_update raise an exception
async def failing_wait(task_id: str) -> None:
raise RuntimeError("Store error")
store.wait_for_update = failing_wait # type: ignore[method-assign]
# Queue a message to unblock the race via the queue path
async def enqueue_later() -> None:
# Wait for queue to start waiting (event gets created when wait starts)
while task.taskId not in queue._events:
await anyio.sleep(0)
await queue.enqueue(
task.taskId,
QueuedMessage(
type="notification",
message=JSONRPCRequest(
jsonrpc="2.0",
id="notif-1",
method="test/notification",
params={},
),
),
)
async with anyio.create_task_group() as tg:
tg.start_soon(enqueue_later)
# This should complete via the queue path even though store raises
await handler._wait_for_task_update(task.taskId)
@pytest.mark.anyio
async def test_wait_for_task_update_handles_queue_exception(
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
) -> None:
"""Test that _wait_for_task_update handles queue exception gracefully."""
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
# Make wait_for_message raise an exception
async def failing_wait(task_id: str) -> None:
raise RuntimeError("Queue error")
queue.wait_for_message = failing_wait # type: ignore[method-assign]
# Update the store to unblock the race via the store path
async def update_later() -> None:
# Wait for store to start waiting (event gets created when wait starts)
while task.taskId not in store._update_events:
await anyio.sleep(0)
await store.update_task(task.taskId, status="completed")
async with anyio.create_task_group() as tg:
tg.start_soon(update_later)
# This should complete via the store path even though queue raises
await handler._wait_for_task_update(task.taskId)