Max Isbey
543 lines
20 KiB
Python
543 lines
20 KiB
Python
"""
|
|
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
|
|
# Verify nested call capability is present
|
|
assert caps_after.tasks.requests is not None
|
|
assert caps_after.tasks.requests.tools is not None
|
|
assert caps_after.tasks.requests.tools.call is not None
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_enable_tasks_with_custom_store_and_queue() -> None:
|
|
"""Test that enable_tasks() uses provided store and queue instead of defaults."""
|
|
server = Server("test-custom-store-queue")
|
|
|
|
# Create custom store and queue
|
|
custom_store = InMemoryTaskStore()
|
|
custom_queue = InMemoryTaskMessageQueue()
|
|
|
|
# Enable tasks with custom implementations
|
|
task_support = server.experimental.enable_tasks(store=custom_store, queue=custom_queue)
|
|
|
|
# Verify our custom implementations are used
|
|
assert task_support.store is custom_store
|
|
assert task_support.queue is custom_queue
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_enable_tasks_skips_default_handlers_when_custom_registered() -> None:
|
|
"""Test that enable_tasks() doesn't override already-registered handlers."""
|
|
server = Server("test-custom-handlers")
|
|
|
|
# Register custom handlers BEFORE enable_tasks (never called, just for registration)
|
|
@server.experimental.get_task()
|
|
async def custom_get_task(req: GetTaskRequest) -> GetTaskResult:
|
|
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)
|