Max Isbey
368 lines
13 KiB
Python
368 lines
13 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 unittest.mock import Mock
|
|
|
|
import anyio
|
|
import pytest
|
|
from anyio import Event
|
|
|
|
from mcp import Client
|
|
from mcp.server import Server, ServerRequestContext
|
|
from mcp.server.experimental.request_context import Experimental
|
|
from mcp.server.experimental.task_context import ServerTaskContext
|
|
from mcp.server.experimental.task_support import TaskSupport
|
|
from mcp.server.lowlevel import NotificationOptions
|
|
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
|
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
|
|
from mcp.types import (
|
|
TASK_REQUIRED,
|
|
CallToolRequestParams,
|
|
CallToolResult,
|
|
CreateTaskResult,
|
|
GetTaskRequestParams,
|
|
GetTaskResult,
|
|
ListToolsResult,
|
|
PaginatedRequestParams,
|
|
TextContent,
|
|
)
|
|
|
|
pytestmark = pytest.mark.anyio
|
|
|
|
|
|
async def _handle_list_tools_simple_task(
|
|
ctx: ServerRequestContext, params: PaginatedRequestParams | None
|
|
) -> ListToolsResult:
|
|
raise NotImplementedError
|
|
|
|
|
|
async def test_run_task_basic_flow() -> None:
|
|
"""Test the basic run_task flow without elicitation."""
|
|
work_completed = Event()
|
|
received_meta: list[str | None] = [None]
|
|
|
|
async def handle_call_tool(
|
|
ctx: ServerRequestContext, params: CallToolRequestParams
|
|
) -> CallToolResult | CreateTaskResult:
|
|
ctx.experimental.validate_task_mode(TASK_REQUIRED)
|
|
|
|
if ctx.meta is not None: # pragma: no branch
|
|
received_meta[0] = ctx.meta.get("custom_field")
|
|
|
|
async def work(task: ServerTaskContext) -> CallToolResult:
|
|
await task.update_status("Working...")
|
|
input_val = (params.arguments or {}).get("input", "default")
|
|
result = CallToolResult(content=[TextContent(type="text", text=f"Processed: {input_val}")])
|
|
work_completed.set()
|
|
return result
|
|
|
|
return await ctx.experimental.run_task(work)
|
|
|
|
server = Server(
|
|
"test-run-task",
|
|
on_list_tools=_handle_list_tools_simple_task,
|
|
on_call_tool=handle_call_tool,
|
|
)
|
|
server.experimental.enable_tasks()
|
|
|
|
async with Client(server) as client:
|
|
result = await client.session.experimental.call_tool_as_task(
|
|
"simple_task",
|
|
{"input": "hello"},
|
|
meta={"custom_field": "test_value"},
|
|
)
|
|
|
|
task_id = result.task.task_id
|
|
assert result.task.status == "working"
|
|
|
|
with anyio.fail_after(5):
|
|
await work_completed.wait()
|
|
|
|
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
|
|
|
|
assert received_meta[0] == "test_value"
|
|
|
|
|
|
async def test_run_task_auto_fails_on_exception() -> None:
|
|
"""Test that run_task automatically fails the task when work raises."""
|
|
work_failed = Event()
|
|
|
|
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
|
|
raise NotImplementedError
|
|
|
|
async def handle_call_tool(
|
|
ctx: ServerRequestContext, params: CallToolRequestParams
|
|
) -> CallToolResult | CreateTaskResult:
|
|
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 = Server(
|
|
"test-run-task-fail",
|
|
on_list_tools=handle_list_tools,
|
|
on_call_tool=handle_call_tool,
|
|
)
|
|
server.experimental.enable_tasks()
|
|
|
|
async with Client(server) as client:
|
|
result = await client.session.experimental.call_tool_as_task("failing_task", {})
|
|
task_id = result.task.task_id
|
|
|
|
with anyio.fail_after(5):
|
|
await work_failed.wait()
|
|
|
|
with anyio.fail_after(5):
|
|
while True:
|
|
task_status = await client.session.experimental.get_task(task_id)
|
|
if task_status.status == "failed": # pragma: no branch
|
|
break
|
|
|
|
assert "Something went wrong" in (task_status.status_message or "")
|
|
|
|
|
|
async 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
|
|
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
|
|
|
|
|
|
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")
|
|
|
|
custom_store = InMemoryTaskStore()
|
|
custom_queue = InMemoryTaskMessageQueue()
|
|
|
|
task_support = server.experimental.enable_tasks(store=custom_store, queue=custom_queue)
|
|
|
|
assert task_support.store is custom_store
|
|
assert task_support.queue is custom_queue
|
|
|
|
|
|
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 via enable_tasks kwargs
|
|
async def custom_get_task(ctx: ServerRequestContext, params: GetTaskRequestParams) -> GetTaskResult:
|
|
raise NotImplementedError
|
|
|
|
server.experimental.enable_tasks(on_get_task=custom_get_task)
|
|
|
|
# Verify handler is registered
|
|
assert server._has_handler("tasks/get")
|
|
assert server._has_handler("tasks/list")
|
|
assert server._has_handler("tasks/cancel")
|
|
assert server._has_handler("tasks/result")
|
|
|
|
|
|
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)
|
|
|
|
|
|
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
|
|
|
|
|
|
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)
|
|
|
|
|
|
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)
|
|
|
|
|
|
async def test_run_task_with_model_immediate_response() -> None:
|
|
"""Test that run_task includes model_immediate_response in CreateTaskResult._meta."""
|
|
work_completed = Event()
|
|
immediate_response_text = "Processing your request..."
|
|
|
|
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
|
|
raise NotImplementedError
|
|
|
|
async def handle_call_tool(
|
|
ctx: ServerRequestContext, params: CallToolRequestParams
|
|
) -> CallToolResult | CreateTaskResult:
|
|
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 = Server(
|
|
"test-run-task-immediate",
|
|
on_list_tools=handle_list_tools,
|
|
on_call_tool=handle_call_tool,
|
|
)
|
|
server.experimental.enable_tasks()
|
|
|
|
async with Client(server) as client:
|
|
result = await client.session.experimental.call_tool_as_task("task_with_immediate", {})
|
|
|
|
assert result.meta is not None
|
|
assert "io.modelcontextprotocol/model-immediate-response" in result.meta
|
|
assert result.meta["io.modelcontextprotocol/model-immediate-response"] == immediate_response_text
|
|
|
|
with anyio.fail_after(5):
|
|
await work_completed.wait()
|
|
|
|
|
|
async def test_run_task_doesnt_complete_if_already_terminal() -> None:
|
|
"""Test that run_task doesn't auto-complete if work manually completed the task."""
|
|
work_completed = Event()
|
|
|
|
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
|
|
raise NotImplementedError
|
|
|
|
async def handle_call_tool(
|
|
ctx: ServerRequestContext, params: CallToolRequestParams
|
|
) -> CallToolResult | CreateTaskResult:
|
|
ctx.experimental.validate_task_mode(TASK_REQUIRED)
|
|
|
|
async def work(task: ServerTaskContext) -> CallToolResult:
|
|
manual_result = CallToolResult(content=[TextContent(type="text", text="Manually completed")])
|
|
await task.complete(manual_result, notify=False)
|
|
work_completed.set()
|
|
return CallToolResult(content=[TextContent(type="text", text="This should be ignored")])
|
|
|
|
return await ctx.experimental.run_task(work)
|
|
|
|
server = Server(
|
|
"test-already-complete",
|
|
on_list_tools=handle_list_tools,
|
|
on_call_tool=handle_call_tool,
|
|
)
|
|
server.experimental.enable_tasks()
|
|
|
|
async with Client(server) as client:
|
|
result = await client.session.experimental.call_tool_as_task("manual_complete_task", {})
|
|
task_id = result.task.task_id
|
|
|
|
with anyio.fail_after(5):
|
|
await work_completed.wait()
|
|
|
|
with anyio.fail_after(5):
|
|
while True:
|
|
status = await client.session.experimental.get_task(task_id)
|
|
if status.status == "completed": # pragma: no branch
|
|
break
|
|
|
|
|
|
async def test_run_task_doesnt_fail_if_already_terminal() -> None:
|
|
"""Test that run_task doesn't auto-fail if work manually failed/cancelled the task."""
|
|
work_completed = Event()
|
|
|
|
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
|
|
raise NotImplementedError
|
|
|
|
async def handle_call_tool(
|
|
ctx: ServerRequestContext, params: CallToolRequestParams
|
|
) -> CallToolResult | CreateTaskResult:
|
|
ctx.experimental.validate_task_mode(TASK_REQUIRED)
|
|
|
|
async def work(task: ServerTaskContext) -> CallToolResult:
|
|
await task.fail("Manually failed", notify=False)
|
|
work_completed.set()
|
|
raise RuntimeError("This error should not change status")
|
|
|
|
return await ctx.experimental.run_task(work)
|
|
|
|
server = Server(
|
|
"test-already-failed",
|
|
on_list_tools=handle_list_tools,
|
|
on_call_tool=handle_call_tool,
|
|
)
|
|
server.experimental.enable_tasks()
|
|
|
|
async with Client(server) as client:
|
|
result = await client.session.experimental.call_tool_as_task("manual_cancel_task", {})
|
|
task_id = result.task.task_id
|
|
|
|
with anyio.fail_after(5):
|
|
await work_completed.wait()
|
|
|
|
with anyio.fail_after(5):
|
|
while True:
|
|
status = await client.session.experimental.get_task(task_id)
|
|
if status.status == "failed": # pragma: no branch
|
|
break
|
|
|
|
assert status.status_message == "Manually failed"
|