Max Isbey
358 lines
14 KiB
Python
358 lines
14 KiB
Python
"""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()
|