diff --git a/README.md b/README.md index ca0655f579..bb20a19d13 100644 --- a/README.md +++ b/README.md @@ -2512,6 +2512,7 @@ MCP servers declare capabilities during initialization: ## Documentation - [API Reference](https://modelcontextprotocol.github.io/python-sdk/api/) +- [Experimental Features (Tasks)](https://modelcontextprotocol.github.io/python-sdk/experimental/tasks/) - [Model Context Protocol documentation](https://modelcontextprotocol.io) - [Model Context Protocol specification](https://modelcontextprotocol.io/specification/latest) - [Officially supported servers](https://github.com/modelcontextprotocol/servers) diff --git a/docs/experimental/index.md b/docs/experimental/index.md new file mode 100644 index 0000000000..1d496b3f10 --- /dev/null +++ b/docs/experimental/index.md @@ -0,0 +1,43 @@ +# Experimental Features + +!!! warning "Experimental APIs" + + The features in this section are experimental and may change without notice. + They track the evolving MCP specification and are not yet stable. + +This section documents experimental features in the MCP Python SDK. These features +implement draft specifications that are still being refined. + +## Available Experimental Features + +### [Tasks](tasks.md) + +Tasks enable asynchronous execution of MCP operations. Instead of waiting for a +long-running operation to complete, the server returns a task reference immediately. +Clients can then poll for status updates and retrieve results when ready. + +Tasks are useful for: + +- **Long-running computations** that would otherwise block +- **Batch operations** that process many items +- **Interactive workflows** that require user input (elicitation) or LLM assistance (sampling) + +## Using Experimental APIs + +Experimental features are accessed via the `.experimental` property: + +```python +# Server-side +@server.experimental.get_task() +async def handle_get_task(request: GetTaskRequest) -> GetTaskResult: + ... + +# Client-side +result = await session.experimental.call_tool_as_task("tool_name", {"arg": "value"}) +``` + +## Providing Feedback + +Since these features are experimental, feedback is especially valuable. If you encounter +issues or have suggestions, please open an issue on the +[python-sdk repository](https://github.com/modelcontextprotocol/python-sdk/issues). diff --git a/docs/experimental/tasks-client.md b/docs/experimental/tasks-client.md new file mode 100644 index 0000000000..cfd23e4e14 --- /dev/null +++ b/docs/experimental/tasks-client.md @@ -0,0 +1,361 @@ +# Client Task Usage + +!!! warning "Experimental" + + Tasks are an experimental feature. The API may change without notice. + +This guide covers calling task-augmented tools from clients, handling the `input_required` status, and advanced patterns like receiving task requests from servers. + +## Quick Start + +Call a tool as a task and poll for the result: + +```python +from mcp.client.session import ClientSession +from mcp.types import CallToolResult + +async with ClientSession(read, write) as session: + await session.initialize() + + # Call tool as task + result = await session.experimental.call_tool_as_task( + "process_data", + {"input": "hello"}, + ttl=60000, + ) + task_id = result.task.taskId + + # Poll until complete + async for status in session.experimental.poll_task(task_id): + print(f"Status: {status.status} - {status.statusMessage or ''}") + + # Get result + final = await session.experimental.get_task_result(task_id, CallToolResult) + print(f"Result: {final.content[0].text}") +``` + +## Calling Tools as Tasks + +Use `call_tool_as_task()` to invoke a tool with task augmentation: + +```python +result = await session.experimental.call_tool_as_task( + "my_tool", # Tool name + {"arg": "value"}, # Arguments + ttl=60000, # Time-to-live in milliseconds + meta={"key": "val"}, # Optional metadata +) + +task_id = result.task.taskId +print(f"Task: {task_id}, Status: {result.task.status}") +``` + +The response is a `CreateTaskResult` containing: + +- `task.taskId` - Unique identifier for polling +- `task.status` - Initial status (usually `"working"`) +- `task.pollInterval` - Suggested polling interval (milliseconds) +- `task.ttl` - Time-to-live for results +- `task.createdAt` - Creation timestamp + +## Polling with poll_task + +The `poll_task()` async iterator polls until the task reaches a terminal state: + +```python +async for status in session.experimental.poll_task(task_id): + print(f"Status: {status.status}") + if status.statusMessage: + print(f"Progress: {status.statusMessage}") +``` + +It automatically: + +- Respects the server's suggested `pollInterval` +- Stops when status is `completed`, `failed`, or `cancelled` +- Yields each status for progress display + +### Handling input_required + +When a task needs user input (elicitation), it transitions to `input_required`. You must call `get_task_result()` to receive and respond to the elicitation: + +```python +async for status in session.experimental.poll_task(task_id): + print(f"Status: {status.status}") + + if status.status == "input_required": + # This delivers the elicitation and waits for completion + final = await session.experimental.get_task_result(task_id, CallToolResult) + break +``` + +The elicitation callback (set during session creation) handles the actual user interaction. + +## Elicitation Callbacks + +To handle elicitation requests from the server, provide a callback when creating the session: + +```python +from mcp.types import ElicitRequestParams, ElicitResult + +async def handle_elicitation(context, params: ElicitRequestParams) -> ElicitResult: + # Display the message to the user + print(f"Server asks: {params.message}") + + # Collect user input (this is a simplified example) + response = input("Your response (y/n): ") + confirmed = response.lower() == "y" + + return ElicitResult( + action="accept", + content={"confirm": confirmed}, + ) + +async with ClientSession( + read, + write, + elicitation_callback=handle_elicitation, +) as session: + await session.initialize() + # ... call tasks that may require elicitation +``` + +## Sampling Callbacks + +Similarly, handle sampling requests with a callback: + +```python +from mcp.types import CreateMessageRequestParams, CreateMessageResult, TextContent + +async def handle_sampling(context, params: CreateMessageRequestParams) -> CreateMessageResult: + # In a real implementation, call your LLM here + prompt = params.messages[-1].content.text if params.messages else "" + + # Return a mock response + return CreateMessageResult( + role="assistant", + content=TextContent(type="text", text=f"Response to: {prompt}"), + model="my-model", + ) + +async with ClientSession( + read, + write, + sampling_callback=handle_sampling, +) as session: + # ... +``` + +## Retrieving Results + +Once a task completes, retrieve the result: + +```python +if status.status == "completed": + result = await session.experimental.get_task_result(task_id, CallToolResult) + for content in result.content: + if hasattr(content, "text"): + print(content.text) + +elif status.status == "failed": + print(f"Task failed: {status.statusMessage}") + +elif status.status == "cancelled": + print("Task was cancelled") +``` + +The result type matches the original request: + +- `tools/call` → `CallToolResult` +- `sampling/createMessage` → `CreateMessageResult` +- `elicitation/create` → `ElicitResult` + +## Cancellation + +Cancel a running task: + +```python +cancel_result = await session.experimental.cancel_task(task_id) +print(f"Cancelled, status: {cancel_result.status}") +``` + +Note: Cancellation is cooperative—the server must check for and handle cancellation. + +## Listing Tasks + +View all tasks on the server: + +```python +result = await session.experimental.list_tasks() +for task in result.tasks: + print(f"{task.taskId}: {task.status}") + +# Handle pagination +while result.nextCursor: + result = await session.experimental.list_tasks(cursor=result.nextCursor) + for task in result.tasks: + print(f"{task.taskId}: {task.status}") +``` + +## Advanced: Client as Task Receiver + +Servers can send task-augmented requests to clients. This is useful when the server needs the client to perform async work (like complex sampling or user interaction). + +### Declaring Client Capabilities + +Register task handlers to declare what task-augmented requests your client accepts: + +```python +from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers +from mcp.types import ( + CreateTaskResult, GetTaskResult, GetTaskPayloadResult, + TaskMetadata, ElicitRequestParams, +) +from mcp.shared.experimental.tasks import InMemoryTaskStore + +# Client-side task store +client_store = InMemoryTaskStore() + +async def handle_augmented_elicitation(context, params: ElicitRequestParams, task_metadata: TaskMetadata): + """Handle task-augmented elicitation from server.""" + # Create a task for this elicitation + task = await client_store.create_task(task_metadata) + + # Start async work (e.g., show UI, wait for user) + async def complete_elicitation(): + # ... do async work ... + result = ElicitResult(action="accept", content={"confirm": True}) + await client_store.store_result(task.taskId, result) + await client_store.update_task(task.taskId, status="completed") + + context.session._task_group.start_soon(complete_elicitation) + + # Return task reference immediately + return CreateTaskResult(task=task) + +async def handle_get_task(context, params): + """Handle tasks/get from server.""" + task = await client_store.get_task(params.taskId) + return GetTaskResult( + taskId=task.taskId, + status=task.status, + statusMessage=task.statusMessage, + createdAt=task.createdAt, + lastUpdatedAt=task.lastUpdatedAt, + ttl=task.ttl, + pollInterval=100, + ) + +async def handle_get_task_result(context, params): + """Handle tasks/result from server.""" + result = await client_store.get_result(params.taskId) + return GetTaskPayloadResult.model_validate(result.model_dump()) + +task_handlers = ExperimentalTaskHandlers( + augmented_elicitation=handle_augmented_elicitation, + get_task=handle_get_task, + get_task_result=handle_get_task_result, +) + +async with ClientSession( + read, + write, + experimental_task_handlers=task_handlers, +) as session: + # Client now accepts task-augmented elicitation from server + await session.initialize() +``` + +This enables flows where: + +1. Client calls a task-augmented tool +2. Server's tool work calls `task.elicit_as_task()` +3. Client receives task-augmented elicitation +4. Client creates its own task, does async work +5. Server polls client's task +6. Eventually both tasks complete + +## Complete Example + +A client that handles all task scenarios: + +```python +import anyio +from mcp.client.session import ClientSession +from mcp.client.stdio import stdio_client +from mcp.types import CallToolResult, ElicitRequestParams, ElicitResult + + +async def elicitation_callback(context, params: ElicitRequestParams) -> ElicitResult: + print(f"\n[Elicitation] {params.message}") + response = input("Confirm? (y/n): ") + return ElicitResult(action="accept", content={"confirm": response.lower() == "y"}) + + +async def main(): + async with stdio_client(command="python", args=["server.py"]) as (read, write): + async with ClientSession( + read, + write, + elicitation_callback=elicitation_callback, + ) as session: + await session.initialize() + + # List available tools + tools = await session.list_tools() + print("Tools:", [t.name for t in tools.tools]) + + # Call a task-augmented tool + print("\nCalling task tool...") + result = await session.experimental.call_tool_as_task( + "confirm_action", + {"action": "delete files"}, + ) + task_id = result.task.taskId + print(f"Task created: {task_id}") + + # Poll and handle input_required + async for status in session.experimental.poll_task(task_id): + print(f"Status: {status.status}") + + if status.status == "input_required": + final = await session.experimental.get_task_result(task_id, CallToolResult) + print(f"Result: {final.content[0].text}") + break + + if status.status == "completed": + final = await session.experimental.get_task_result(task_id, CallToolResult) + print(f"Result: {final.content[0].text}") + + +if __name__ == "__main__": + anyio.run(main) +``` + +## Error Handling + +Handle task errors gracefully: + +```python +from mcp.shared.exceptions import McpError + +try: + result = await session.experimental.call_tool_as_task("my_tool", args) + task_id = result.task.taskId + + async for status in session.experimental.poll_task(task_id): + if status.status == "failed": + raise RuntimeError(f"Task failed: {status.statusMessage}") + + final = await session.experimental.get_task_result(task_id, CallToolResult) + +except McpError as e: + print(f"MCP error: {e.error.message}") +except Exception as e: + print(f"Error: {e}") +``` + +## Next Steps + +- [Server Implementation](tasks-server.md) - Build task-supporting servers +- [Tasks Overview](tasks.md) - Review lifecycle and concepts diff --git a/docs/experimental/tasks-server.md b/docs/experimental/tasks-server.md new file mode 100644 index 0000000000..761dc5de5c --- /dev/null +++ b/docs/experimental/tasks-server.md @@ -0,0 +1,597 @@ +# Server Task Implementation + +!!! warning "Experimental" + + Tasks are an experimental feature. The API may change without notice. + +This guide covers implementing task support in MCP servers, from basic setup to advanced patterns like elicitation and sampling within tasks. + +## Quick Start + +The simplest way to add task support: + +```python +from mcp.server import Server +from mcp.server.experimental.task_context import ServerTaskContext +from mcp.types import CallToolResult, CreateTaskResult, TextContent, Tool, ToolExecution, TASK_REQUIRED + +server = Server("my-server") +server.experimental.enable_tasks() # Registers all task handlers automatically + +@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_tool(name: str, arguments: dict) -> CallToolResult | CreateTaskResult: + if name == "process_data": + return await handle_process_data(arguments) + return CallToolResult(content=[TextContent(type="text", text=f"Unknown: {name}")], isError=True) + +async def handle_process_data(arguments: dict) -> CreateTaskResult: + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + async def work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Processing...") + result = arguments.get("input", "").upper() + return CallToolResult(content=[TextContent(type="text", text=result)]) + + return await ctx.experimental.run_task(work) +``` + +That's it. `enable_tasks()` automatically: + +- Creates an in-memory task store +- Registers handlers for `tasks/get`, `tasks/result`, `tasks/list`, `tasks/cancel` +- Updates server capabilities + +## Tool Declaration + +Tools declare task support via the `execution.taskSupport` field: + +```python +from mcp.types import Tool, ToolExecution, TASK_REQUIRED, TASK_OPTIONAL, TASK_FORBIDDEN + +Tool( + name="my_tool", + inputSchema={"type": "object"}, + execution=ToolExecution(taskSupport=TASK_REQUIRED), # or TASK_OPTIONAL, TASK_FORBIDDEN +) +``` + +| Value | Meaning | +|-------|---------| +| `TASK_REQUIRED` | Tool **must** be called as a task | +| `TASK_OPTIONAL` | Tool supports both sync and task execution | +| `TASK_FORBIDDEN` | Tool **cannot** be called as a task (default) | + +Validate the request matches your tool's requirements: + +```python +@server.call_tool() +async def handle_tool(name: str, arguments: dict): + ctx = server.request_context + + if name == "required_task_tool": + ctx.experimental.validate_task_mode(TASK_REQUIRED) # Raises if not task mode + return await handle_as_task(arguments) + + elif name == "optional_task_tool": + if ctx.experimental.is_task: + return await handle_as_task(arguments) + else: + return handle_sync(arguments) +``` + +## The run_task Pattern + +`run_task()` is the recommended way to execute task work: + +```python +async def handle_my_tool(arguments: dict) -> CreateTaskResult: + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + async def work(task: ServerTaskContext) -> CallToolResult: + # Your work here + return CallToolResult(content=[TextContent(type="text", text="Done")]) + + return await ctx.experimental.run_task(work) +``` + +**What `run_task()` does:** + +1. Creates a task in the store +2. Spawns your work function in the background +3. Returns `CreateTaskResult` immediately +4. Auto-completes the task when your function returns +5. Auto-fails the task if your function raises + +**The `ServerTaskContext` provides:** + +- `task.task_id` - The task identifier +- `task.update_status(message)` - Update progress +- `task.complete(result)` - Explicitly complete (usually automatic) +- `task.fail(error)` - Explicitly fail +- `task.is_cancelled` - Check if cancellation requested + +## Status Updates + +Keep clients informed of progress: + +```python +async def work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Starting...") + + for i, item in enumerate(items): + await task.update_status(f"Processing {i+1}/{len(items)}") + await process_item(item) + + await task.update_status("Finalizing...") + return CallToolResult(content=[TextContent(type="text", text="Complete")]) +``` + +Status messages appear in `tasks/get` responses, letting clients show progress to users. + +## Elicitation Within Tasks + +Tasks can request user input via elicitation. This transitions the task to `input_required` status. + +### Form Elicitation + +Collect structured data from the user: + +```python +async def work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Waiting for confirmation...") + + result = await task.elicit( + message="Delete these files?", + requestedSchema={ + "type": "object", + "properties": { + "confirm": {"type": "boolean"}, + "reason": {"type": "string"}, + }, + "required": ["confirm"], + }, + ) + + if result.action == "accept" and result.content.get("confirm"): + # User confirmed + return CallToolResult(content=[TextContent(type="text", text="Files deleted")]) + else: + # User declined or cancelled + return CallToolResult(content=[TextContent(type="text", text="Cancelled")]) +``` + +### URL Elicitation + +Direct users to external URLs for OAuth, payments, or other out-of-band flows: + +```python +async def work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Waiting for OAuth...") + + result = await task.elicit_url( + message="Please authorize with GitHub", + url="https://github.com/login/oauth/authorize?client_id=...", + elicitation_id="oauth-github-123", + ) + + if result.action == "accept": + # User completed OAuth flow + return CallToolResult(content=[TextContent(type="text", text="Connected to GitHub")]) + else: + return CallToolResult(content=[TextContent(type="text", text="OAuth cancelled")]) +``` + +## Sampling Within Tasks + +Tasks can request LLM completions from the client: + +```python +from mcp.types import SamplingMessage, TextContent + +async def work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Generating response...") + + result = await task.create_message( + messages=[ + SamplingMessage( + role="user", + content=TextContent(type="text", text="Write a haiku about coding"), + ) + ], + max_tokens=100, + ) + + haiku = result.content.text if isinstance(result.content, TextContent) else "Error" + return CallToolResult(content=[TextContent(type="text", text=haiku)]) +``` + +Sampling supports additional parameters: + +```python +result = await task.create_message( + messages=[...], + max_tokens=500, + system_prompt="You are a helpful assistant", + temperature=0.7, + stop_sequences=["\n\n"], + model_preferences=ModelPreferences(hints=[ModelHint(name="claude-3")]), +) +``` + +## Cancellation Support + +Check for cancellation in long-running work: + +```python +async def work(task: ServerTaskContext) -> CallToolResult: + for i in range(1000): + if task.is_cancelled: + # Clean up and exit + return CallToolResult(content=[TextContent(type="text", text="Cancelled")]) + + await task.update_status(f"Step {i}/1000") + await process_step(i) + + return CallToolResult(content=[TextContent(type="text", text="Complete")]) +``` + +The SDK's default cancel handler updates the task status. Your work function should check `is_cancelled` periodically. + +## Custom Task Store + +For production, implement `TaskStore` with persistent storage: + +```python +from mcp.shared.experimental.tasks.store import TaskStore +from mcp.types import Task, TaskMetadata, Result + +class RedisTaskStore(TaskStore): + def __init__(self, redis_client): + self.redis = redis_client + + async def create_task(self, metadata: TaskMetadata, task_id: str | None = None) -> Task: + # Create and persist task + ... + + async def get_task(self, task_id: str) -> Task | None: + # Retrieve task from Redis + ... + + async def update_task(self, task_id: str, status: str | None = None, ...) -> Task: + # Update and persist + ... + + async def store_result(self, task_id: str, result: Result) -> None: + # Store result in Redis + ... + + async def get_result(self, task_id: str) -> Result | None: + # Retrieve result + ... + + # ... implement remaining methods +``` + +Use your custom store: + +```python +store = RedisTaskStore(redis_client) +server.experimental.enable_tasks(store=store) +``` + +## Complete Example + +A server with multiple task-supporting tools: + +```python +from mcp.server import Server +from mcp.server.experimental.task_context import ServerTaskContext +from mcp.types import ( + CallToolResult, CreateTaskResult, TextContent, Tool, ToolExecution, + SamplingMessage, TASK_REQUIRED, +) + +server = Server("task-demo") +server.experimental.enable_tasks() + + +@server.list_tools() +async def list_tools(): + return [ + Tool( + name="confirm_action", + description="Requires user confirmation", + inputSchema={"type": "object", "properties": {"action": {"type": "string"}}}, + execution=ToolExecution(taskSupport=TASK_REQUIRED), + ), + Tool( + name="generate_text", + description="Generate text via LLM", + inputSchema={"type": "object", "properties": {"prompt": {"type": "string"}}}, + execution=ToolExecution(taskSupport=TASK_REQUIRED), + ), + ] + + +async def handle_confirm_action(arguments: dict) -> CreateTaskResult: + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + action = arguments.get("action", "unknown action") + + async def work(task: ServerTaskContext) -> CallToolResult: + result = await task.elicit( + message=f"Confirm: {action}?", + requestedSchema={ + "type": "object", + "properties": {"confirm": {"type": "boolean"}}, + "required": ["confirm"], + }, + ) + + if result.action == "accept" and result.content.get("confirm"): + return CallToolResult(content=[TextContent(type="text", text=f"Executed: {action}")]) + return CallToolResult(content=[TextContent(type="text", text="Cancelled")]) + + return await ctx.experimental.run_task(work) + + +async def handle_generate_text(arguments: dict) -> CreateTaskResult: + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + prompt = arguments.get("prompt", "Hello") + + async def work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Generating...") + + result = await task.create_message( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text=prompt))], + max_tokens=200, + ) + + text = result.content.text if isinstance(result.content, TextContent) else "Error" + return CallToolResult(content=[TextContent(type="text", text=text)]) + + return await ctx.experimental.run_task(work) + + +@server.call_tool() +async def handle_tool(name: str, arguments: dict) -> CallToolResult | CreateTaskResult: + if name == "confirm_action": + return await handle_confirm_action(arguments) + elif name == "generate_text": + return await handle_generate_text(arguments) + return CallToolResult(content=[TextContent(type="text", text=f"Unknown: {name}")], isError=True) +``` + +## Error Handling in Tasks + +Tasks handle errors automatically, but you can also fail explicitly: + +```python +async def work(task: ServerTaskContext) -> CallToolResult: + try: + result = await risky_operation() + return CallToolResult(content=[TextContent(type="text", text=result)]) + except PermissionError: + await task.fail("Access denied - insufficient permissions") + raise + except TimeoutError: + await task.fail("Operation timed out after 30 seconds") + raise +``` + +When `run_task()` catches an exception, it automatically: + +1. Marks the task as `failed` +2. Sets `statusMessage` to the exception message +3. Propagates the exception (which is caught by the task group) + +For custom error messages, call `task.fail()` before raising. + +## HTTP Transport Example + +For web applications, use the Streamable HTTP transport: + +```python +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +import uvicorn +from starlette.applications import Starlette +from starlette.routing import Mount + +from mcp.server import Server +from mcp.server.experimental.task_context import ServerTaskContext +from mcp.server.streamable_http_manager import StreamableHTTPSessionManager +from mcp.types import ( + CallToolResult, CreateTaskResult, TextContent, Tool, ToolExecution, TASK_REQUIRED, +) + + +server = Server("http-task-server") +server.experimental.enable_tasks() + + +@server.list_tools() +async def list_tools(): + return [ + Tool( + name="long_operation", + description="A long-running operation", + inputSchema={"type": "object", "properties": {"duration": {"type": "number"}}}, + execution=ToolExecution(taskSupport=TASK_REQUIRED), + ) + ] + + +async def handle_long_operation(arguments: dict) -> CreateTaskResult: + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + duration = arguments.get("duration", 5) + + async def work(task: ServerTaskContext) -> CallToolResult: + import anyio + for i in range(int(duration)): + await task.update_status(f"Step {i+1}/{int(duration)}") + await anyio.sleep(1) + return CallToolResult(content=[TextContent(type="text", text=f"Completed after {duration}s")]) + + return await ctx.experimental.run_task(work) + + +@server.call_tool() +async def handle_tool(name: str, arguments: dict) -> CallToolResult | CreateTaskResult: + if name == "long_operation": + return await handle_long_operation(arguments) + return CallToolResult(content=[TextContent(type="text", text=f"Unknown: {name}")], isError=True) + + +def create_app(): + session_manager = StreamableHTTPSessionManager(app=server) + + @asynccontextmanager + async def lifespan(app: Starlette) -> AsyncIterator[None]: + async with session_manager.run(): + yield + + return Starlette( + routes=[Mount("/mcp", app=session_manager.handle_request)], + lifespan=lifespan, + ) + + +if __name__ == "__main__": + uvicorn.run(create_app(), host="127.0.0.1", port=8000) +``` + +## Testing Task Servers + +Test task functionality with the SDK's testing utilities: + +```python +import pytest +import anyio +from mcp.client.session import ClientSession +from mcp.types import CallToolResult + + +@pytest.mark.anyio +async def test_task_tool(): + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream(10) + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream(10) + + async def run_server(): + await server.run( + client_to_server_receive, + server_to_client_send, + server.create_initialization_options(), + ) + + async def run_client(): + async with ClientSession(server_to_client_receive, client_to_server_send) as session: + await session.initialize() + + # Call the tool as a task + result = await session.experimental.call_tool_as_task("my_tool", {"arg": "value"}) + task_id = result.task.taskId + assert result.task.status == "working" + + # Poll until complete + async for status in session.experimental.poll_task(task_id): + if status.status in ("completed", "failed"): + break + + # Get result + final = await session.experimental.get_task_result(task_id, CallToolResult) + assert len(final.content) > 0 + + async with anyio.create_task_group() as tg: + tg.start_soon(run_server) + tg.start_soon(run_client) +``` + +## Best Practices + +### Keep Work Functions Focused + +```python +# Good: focused work function +async def work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Validating...") + validate_input(arguments) + + await task.update_status("Processing...") + result = await process_data(arguments) + + return CallToolResult(content=[TextContent(type="text", text=result)]) +``` + +### Check Cancellation in Loops + +```python +async def work(task: ServerTaskContext) -> CallToolResult: + results = [] + for item in large_dataset: + if task.is_cancelled: + return CallToolResult(content=[TextContent(type="text", text="Cancelled")]) + + results.append(await process(item)) + + return CallToolResult(content=[TextContent(type="text", text=str(results))]) +``` + +### Use Meaningful Status Messages + +```python +async def work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Connecting to database...") + db = await connect() + + await task.update_status("Fetching records (0/1000)...") + for i, record in enumerate(records): + if i % 100 == 0: + await task.update_status(f"Processing records ({i}/1000)...") + await process(record) + + await task.update_status("Finalizing results...") + return CallToolResult(content=[TextContent(type="text", text="Done")]) +``` + +### Handle Elicitation Responses + +```python +async def work(task: ServerTaskContext) -> CallToolResult: + result = await task.elicit(message="Continue?", requestedSchema={...}) + + match result.action: + case "accept": + # User accepted, process content + return await process_accepted(result.content) + case "decline": + # User explicitly declined + return CallToolResult(content=[TextContent(type="text", text="User declined")]) + case "cancel": + # User cancelled the elicitation + return CallToolResult(content=[TextContent(type="text", text="Cancelled")]) +``` + +## Next Steps + +- [Client Usage](tasks-client.md) - Learn how clients interact with task servers +- [Tasks Overview](tasks.md) - Review lifecycle and concepts diff --git a/docs/experimental/tasks.md b/docs/experimental/tasks.md new file mode 100644 index 0000000000..2d4d06a025 --- /dev/null +++ b/docs/experimental/tasks.md @@ -0,0 +1,188 @@ +# Tasks + +!!! warning "Experimental" + + Tasks are an experimental feature tracking the draft MCP specification. + The API may change without notice. + +Tasks enable asynchronous request handling in MCP. Instead of blocking until an operation completes, the receiver creates a task, returns immediately, and the requestor polls for the result. + +## When to Use Tasks + +Tasks are designed for operations that: + +- Take significant time (seconds to minutes) +- Need progress updates during execution +- Require user input mid-execution (elicitation, sampling) +- Should run without blocking the requestor + +Common use cases: + +- Long-running data processing +- Multi-step workflows with user confirmation +- LLM-powered operations requiring sampling +- OAuth flows requiring user browser interaction + +## Task Lifecycle + +```text + ┌─────────────┐ + │ working │ + └──────┬──────┘ + │ + ┌────────────┼────────────┐ + │ │ │ + ▼ ▼ ▼ + ┌────────────┐ ┌───────────┐ ┌───────────┐ + │ completed │ │ failed │ │ cancelled │ + └────────────┘ └───────────┘ └───────────┘ + ▲ + │ + ┌────────┴────────┐ + │ input_required │◄──────┐ + └────────┬────────┘ │ + │ │ + └────────────────┘ +``` + +| Status | Description | +|--------|-------------| +| `working` | Task is being processed | +| `input_required` | Receiver needs input from requestor (elicitation/sampling) | +| `completed` | Task finished successfully | +| `failed` | Task encountered an error | +| `cancelled` | Task was cancelled by requestor | + +Terminal states (`completed`, `failed`, `cancelled`) are final—tasks cannot transition out of them. + +## Bidirectional Flow + +Tasks work in both directions: + +**Client → Server** (most common): + +```text +Client Server + │ │ + │── tools/call (task) ──────────────>│ Creates task + │<── CreateTaskResult ───────────────│ + │ │ + │── tasks/get ──────────────────────>│ + │<── status: working ────────────────│ + │ │ ... work continues ... + │── tasks/get ──────────────────────>│ + │<── status: completed ──────────────│ + │ │ + │── tasks/result ───────────────────>│ + │<── CallToolResult ─────────────────│ +``` + +**Server → Client** (for elicitation/sampling): + +```text +Server Client + │ │ + │── elicitation/create (task) ──────>│ Creates task + │<── CreateTaskResult ───────────────│ + │ │ + │── tasks/get ──────────────────────>│ + │<── status: working ────────────────│ + │ │ ... user interaction ... + │── tasks/get ──────────────────────>│ + │<── status: completed ──────────────│ + │ │ + │── tasks/result ───────────────────>│ + │<── ElicitResult ───────────────────│ +``` + +## Key Concepts + +### Task Metadata + +When augmenting a request with task execution, include `TaskMetadata`: + +```python +from mcp.types import TaskMetadata + +task = TaskMetadata(ttl=60000) # TTL in milliseconds +``` + +The `ttl` (time-to-live) specifies how long the task and result are retained after completion. + +### Task Store + +Servers persist task state in a `TaskStore`. The SDK provides `InMemoryTaskStore` for development: + +```python +from mcp.shared.experimental.tasks import InMemoryTaskStore + +store = InMemoryTaskStore() +``` + +For production, implement `TaskStore` with a database or distributed cache. + +### Capabilities + +Both servers and clients declare task support through capabilities: + +**Server capabilities:** + +- `tasks.requests.tools.call` - Server accepts task-augmented tool calls + +**Client capabilities:** + +- `tasks.requests.sampling.createMessage` - Client accepts task-augmented sampling +- `tasks.requests.elicitation.create` - Client accepts task-augmented elicitation + +The SDK manages these automatically when you enable task support. + +## Quick Example + +**Server** (simplified API): + +```python +from mcp.server import Server +from mcp.server.experimental.task_context import ServerTaskContext +from mcp.types import CallToolResult, TextContent, TASK_REQUIRED + +server = Server("my-server") +server.experimental.enable_tasks() # One-line setup + +@server.call_tool() +async def handle_tool(name: str, arguments: dict): + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + async def work(task: ServerTaskContext): + await task.update_status("Processing...") + # ... do work ... + return CallToolResult(content=[TextContent(type="text", text="Done!")]) + + return await ctx.experimental.run_task(work) +``` + +**Client:** + +```python +from mcp.client.session import ClientSession +from mcp.types import CallToolResult + +async with ClientSession(read, write) as session: + await session.initialize() + + # Call tool as task + result = await session.experimental.call_tool_as_task("my_tool", {"arg": "value"}) + task_id = result.task.taskId + + # Poll until done + async for status in session.experimental.poll_task(task_id): + print(f"Status: {status.status}") + + # Get result + final = await session.experimental.get_task_result(task_id, CallToolResult) +``` + +## Next Steps + +- [Server Implementation](tasks-server.md) - Build task-supporting servers +- [Client Usage](tasks-client.md) - Call and poll tasks from clients diff --git a/examples/clients/simple-task-client/README.md b/examples/clients/simple-task-client/README.md new file mode 100644 index 0000000000..103be0f1fb --- /dev/null +++ b/examples/clients/simple-task-client/README.md @@ -0,0 +1,43 @@ +# Simple Task Client + +A minimal MCP client demonstrating polling for task results over streamable HTTP. + +## Running + +First, start the simple-task server in another terminal: + +```bash +cd examples/servers/simple-task +uv run mcp-simple-task +``` + +Then run the client: + +```bash +cd examples/clients/simple-task-client +uv run mcp-simple-task-client +``` + +Use `--url` to connect to a different server. + +## What it does + +1. Connects to the server via streamable HTTP +2. Calls the `long_running_task` tool as a task +3. Polls the task status until completion +4. Retrieves and prints the result + +## Expected output + +```text +Available tools: ['long_running_task'] + +Calling tool as a task... +Task created: + Status: working - Starting work... + Status: working - Processing step 1... + Status: working - Processing step 2... + Status: completed - + +Result: Task completed! +``` diff --git a/examples/clients/simple-task-client/mcp_simple_task_client/__init__.py b/examples/clients/simple-task-client/mcp_simple_task_client/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/clients/simple-task-client/mcp_simple_task_client/__main__.py b/examples/clients/simple-task-client/mcp_simple_task_client/__main__.py new file mode 100644 index 0000000000..2fc2cda8d9 --- /dev/null +++ b/examples/clients/simple-task-client/mcp_simple_task_client/__main__.py @@ -0,0 +1,5 @@ +import sys + +from .main import main + +sys.exit(main()) # type: ignore[call-arg] diff --git a/examples/clients/simple-task-client/mcp_simple_task_client/main.py b/examples/clients/simple-task-client/mcp_simple_task_client/main.py new file mode 100644 index 0000000000..12691162ab --- /dev/null +++ b/examples/clients/simple-task-client/mcp_simple_task_client/main.py @@ -0,0 +1,55 @@ +"""Simple task client demonstrating MCP tasks polling over streamable HTTP.""" + +import asyncio + +import click +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client +from mcp.types import CallToolResult, TextContent + + +async def run(url: str) -> None: + async with streamablehttp_client(url) as (read, write, _): + async with ClientSession(read, write) as session: + await session.initialize() + + # List tools + tools = await session.list_tools() + print(f"Available tools: {[t.name for t in tools.tools]}") + + # Call the tool as a task + print("\nCalling tool as a task...") + + result = await session.experimental.call_tool_as_task( + "long_running_task", + arguments={}, + ttl=60000, + ) + task_id = result.task.taskId + print(f"Task created: {task_id}") + + # Poll until done (respects server's pollInterval hint) + async for status in session.experimental.poll_task(task_id): + print(f" Status: {status.status} - {status.statusMessage or ''}") + + # Check final status + if status.status != "completed": + print(f"Task ended with status: {status.status}") + return + + # Get the result + task_result = await session.experimental.get_task_result(task_id, CallToolResult) + content = task_result.content[0] + if isinstance(content, TextContent): + print(f"\nResult: {content.text}") + + +@click.command() +@click.option("--url", default="http://localhost:8000/mcp", help="Server URL") +def main(url: str) -> int: + asyncio.run(run(url)) + return 0 + + +if __name__ == "__main__": + main() diff --git a/examples/clients/simple-task-client/pyproject.toml b/examples/clients/simple-task-client/pyproject.toml new file mode 100644 index 0000000000..da10392e3c --- /dev/null +++ b/examples/clients/simple-task-client/pyproject.toml @@ -0,0 +1,43 @@ +[project] +name = "mcp-simple-task-client" +version = "0.1.0" +description = "A simple MCP client demonstrating task polling" +readme = "README.md" +requires-python = ">=3.10" +authors = [{ name = "Anthropic, PBC." }] +keywords = ["mcp", "llm", "tasks", "client"] +license = { text = "MIT" } +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", +] +dependencies = ["click>=8.0", "mcp"] + +[project.scripts] +mcp-simple-task-client = "mcp_simple_task_client.main:main" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["mcp_simple_task_client"] + +[tool.pyright] +include = ["mcp_simple_task_client"] +venvPath = "." +venv = ".venv" + +[tool.ruff.lint] +select = ["E", "F", "I"] +ignore = [] + +[tool.ruff] +line-length = 120 +target-version = "py310" + +[dependency-groups] +dev = ["pyright>=1.1.378", "ruff>=0.6.9"] diff --git a/examples/clients/simple-task-interactive-client/README.md b/examples/clients/simple-task-interactive-client/README.md new file mode 100644 index 0000000000..ac73d2bc12 --- /dev/null +++ b/examples/clients/simple-task-interactive-client/README.md @@ -0,0 +1,87 @@ +# Simple Interactive Task Client + +A minimal MCP client demonstrating responses to interactive tasks (elicitation and sampling). + +## Running + +First, start the interactive task server in another terminal: + +```bash +cd examples/servers/simple-task-interactive +uv run mcp-simple-task-interactive +``` + +Then run the client: + +```bash +cd examples/clients/simple-task-interactive-client +uv run mcp-simple-task-interactive-client +``` + +Use `--url` to connect to a different server. + +## What it does + +1. Connects to the server via streamable HTTP +2. Calls `confirm_delete` - server asks for confirmation, client responds via terminal +3. Calls `write_haiku` - server requests LLM completion, client returns a hardcoded haiku + +## Key concepts + +### Elicitation callback + +```python +async def elicitation_callback(context, params) -> ElicitResult: + # Handle user input request from server + return ElicitResult(action="accept", content={"confirm": True}) +``` + +### Sampling callback + +```python +async def sampling_callback(context, params) -> CreateMessageResult: + # Handle LLM completion request from server + return CreateMessageResult(model="...", role="assistant", content=...) +``` + +### Using call_tool_as_task + +```python +# Call a tool as a task (returns immediately with task reference) +result = await session.experimental.call_tool_as_task("tool_name", {"arg": "value"}) +task_id = result.task.taskId + +# Get result - this delivers elicitation/sampling requests and blocks until complete +final = await session.experimental.get_task_result(task_id, CallToolResult) +``` + +**Important**: The `get_task_result()` call is what triggers the delivery of elicitation +and sampling requests to your callbacks. It blocks until the task completes and returns +the final result. + +## Expected output + +```text +Available tools: ['confirm_delete', 'write_haiku'] + +--- Demo 1: Elicitation --- +Calling confirm_delete tool... +Task created: + +[Elicitation] Server asks: Are you sure you want to delete 'important.txt'? +Your response (y/n): y +[Elicitation] Responding with: confirm=True +Result: Deleted 'important.txt' + +--- Demo 2: Sampling --- +Calling write_haiku tool... +Task created: + +[Sampling] Server requests LLM completion for: Write a haiku about autumn leaves +[Sampling] Responding with haiku +Result: +Haiku: +Cherry blossoms fall +Softly on the quiet pond +Spring whispers goodbye +``` diff --git a/examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/__init__.py b/examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/__main__.py b/examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/__main__.py new file mode 100644 index 0000000000..2fc2cda8d9 --- /dev/null +++ b/examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/__main__.py @@ -0,0 +1,5 @@ +import sys + +from .main import main + +sys.exit(main()) # type: ignore[call-arg] diff --git a/examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/main.py b/examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/main.py new file mode 100644 index 0000000000..a8a47dc57c --- /dev/null +++ b/examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/main.py @@ -0,0 +1,138 @@ +"""Simple interactive task client demonstrating elicitation and sampling responses. + +This example demonstrates the spec-compliant polling pattern: +1. Poll tasks/get watching for status changes +2. On input_required, call tasks/result to receive elicitation/sampling requests +3. Continue until terminal status, then retrieve final result +""" + +import asyncio +from typing import Any + +import click +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client +from mcp.shared.context import RequestContext +from mcp.types import ( + CallToolResult, + CreateMessageRequestParams, + CreateMessageResult, + ElicitRequestParams, + ElicitResult, + TextContent, +) + + +async def elicitation_callback( + context: RequestContext[ClientSession, Any], + params: ElicitRequestParams, +) -> ElicitResult: + """Handle elicitation requests from the server.""" + print(f"\n[Elicitation] Server asks: {params.message}") + + # Simple terminal prompt + response = input("Your response (y/n): ").strip().lower() + confirmed = response in ("y", "yes", "true", "1") + + print(f"[Elicitation] Responding with: confirm={confirmed}") + return ElicitResult(action="accept", content={"confirm": confirmed}) + + +async def sampling_callback( + context: RequestContext[ClientSession, Any], + params: CreateMessageRequestParams, +) -> CreateMessageResult: + """Handle sampling requests from the server.""" + # Get the prompt from the first message + prompt = "unknown" + if params.messages: + content = params.messages[0].content + if isinstance(content, TextContent): + prompt = content.text + + print(f"\n[Sampling] Server requests LLM completion for: {prompt}") + + # Return a hardcoded haiku (in real use, call your LLM here) + haiku = """Cherry blossoms fall +Softly on the quiet pond +Spring whispers goodbye""" + + print("[Sampling] Responding with haiku") + return CreateMessageResult( + model="mock-haiku-model", + role="assistant", + content=TextContent(type="text", text=haiku), + ) + + +def get_text(result: CallToolResult) -> str: + """Extract text from a CallToolResult.""" + if result.content and isinstance(result.content[0], TextContent): + return result.content[0].text + return "(no text)" + + +async def run(url: str) -> None: + async with streamablehttp_client(url) as (read, write, _): + async with ClientSession( + read, + write, + elicitation_callback=elicitation_callback, + sampling_callback=sampling_callback, + ) as session: + await session.initialize() + + # List tools + tools = await session.list_tools() + print(f"Available tools: {[t.name for t in tools.tools]}") + + # Demo 1: Elicitation (confirm_delete) + print("\n--- Demo 1: Elicitation ---") + print("Calling confirm_delete tool...") + + elicit_task = await session.experimental.call_tool_as_task("confirm_delete", {"filename": "important.txt"}) + elicit_task_id = elicit_task.task.taskId + print(f"Task created: {elicit_task_id}") + + # Poll until terminal, calling tasks/result on input_required + async for status in session.experimental.poll_task(elicit_task_id): + print(f"[Poll] Status: {status.status}") + if status.status == "input_required": + # Server needs input - tasks/result delivers the elicitation request + elicit_result = await session.experimental.get_task_result(elicit_task_id, CallToolResult) + break + else: + # poll_task exited due to terminal status + elicit_result = await session.experimental.get_task_result(elicit_task_id, CallToolResult) + + print(f"Result: {get_text(elicit_result)}") + + # Demo 2: Sampling (write_haiku) + print("\n--- Demo 2: Sampling ---") + print("Calling write_haiku tool...") + + sampling_task = await session.experimental.call_tool_as_task("write_haiku", {"topic": "autumn leaves"}) + sampling_task_id = sampling_task.task.taskId + print(f"Task created: {sampling_task_id}") + + # Poll until terminal, calling tasks/result on input_required + async for status in session.experimental.poll_task(sampling_task_id): + print(f"[Poll] Status: {status.status}") + if status.status == "input_required": + sampling_result = await session.experimental.get_task_result(sampling_task_id, CallToolResult) + break + else: + sampling_result = await session.experimental.get_task_result(sampling_task_id, CallToolResult) + + print(f"Result:\n{get_text(sampling_result)}") + + +@click.command() +@click.option("--url", default="http://localhost:8000/mcp", help="Server URL") +def main(url: str) -> int: + asyncio.run(run(url)) + return 0 + + +if __name__ == "__main__": + main() diff --git a/examples/clients/simple-task-interactive-client/pyproject.toml b/examples/clients/simple-task-interactive-client/pyproject.toml new file mode 100644 index 0000000000..224bbc5917 --- /dev/null +++ b/examples/clients/simple-task-interactive-client/pyproject.toml @@ -0,0 +1,43 @@ +[project] +name = "mcp-simple-task-interactive-client" +version = "0.1.0" +description = "A simple MCP client demonstrating interactive task responses" +readme = "README.md" +requires-python = ">=3.10" +authors = [{ name = "Anthropic, PBC." }] +keywords = ["mcp", "llm", "tasks", "client", "elicitation", "sampling"] +license = { text = "MIT" } +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", +] +dependencies = ["click>=8.0", "mcp"] + +[project.scripts] +mcp-simple-task-interactive-client = "mcp_simple_task_interactive_client.main:main" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["mcp_simple_task_interactive_client"] + +[tool.pyright] +include = ["mcp_simple_task_interactive_client"] +venvPath = "." +venv = ".venv" + +[tool.ruff.lint] +select = ["E", "F", "I"] +ignore = [] + +[tool.ruff] +line-length = 120 +target-version = "py310" + +[dependency-groups] +dev = ["pyright>=1.1.378", "ruff>=0.6.9"] diff --git a/examples/servers/simple-task-interactive/README.md b/examples/servers/simple-task-interactive/README.md new file mode 100644 index 0000000000..b8f384cb48 --- /dev/null +++ b/examples/servers/simple-task-interactive/README.md @@ -0,0 +1,74 @@ +# Simple Interactive Task Server + +A minimal MCP server demonstrating interactive tasks with elicitation and sampling. + +## Running + +```bash +cd examples/servers/simple-task-interactive +uv run mcp-simple-task-interactive +``` + +The server starts on `http://localhost:8000/mcp` by default. Use `--port` to change. + +## What it does + +This server exposes two tools: + +### `confirm_delete` (demonstrates elicitation) + +Asks the user for confirmation before "deleting" a file. + +- Uses `task.elicit()` to request user input +- Shows the elicitation flow: task -> input_required -> response -> complete + +### `write_haiku` (demonstrates sampling) + +Asks the LLM to write a haiku about a topic. + +- Uses `task.create_message()` to request LLM completion +- Shows the sampling flow: task -> input_required -> response -> complete + +## Usage with the client + +In one terminal, start the server: + +```bash +cd examples/servers/simple-task-interactive +uv run mcp-simple-task-interactive +``` + +In another terminal, run the interactive client: + +```bash +cd examples/clients/simple-task-interactive-client +uv run mcp-simple-task-interactive-client +``` + +## Expected server output + +When a client connects and calls the tools, you'll see: + +```text +Starting server on http://localhost:8000/mcp + +[Server] confirm_delete called for 'important.txt' +[Server] Task created: +[Server] Sending elicitation request to client... +[Server] Received elicitation response: action=accept, content={'confirm': True} +[Server] Completing task with result: Deleted 'important.txt' + +[Server] write_haiku called for topic 'autumn leaves' +[Server] Task created: +[Server] Sending sampling request to client... +[Server] Received sampling response: Cherry blossoms fall +Softly on the quiet pon... +[Server] Completing task with haiku +``` + +## Key concepts + +1. **ServerTaskContext**: Provides `elicit()` and `create_message()` for user interaction +2. **run_task()**: Spawns background work, auto-completes/fails, returns immediately +3. **TaskResultHandler**: Delivers queued messages and routes responses +4. **Response routing**: Responses are routed back to waiting resolvers diff --git a/examples/servers/simple-task-interactive/mcp_simple_task_interactive/__init__.py b/examples/servers/simple-task-interactive/mcp_simple_task_interactive/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/servers/simple-task-interactive/mcp_simple_task_interactive/__main__.py b/examples/servers/simple-task-interactive/mcp_simple_task_interactive/__main__.py new file mode 100644 index 0000000000..e7ef16530b --- /dev/null +++ b/examples/servers/simple-task-interactive/mcp_simple_task_interactive/__main__.py @@ -0,0 +1,5 @@ +import sys + +from .server import main + +sys.exit(main()) # type: ignore[call-arg] diff --git a/examples/servers/simple-task-interactive/mcp_simple_task_interactive/server.py b/examples/servers/simple-task-interactive/mcp_simple_task_interactive/server.py new file mode 100644 index 0000000000..4d35ca8094 --- /dev/null +++ b/examples/servers/simple-task-interactive/mcp_simple_task_interactive/server.py @@ -0,0 +1,147 @@ +"""Simple interactive task server demonstrating elicitation and sampling. + +This example shows the simplified task API where: +- server.experimental.enable_tasks() sets up all infrastructure +- ctx.experimental.run_task() handles task lifecycle automatically +- ServerTaskContext.elicit() and ServerTaskContext.create_message() queue requests properly +""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any + +import click +import mcp.types as types +import uvicorn +from mcp.server.experimental.task_context import ServerTaskContext +from mcp.server.lowlevel import Server +from mcp.server.streamable_http_manager import StreamableHTTPSessionManager +from starlette.applications import Starlette +from starlette.routing import Mount + +server = Server("simple-task-interactive") + +# Enable task support - this auto-registers all handlers +server.experimental.enable_tasks() + + +@server.list_tools() +async def list_tools() -> list[types.Tool]: + return [ + types.Tool( + name="confirm_delete", + description="Asks for confirmation before deleting (demonstrates elicitation)", + inputSchema={ + "type": "object", + "properties": {"filename": {"type": "string"}}, + }, + execution=types.ToolExecution(taskSupport=types.TASK_REQUIRED), + ), + types.Tool( + name="write_haiku", + description="Asks LLM to write a haiku (demonstrates sampling)", + inputSchema={"type": "object", "properties": {"topic": {"type": "string"}}}, + execution=types.ToolExecution(taskSupport=types.TASK_REQUIRED), + ), + ] + + +async def handle_confirm_delete(arguments: dict[str, Any]) -> types.CreateTaskResult: + """Handle the confirm_delete tool - demonstrates elicitation.""" + ctx = server.request_context + ctx.experimental.validate_task_mode(types.TASK_REQUIRED) + + filename = arguments.get("filename", "unknown.txt") + print(f"\n[Server] confirm_delete called for '{filename}'") + + async def work(task: ServerTaskContext) -> types.CallToolResult: + print(f"[Server] Task {task.task_id} starting elicitation...") + + result = await task.elicit( + message=f"Are you sure you want to delete '{filename}'?", + requestedSchema={ + "type": "object", + "properties": {"confirm": {"type": "boolean"}}, + "required": ["confirm"], + }, + ) + + print(f"[Server] Received elicitation response: action={result.action}, content={result.content}") + + if result.action == "accept" and result.content: + confirmed = result.content.get("confirm", False) + text = f"Deleted '{filename}'" if confirmed else "Deletion cancelled" + else: + text = "Deletion cancelled" + + print(f"[Server] Completing task with result: {text}") + return types.CallToolResult(content=[types.TextContent(type="text", text=text)]) + + return await ctx.experimental.run_task(work) + + +async def handle_write_haiku(arguments: dict[str, Any]) -> types.CreateTaskResult: + """Handle the write_haiku tool - demonstrates sampling.""" + ctx = server.request_context + ctx.experimental.validate_task_mode(types.TASK_REQUIRED) + + topic = arguments.get("topic", "nature") + print(f"\n[Server] write_haiku called for topic '{topic}'") + + async def work(task: ServerTaskContext) -> types.CallToolResult: + print(f"[Server] Task {task.task_id} starting sampling...") + + result = await task.create_message( + messages=[ + types.SamplingMessage( + role="user", + content=types.TextContent(type="text", text=f"Write a haiku about {topic}"), + ) + ], + max_tokens=50, + ) + + haiku = "No response" + if isinstance(result.content, types.TextContent): + haiku = result.content.text + + print(f"[Server] Received sampling response: {haiku[:50]}...") + return types.CallToolResult(content=[types.TextContent(type="text", text=f"Haiku:\n{haiku}")]) + + return await ctx.experimental.run_task(work) + + +@server.call_tool() +async def handle_call_tool(name: str, arguments: dict[str, Any]) -> types.CallToolResult | types.CreateTaskResult: + """Dispatch tool calls to their handlers.""" + if name == "confirm_delete": + return await handle_confirm_delete(arguments) + elif name == "write_haiku": + return await handle_write_haiku(arguments) + else: + return types.CallToolResult( + content=[types.TextContent(type="text", text=f"Unknown tool: {name}")], + isError=True, + ) + + +def create_app(session_manager: StreamableHTTPSessionManager) -> Starlette: + @asynccontextmanager + async def app_lifespan(app: Starlette) -> AsyncIterator[None]: + async with session_manager.run(): + yield + + return Starlette( + routes=[Mount("/mcp", app=session_manager.handle_request)], + lifespan=app_lifespan, + ) + + +@click.command() +@click.option("--port", default=8000, help="Port to listen on") +def main(port: int) -> int: + session_manager = StreamableHTTPSessionManager(app=server) + starlette_app = create_app(session_manager) + print(f"Starting server on http://localhost:{port}/mcp") + uvicorn.run(starlette_app, host="127.0.0.1", port=port) + return 0 diff --git a/examples/servers/simple-task-interactive/pyproject.toml b/examples/servers/simple-task-interactive/pyproject.toml new file mode 100644 index 0000000000..492345ff52 --- /dev/null +++ b/examples/servers/simple-task-interactive/pyproject.toml @@ -0,0 +1,43 @@ +[project] +name = "mcp-simple-task-interactive" +version = "0.1.0" +description = "A simple MCP server demonstrating interactive tasks (elicitation & sampling)" +readme = "README.md" +requires-python = ">=3.10" +authors = [{ name = "Anthropic, PBC." }] +keywords = ["mcp", "llm", "tasks", "elicitation", "sampling"] +license = { text = "MIT" } +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", +] +dependencies = ["anyio>=4.5", "click>=8.0", "mcp", "starlette", "uvicorn"] + +[project.scripts] +mcp-simple-task-interactive = "mcp_simple_task_interactive.server:main" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["mcp_simple_task_interactive"] + +[tool.pyright] +include = ["mcp_simple_task_interactive"] +venvPath = "." +venv = ".venv" + +[tool.ruff.lint] +select = ["E", "F", "I"] +ignore = [] + +[tool.ruff] +line-length = 120 +target-version = "py310" + +[dependency-groups] +dev = ["pyright>=1.1.378", "ruff>=0.6.9"] diff --git a/examples/servers/simple-task/README.md b/examples/servers/simple-task/README.md new file mode 100644 index 0000000000..6914e0414f --- /dev/null +++ b/examples/servers/simple-task/README.md @@ -0,0 +1,37 @@ +# Simple Task Server + +A minimal MCP server demonstrating the experimental tasks feature over streamable HTTP. + +## Running + +```bash +cd examples/servers/simple-task +uv run mcp-simple-task +``` + +The server starts on `http://localhost:8000/mcp` by default. Use `--port` to change. + +## What it does + +This server exposes a single tool `long_running_task` that: + +1. Must be called as a task (with `task` metadata in the request) +2. Takes ~3 seconds to complete +3. Sends status updates during execution +4. Returns a result when complete + +## Usage with the client + +In one terminal, start the server: + +```bash +cd examples/servers/simple-task +uv run mcp-simple-task +``` + +In another terminal, run the client: + +```bash +cd examples/clients/simple-task-client +uv run mcp-simple-task-client +``` diff --git a/examples/servers/simple-task/mcp_simple_task/__init__.py b/examples/servers/simple-task/mcp_simple_task/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/servers/simple-task/mcp_simple_task/__main__.py b/examples/servers/simple-task/mcp_simple_task/__main__.py new file mode 100644 index 0000000000..e7ef16530b --- /dev/null +++ b/examples/servers/simple-task/mcp_simple_task/__main__.py @@ -0,0 +1,5 @@ +import sys + +from .server import main + +sys.exit(main()) # type: ignore[call-arg] diff --git a/examples/servers/simple-task/mcp_simple_task/server.py b/examples/servers/simple-task/mcp_simple_task/server.py new file mode 100644 index 0000000000..d0681b8423 --- /dev/null +++ b/examples/servers/simple-task/mcp_simple_task/server.py @@ -0,0 +1,84 @@ +"""Simple task server demonstrating MCP tasks over streamable HTTP.""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any + +import anyio +import click +import mcp.types as types +import uvicorn +from mcp.server.experimental.task_context import ServerTaskContext +from mcp.server.lowlevel import Server +from mcp.server.streamable_http_manager import StreamableHTTPSessionManager +from starlette.applications import Starlette +from starlette.routing import Mount + +server = Server("simple-task-server") + +# One-line setup: auto-registers get_task, get_task_result, list_tasks, cancel_task +server.experimental.enable_tasks() + + +@server.list_tools() +async def list_tools() -> list[types.Tool]: + return [ + types.Tool( + name="long_running_task", + description="A task that takes a few seconds to complete with status updates", + inputSchema={"type": "object", "properties": {}}, + execution=types.ToolExecution(taskSupport=types.TASK_REQUIRED), + ) + ] + + +async def handle_long_running_task(arguments: dict[str, Any]) -> types.CreateTaskResult: + """Handle the long_running_task tool - demonstrates status updates.""" + ctx = server.request_context + ctx.experimental.validate_task_mode(types.TASK_REQUIRED) + + async def work(task: ServerTaskContext) -> types.CallToolResult: + await task.update_status("Starting work...") + await anyio.sleep(1) + + await task.update_status("Processing step 1...") + await anyio.sleep(1) + + await task.update_status("Processing step 2...") + await anyio.sleep(1) + + return types.CallToolResult(content=[types.TextContent(type="text", text="Task completed!")]) + + return await ctx.experimental.run_task(work) + + +@server.call_tool() +async def handle_call_tool(name: str, arguments: dict[str, Any]) -> types.CallToolResult | types.CreateTaskResult: + """Dispatch tool calls to their handlers.""" + if name == "long_running_task": + return await handle_long_running_task(arguments) + else: + return types.CallToolResult( + content=[types.TextContent(type="text", text=f"Unknown tool: {name}")], + isError=True, + ) + + +@click.command() +@click.option("--port", default=8000, help="Port to listen on") +def main(port: int) -> int: + session_manager = StreamableHTTPSessionManager(app=server) + + @asynccontextmanager + async def app_lifespan(app: Starlette) -> AsyncIterator[None]: + async with session_manager.run(): + yield + + starlette_app = Starlette( + routes=[Mount("/mcp", app=session_manager.handle_request)], + lifespan=app_lifespan, + ) + + print(f"Starting server on http://localhost:{port}/mcp") + uvicorn.run(starlette_app, host="127.0.0.1", port=port) + return 0 diff --git a/examples/servers/simple-task/pyproject.toml b/examples/servers/simple-task/pyproject.toml new file mode 100644 index 0000000000..a8fba8bdc1 --- /dev/null +++ b/examples/servers/simple-task/pyproject.toml @@ -0,0 +1,43 @@ +[project] +name = "mcp-simple-task" +version = "0.1.0" +description = "A simple MCP server demonstrating tasks" +readme = "README.md" +requires-python = ">=3.10" +authors = [{ name = "Anthropic, PBC." }] +keywords = ["mcp", "llm", "tasks"] +license = { text = "MIT" } +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", +] +dependencies = ["anyio>=4.5", "click>=8.0", "mcp", "starlette", "uvicorn"] + +[project.scripts] +mcp-simple-task = "mcp_simple_task.server:main" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["mcp_simple_task"] + +[tool.pyright] +include = ["mcp_simple_task"] +venvPath = "." +venv = ".venv" + +[tool.ruff.lint] +select = ["E", "F", "I"] +ignore = [] + +[tool.ruff] +line-length = 120 +target-version = "py310" + +[dependency-groups] +dev = ["pyright>=1.1.378", "ruff>=0.6.9"] diff --git a/mkdocs.yml b/mkdocs.yml index 18cbb034bb..22c323d9d4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -18,6 +18,12 @@ nav: - Low-Level Server: low-level-server.md - Authorization: authorization.md - Testing: testing.md + - Experimental: + - Overview: experimental/index.md + - Tasks: + - Introduction: experimental/tasks.md + - Server Implementation: experimental/tasks-server.md + - Client Usage: experimental/tasks-client.md - API Reference: api.md theme: diff --git a/src/mcp/client/experimental/__init__.py b/src/mcp/client/experimental/__init__.py new file mode 100644 index 0000000000..b6579b191e --- /dev/null +++ b/src/mcp/client/experimental/__init__.py @@ -0,0 +1,9 @@ +""" +Experimental client features. + +WARNING: These APIs are experimental and may change without notice. +""" + +from mcp.client.experimental.tasks import ExperimentalClientFeatures + +__all__ = ["ExperimentalClientFeatures"] diff --git a/src/mcp/client/experimental/task_handlers.py b/src/mcp/client/experimental/task_handlers.py new file mode 100644 index 0000000000..a47508674b --- /dev/null +++ b/src/mcp/client/experimental/task_handlers.py @@ -0,0 +1,290 @@ +""" +Experimental task handler protocols for server -> client requests. + +This module provides Protocol types and default handlers for when servers +send task-related requests to clients (the reverse of normal client -> server flow). + +WARNING: These APIs are experimental and may change without notice. + +Use cases: +- Server sends task-augmented sampling/elicitation request to client +- Client creates a local task, spawns background work, returns CreateTaskResult +- Server polls client's task status via tasks/get, tasks/result, etc. +""" + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Protocol + +from pydantic import TypeAdapter + +import mcp.types as types +from mcp.shared.context import RequestContext +from mcp.shared.session import RequestResponder + +if TYPE_CHECKING: + from mcp.client.session import ClientSession + + +class GetTaskHandlerFnT(Protocol): + """Handler for tasks/get requests from server. + + WARNING: This is experimental and may change without notice. + """ + + async def __call__( + self, + context: RequestContext["ClientSession", Any], + params: types.GetTaskRequestParams, + ) -> types.GetTaskResult | types.ErrorData: ... # pragma: no branch + + +class GetTaskResultHandlerFnT(Protocol): + """Handler for tasks/result requests from server. + + WARNING: This is experimental and may change without notice. + """ + + async def __call__( + self, + context: RequestContext["ClientSession", Any], + params: types.GetTaskPayloadRequestParams, + ) -> types.GetTaskPayloadResult | types.ErrorData: ... # pragma: no branch + + +class ListTasksHandlerFnT(Protocol): + """Handler for tasks/list requests from server. + + WARNING: This is experimental and may change without notice. + """ + + async def __call__( + self, + context: RequestContext["ClientSession", Any], + params: types.PaginatedRequestParams | None, + ) -> types.ListTasksResult | types.ErrorData: ... # pragma: no branch + + +class CancelTaskHandlerFnT(Protocol): + """Handler for tasks/cancel requests from server. + + WARNING: This is experimental and may change without notice. + """ + + async def __call__( + self, + context: RequestContext["ClientSession", Any], + params: types.CancelTaskRequestParams, + ) -> types.CancelTaskResult | types.ErrorData: ... # pragma: no branch + + +class TaskAugmentedSamplingFnT(Protocol): + """Handler for task-augmented sampling/createMessage requests from server. + + When server sends a CreateMessageRequest with task field, this callback + is invoked. The callback should create a task, spawn background work, + and return CreateTaskResult immediately. + + WARNING: This is experimental and may change without notice. + """ + + async def __call__( + self, + context: RequestContext["ClientSession", Any], + params: types.CreateMessageRequestParams, + task_metadata: types.TaskMetadata, + ) -> types.CreateTaskResult | types.ErrorData: ... # pragma: no branch + + +class TaskAugmentedElicitationFnT(Protocol): + """Handler for task-augmented elicitation/create requests from server. + + When server sends an ElicitRequest with task field, this callback + is invoked. The callback should create a task, spawn background work, + and return CreateTaskResult immediately. + + WARNING: This is experimental and may change without notice. + """ + + async def __call__( + self, + context: RequestContext["ClientSession", Any], + params: types.ElicitRequestParams, + task_metadata: types.TaskMetadata, + ) -> types.CreateTaskResult | types.ErrorData: ... # pragma: no branch + + +async def default_get_task_handler( + context: RequestContext["ClientSession", Any], + params: types.GetTaskRequestParams, +) -> types.GetTaskResult | types.ErrorData: + return types.ErrorData( + code=types.METHOD_NOT_FOUND, + message="tasks/get not supported", + ) + + +async def default_get_task_result_handler( + context: RequestContext["ClientSession", Any], + params: types.GetTaskPayloadRequestParams, +) -> types.GetTaskPayloadResult | types.ErrorData: + return types.ErrorData( + code=types.METHOD_NOT_FOUND, + message="tasks/result not supported", + ) + + +async def default_list_tasks_handler( + context: RequestContext["ClientSession", Any], + params: types.PaginatedRequestParams | None, +) -> types.ListTasksResult | types.ErrorData: + return types.ErrorData( + code=types.METHOD_NOT_FOUND, + message="tasks/list not supported", + ) + + +async def default_cancel_task_handler( + context: RequestContext["ClientSession", Any], + params: types.CancelTaskRequestParams, +) -> types.CancelTaskResult | types.ErrorData: + return types.ErrorData( + code=types.METHOD_NOT_FOUND, + message="tasks/cancel not supported", + ) + + +async def default_task_augmented_sampling( + context: RequestContext["ClientSession", Any], + params: types.CreateMessageRequestParams, + task_metadata: types.TaskMetadata, +) -> types.CreateTaskResult | types.ErrorData: + return types.ErrorData( + code=types.INVALID_REQUEST, + message="Task-augmented sampling not supported", + ) + + +async def default_task_augmented_elicitation( + context: RequestContext["ClientSession", Any], + params: types.ElicitRequestParams, + task_metadata: types.TaskMetadata, +) -> types.CreateTaskResult | types.ErrorData: + return types.ErrorData( + code=types.INVALID_REQUEST, + message="Task-augmented elicitation not supported", + ) + + +@dataclass +class ExperimentalTaskHandlers: + """Container for experimental task handlers. + + Groups all task-related handlers that handle server -> client requests. + This includes both pure task requests (get, list, cancel, result) and + task-augmented request handlers (sampling, elicitation with task field). + + WARNING: These APIs are experimental and may change without notice. + + Example: + handlers = ExperimentalTaskHandlers( + get_task=my_get_task_handler, + list_tasks=my_list_tasks_handler, + ) + session = ClientSession(..., experimental_task_handlers=handlers) + """ + + # Pure task request handlers + get_task: GetTaskHandlerFnT = field(default=default_get_task_handler) + get_task_result: GetTaskResultHandlerFnT = field(default=default_get_task_result_handler) + list_tasks: ListTasksHandlerFnT = field(default=default_list_tasks_handler) + cancel_task: CancelTaskHandlerFnT = field(default=default_cancel_task_handler) + + # Task-augmented request handlers + augmented_sampling: TaskAugmentedSamplingFnT = field(default=default_task_augmented_sampling) + augmented_elicitation: TaskAugmentedElicitationFnT = field(default=default_task_augmented_elicitation) + + def build_capability(self) -> types.ClientTasksCapability | None: + """Build ClientTasksCapability from the configured handlers. + + Returns a capability object that reflects which handlers are configured + (i.e., not using the default "not supported" handlers). + + Returns: + ClientTasksCapability if any handlers are provided, None otherwise + """ + has_list = self.list_tasks is not default_list_tasks_handler + has_cancel = self.cancel_task is not default_cancel_task_handler + has_sampling = self.augmented_sampling is not default_task_augmented_sampling + has_elicitation = self.augmented_elicitation is not default_task_augmented_elicitation + + # If no handlers are provided, return None + if not any([has_list, has_cancel, has_sampling, has_elicitation]): + return None + + # Build requests capability if any request handlers are provided + requests_capability: types.ClientTasksRequestsCapability | None = None + if has_sampling or has_elicitation: + requests_capability = types.ClientTasksRequestsCapability( + sampling=types.TasksSamplingCapability(createMessage=types.TasksCreateMessageCapability()) + if has_sampling + else None, + elicitation=types.TasksElicitationCapability(create=types.TasksCreateElicitationCapability()) + if has_elicitation + else None, + ) + + return types.ClientTasksCapability( + list=types.TasksListCapability() if has_list else None, + cancel=types.TasksCancelCapability() if has_cancel else None, + requests=requests_capability, + ) + + @staticmethod + def handles_request(request: types.ServerRequest) -> bool: + """Check if this handler handles the given request type.""" + return isinstance( + request.root, + types.GetTaskRequest | types.GetTaskPayloadRequest | types.ListTasksRequest | types.CancelTaskRequest, + ) + + async def handle_request( + self, + ctx: RequestContext["ClientSession", Any], + responder: RequestResponder[types.ServerRequest, types.ClientResult], + ) -> None: + """Handle a task-related request from the server. + + Call handles_request() first to check if this handler can handle the request. + """ + client_response_type: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter( + types.ClientResult | types.ErrorData + ) + + match responder.request.root: + case types.GetTaskRequest(params=params): + response = await self.get_task(ctx, params) + client_response = client_response_type.validate_python(response) + await responder.respond(client_response) + + case types.GetTaskPayloadRequest(params=params): + response = await self.get_task_result(ctx, params) + client_response = client_response_type.validate_python(response) + await responder.respond(client_response) + + case types.ListTasksRequest(params=params): + response = await self.list_tasks(ctx, params) + client_response = client_response_type.validate_python(response) + await responder.respond(client_response) + + case types.CancelTaskRequest(params=params): + response = await self.cancel_task(ctx, params) + client_response = client_response_type.validate_python(response) + await responder.respond(client_response) + + case _: # pragma: no cover + raise ValueError(f"Unhandled request type: {type(responder.request.root)}") + + +# Backwards compatibility aliases +default_task_augmented_sampling_callback = default_task_augmented_sampling +default_task_augmented_elicitation_callback = default_task_augmented_elicitation diff --git a/src/mcp/client/experimental/tasks.py b/src/mcp/client/experimental/tasks.py new file mode 100644 index 0000000000..ce9c387462 --- /dev/null +++ b/src/mcp/client/experimental/tasks.py @@ -0,0 +1,224 @@ +""" +Experimental client-side task support. + +This module provides client methods for interacting with MCP tasks. + +WARNING: These APIs are experimental and may change without notice. + +Example: + # Call a tool as a task + result = await session.experimental.call_tool_as_task("tool_name", {"arg": "value"}) + task_id = result.task.taskId + + # Get task status + status = await session.experimental.get_task(task_id) + + # Get task result when complete + if status.status == "completed": + result = await session.experimental.get_task_result(task_id, CallToolResult) + + # List all tasks + tasks = await session.experimental.list_tasks() + + # Cancel a task + await session.experimental.cancel_task(task_id) +""" + +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Any, TypeVar + +import mcp.types as types +from mcp.shared.experimental.tasks.polling import poll_until_terminal + +if TYPE_CHECKING: + from mcp.client.session import ClientSession + +ResultT = TypeVar("ResultT", bound=types.Result) + + +class ExperimentalClientFeatures: + """ + Experimental client features for tasks and other experimental APIs. + + WARNING: These APIs are experimental and may change without notice. + + Access via session.experimental: + status = await session.experimental.get_task(task_id) + """ + + def __init__(self, session: "ClientSession") -> None: + self._session = session + + async def call_tool_as_task( + self, + name: str, + arguments: dict[str, Any] | None = None, + *, + ttl: int = 60000, + meta: dict[str, Any] | None = None, + ) -> types.CreateTaskResult: + """Call a tool as a task, returning a CreateTaskResult for polling. + + This is a convenience method for calling tools that support task execution. + The server will return a task reference instead of the immediate result, + which can then be polled via `get_task()` and retrieved via `get_task_result()`. + + Args: + name: The tool name + arguments: Tool arguments + ttl: Task time-to-live in milliseconds (default: 60000 = 1 minute) + meta: Optional metadata to include in the request + + Returns: + CreateTaskResult containing the task reference + + Example: + # Create task + result = await session.experimental.call_tool_as_task( + "long_running_tool", {"input": "data"} + ) + task_id = result.task.taskId + + # Poll for completion + while True: + status = await session.experimental.get_task(task_id) + if status.status == "completed": + break + await asyncio.sleep(0.5) + + # Get result + final = await session.experimental.get_task_result(task_id, CallToolResult) + """ + _meta: types.RequestParams.Meta | None = None + if meta is not None: + _meta = types.RequestParams.Meta(**meta) + + return await self._session.send_request( + types.ClientRequest( + types.CallToolRequest( + params=types.CallToolRequestParams( + name=name, + arguments=arguments, + task=types.TaskMetadata(ttl=ttl), + _meta=_meta, + ), + ) + ), + types.CreateTaskResult, + ) + + async def get_task(self, task_id: str) -> types.GetTaskResult: + """ + Get the current status of a task. + + Args: + task_id: The task identifier + + Returns: + GetTaskResult containing the task status and metadata + """ + return await self._session.send_request( + types.ClientRequest( + types.GetTaskRequest( + params=types.GetTaskRequestParams(taskId=task_id), + ) + ), + types.GetTaskResult, + ) + + async def get_task_result( + self, + task_id: str, + result_type: type[ResultT], + ) -> ResultT: + """ + Get the result of a completed task. + + The result type depends on the original request type: + - tools/call tasks return CallToolResult + - Other request types return their corresponding result type + + Args: + task_id: The task identifier + result_type: The expected result type (e.g., CallToolResult) + + Returns: + The task result, validated against result_type + """ + return await self._session.send_request( + types.ClientRequest( + types.GetTaskPayloadRequest( + params=types.GetTaskPayloadRequestParams(taskId=task_id), + ) + ), + result_type, + ) + + async def list_tasks( + self, + cursor: str | None = None, + ) -> types.ListTasksResult: + """ + List all tasks. + + Args: + cursor: Optional pagination cursor + + Returns: + ListTasksResult containing tasks and optional next cursor + """ + params = types.PaginatedRequestParams(cursor=cursor) if cursor else None + return await self._session.send_request( + types.ClientRequest( + types.ListTasksRequest(params=params), + ), + types.ListTasksResult, + ) + + async def cancel_task(self, task_id: str) -> types.CancelTaskResult: + """ + Cancel a running task. + + Args: + task_id: The task identifier + + Returns: + CancelTaskResult with the updated task state + """ + return await self._session.send_request( + types.ClientRequest( + types.CancelTaskRequest( + params=types.CancelTaskRequestParams(taskId=task_id), + ) + ), + types.CancelTaskResult, + ) + + async def poll_task(self, task_id: str) -> AsyncIterator[types.GetTaskResult]: + """ + Poll a task until it reaches a terminal status. + + Yields GetTaskResult for each poll, allowing the caller to react to + status changes (e.g., handle input_required). Exits when task reaches + a terminal status (completed, failed, cancelled). + + Respects the pollInterval hint from the server. + + Args: + task_id: The task identifier + + Yields: + GetTaskResult for each poll + + Example: + async for status in session.experimental.poll_task(task_id): + print(f"Status: {status.status}") + if status.status == "input_required": + # Handle elicitation request via tasks/result + pass + + # Task is now terminal, get the result + result = await session.experimental.get_task_result(task_id, CallToolResult) + """ + async for status in poll_until_terminal(self.get_task, task_id): + yield status diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index be47d681fb..53fc53a1f3 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -8,6 +8,8 @@ from typing_extensions import deprecated import mcp.types as types +from mcp.client.experimental import ExperimentalClientFeatures +from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers from mcp.shared.context import RequestContext from mcp.shared.message import SessionMessage from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder @@ -118,6 +120,8 @@ def __init__( logging_callback: LoggingFnT | None = None, message_handler: MessageHandlerFnT | None = None, client_info: types.Implementation | None = None, + *, + experimental_task_handlers: ExperimentalTaskHandlers | None = None, ) -> None: super().__init__( read_stream, @@ -134,6 +138,10 @@ def __init__( self._message_handler = message_handler or _default_message_handler self._tool_output_schemas: dict[str, dict[str, Any] | None] = {} self._server_capabilities: types.ServerCapabilities | None = None + self._experimental_features: ExperimentalClientFeatures | None = None + + # Experimental: Task handlers (use defaults if not provided) + self._task_handlers = experimental_task_handlers or ExperimentalTaskHandlers() async def initialize(self) -> types.InitializeResult: sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None @@ -164,6 +172,7 @@ async def initialize(self) -> types.InitializeResult: elicitation=elicitation, experimental=None, roots=roots, + tasks=self._task_handlers.build_capability(), ), clientInfo=self._client_info, ), @@ -188,6 +197,20 @@ def get_server_capabilities(self) -> types.ServerCapabilities | None: """ return self._server_capabilities + @property + def experimental(self) -> ExperimentalClientFeatures: + """Experimental APIs for tasks and other features. + + WARNING: These APIs are experimental and may change without notice. + + Example: + status = await session.experimental.get_task(task_id) + result = await session.experimental.get_task_result(task_id, CallToolResult) + """ + if self._experimental_features is None: + self._experimental_features = ExperimentalClientFeatures(self) + return self._experimental_features + async def send_ping(self) -> types.EmptyResult: """Send a ping request.""" return await self.send_request( @@ -521,16 +544,31 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques lifespan_context=None, ) + # Delegate to experimental task handler if applicable + if self._task_handlers.handles_request(responder.request): + with responder: + await self._task_handlers.handle_request(ctx, responder) + return None + + # Core request handling match responder.request.root: case types.CreateMessageRequest(params=params): with responder: - response = await self._sampling_callback(ctx, params) + # Check if this is a task-augmented request + if params.task is not None: + response = await self._task_handlers.augmented_sampling(ctx, params, params.task) + else: + response = await self._sampling_callback(ctx, params) client_response = ClientResponse.validate_python(response) await responder.respond(client_response) case types.ElicitRequest(params=params): with responder: - response = await self._elicitation_callback(ctx, params) + # Check if this is a task-augmented request + if params.task is not None: + response = await self._task_handlers.augmented_elicitation(ctx, params, params.task) + else: + response = await self._elicitation_callback(ctx, params) client_response = ClientResponse.validate_python(response) await responder.respond(client_response) @@ -544,6 +582,11 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques with responder: return await responder.respond(types.ClientResult(root=types.EmptyResult())) + case _: # pragma: no cover + pass # Task requests handled above by _task_handlers + + return None + async def _handle_incoming( self, req: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception, diff --git a/src/mcp/server/experimental/__init__.py b/src/mcp/server/experimental/__init__.py new file mode 100644 index 0000000000..824bb8b8be --- /dev/null +++ b/src/mcp/server/experimental/__init__.py @@ -0,0 +1,11 @@ +""" +Server-side experimental features. + +WARNING: These APIs are experimental and may change without notice. + +Import directly from submodules: +- mcp.server.experimental.task_context.ServerTaskContext +- mcp.server.experimental.task_support.TaskSupport +- mcp.server.experimental.task_result_handler.TaskResultHandler +- mcp.server.experimental.request_context.Experimental +""" diff --git a/src/mcp/server/experimental/request_context.py b/src/mcp/server/experimental/request_context.py new file mode 100644 index 0000000000..78e75beb6a --- /dev/null +++ b/src/mcp/server/experimental/request_context.py @@ -0,0 +1,238 @@ +""" +Experimental request context features. + +This module provides the Experimental class which gives access to experimental +features within a request context, such as task-augmented request handling. + +WARNING: These APIs are experimental and may change without notice. +""" + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any + +from mcp.server.experimental.task_context import ServerTaskContext +from mcp.server.experimental.task_support import TaskSupport +from mcp.server.session import ServerSession +from mcp.shared.exceptions import McpError +from mcp.shared.experimental.tasks.helpers import MODEL_IMMEDIATE_RESPONSE_KEY, is_terminal +from mcp.types import ( + METHOD_NOT_FOUND, + TASK_FORBIDDEN, + TASK_REQUIRED, + ClientCapabilities, + CreateTaskResult, + ErrorData, + Result, + TaskExecutionMode, + TaskMetadata, + Tool, +) + + +@dataclass +class Experimental: + """ + Experimental features context for task-augmented requests. + + Provides helpers for validating task execution compatibility and + running tasks with automatic lifecycle management. + + WARNING: This API is experimental and may change without notice. + """ + + task_metadata: TaskMetadata | None = None + _client_capabilities: ClientCapabilities | None = field(default=None, repr=False) + _session: ServerSession | None = field(default=None, repr=False) + _task_support: TaskSupport | None = field(default=None, repr=False) + + @property + def is_task(self) -> bool: + """Check if this request is task-augmented.""" + return self.task_metadata is not None + + @property + def client_supports_tasks(self) -> bool: + """Check if the client declared task support.""" + if self._client_capabilities is None: + return False + return self._client_capabilities.tasks is not None + + def validate_task_mode( + self, + tool_task_mode: TaskExecutionMode | None, + *, + raise_error: bool = True, + ) -> ErrorData | None: + """ + Validate that the request is compatible with the tool's task execution mode. + + Per MCP spec: + - "required": Clients MUST invoke as task. Server returns -32601 if not. + - "forbidden" (or None): Clients MUST NOT invoke as task. Server returns -32601 if they do. + - "optional": Either is acceptable. + + Args: + tool_task_mode: The tool's execution.taskSupport value + ("forbidden", "optional", "required", or None) + raise_error: If True, raises McpError on validation failure. If False, returns ErrorData. + + Returns: + None if valid, ErrorData if invalid and raise_error=False + + Raises: + McpError: If invalid and raise_error=True + """ + + mode = tool_task_mode or TASK_FORBIDDEN + + error: ErrorData | None = None + + if mode == TASK_REQUIRED and not self.is_task: + error = ErrorData( + code=METHOD_NOT_FOUND, + message="This tool requires task-augmented invocation", + ) + elif mode == TASK_FORBIDDEN and self.is_task: + error = ErrorData( + code=METHOD_NOT_FOUND, + message="This tool does not support task-augmented invocation", + ) + + if error is not None and raise_error: + raise McpError(error) + + return error + + def validate_for_tool( + self, + tool: Tool, + *, + raise_error: bool = True, + ) -> ErrorData | None: + """ + Validate that the request is compatible with the given tool. + + Convenience wrapper around validate_task_mode that extracts the mode from a Tool. + + Args: + tool: The Tool definition + raise_error: If True, raises McpError on validation failure. + + Returns: + None if valid, ErrorData if invalid and raise_error=False + """ + mode = tool.execution.taskSupport if tool.execution else None + return self.validate_task_mode(mode, raise_error=raise_error) + + def can_use_tool(self, tool_task_mode: TaskExecutionMode | None) -> bool: + """ + Check if this client can use a tool with the given task mode. + + Useful for filtering tool lists or providing warnings. + Returns False if tool requires "required" but client doesn't support tasks. + + Args: + tool_task_mode: The tool's execution.taskSupport value + + Returns: + True if the client can use this tool, False otherwise + """ + mode = tool_task_mode or TASK_FORBIDDEN + if mode == TASK_REQUIRED and not self.client_supports_tasks: + return False + return True + + async def run_task( + self, + work: Callable[[ServerTaskContext], Awaitable[Result]], + *, + task_id: str | None = None, + model_immediate_response: str | None = None, + ) -> CreateTaskResult: + """ + Create a task, spawn background work, and return CreateTaskResult immediately. + + This is the recommended way to handle task-augmented tool calls. It: + 1. Creates a task in the store + 2. Spawns the work function in a background task + 3. Returns CreateTaskResult immediately + + The work function receives a ServerTaskContext with: + - elicit() for sending elicitation requests + - create_message() for sampling requests + - update_status() for progress updates + - complete()/fail() for finishing the task + + When work() returns a Result, the task is auto-completed with that result. + If work() raises an exception, the task is auto-failed. + + Args: + work: Async function that does the actual work + task_id: Optional task ID (generated if not provided) + model_immediate_response: Optional string to include in _meta as + io.modelcontextprotocol/model-immediate-response + + Returns: + CreateTaskResult to return to the client + + Raises: + RuntimeError: If task support is not enabled or task_metadata is missing + + Example: + @server.call_tool() + async def handle_tool(name: str, args: dict): + ctx = server.request_context + + async def work(task: ServerTaskContext) -> CallToolResult: + result = await task.elicit( + message="Are you sure?", + requestedSchema={"type": "object", ...} + ) + confirmed = result.content.get("confirm", False) + return CallToolResult(content=[TextContent(text="Done" if confirmed else "Cancelled")]) + + return await ctx.experimental.run_task(work) + + WARNING: This API is experimental and may change without notice. + """ + if self._task_support is None: + raise RuntimeError("Task support not enabled. Call server.experimental.enable_tasks() first.") + if self._session is None: + raise RuntimeError("Session not available.") + if self.task_metadata is None: + raise RuntimeError( + "Request is not task-augmented (no task field in params). " + "The client must send a task-augmented request." + ) + + support = self._task_support + # Access task_group via TaskSupport - raises if not in run() context + task_group = support.task_group + + task = await support.store.create_task(self.task_metadata, task_id) + + task_ctx = ServerTaskContext( + task=task, + store=support.store, + session=self._session, + queue=support.queue, + handler=support.handler, + ) + + async def execute() -> None: + try: + result = await work(task_ctx) + if not is_terminal(task_ctx.task.status): + await task_ctx.complete(result) + except Exception as e: + if not is_terminal(task_ctx.task.status): + await task_ctx.fail(str(e)) + + task_group.start_soon(execute) + + meta: dict[str, Any] | None = None + if model_immediate_response is not None: + meta = {MODEL_IMMEDIATE_RESPONSE_KEY: model_immediate_response} + + return CreateTaskResult(task=task, **{"_meta": meta} if meta else {}) diff --git a/src/mcp/server/experimental/session_features.py b/src/mcp/server/experimental/session_features.py new file mode 100644 index 0000000000..4842da5175 --- /dev/null +++ b/src/mcp/server/experimental/session_features.py @@ -0,0 +1,220 @@ +""" +Experimental server session features for server→client task operations. + +This module provides the server-side equivalent of ExperimentalClientFeatures, +allowing the server to send task-augmented requests to the client and poll for results. + +WARNING: These APIs are experimental and may change without notice. +""" + +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Any, TypeVar + +import mcp.types as types +from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages +from mcp.shared.experimental.tasks.capabilities import ( + require_task_augmented_elicitation, + require_task_augmented_sampling, +) +from mcp.shared.experimental.tasks.polling import poll_until_terminal + +if TYPE_CHECKING: + from mcp.server.session import ServerSession + +ResultT = TypeVar("ResultT", bound=types.Result) + + +class ExperimentalServerSessionFeatures: + """ + Experimental server session features for server→client task operations. + + This provides the server-side equivalent of ExperimentalClientFeatures, + allowing the server to send task-augmented requests to the client and + poll for results. + + WARNING: These APIs are experimental and may change without notice. + + Access via session.experimental: + result = await session.experimental.elicit_as_task(...) + """ + + def __init__(self, session: "ServerSession") -> None: + self._session = session + + async def get_task(self, task_id: str) -> types.GetTaskResult: + """ + Send tasks/get to the client to get task status. + + Args: + task_id: The task identifier + + Returns: + GetTaskResult containing the task status + """ + return await self._session.send_request( + types.ServerRequest(types.GetTaskRequest(params=types.GetTaskRequestParams(taskId=task_id))), + types.GetTaskResult, + ) + + async def get_task_result( + self, + task_id: str, + result_type: type[ResultT], + ) -> ResultT: + """ + Send tasks/result to the client to retrieve the final result. + + Args: + task_id: The task identifier + result_type: The expected result type + + Returns: + The task result, validated against result_type + """ + return await self._session.send_request( + types.ServerRequest(types.GetTaskPayloadRequest(params=types.GetTaskPayloadRequestParams(taskId=task_id))), + result_type, + ) + + async def poll_task(self, task_id: str) -> AsyncIterator[types.GetTaskResult]: + """ + Poll a client task until it reaches terminal status. + + Yields GetTaskResult for each poll, allowing the caller to react to + status changes. Exits when task reaches a terminal status. + + Respects the pollInterval hint from the client. + + Args: + task_id: The task identifier + + Yields: + GetTaskResult for each poll + """ + async for status in poll_until_terminal(self.get_task, task_id): + yield status + + async def elicit_as_task( + self, + message: str, + requestedSchema: types.ElicitRequestedSchema, + *, + ttl: int = 60000, + ) -> types.ElicitResult: + """ + Send a task-augmented elicitation to the client and poll until complete. + + The client will create a local task, process the elicitation asynchronously, + and return the result when ready. This method handles the full flow: + 1. Send elicitation with task field + 2. Receive CreateTaskResult from client + 3. Poll client's task until terminal + 4. Retrieve and return the final ElicitResult + + Args: + message: The message to present to the user + requestedSchema: Schema defining the expected response + ttl: Task time-to-live in milliseconds + + Returns: + The client's elicitation response + + Raises: + McpError: If client doesn't support task-augmented elicitation + """ + client_caps = self._session.client_params.capabilities if self._session.client_params else None + require_task_augmented_elicitation(client_caps) + + create_result = await self._session.send_request( + types.ServerRequest( + types.ElicitRequest( + params=types.ElicitRequestFormParams( + message=message, + requestedSchema=requestedSchema, + task=types.TaskMetadata(ttl=ttl), + ) + ) + ), + types.CreateTaskResult, + ) + + task_id = create_result.task.taskId + + async for _ in self.poll_task(task_id): + pass + + return await self.get_task_result(task_id, types.ElicitResult) + + async def create_message_as_task( + self, + messages: list[types.SamplingMessage], + *, + max_tokens: int, + ttl: int = 60000, + system_prompt: str | None = None, + include_context: types.IncludeContext | None = None, + temperature: float | None = None, + stop_sequences: list[str] | None = None, + metadata: dict[str, Any] | None = None, + model_preferences: types.ModelPreferences | None = None, + tools: list[types.Tool] | None = None, + tool_choice: types.ToolChoice | None = None, + ) -> types.CreateMessageResult: + """ + Send a task-augmented sampling request and poll until complete. + + The client will create a local task, process the sampling request + asynchronously, and return the result when ready. + + Args: + messages: The conversation messages for sampling + max_tokens: Maximum tokens in the response + ttl: Task time-to-live in milliseconds + system_prompt: Optional system prompt + include_context: Context inclusion strategy + temperature: Sampling temperature + stop_sequences: Stop sequences + metadata: Additional metadata + model_preferences: Model selection preferences + tools: Optional list of tools the LLM can use during sampling + tool_choice: Optional control over tool usage behavior + + Returns: + The sampling result from the client + + Raises: + McpError: If client doesn't support task-augmented sampling or tools + ValueError: If tool_use or tool_result message structure is invalid + """ + client_caps = self._session.client_params.capabilities if self._session.client_params else None + require_task_augmented_sampling(client_caps) + validate_sampling_tools(client_caps, tools, tool_choice) + validate_tool_use_result_messages(messages) + + create_result = await self._session.send_request( + types.ServerRequest( + types.CreateMessageRequest( + params=types.CreateMessageRequestParams( + messages=messages, + maxTokens=max_tokens, + systemPrompt=system_prompt, + includeContext=include_context, + temperature=temperature, + stopSequences=stop_sequences, + metadata=metadata, + modelPreferences=model_preferences, + tools=tools, + toolChoice=tool_choice, + task=types.TaskMetadata(ttl=ttl), + ) + ) + ), + types.CreateTaskResult, + ) + + task_id = create_result.task.taskId + + async for _ in self.poll_task(task_id): + pass + + return await self.get_task_result(task_id, types.CreateMessageResult) diff --git a/src/mcp/server/experimental/task_context.py b/src/mcp/server/experimental/task_context.py new file mode 100644 index 0000000000..e6e14fc938 --- /dev/null +++ b/src/mcp/server/experimental/task_context.py @@ -0,0 +1,612 @@ +""" +ServerTaskContext - Server-integrated task context with elicitation and sampling. + +This wraps the pure TaskContext and adds server-specific functionality: +- Elicitation (task.elicit()) +- Sampling (task.create_message()) +- Status notifications +""" + +from typing import Any + +import anyio + +from mcp.server.experimental.task_result_handler import TaskResultHandler +from mcp.server.session import ServerSession +from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages +from mcp.shared.exceptions import McpError +from mcp.shared.experimental.tasks.capabilities import ( + require_task_augmented_elicitation, + require_task_augmented_sampling, +) +from mcp.shared.experimental.tasks.context import TaskContext +from mcp.shared.experimental.tasks.message_queue import QueuedMessage, TaskMessageQueue +from mcp.shared.experimental.tasks.resolver import Resolver +from mcp.shared.experimental.tasks.store import TaskStore +from mcp.types import ( + INVALID_REQUEST, + TASK_STATUS_INPUT_REQUIRED, + TASK_STATUS_WORKING, + ClientCapabilities, + CreateMessageResult, + CreateTaskResult, + ElicitationCapability, + ElicitRequestedSchema, + ElicitResult, + ErrorData, + IncludeContext, + ModelPreferences, + RequestId, + Result, + SamplingCapability, + SamplingMessage, + ServerNotification, + Task, + TaskMetadata, + TaskStatusNotification, + TaskStatusNotificationParams, + Tool, + ToolChoice, +) + + +class ServerTaskContext: + """ + Server-integrated task context with elicitation and sampling. + + This wraps a pure TaskContext and adds server-specific functionality: + - elicit() for sending elicitation requests to the client + - create_message() for sampling requests + - Status notifications via the session + + Example: + async def my_task_work(task: ServerTaskContext) -> CallToolResult: + await task.update_status("Starting...") + + result = await task.elicit( + message="Continue?", + requestedSchema={"type": "object", "properties": {"ok": {"type": "boolean"}}} + ) + + if result.content.get("ok"): + return CallToolResult(content=[TextContent(text="Done!")]) + else: + return CallToolResult(content=[TextContent(text="Cancelled")]) + """ + + def __init__( + self, + *, + task: Task, + store: TaskStore, + session: ServerSession, + queue: TaskMessageQueue, + handler: TaskResultHandler | None = None, + ): + """ + Create a ServerTaskContext. + + Args: + task: The Task object + store: The task store + session: The server session + queue: The message queue for elicitation/sampling + handler: The result handler for response routing (required for elicit/create_message) + """ + self._ctx = TaskContext(task=task, store=store) + self._session = session + self._queue = queue + self._handler = handler + self._store = store + + # Delegate pure properties to inner context + + @property + def task_id(self) -> str: + """The task identifier.""" + return self._ctx.task_id + + @property + def task(self) -> Task: + """The current task state.""" + return self._ctx.task + + @property + def is_cancelled(self) -> bool: + """Whether cancellation has been requested.""" + return self._ctx.is_cancelled + + def request_cancellation(self) -> None: + """Request cancellation of this task.""" + self._ctx.request_cancellation() + + # Enhanced methods with notifications + + async def update_status(self, message: str, *, notify: bool = True) -> None: + """ + Update the task's status message. + + Args: + message: The new status message + notify: Whether to send a notification to the client + """ + await self._ctx.update_status(message) + if notify: + await self._send_notification() + + async def complete(self, result: Result, *, notify: bool = True) -> None: + """ + Mark the task as completed with the given result. + + Args: + result: The task result + notify: Whether to send a notification to the client + """ + await self._ctx.complete(result) + if notify: + await self._send_notification() + + async def fail(self, error: str, *, notify: bool = True) -> None: + """ + Mark the task as failed with an error message. + + Args: + error: The error message + notify: Whether to send a notification to the client + """ + await self._ctx.fail(error) + if notify: + await self._send_notification() + + async def _send_notification(self) -> None: + """Send a task status notification to the client.""" + task = self._ctx.task + await self._session.send_notification( + ServerNotification( + TaskStatusNotification( + params=TaskStatusNotificationParams( + taskId=task.taskId, + status=task.status, + statusMessage=task.statusMessage, + createdAt=task.createdAt, + lastUpdatedAt=task.lastUpdatedAt, + ttl=task.ttl, + pollInterval=task.pollInterval, + ) + ) + ) + ) + + # Server-specific methods: elicitation and sampling + + def _check_elicitation_capability(self) -> None: + """Check if the client supports elicitation.""" + if not self._session.check_client_capability(ClientCapabilities(elicitation=ElicitationCapability())): + raise McpError( + ErrorData( + code=INVALID_REQUEST, + message="Client does not support elicitation capability", + ) + ) + + def _check_sampling_capability(self) -> None: + """Check if the client supports sampling.""" + if not self._session.check_client_capability(ClientCapabilities(sampling=SamplingCapability())): + raise McpError( + ErrorData( + code=INVALID_REQUEST, + message="Client does not support sampling capability", + ) + ) + + async def elicit( + self, + message: str, + requestedSchema: ElicitRequestedSchema, + ) -> ElicitResult: + """ + Send an elicitation request via the task message queue. + + This method: + 1. Checks client capability + 2. Updates task status to "input_required" + 3. Queues the elicitation request + 4. Waits for the response (delivered via tasks/result round-trip) + 5. Updates task status back to "working" + 6. Returns the result + + Args: + message: The message to present to the user + requestedSchema: Schema defining the expected response structure + + Returns: + The client's response + + Raises: + McpError: If client doesn't support elicitation capability + """ + self._check_elicitation_capability() + + if self._handler is None: + raise RuntimeError("handler is required for elicit(). Pass handler= to ServerTaskContext.") + + # Update status to input_required + await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED) + + # Build the request using session's helper + request = self._session._build_elicit_form_request( # pyright: ignore[reportPrivateUsage] + message=message, + requestedSchema=requestedSchema, + related_task_id=self.task_id, + ) + request_id: RequestId = request.id + + resolver: Resolver[dict[str, Any]] = Resolver() + self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage] + + queued = QueuedMessage( + type="request", + message=request, + resolver=resolver, + original_request_id=request_id, + ) + await self._queue.enqueue(self.task_id, queued) + + try: + # Wait for response (routed back via TaskResultHandler) + response_data = await resolver.wait() + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + return ElicitResult.model_validate(response_data) + except anyio.get_cancelled_exc_class(): # pragma: no cover + # Coverage can't track async exception handlers reliably. + # This path is tested in test_elicit_restores_status_on_cancellation + # which verifies status is restored to "working" after cancellation. + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + raise + + async def elicit_url( + self, + message: str, + url: str, + elicitation_id: str, + ) -> ElicitResult: + """ + Send a URL mode elicitation request via the task message queue. + + This directs the user to an external URL for out-of-band interactions + like OAuth flows, credential collection, or payment processing. + + This method: + 1. Checks client capability + 2. Updates task status to "input_required" + 3. Queues the elicitation request + 4. Waits for the response (delivered via tasks/result round-trip) + 5. Updates task status back to "working" + 6. Returns the result + + Args: + message: Human-readable explanation of why the interaction is needed + url: The URL the user should navigate to + elicitation_id: Unique identifier for tracking this elicitation + + Returns: + The client's response indicating acceptance, decline, or cancellation + + Raises: + McpError: If client doesn't support elicitation capability + RuntimeError: If handler is not configured + """ + self._check_elicitation_capability() + + if self._handler is None: + raise RuntimeError("handler is required for elicit_url(). Pass handler= to ServerTaskContext.") + + # Update status to input_required + await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED) + + # Build the request using session's helper + request = self._session._build_elicit_url_request( # pyright: ignore[reportPrivateUsage] + message=message, + url=url, + elicitation_id=elicitation_id, + related_task_id=self.task_id, + ) + request_id: RequestId = request.id + + resolver: Resolver[dict[str, Any]] = Resolver() + self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage] + + queued = QueuedMessage( + type="request", + message=request, + resolver=resolver, + original_request_id=request_id, + ) + await self._queue.enqueue(self.task_id, queued) + + try: + # Wait for response (routed back via TaskResultHandler) + response_data = await resolver.wait() + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + return ElicitResult.model_validate(response_data) + except anyio.get_cancelled_exc_class(): # pragma: no cover + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + raise + + async def create_message( + self, + messages: list[SamplingMessage], + *, + max_tokens: int, + system_prompt: str | None = None, + include_context: IncludeContext | None = None, + temperature: float | None = None, + stop_sequences: list[str] | None = None, + metadata: dict[str, Any] | None = None, + model_preferences: ModelPreferences | None = None, + tools: list[Tool] | None = None, + tool_choice: ToolChoice | None = None, + ) -> CreateMessageResult: + """ + Send a sampling request via the task message queue. + + This method: + 1. Checks client capability + 2. Updates task status to "input_required" + 3. Queues the sampling request + 4. Waits for the response (delivered via tasks/result round-trip) + 5. Updates task status back to "working" + 6. Returns the result + + Args: + messages: The conversation messages for sampling + max_tokens: Maximum tokens in the response + system_prompt: Optional system prompt + include_context: Context inclusion strategy + temperature: Sampling temperature + stop_sequences: Stop sequences + metadata: Additional metadata + model_preferences: Model selection preferences + tools: Optional list of tools the LLM can use during sampling + tool_choice: Optional control over tool usage behavior + + Returns: + The sampling result from the client + + Raises: + McpError: If client doesn't support sampling capability or tools + ValueError: If tool_use or tool_result message structure is invalid + """ + self._check_sampling_capability() + client_caps = self._session.client_params.capabilities if self._session.client_params else None + validate_sampling_tools(client_caps, tools, tool_choice) + validate_tool_use_result_messages(messages) + + if self._handler is None: + raise RuntimeError("handler is required for create_message(). Pass handler= to ServerTaskContext.") + + # Update status to input_required + await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED) + + # Build the request using session's helper + request = self._session._build_create_message_request( # pyright: ignore[reportPrivateUsage] + messages=messages, + max_tokens=max_tokens, + system_prompt=system_prompt, + include_context=include_context, + temperature=temperature, + stop_sequences=stop_sequences, + metadata=metadata, + model_preferences=model_preferences, + tools=tools, + tool_choice=tool_choice, + related_task_id=self.task_id, + ) + request_id: RequestId = request.id + + resolver: Resolver[dict[str, Any]] = Resolver() + self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage] + + queued = QueuedMessage( + type="request", + message=request, + resolver=resolver, + original_request_id=request_id, + ) + await self._queue.enqueue(self.task_id, queued) + + try: + # Wait for response (routed back via TaskResultHandler) + response_data = await resolver.wait() + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + return CreateMessageResult.model_validate(response_data) + except anyio.get_cancelled_exc_class(): # pragma: no cover + # Coverage can't track async exception handlers reliably. + # This path is tested in test_create_message_restores_status_on_cancellation + # which verifies status is restored to "working" after cancellation. + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + raise + + async def elicit_as_task( + self, + message: str, + requestedSchema: ElicitRequestedSchema, + *, + ttl: int = 60000, + ) -> ElicitResult: + """ + Send a task-augmented elicitation via the queue, then poll client. + + This is for use inside a task-augmented tool call when you want the client + to handle the elicitation as its own task. The elicitation request is queued + and delivered when the client calls tasks/result. After the client responds + with CreateTaskResult, we poll the client's task until complete. + + Args: + message: The message to present to the user + requestedSchema: Schema defining the expected response structure + ttl: Task time-to-live in milliseconds for the client's task + + Returns: + The client's elicitation response + + Raises: + McpError: If client doesn't support task-augmented elicitation + RuntimeError: If handler is not configured + """ + client_caps = self._session.client_params.capabilities if self._session.client_params else None + require_task_augmented_elicitation(client_caps) + + if self._handler is None: + raise RuntimeError("handler is required for elicit_as_task()") + + # Update status to input_required + await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED) + + request = self._session._build_elicit_form_request( # pyright: ignore[reportPrivateUsage] + message=message, + requestedSchema=requestedSchema, + related_task_id=self.task_id, + task=TaskMetadata(ttl=ttl), + ) + request_id: RequestId = request.id + + resolver: Resolver[dict[str, Any]] = Resolver() + self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage] + + queued = QueuedMessage( + type="request", + message=request, + resolver=resolver, + original_request_id=request_id, + ) + await self._queue.enqueue(self.task_id, queued) + + try: + # Wait for initial response (CreateTaskResult from client) + response_data = await resolver.wait() + create_result = CreateTaskResult.model_validate(response_data) + client_task_id = create_result.task.taskId + + # Poll the client's task using session.experimental + async for _ in self._session.experimental.poll_task(client_task_id): + pass + + # Get final result from client + result = await self._session.experimental.get_task_result( + client_task_id, + ElicitResult, + ) + + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + return result + + except anyio.get_cancelled_exc_class(): # pragma: no cover + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + raise + + async def create_message_as_task( + self, + messages: list[SamplingMessage], + *, + max_tokens: int, + ttl: int = 60000, + system_prompt: str | None = None, + include_context: IncludeContext | None = None, + temperature: float | None = None, + stop_sequences: list[str] | None = None, + metadata: dict[str, Any] | None = None, + model_preferences: ModelPreferences | None = None, + tools: list[Tool] | None = None, + tool_choice: ToolChoice | None = None, + ) -> CreateMessageResult: + """ + Send a task-augmented sampling request via the queue, then poll client. + + This is for use inside a task-augmented tool call when you want the client + to handle the sampling as its own task. The request is queued and delivered + when the client calls tasks/result. After the client responds with + CreateTaskResult, we poll the client's task until complete. + + Args: + messages: The conversation messages for sampling + max_tokens: Maximum tokens in the response + ttl: Task time-to-live in milliseconds for the client's task + system_prompt: Optional system prompt + include_context: Context inclusion strategy + temperature: Sampling temperature + stop_sequences: Stop sequences + metadata: Additional metadata + model_preferences: Model selection preferences + tools: Optional list of tools the LLM can use during sampling + tool_choice: Optional control over tool usage behavior + + Returns: + The sampling result from the client + + Raises: + McpError: If client doesn't support task-augmented sampling or tools + ValueError: If tool_use or tool_result message structure is invalid + RuntimeError: If handler is not configured + """ + client_caps = self._session.client_params.capabilities if self._session.client_params else None + require_task_augmented_sampling(client_caps) + validate_sampling_tools(client_caps, tools, tool_choice) + validate_tool_use_result_messages(messages) + + if self._handler is None: + raise RuntimeError("handler is required for create_message_as_task()") + + # Update status to input_required + await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED) + + # Build request WITH task field for task-augmented sampling + request = self._session._build_create_message_request( # pyright: ignore[reportPrivateUsage] + messages=messages, + max_tokens=max_tokens, + system_prompt=system_prompt, + include_context=include_context, + temperature=temperature, + stop_sequences=stop_sequences, + metadata=metadata, + model_preferences=model_preferences, + tools=tools, + tool_choice=tool_choice, + related_task_id=self.task_id, + task=TaskMetadata(ttl=ttl), + ) + request_id: RequestId = request.id + + resolver: Resolver[dict[str, Any]] = Resolver() + self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage] + + queued = QueuedMessage( + type="request", + message=request, + resolver=resolver, + original_request_id=request_id, + ) + await self._queue.enqueue(self.task_id, queued) + + try: + # Wait for initial response (CreateTaskResult from client) + response_data = await resolver.wait() + create_result = CreateTaskResult.model_validate(response_data) + client_task_id = create_result.task.taskId + + # Poll the client's task using session.experimental + async for _ in self._session.experimental.poll_task(client_task_id): + pass + + # Get final result from client + result = await self._session.experimental.get_task_result( + client_task_id, + CreateMessageResult, + ) + + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + return result + + except anyio.get_cancelled_exc_class(): # pragma: no cover + await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING) + raise diff --git a/src/mcp/server/experimental/task_result_handler.py b/src/mcp/server/experimental/task_result_handler.py new file mode 100644 index 0000000000..0b869216e8 --- /dev/null +++ b/src/mcp/server/experimental/task_result_handler.py @@ -0,0 +1,235 @@ +""" +TaskResultHandler - Integrated handler for tasks/result endpoint. + +This implements the dequeue-send-wait pattern from the MCP Tasks spec: +1. Dequeue all pending messages for the task +2. Send them to the client via transport with relatedRequestId routing +3. Wait if task is not in terminal state +4. Return final result when task completes + +This is the core of the task message queue pattern. +""" + +import logging +from typing import Any + +import anyio + +from mcp.server.session import ServerSession +from mcp.shared.exceptions import McpError +from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY, is_terminal +from mcp.shared.experimental.tasks.message_queue import TaskMessageQueue +from mcp.shared.experimental.tasks.resolver import Resolver +from mcp.shared.experimental.tasks.store import TaskStore +from mcp.shared.message import ServerMessageMetadata, SessionMessage +from mcp.types import ( + INVALID_PARAMS, + ErrorData, + GetTaskPayloadRequest, + GetTaskPayloadResult, + JSONRPCMessage, + RelatedTaskMetadata, + RequestId, +) + +logger = logging.getLogger(__name__) + + +class TaskResultHandler: + """ + Handler for tasks/result that implements the message queue pattern. + + This handler: + 1. Dequeues pending messages (elicitations, notifications) for the task + 2. Sends them to the client via the response stream + 3. Waits for responses and resolves them back to callers + 4. Blocks until task reaches terminal state + 5. Returns the final result + + Usage: + # Create handler with store and queue + handler = TaskResultHandler(task_store, message_queue) + + # Register it with the server + @server.experimental.get_task_result() + async def handle_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult: + ctx = server.request_context + return await handler.handle(req, ctx.session, ctx.request_id) + + # Or use the convenience method + handler.register(server) + """ + + def __init__( + self, + store: TaskStore, + queue: TaskMessageQueue, + ): + self._store = store + self._queue = queue + # Map from internal request ID to resolver for routing responses + self._pending_requests: dict[RequestId, Resolver[dict[str, Any]]] = {} + + async def send_message( + self, + session: ServerSession, + message: SessionMessage, + ) -> None: + """ + Send a message via the session. + + This is a helper for delivering queued task messages. + """ + await session.send_message(message) + + async def handle( + self, + request: GetTaskPayloadRequest, + session: ServerSession, + request_id: RequestId, + ) -> GetTaskPayloadResult: + """ + Handle a tasks/result request. + + This implements the dequeue-send-wait loop: + 1. Dequeue all pending messages + 2. Send each via transport with relatedRequestId = this request's ID + 3. If task not terminal, wait for status change + 4. Loop until task is terminal + 5. Return final result + + Args: + request: The GetTaskPayloadRequest + session: The server session for sending messages + request_id: The request ID for relatedRequestId routing + + Returns: + GetTaskPayloadResult with the task's final payload + """ + task_id = request.params.taskId + + while True: + task = await self._store.get_task(task_id) + if task is None: + raise McpError( + ErrorData( + code=INVALID_PARAMS, + message=f"Task not found: {task_id}", + ) + ) + + await self._deliver_queued_messages(task_id, session, request_id) + + # If task is terminal, return result + if is_terminal(task.status): + result = await self._store.get_result(task_id) + # GetTaskPayloadResult is a Result with extra="allow" + # The stored result contains the actual payload data + # Per spec: tasks/result MUST include _meta with related-task metadata + related_task = RelatedTaskMetadata(taskId=task_id) + related_task_meta: dict[str, Any] = {RELATED_TASK_METADATA_KEY: related_task.model_dump(by_alias=True)} + if result is not None: + result_data = result.model_dump(by_alias=True) + existing_meta: dict[str, Any] = result_data.get("_meta") or {} + result_data["_meta"] = {**existing_meta, **related_task_meta} + return GetTaskPayloadResult.model_validate(result_data) + return GetTaskPayloadResult.model_validate({"_meta": related_task_meta}) + + # Wait for task update (status change or new messages) + await self._wait_for_task_update(task_id) + + async def _deliver_queued_messages( + self, + task_id: str, + session: ServerSession, + request_id: RequestId, + ) -> None: + """ + Dequeue and send all pending messages for a task. + + Each message is sent via the session's write stream with + relatedRequestId set so responses route back to this stream. + """ + while True: + message = await self._queue.dequeue(task_id) + if message is None: + break + + # If this is a request (not notification), wait for response + if message.type == "request" and message.resolver is not None: + # Store the resolver so we can route the response back + original_id = message.original_request_id + if original_id is not None: + self._pending_requests[original_id] = message.resolver + + logger.debug("Delivering queued message for task %s: %s", task_id, message.type) + + # Send the message with relatedRequestId for routing + session_message = SessionMessage( + message=JSONRPCMessage(message.message), + metadata=ServerMessageMetadata(related_request_id=request_id), + ) + await self.send_message(session, session_message) + + async def _wait_for_task_update(self, task_id: str) -> None: + """ + Wait for task to be updated (status change or new message). + + Races between store update and queue message - first one wins. + """ + async with anyio.create_task_group() as tg: + + async def wait_for_store() -> None: + try: + await self._store.wait_for_update(task_id) + except Exception: + pass + finally: + tg.cancel_scope.cancel() + + async def wait_for_queue() -> None: + try: + await self._queue.wait_for_message(task_id) + except Exception: + pass + finally: + tg.cancel_scope.cancel() + + tg.start_soon(wait_for_store) + tg.start_soon(wait_for_queue) + + def route_response(self, request_id: RequestId, response: dict[str, Any]) -> bool: + """ + Route a response back to the waiting resolver. + + This is called when a response arrives for a queued request. + + Args: + request_id: The request ID from the response + response: The response data + + Returns: + True if response was routed, False if no pending request + """ + resolver = self._pending_requests.pop(request_id, None) + if resolver is not None and not resolver.done(): + resolver.set_result(response) + return True + return False + + def route_error(self, request_id: RequestId, error: ErrorData) -> bool: + """ + Route an error back to the waiting resolver. + + Args: + request_id: The request ID from the error response + error: The error data + + Returns: + True if error was routed, False if no pending request + """ + resolver = self._pending_requests.pop(request_id, None) + if resolver is not None and not resolver.done(): + resolver.set_exception(McpError(error)) + return True + return False diff --git a/src/mcp/server/experimental/task_support.py b/src/mcp/server/experimental/task_support.py new file mode 100644 index 0000000000..dbb2ed6d2b --- /dev/null +++ b/src/mcp/server/experimental/task_support.py @@ -0,0 +1,115 @@ +""" +TaskSupport - Configuration for experimental task support. + +This module provides the TaskSupport class which encapsulates all the +infrastructure needed for task-augmented requests: store, queue, and handler. +""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass, field + +import anyio +from anyio.abc import TaskGroup + +from mcp.server.experimental.task_result_handler import TaskResultHandler +from mcp.server.session import ServerSession +from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore +from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue +from mcp.shared.experimental.tasks.store import TaskStore + + +@dataclass +class TaskSupport: + """ + Configuration for experimental task support. + + Encapsulates the task store, message queue, result handler, and task group + for spawning background work. + + When enabled on a server, this automatically: + - Configures response routing for each session + - Provides default handlers for task operations + - Manages a task group for background task execution + + Example: + # Simple in-memory setup + server.experimental.enable_tasks() + + # Custom store/queue for distributed systems + server.experimental.enable_tasks( + store=RedisTaskStore(redis_url), + queue=RedisTaskMessageQueue(redis_url), + ) + """ + + store: TaskStore + queue: TaskMessageQueue + handler: TaskResultHandler = field(init=False) + _task_group: TaskGroup | None = field(init=False, default=None) + + def __post_init__(self) -> None: + """Create the result handler from store and queue.""" + self.handler = TaskResultHandler(self.store, self.queue) + + @property + def task_group(self) -> TaskGroup: + """Get the task group for spawning background work. + + Raises: + RuntimeError: If not within a run() context + """ + if self._task_group is None: + raise RuntimeError("TaskSupport not running. Ensure Server.run() is active.") + return self._task_group + + @asynccontextmanager + async def run(self) -> AsyncIterator[None]: + """ + Run the task support lifecycle. + + This creates a task group for spawning background task work. + Called automatically by Server.run(). + + Usage: + async with task_support.run(): + # Task group is now available + ... + """ + async with anyio.create_task_group() as tg: + self._task_group = tg + try: + yield + finally: + self._task_group = None + + def configure_session(self, session: ServerSession) -> None: + """ + Configure a session for task support. + + This registers the result handler as a response router so that + responses to queued requests (elicitation, sampling) are routed + back to the waiting resolvers. + + Called automatically by Server.run() for each new session. + + Args: + session: The session to configure + """ + session.add_response_router(self.handler) + + @classmethod + def in_memory(cls) -> "TaskSupport": + """ + Create in-memory task support. + + Suitable for development, testing, and single-process servers. + For distributed systems, provide custom store and queue implementations. + + Returns: + TaskSupport configured with in-memory store and queue + """ + return cls( + store=InMemoryTaskStore(), + queue=InMemoryTaskMessageQueue(), + ) diff --git a/src/mcp/server/lowlevel/experimental.py b/src/mcp/server/lowlevel/experimental.py new file mode 100644 index 0000000000..0e6655b3de --- /dev/null +++ b/src/mcp/server/lowlevel/experimental.py @@ -0,0 +1,288 @@ +"""Experimental handlers for the low-level MCP server. + +WARNING: These APIs are experimental and may change without notice. +""" + +from __future__ import annotations + +import logging +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING + +from mcp.server.experimental.task_support import TaskSupport +from mcp.server.lowlevel.func_inspection import create_call_wrapper +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.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue +from mcp.shared.experimental.tasks.store import TaskStore +from mcp.types import ( + INVALID_PARAMS, + CancelTaskRequest, + CancelTaskResult, + ErrorData, + GetTaskPayloadRequest, + GetTaskPayloadResult, + GetTaskRequest, + GetTaskResult, + ListTasksRequest, + ListTasksResult, + ServerCapabilities, + ServerResult, + ServerTasksCapability, + ServerTasksRequestsCapability, + TasksCancelCapability, + TasksListCapability, + TasksToolsCapability, +) + +if TYPE_CHECKING: + from mcp.server.lowlevel.server import Server + +logger = logging.getLogger(__name__) + + +class ExperimentalHandlers: + """Experimental request/notification handlers. + + WARNING: These APIs are experimental and may change without notice. + """ + + def __init__( + self, + server: Server, + request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]], + notification_handlers: dict[type, Callable[..., Awaitable[None]]], + ): + self._server = server + self._request_handlers = request_handlers + self._notification_handlers = notification_handlers + self._task_support: TaskSupport | None = None + + @property + def task_support(self) -> TaskSupport | None: + """Get the task support configuration, if enabled.""" + return self._task_support + + def update_capabilities(self, capabilities: ServerCapabilities) -> None: + # Only add tasks capability if handlers are registered + if not any( + req_type in self._request_handlers + for req_type in [GetTaskRequest, ListTasksRequest, CancelTaskRequest, GetTaskPayloadRequest] + ): + return + + capabilities.tasks = ServerTasksCapability() + if ListTasksRequest in self._request_handlers: + capabilities.tasks.list = TasksListCapability() + if CancelTaskRequest in self._request_handlers: + capabilities.tasks.cancel = TasksCancelCapability() + + capabilities.tasks.requests = ServerTasksRequestsCapability( + tools=TasksToolsCapability() + ) # assuming always supported for now + + def enable_tasks( + self, + store: TaskStore | None = None, + queue: TaskMessageQueue | None = None, + ) -> TaskSupport: + """ + Enable experimental task support. + + This sets up the task infrastructure and auto-registers default handlers + for tasks/get, tasks/result, tasks/list, and tasks/cancel. + + Args: + store: Custom TaskStore implementation (defaults to InMemoryTaskStore) + queue: Custom TaskMessageQueue implementation (defaults to InMemoryTaskMessageQueue) + + Returns: + The TaskSupport configuration object + + Example: + # Simple in-memory setup + server.experimental.enable_tasks() + + # Custom store/queue for distributed systems + server.experimental.enable_tasks( + store=RedisTaskStore(redis_url), + queue=RedisTaskMessageQueue(redis_url), + ) + + WARNING: This API is experimental and may change without notice. + """ + if store is None: + store = InMemoryTaskStore() + if queue is None: + queue = InMemoryTaskMessageQueue() + + self._task_support = TaskSupport(store=store, queue=queue) + + # Auto-register default handlers + self._register_default_task_handlers() + + return self._task_support + + def _register_default_task_handlers(self) -> None: + """Register default handlers for task operations.""" + assert self._task_support is not None + support = self._task_support + + # Register get_task handler if not already registered + if GetTaskRequest not in self._request_handlers: + + async def _default_get_task(req: GetTaskRequest) -> ServerResult: + task = await support.store.get_task(req.params.taskId) + if task is None: + raise McpError( + ErrorData( + code=INVALID_PARAMS, + message=f"Task not found: {req.params.taskId}", + ) + ) + return ServerResult( + GetTaskResult( + taskId=task.taskId, + status=task.status, + statusMessage=task.statusMessage, + createdAt=task.createdAt, + lastUpdatedAt=task.lastUpdatedAt, + ttl=task.ttl, + pollInterval=task.pollInterval, + ) + ) + + self._request_handlers[GetTaskRequest] = _default_get_task + + # Register get_task_result handler if not already registered + if GetTaskPayloadRequest not in self._request_handlers: + + async def _default_get_task_result(req: GetTaskPayloadRequest) -> ServerResult: + ctx = self._server.request_context + result = await support.handler.handle(req, ctx.session, ctx.request_id) + return ServerResult(result) + + self._request_handlers[GetTaskPayloadRequest] = _default_get_task_result + + # Register list_tasks handler if not already registered + if ListTasksRequest not in self._request_handlers: + + async def _default_list_tasks(req: ListTasksRequest) -> ServerResult: + cursor = req.params.cursor if req.params else None + tasks, next_cursor = await support.store.list_tasks(cursor) + return ServerResult(ListTasksResult(tasks=tasks, nextCursor=next_cursor)) + + self._request_handlers[ListTasksRequest] = _default_list_tasks + + # Register cancel_task handler if not already registered + if CancelTaskRequest not in self._request_handlers: + + async def _default_cancel_task(req: CancelTaskRequest) -> ServerResult: + result = await cancel_task(support.store, req.params.taskId) + return ServerResult(result) + + self._request_handlers[CancelTaskRequest] = _default_cancel_task + + def list_tasks( + self, + ) -> Callable[ + [Callable[[ListTasksRequest], Awaitable[ListTasksResult]]], + Callable[[ListTasksRequest], Awaitable[ListTasksResult]], + ]: + """Register a handler for listing tasks. + + WARNING: This API is experimental and may change without notice. + """ + + def decorator( + func: Callable[[ListTasksRequest], Awaitable[ListTasksResult]], + ) -> Callable[[ListTasksRequest], Awaitable[ListTasksResult]]: + logger.debug("Registering handler for ListTasksRequest") + wrapper = create_call_wrapper(func, ListTasksRequest) + + async def handler(req: ListTasksRequest) -> ServerResult: + result = await wrapper(req) + return ServerResult(result) + + self._request_handlers[ListTasksRequest] = handler + return func + + return decorator + + def get_task( + self, + ) -> Callable[ + [Callable[[GetTaskRequest], Awaitable[GetTaskResult]]], Callable[[GetTaskRequest], Awaitable[GetTaskResult]] + ]: + """Register a handler for getting task status. + + WARNING: This API is experimental and may change without notice. + """ + + def decorator( + func: Callable[[GetTaskRequest], Awaitable[GetTaskResult]], + ) -> Callable[[GetTaskRequest], Awaitable[GetTaskResult]]: + logger.debug("Registering handler for GetTaskRequest") + wrapper = create_call_wrapper(func, GetTaskRequest) + + async def handler(req: GetTaskRequest) -> ServerResult: + result = await wrapper(req) + return ServerResult(result) + + self._request_handlers[GetTaskRequest] = handler + return func + + return decorator + + def get_task_result( + self, + ) -> Callable[ + [Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]], + Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]], + ]: + """Register a handler for getting task results/payload. + + WARNING: This API is experimental and may change without notice. + """ + + def decorator( + func: Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]], + ) -> Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]: + logger.debug("Registering handler for GetTaskPayloadRequest") + wrapper = create_call_wrapper(func, GetTaskPayloadRequest) + + async def handler(req: GetTaskPayloadRequest) -> ServerResult: + result = await wrapper(req) + return ServerResult(result) + + self._request_handlers[GetTaskPayloadRequest] = handler + return func + + return decorator + + def cancel_task( + self, + ) -> Callable[ + [Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]], + Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]], + ]: + """Register a handler for cancelling tasks. + + WARNING: This API is experimental and may change without notice. + """ + + def decorator( + func: Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]], + ) -> Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]: + logger.debug("Registering handler for CancelTaskRequest") + wrapper = create_call_wrapper(func, CancelTaskRequest) + + async def handler(req: CancelTaskRequest) -> ServerResult: + result = await wrapper(req) + return ServerResult(result) + + self._request_handlers[CancelTaskRequest] = handler + return func + + return decorator diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index a0617036f9..71cee3154b 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -67,6 +67,7 @@ async def main(): from __future__ import annotations as _annotations +import base64 import contextvars import json import logging @@ -82,6 +83,8 @@ async def main(): from typing_extensions import TypeVar import mcp.types as types +from mcp.server.experimental.request_context import Experimental +from mcp.server.lowlevel.experimental import ExperimentalHandlers from mcp.server.lowlevel.func_inspection import create_call_wrapper from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.server.models import InitializationOptions @@ -155,6 +158,7 @@ def __init__( } self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {} self._tool_cache: dict[str, types.Tool] = {} + self._experimental_handlers: ExperimentalHandlers | None = None logger.debug("Initializing server %r", name) def create_initialization_options( @@ -220,7 +224,7 @@ def get_capabilities( if types.CompleteRequest in self.request_handlers: completions_capability = types.CompletionsCapability() - return types.ServerCapabilities( + capabilities = types.ServerCapabilities( prompts=prompts_capability, resources=resources_capability, tools=tools_capability, @@ -228,6 +232,9 @@ def get_capabilities( experimental=experimental_capabilities, completions=completions_capability, ) + if self._experimental_handlers: + self._experimental_handlers.update_capabilities(capabilities) + return capabilities @property def request_context( @@ -236,6 +243,18 @@ def request_context( """If called outside of a request context, this will raise a LookupError.""" return request_ctx.get() + @property + def experimental(self) -> ExperimentalHandlers: + """Experimental APIs for tasks and other features. + + WARNING: These APIs are experimental and may change without notice. + """ + + # We create this inline so we only add these capabilities _if_ they're actually used + if self._experimental_handlers is None: + self._experimental_handlers = ExperimentalHandlers(self, self.request_handlers, self.notification_handlers) + return self._experimental_handlers + def list_prompts(self): def decorator( func: Callable[[], Awaitable[list[types.Prompt]]] @@ -328,8 +347,6 @@ def create_content(data: str | bytes, mime_type: str | None): mimeType=mime_type or "text/plain", ) case bytes() as data: # pragma: no cover - import base64 - return types.BlobResourceContents( uri=req.params.uri, blob=base64.b64encode(data).decode(), @@ -483,7 +500,13 @@ def call_tool(self, *, validate_input: bool = True): def decorator( func: Callable[ ..., - Awaitable[UnstructuredContent | StructuredContent | CombinationContent | types.CallToolResult], + Awaitable[ + UnstructuredContent + | StructuredContent + | CombinationContent + | types.CallToolResult + | types.CreateTaskResult + ], ], ): logger.debug("Registering handler for CallToolRequest") @@ -509,6 +532,9 @@ async def handler(req: types.CallToolRequest): maybe_structured_content: StructuredContent | None if isinstance(results, types.CallToolResult): return types.ServerResult(results) + elif isinstance(results, types.CreateTaskResult): + # Task-augmented execution returns task info instead of result + return types.ServerResult(results) elif isinstance(results, tuple) and len(results) == 2: # tool returned both structured and unstructured content unstructured_content, maybe_structured_content = cast(CombinationContent, results) @@ -627,6 +653,12 @@ async def run( ) ) + # Configure task support for this session if enabled + task_support = self._experimental_handlers.task_support if self._experimental_handlers else None + if task_support is not None: + task_support.configure_session(session) + await stack.enter_async_context(task_support.run()) + async with anyio.create_task_group() as tg: async for message in session.incoming_messages: logger.debug("Received message: %s", message) @@ -669,13 +701,14 @@ async def _handle_message( async def _handle_request( self, message: RequestResponder[types.ClientRequest, types.ServerResult], - req: Any, + req: types.ClientRequestType, session: ServerSession, lifespan_context: LifespanResultT, raise_exceptions: bool, ): logger.info("Processing request of type %s", type(req).__name__) - if handler := self.request_handlers.get(type(req)): # type: ignore + + if handler := self.request_handlers.get(type(req)): logger.debug("Dispatching request of type %s", type(req).__name__) token = None @@ -689,12 +722,24 @@ async def _handle_request( # Set our global state that can be retrieved via # app.get_request_context() + client_capabilities = session.client_params.capabilities if session.client_params else None + task_support = self._experimental_handlers.task_support if self._experimental_handlers else None + # Get task metadata from request params if present + task_metadata = None + if hasattr(req, "params") and req.params is not None: + task_metadata = getattr(req.params, "task", None) token = request_ctx.set( RequestContext( message.request_id, message.request_meta, session, lifespan_context, + Experimental( + task_metadata=task_metadata, + _client_capabilities=client_capabilities, + _session=session, + _task_support=task_support, + ), request=request_data, ) ) diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py index b116fbe384..be8eca8fb1 100644 --- a/src/mcp/server/session.py +++ b/src/mcp/server/session.py @@ -46,8 +46,11 @@ async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]: from pydantic import AnyUrl import mcp.types as types +from mcp.server.experimental.session_features import ExperimentalServerSessionFeatures from mcp.server.models import InitializationOptions -from mcp.shared.exceptions import McpError +from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages +from mcp.shared.experimental.tasks.capabilities import check_tasks_capability +from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY from mcp.shared.message import ServerMessageMetadata, SessionMessage from mcp.shared.session import ( BaseSession, @@ -80,6 +83,7 @@ class ServerSession( ): _initialized: InitializationState = InitializationState.NotInitialized _client_params: types.InitializeRequestParams | None = None + _experimental_features: ExperimentalServerSessionFeatures | None = None def __init__( self, @@ -103,15 +107,23 @@ def __init__( def client_params(self) -> types.InitializeRequestParams | None: return self._client_params # pragma: no cover + @property + def experimental(self) -> ExperimentalServerSessionFeatures: + """Experimental APIs for server→client task operations. + + WARNING: These APIs are experimental and may change without notice. + """ + if self._experimental_features is None: + self._experimental_features = ExperimentalServerSessionFeatures(self) + return self._experimental_features + def check_client_capability(self, capability: types.ClientCapabilities) -> bool: # pragma: no cover """Check if the client supports a specific capability.""" if self._client_params is None: return False - # Get client capabilities from initialization params client_caps = self._client_params.capabilities - # Check each specified capability in the passed in capability object if capability.roots is not None: if client_caps.roots is None: return False @@ -121,25 +133,27 @@ def check_client_capability(self, capability: types.ClientCapabilities) -> bool: if capability.sampling is not None: if client_caps.sampling is None: return False - if capability.sampling.context is not None: - if client_caps.sampling.context is None: - return False - if capability.sampling.tools is not None: - if client_caps.sampling.tools is None: - return False - - if capability.elicitation is not None: - if client_caps.elicitation is None: + if capability.sampling.context is not None and client_caps.sampling.context is None: return False + if capability.sampling.tools is not None and client_caps.sampling.tools is None: + return False + + if capability.elicitation is not None and client_caps.elicitation is None: + return False if capability.experimental is not None: if client_caps.experimental is None: return False - # Check each experimental capability for exp_key, exp_value in capability.experimental.items(): if exp_key not in client_caps.experimental or client_caps.experimental[exp_key] != exp_value: return False + if capability.tasks is not None: + if client_caps.tasks is None: + return False + if not check_tasks_capability(capability.tasks, client_caps.tasks): + return False + return True async def _receive_loop(self) -> None: @@ -257,47 +271,12 @@ async def create_message( The sampling result from the client. Raises: - McpError: If tool_use or tool_result blocks are misused when tools are provided. + McpError: If tools are provided but client doesn't support them. + ValueError: If tool_use or tool_result message structure is invalid. """ - - if tools is not None or tool_choice is not None: - has_tools_cap = self.check_client_capability( - types.ClientCapabilities(sampling=types.SamplingCapability(tools=types.SamplingToolsCapability())) - ) - if not has_tools_cap: - raise McpError( - types.ErrorData( - code=types.INVALID_PARAMS, - message="Client does not support sampling tools capability", - ) - ) - - # Validate tool_use/tool_result message structure per SEP-1577: - # https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1577 - # This validation runs regardless of whether `tools` is in this request, - # since a tool loop continuation may omit `tools` while still containing - # tool_result content that must match previous tool_use. - if messages: - last_content = messages[-1].content_as_list - has_tool_results = any(c.type == "tool_result" for c in last_content) - - previous_content = messages[-2].content_as_list if len(messages) >= 2 else None - has_previous_tool_use = previous_content and any(c.type == "tool_use" for c in previous_content) - - if has_tool_results: - # Per spec: "SamplingMessage with tool result content blocks - # MUST NOT contain other content types." - if any(c.type != "tool_result" for c in last_content): - raise ValueError("The last message must contain only tool_result content if any is present") - if previous_content is None: - raise ValueError("tool_result requires a previous message containing tool_use") - if not has_previous_tool_use: - raise ValueError("tool_result blocks do not match any tool_use in the previous message") - if has_previous_tool_use and previous_content: - tool_use_ids = {c.id for c in previous_content if c.type == "tool_use"} - tool_result_ids = {c.toolUseId for c in last_content if c.type == "tool_result"} - if tool_use_ids != tool_result_ids: - raise ValueError("ids of tool_result blocks and tool_use blocks from previous message do not match") + client_caps = self._client_params.capabilities if self._client_params else None + validate_sampling_tools(client_caps, tools, tool_choice) + validate_tool_use_result_messages(messages) return await self.send_request( request=types.ServerRequest( @@ -481,6 +460,181 @@ async def send_elicit_complete( related_request_id, ) + def _build_elicit_form_request( + self, + message: str, + requestedSchema: types.ElicitRequestedSchema, + related_task_id: str | None = None, + task: types.TaskMetadata | None = None, + ) -> types.JSONRPCRequest: + """Build a form mode elicitation request without sending it. + + Args: + message: The message to present to the user + requestedSchema: Schema defining the expected response structure + related_task_id: If provided, adds io.modelcontextprotocol/related-task metadata + task: If provided, makes this a task-augmented request + + Returns: + A JSONRPCRequest ready to be sent or queued + """ + params = types.ElicitRequestFormParams( + message=message, + requestedSchema=requestedSchema, + task=task, + ) + params_data = params.model_dump(by_alias=True, mode="json", exclude_none=True) + + # Add related-task metadata if associated with a parent task + if related_task_id is not None: + # Defensive: model_dump() never includes _meta, but guard against future changes + if "_meta" not in params_data: # pragma: no cover + params_data["_meta"] = {} + params_data["_meta"][RELATED_TASK_METADATA_KEY] = types.RelatedTaskMetadata( + taskId=related_task_id + ).model_dump(by_alias=True) + + request_id = f"task-{related_task_id}-{id(params)}" if related_task_id else self._request_id + if related_task_id is None: + self._request_id += 1 + + return types.JSONRPCRequest( + jsonrpc="2.0", + id=request_id, + method="elicitation/create", + params=params_data, + ) + + def _build_elicit_url_request( + self, + message: str, + url: str, + elicitation_id: str, + related_task_id: str | None = None, + ) -> types.JSONRPCRequest: + """Build a URL mode elicitation request without sending it. + + Args: + message: Human-readable explanation of why the interaction is needed + url: The URL the user should navigate to + elicitation_id: Unique identifier for tracking this elicitation + related_task_id: If provided, adds io.modelcontextprotocol/related-task metadata + + Returns: + A JSONRPCRequest ready to be sent or queued + """ + params = types.ElicitRequestURLParams( + message=message, + url=url, + elicitationId=elicitation_id, + ) + params_data = params.model_dump(by_alias=True, mode="json", exclude_none=True) + + # Add related-task metadata if associated with a parent task + if related_task_id is not None: + # Defensive: model_dump() never includes _meta, but guard against future changes + if "_meta" not in params_data: # pragma: no cover + params_data["_meta"] = {} + params_data["_meta"][RELATED_TASK_METADATA_KEY] = types.RelatedTaskMetadata( + taskId=related_task_id + ).model_dump(by_alias=True) + + request_id = f"task-{related_task_id}-{id(params)}" if related_task_id else self._request_id + if related_task_id is None: + self._request_id += 1 + + return types.JSONRPCRequest( + jsonrpc="2.0", + id=request_id, + method="elicitation/create", + params=params_data, + ) + + def _build_create_message_request( + self, + messages: list[types.SamplingMessage], + *, + max_tokens: int, + system_prompt: str | None = None, + include_context: types.IncludeContext | None = None, + temperature: float | None = None, + stop_sequences: list[str] | None = None, + metadata: dict[str, Any] | None = None, + model_preferences: types.ModelPreferences | None = None, + tools: list[types.Tool] | None = None, + tool_choice: types.ToolChoice | None = None, + related_task_id: str | None = None, + task: types.TaskMetadata | None = None, + ) -> types.JSONRPCRequest: + """Build a sampling/createMessage request without sending it. + + Args: + messages: The conversation messages to send + max_tokens: Maximum number of tokens to generate + system_prompt: Optional system prompt + include_context: Optional context inclusion setting + temperature: Optional sampling temperature + stop_sequences: Optional stop sequences + metadata: Optional metadata to pass through to the LLM provider + model_preferences: Optional model selection preferences + tools: Optional list of tools the LLM can use during sampling + tool_choice: Optional control over tool usage behavior + related_task_id: If provided, adds io.modelcontextprotocol/related-task metadata + task: If provided, makes this a task-augmented request + + Returns: + A JSONRPCRequest ready to be sent or queued + """ + params = types.CreateMessageRequestParams( + messages=messages, + systemPrompt=system_prompt, + includeContext=include_context, + temperature=temperature, + maxTokens=max_tokens, + stopSequences=stop_sequences, + metadata=metadata, + modelPreferences=model_preferences, + tools=tools, + toolChoice=tool_choice, + task=task, + ) + params_data = params.model_dump(by_alias=True, mode="json", exclude_none=True) + + # Add related-task metadata if associated with a parent task + if related_task_id is not None: + # Defensive: model_dump() never includes _meta, but guard against future changes + if "_meta" not in params_data: # pragma: no cover + params_data["_meta"] = {} + params_data["_meta"][RELATED_TASK_METADATA_KEY] = types.RelatedTaskMetadata( + taskId=related_task_id + ).model_dump(by_alias=True) + + request_id = f"task-{related_task_id}-{id(params)}" if related_task_id else self._request_id + if related_task_id is None: + self._request_id += 1 + + return types.JSONRPCRequest( + jsonrpc="2.0", + id=request_id, + method="sampling/createMessage", + params=params_data, + ) + + async def send_message(self, message: SessionMessage) -> None: + """Send a raw session message. + + This is primarily used by TaskResultHandler to deliver queued messages + (elicitation/sampling requests) to the client during task execution. + + WARNING: This is a low-level experimental method that may change without + notice. Prefer using higher-level methods like send_notification() or + send_request() for normal operations. + + Args: + message: The session message to send + """ + await self._write_stream.send(message) + async def _handle_incoming(self, req: ServerRequestResponder) -> None: await self._incoming_message_stream_writer.send(req) diff --git a/src/mcp/server/validation.py b/src/mcp/server/validation.py new file mode 100644 index 0000000000..2ccd7056bd --- /dev/null +++ b/src/mcp/server/validation.py @@ -0,0 +1,104 @@ +""" +Shared validation functions for server requests. + +This module provides validation logic for sampling and elicitation requests +that is shared across normal and task-augmented code paths. +""" + +from mcp.shared.exceptions import McpError +from mcp.types import ( + INVALID_PARAMS, + ClientCapabilities, + ErrorData, + SamplingMessage, + Tool, + ToolChoice, +) + + +def check_sampling_tools_capability(client_caps: ClientCapabilities | None) -> bool: + """ + Check if the client supports sampling tools capability. + + Args: + client_caps: The client's declared capabilities + + Returns: + True if client supports sampling.tools, False otherwise + """ + if client_caps is None: + return False + if client_caps.sampling is None: + return False + if client_caps.sampling.tools is None: + return False + return True + + +def validate_sampling_tools( + client_caps: ClientCapabilities | None, + tools: list[Tool] | None, + tool_choice: ToolChoice | None, +) -> None: + """ + Validate that the client supports sampling tools if tools are being used. + + Args: + client_caps: The client's declared capabilities + tools: The tools list, if provided + tool_choice: The tool choice setting, if provided + + Raises: + McpError: If tools/tool_choice are provided but client doesn't support them + """ + if tools is not None or tool_choice is not None: + if not check_sampling_tools_capability(client_caps): + raise McpError( + ErrorData( + code=INVALID_PARAMS, + message="Client does not support sampling tools capability", + ) + ) + + +def validate_tool_use_result_messages(messages: list[SamplingMessage]) -> None: + """ + Validate tool_use/tool_result message structure per SEP-1577. + + This validation ensures: + 1. Messages with tool_result content contain ONLY tool_result content + 2. tool_result messages are preceded by a message with tool_use + 3. tool_result IDs match the tool_use IDs from the previous message + + See: https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1577 + + Args: + messages: The list of sampling messages to validate + + Raises: + ValueError: If the message structure is invalid + """ + if not messages: + return + + last_content = messages[-1].content_as_list + has_tool_results = any(c.type == "tool_result" for c in last_content) + + previous_content = messages[-2].content_as_list if len(messages) >= 2 else None + has_previous_tool_use = previous_content and any(c.type == "tool_use" for c in previous_content) + + if has_tool_results: + # Per spec: "SamplingMessage with tool result content blocks + # MUST NOT contain other content types." + if any(c.type != "tool_result" for c in last_content): + raise ValueError("The last message must contain only tool_result content if any is present") + if previous_content is None: + raise ValueError("tool_result requires a previous message containing tool_use") + if not has_previous_tool_use: + raise ValueError("tool_result blocks do not match any tool_use in the previous message") + + if has_previous_tool_use and previous_content: + tool_use_ids = {c.id for c in previous_content if c.type == "tool_use"} + tool_result_ids = {c.toolUseId for c in last_content if c.type == "tool_result"} + if tool_use_ids != tool_result_ids: + raise ValueError("ids of tool_result blocks and tool_use blocks from previous message do not match") diff --git a/src/mcp/shared/context.py b/src/mcp/shared/context.py index f3006e7d5f..a0a0e40dca 100644 --- a/src/mcp/shared/context.py +++ b/src/mcp/shared/context.py @@ -1,4 +1,8 @@ -from dataclasses import dataclass +""" +Request context for MCP handlers. +""" + +from dataclasses import dataclass, field from typing import Any, Generic from typing_extensions import TypeVar @@ -17,4 +21,9 @@ class RequestContext(Generic[SessionT, LifespanContextT, RequestT]): meta: RequestParams.Meta | None session: SessionT lifespan_context: LifespanContextT + # NOTE: This is typed as Any to avoid circular imports. The actual type is + # mcp.server.experimental.request_context.Experimental, but importing it here + # triggers mcp.server.__init__ -> fastmcp -> tools -> back to this module. + # The Server sets this to an Experimental instance at runtime. + experimental: Any = field(default=None) request: RequestT | None = None diff --git a/src/mcp/shared/experimental/__init__.py b/src/mcp/shared/experimental/__init__.py new file mode 100644 index 0000000000..9b1b1479cb --- /dev/null +++ b/src/mcp/shared/experimental/__init__.py @@ -0,0 +1,7 @@ +""" +Pure experimental MCP features (no server dependencies). + +WARNING: These APIs are experimental and may change without notice. + +For server-integrated experimental features, use mcp.server.experimental. +""" diff --git a/src/mcp/shared/experimental/tasks/__init__.py b/src/mcp/shared/experimental/tasks/__init__.py new file mode 100644 index 0000000000..37d81af50b --- /dev/null +++ b/src/mcp/shared/experimental/tasks/__init__.py @@ -0,0 +1,12 @@ +""" +Pure task state management for MCP. + +WARNING: These APIs are experimental and may change without notice. + +Import directly from submodules: +- mcp.shared.experimental.tasks.store.TaskStore +- mcp.shared.experimental.tasks.context.TaskContext +- mcp.shared.experimental.tasks.in_memory_task_store.InMemoryTaskStore +- mcp.shared.experimental.tasks.message_queue.TaskMessageQueue +- mcp.shared.experimental.tasks.helpers.is_terminal +""" diff --git a/src/mcp/shared/experimental/tasks/capabilities.py b/src/mcp/shared/experimental/tasks/capabilities.py new file mode 100644 index 0000000000..307fcdd6e5 --- /dev/null +++ b/src/mcp/shared/experimental/tasks/capabilities.py @@ -0,0 +1,115 @@ +""" +Tasks capability checking utilities. + +This module provides functions for checking and requiring task-related +capabilities. All tasks capability logic is centralized here to keep +the main session code clean. + +WARNING: These APIs are experimental and may change without notice. +""" + +from mcp.shared.exceptions import McpError +from mcp.types import ( + INVALID_REQUEST, + ClientCapabilities, + ClientTasksCapability, + ErrorData, +) + + +def check_tasks_capability( + required: ClientTasksCapability, + client: ClientTasksCapability, +) -> bool: + """ + Check if client's tasks capability matches the required capability. + + Args: + required: The capability being checked for + client: The client's declared capabilities + + Returns: + True if client has the required capability, False otherwise + """ + if required.requests is None: + return True + if client.requests is None: + return False + + # Check elicitation.create + if required.requests.elicitation is not None: + if client.requests.elicitation is None: + return False + if required.requests.elicitation.create is not None: + if client.requests.elicitation.create is None: + return False + + # Check sampling.createMessage + if required.requests.sampling is not None: + if client.requests.sampling is None: + return False + if required.requests.sampling.createMessage is not None: + if client.requests.sampling.createMessage is None: + return False + + return True + + +def has_task_augmented_elicitation(caps: ClientCapabilities) -> bool: + """Check if capabilities include task-augmented elicitation support.""" + if caps.tasks is None: + return False + if caps.tasks.requests is None: + return False + if caps.tasks.requests.elicitation is None: + return False + return caps.tasks.requests.elicitation.create is not None + + +def has_task_augmented_sampling(caps: ClientCapabilities) -> bool: + """Check if capabilities include task-augmented sampling support.""" + if caps.tasks is None: + return False + if caps.tasks.requests is None: + return False + if caps.tasks.requests.sampling is None: + return False + return caps.tasks.requests.sampling.createMessage is not None + + +def require_task_augmented_elicitation(client_caps: ClientCapabilities | None) -> None: + """ + Raise McpError if client doesn't support task-augmented elicitation. + + Args: + client_caps: The client's declared capabilities, or None if not initialized + + Raises: + McpError: If client doesn't support task-augmented elicitation + """ + if client_caps is None or not has_task_augmented_elicitation(client_caps): + raise McpError( + ErrorData( + code=INVALID_REQUEST, + message="Client does not support task-augmented elicitation", + ) + ) + + +def require_task_augmented_sampling(client_caps: ClientCapabilities | None) -> None: + """ + Raise McpError if client doesn't support task-augmented sampling. + + Args: + client_caps: The client's declared capabilities, or None if not initialized + + Raises: + McpError: If client doesn't support task-augmented sampling + """ + if client_caps is None or not has_task_augmented_sampling(client_caps): + raise McpError( + ErrorData( + code=INVALID_REQUEST, + message="Client does not support task-augmented sampling", + ) + ) diff --git a/src/mcp/shared/experimental/tasks/context.py b/src/mcp/shared/experimental/tasks/context.py new file mode 100644 index 0000000000..12d159515c --- /dev/null +++ b/src/mcp/shared/experimental/tasks/context.py @@ -0,0 +1,101 @@ +""" +TaskContext - Pure task state management. + +This module provides TaskContext, which manages task state without any +server/session dependencies. It can be used standalone for distributed +workers or wrapped by ServerTaskContext for full server integration. +""" + +from mcp.shared.experimental.tasks.store import TaskStore +from mcp.types import TASK_STATUS_COMPLETED, TASK_STATUS_FAILED, Result, Task + + +class TaskContext: + """ + Pure task state management - no session dependencies. + + This class handles: + - Task state (status, result) + - Cancellation tracking + - Store interactions + + For server-integrated features (elicit, create_message, notifications), + use ServerTaskContext from mcp.server.experimental. + + Example (distributed worker): + async def worker_job(task_id: str): + store = RedisTaskStore(redis_url) + task = await store.get_task(task_id) + ctx = TaskContext(task=task, store=store) + + await ctx.update_status("Working...") + result = await do_work() + await ctx.complete(result) + """ + + def __init__(self, task: Task, store: TaskStore): + self._task = task + self._store = store + self._cancelled = False + + @property + def task_id(self) -> str: + """The task identifier.""" + return self._task.taskId + + @property + def task(self) -> Task: + """The current task state.""" + return self._task + + @property + def is_cancelled(self) -> bool: + """Whether cancellation has been requested.""" + return self._cancelled + + def request_cancellation(self) -> None: + """ + Request cancellation of this task. + + This sets is_cancelled=True. Task work should check this + periodically and exit gracefully if set. + """ + self._cancelled = True + + async def update_status(self, message: str) -> None: + """ + Update the task's status message. + + Args: + message: The new status message + """ + self._task = await self._store.update_task( + self.task_id, + status_message=message, + ) + + async def complete(self, result: Result) -> None: + """ + Mark the task as completed with the given result. + + Args: + result: The task result + """ + await self._store.store_result(self.task_id, result) + self._task = await self._store.update_task( + self.task_id, + status=TASK_STATUS_COMPLETED, + ) + + async def fail(self, error: str) -> None: + """ + Mark the task as failed with an error message. + + Args: + error: The error message + """ + self._task = await self._store.update_task( + self.task_id, + status=TASK_STATUS_FAILED, + status_message=error, + ) diff --git a/src/mcp/shared/experimental/tasks/helpers.py b/src/mcp/shared/experimental/tasks/helpers.py new file mode 100644 index 0000000000..5c87f9ef87 --- /dev/null +++ b/src/mcp/shared/experimental/tasks/helpers.py @@ -0,0 +1,181 @@ +""" +Helper functions for pure task management. + +These helpers work with pure TaskContext and don't require server dependencies. +For server-integrated task helpers, use mcp.server.experimental. +""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from uuid import uuid4 + +from mcp.shared.exceptions import McpError +from mcp.shared.experimental.tasks.context import TaskContext +from mcp.shared.experimental.tasks.store import TaskStore +from mcp.types import ( + INVALID_PARAMS, + TASK_STATUS_CANCELLED, + TASK_STATUS_COMPLETED, + TASK_STATUS_FAILED, + TASK_STATUS_WORKING, + CancelTaskResult, + ErrorData, + Task, + TaskMetadata, + TaskStatus, +) + +# Metadata key for model-immediate-response (per MCP spec) +# Servers MAY include this in CreateTaskResult._meta to provide an immediate +# response string while the task executes in the background. +MODEL_IMMEDIATE_RESPONSE_KEY = "io.modelcontextprotocol/model-immediate-response" + +# Metadata key for associating requests with a task (per MCP spec) +RELATED_TASK_METADATA_KEY = "io.modelcontextprotocol/related-task" + + +def is_terminal(status: TaskStatus) -> bool: + """ + Check if a task status represents a terminal state. + + Terminal states are those where the task has finished and will not change. + + Args: + status: The task status to check + + Returns: + True if the status is terminal (completed, failed, or cancelled) + """ + return status in (TASK_STATUS_COMPLETED, TASK_STATUS_FAILED, TASK_STATUS_CANCELLED) + + +async def cancel_task( + store: TaskStore, + task_id: str, +) -> CancelTaskResult: + """ + Cancel a task with spec-compliant validation. + + Per spec: "Receivers MUST reject cancellation of terminal status tasks + with -32602 (Invalid params)" + + This helper validates that the task exists and is not in a terminal state + before setting it to "cancelled". + + Args: + store: The task store + task_id: The task identifier to cancel + + Returns: + CancelTaskResult with the cancelled task state + + Raises: + McpError: With INVALID_PARAMS (-32602) if: + - Task does not exist + - Task is already in a terminal state (completed, failed, cancelled) + + Example: + @server.experimental.cancel_task() + async def handle_cancel(request: CancelTaskRequest) -> CancelTaskResult: + return await cancel_task(store, request.params.taskId) + """ + task = await store.get_task(task_id) + if task is None: + raise McpError( + ErrorData( + code=INVALID_PARAMS, + message=f"Task not found: {task_id}", + ) + ) + + if is_terminal(task.status): + raise McpError( + ErrorData( + code=INVALID_PARAMS, + message=f"Cannot cancel task in terminal state '{task.status}'", + ) + ) + + # Update task to cancelled status + cancelled_task = await store.update_task(task_id, status=TASK_STATUS_CANCELLED) + return CancelTaskResult(**cancelled_task.model_dump()) + + +def generate_task_id() -> str: + """Generate a unique task ID.""" + return str(uuid4()) + + +def create_task_state( + metadata: TaskMetadata, + task_id: str | None = None, +) -> Task: + """ + Create a Task object with initial state. + + This is a helper for TaskStore implementations. + + Args: + metadata: Task metadata + task_id: Optional task ID (generated if not provided) + + Returns: + A new Task in "working" status + """ + now = datetime.now(timezone.utc) + return Task( + taskId=task_id or generate_task_id(), + status=TASK_STATUS_WORKING, + createdAt=now, + lastUpdatedAt=now, + ttl=metadata.ttl, + pollInterval=500, # Default 500ms poll interval + ) + + +@asynccontextmanager +async def task_execution( + task_id: str, + store: TaskStore, +) -> AsyncIterator[TaskContext]: + """ + Context manager for safe task execution (pure, no server dependencies). + + Loads a task from the store and provides a TaskContext for the work. + If an unhandled exception occurs, the task is automatically marked as failed + and the exception is suppressed (since the failure is captured in task state). + + This is useful for distributed workers that don't have a server session. + + Args: + task_id: The task identifier to execute + store: The task store (must be accessible by the worker) + + Yields: + TaskContext for updating status and completing/failing the task + + Raises: + ValueError: If the task is not found in the store + + Example (distributed worker): + async def worker_process(task_id: str): + store = RedisTaskStore(redis_url) + async with task_execution(task_id, store) as ctx: + await ctx.update_status("Working...") + result = await do_work() + await ctx.complete(result) + """ + task = await store.get_task(task_id) + if task is None: + raise ValueError(f"Task {task_id} not found") + + ctx = TaskContext(task, store) + try: + yield ctx + except Exception as e: + # Auto-fail the task if an exception occurs and task isn't already terminal + # Exception is suppressed since failure is captured in task state + if not is_terminal(ctx.task.status): + await ctx.fail(str(e)) + # Don't re-raise - the failure is recorded in task state diff --git a/src/mcp/shared/experimental/tasks/in_memory_task_store.py b/src/mcp/shared/experimental/tasks/in_memory_task_store.py new file mode 100644 index 0000000000..7b630ce6e2 --- /dev/null +++ b/src/mcp/shared/experimental/tasks/in_memory_task_store.py @@ -0,0 +1,219 @@ +""" +In-memory implementation of TaskStore for demonstration purposes. + +This implementation stores all tasks in memory and provides automatic cleanup +based on the TTL duration specified in the task metadata using lazy expiration. + +Note: This is not suitable for production use as all data is lost on restart. +For production, consider implementing TaskStore with a database or distributed cache. +""" + +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone + +import anyio + +from mcp.shared.experimental.tasks.helpers import create_task_state, is_terminal +from mcp.shared.experimental.tasks.store import TaskStore +from mcp.types import Result, Task, TaskMetadata, TaskStatus + + +@dataclass +class StoredTask: + """Internal storage representation of a task.""" + + task: Task + result: Result | None = None + # Time when this task should be removed (None = never) + expires_at: datetime | None = field(default=None) + + +class InMemoryTaskStore(TaskStore): + """ + A simple in-memory implementation of TaskStore. + + Features: + - Automatic TTL-based cleanup (lazy expiration) + - Thread-safe for single-process async use + - Pagination support for list_tasks + + Limitations: + - All data lost on restart + - Not suitable for distributed systems + - No persistence + + For production, implement TaskStore with Redis, PostgreSQL, etc. + """ + + def __init__(self, page_size: int = 10) -> None: + self._tasks: dict[str, StoredTask] = {} + self._page_size = page_size + self._update_events: dict[str, anyio.Event] = {} + + def _calculate_expiry(self, ttl_ms: int | None) -> datetime | None: + """Calculate expiry time from TTL in milliseconds.""" + if ttl_ms is None: + return None + return datetime.now(timezone.utc) + timedelta(milliseconds=ttl_ms) + + def _is_expired(self, stored: StoredTask) -> bool: + """Check if a task has expired.""" + if stored.expires_at is None: + return False + return datetime.now(timezone.utc) >= stored.expires_at + + def _cleanup_expired(self) -> None: + """Remove all expired tasks. Called lazily during access operations.""" + expired_ids = [task_id for task_id, stored in self._tasks.items() if self._is_expired(stored)] + for task_id in expired_ids: + del self._tasks[task_id] + + async def create_task( + self, + metadata: TaskMetadata, + task_id: str | None = None, + ) -> Task: + """Create a new task with the given metadata.""" + # Cleanup expired tasks on access + self._cleanup_expired() + + task = create_task_state(metadata, task_id) + + if task.taskId in self._tasks: + raise ValueError(f"Task with ID {task.taskId} already exists") + + stored = StoredTask( + task=task, + expires_at=self._calculate_expiry(metadata.ttl), + ) + self._tasks[task.taskId] = stored + + # Return a copy to prevent external modification + return Task(**task.model_dump()) + + async def get_task(self, task_id: str) -> Task | None: + """Get a task by ID.""" + # Cleanup expired tasks on access + self._cleanup_expired() + + stored = self._tasks.get(task_id) + if stored is None: + return None + + # Return a copy to prevent external modification + return Task(**stored.task.model_dump()) + + async def update_task( + self, + task_id: str, + status: TaskStatus | None = None, + status_message: str | None = None, + ) -> Task: + """Update a task's status and/or message.""" + stored = self._tasks.get(task_id) + if stored is None: + raise ValueError(f"Task with ID {task_id} not found") + + # Per spec: Terminal states MUST NOT transition to any other status + if status is not None and status != stored.task.status and is_terminal(stored.task.status): + raise ValueError(f"Cannot transition from terminal status '{stored.task.status}'") + + status_changed = False + if status is not None and stored.task.status != status: + stored.task.status = status + status_changed = True + + if status_message is not None: + stored.task.statusMessage = status_message + + # Update lastUpdatedAt on any change + stored.task.lastUpdatedAt = datetime.now(timezone.utc) + + # If task is now terminal and has TTL, reset expiry timer + if status is not None and is_terminal(status) and stored.task.ttl is not None: + stored.expires_at = self._calculate_expiry(stored.task.ttl) + + # Notify waiters if status changed + if status_changed: + await self.notify_update(task_id) + + return Task(**stored.task.model_dump()) + + async def store_result(self, task_id: str, result: Result) -> None: + """Store the result for a task.""" + stored = self._tasks.get(task_id) + if stored is None: + raise ValueError(f"Task with ID {task_id} not found") + + stored.result = result + + async def get_result(self, task_id: str) -> Result | None: + """Get the stored result for a task.""" + stored = self._tasks.get(task_id) + if stored is None: + return None + + return stored.result + + async def list_tasks( + self, + cursor: str | None = None, + ) -> tuple[list[Task], str | None]: + """List tasks with pagination.""" + # Cleanup expired tasks on access + self._cleanup_expired() + + all_task_ids = list(self._tasks.keys()) + + start_index = 0 + if cursor is not None: + try: + cursor_index = all_task_ids.index(cursor) + start_index = cursor_index + 1 + except ValueError: + raise ValueError(f"Invalid cursor: {cursor}") + + page_task_ids = all_task_ids[start_index : start_index + self._page_size] + tasks = [Task(**self._tasks[tid].task.model_dump()) for tid in page_task_ids] + + # Determine next cursor + next_cursor = None + if start_index + self._page_size < len(all_task_ids) and page_task_ids: + next_cursor = page_task_ids[-1] + + return tasks, next_cursor + + async def delete_task(self, task_id: str) -> bool: + """Delete a task.""" + if task_id not in self._tasks: + return False + + del self._tasks[task_id] + return True + + async def wait_for_update(self, task_id: str) -> None: + """Wait until the task status changes.""" + if task_id not in self._tasks: + raise ValueError(f"Task with ID {task_id} not found") + + # Create a fresh event for waiting (anyio.Event can't be cleared) + self._update_events[task_id] = anyio.Event() + event = self._update_events[task_id] + await event.wait() + + async def notify_update(self, task_id: str) -> None: + """Signal that a task has been updated.""" + if task_id in self._update_events: + self._update_events[task_id].set() + + # --- Testing/debugging helpers --- + + def cleanup(self) -> None: + """Cleanup all tasks (useful for testing or graceful shutdown).""" + self._tasks.clear() + self._update_events.clear() + + def get_all_tasks(self) -> list[Task]: + """Get all tasks (useful for debugging). Returns copies to prevent modification.""" + self._cleanup_expired() + return [Task(**stored.task.model_dump()) for stored in self._tasks.values()] diff --git a/src/mcp/shared/experimental/tasks/message_queue.py b/src/mcp/shared/experimental/tasks/message_queue.py new file mode 100644 index 0000000000..69b6609887 --- /dev/null +++ b/src/mcp/shared/experimental/tasks/message_queue.py @@ -0,0 +1,241 @@ +""" +TaskMessageQueue - FIFO queue for task-related messages. + +This implements the core message queue pattern from the MCP Tasks spec. +When a handler needs to send a request (like elicitation) during a task-augmented +request, the message is enqueued instead of sent directly. Messages are delivered +to the client only through the `tasks/result` endpoint. + +This pattern enables: +1. Decoupling request handling from message delivery +2. Proper bidirectional communication via the tasks/result stream +3. Automatic status management (working <-> input_required) +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Literal + +import anyio + +from mcp.shared.experimental.tasks.resolver import Resolver +from mcp.types import JSONRPCNotification, JSONRPCRequest, RequestId + + +@dataclass +class QueuedMessage: + """ + A message queued for delivery via tasks/result. + + Messages are stored with their type and a resolver for requests + that expect responses. + """ + + type: Literal["request", "notification"] + """Whether this is a request (expects response) or notification (one-way).""" + + message: JSONRPCRequest | JSONRPCNotification + """The JSON-RPC message to send.""" + + timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + """When the message was enqueued.""" + + resolver: Resolver[dict[str, Any]] | None = None + """Resolver to set when response arrives (only for requests).""" + + original_request_id: RequestId | None = None + """The original request ID used internally, for routing responses back.""" + + +class TaskMessageQueue(ABC): + """ + Abstract interface for task message queuing. + + This is a FIFO queue that stores messages to be delivered via `tasks/result`. + When a task-augmented handler calls elicit() or sends a notification, the + message is enqueued here instead of being sent directly to the client. + + The `tasks/result` handler then dequeues and sends these messages through + the transport, with `relatedRequestId` set to the tasks/result request ID + so responses are routed correctly. + + Implementations can use in-memory storage, Redis, etc. + """ + + @abstractmethod + async def enqueue(self, task_id: str, message: QueuedMessage) -> None: + """ + Add a message to the queue for a task. + + Args: + task_id: The task identifier + message: The message to enqueue + """ + + @abstractmethod + async def dequeue(self, task_id: str) -> QueuedMessage | None: + """ + Remove and return the next message from the queue. + + Args: + task_id: The task identifier + + Returns: + The next message, or None if queue is empty + """ + + @abstractmethod + async def peek(self, task_id: str) -> QueuedMessage | None: + """ + Return the next message without removing it. + + Args: + task_id: The task identifier + + Returns: + The next message, or None if queue is empty + """ + + @abstractmethod + async def is_empty(self, task_id: str) -> bool: + """ + Check if the queue is empty for a task. + + Args: + task_id: The task identifier + + Returns: + True if no messages are queued + """ + + @abstractmethod + async def clear(self, task_id: str) -> list[QueuedMessage]: + """ + Remove and return all messages from the queue. + + This is useful for cleanup when a task is cancelled or completed. + + Args: + task_id: The task identifier + + Returns: + All queued messages (may be empty) + """ + + @abstractmethod + async def wait_for_message(self, task_id: str) -> None: + """ + Wait until a message is available in the queue. + + This blocks until either: + 1. A message is enqueued for this task + 2. The wait is cancelled + + Args: + task_id: The task identifier + """ + + @abstractmethod + async def notify_message_available(self, task_id: str) -> None: + """ + Signal that a message is available for a task. + + This wakes up any coroutines waiting in wait_for_message(). + + Args: + task_id: The task identifier + """ + + +class InMemoryTaskMessageQueue(TaskMessageQueue): + """ + In-memory implementation of TaskMessageQueue. + + This is suitable for single-process servers. For distributed systems, + implement TaskMessageQueue with Redis, RabbitMQ, etc. + + Features: + - FIFO ordering per task + - Async wait for message availability + - Thread-safe for single-process async use + """ + + def __init__(self) -> None: + self._queues: dict[str, list[QueuedMessage]] = {} + self._events: dict[str, anyio.Event] = {} + + def _get_queue(self, task_id: str) -> list[QueuedMessage]: + """Get or create the queue for a task.""" + if task_id not in self._queues: + self._queues[task_id] = [] + return self._queues[task_id] + + async def enqueue(self, task_id: str, message: QueuedMessage) -> None: + """Add a message to the queue.""" + queue = self._get_queue(task_id) + queue.append(message) + # Signal that a message is available + await self.notify_message_available(task_id) + + async def dequeue(self, task_id: str) -> QueuedMessage | None: + """Remove and return the next message.""" + queue = self._get_queue(task_id) + if not queue: + return None + return queue.pop(0) + + async def peek(self, task_id: str) -> QueuedMessage | None: + """Return the next message without removing it.""" + queue = self._get_queue(task_id) + if not queue: + return None + return queue[0] + + async def is_empty(self, task_id: str) -> bool: + """Check if the queue is empty.""" + queue = self._get_queue(task_id) + return len(queue) == 0 + + async def clear(self, task_id: str) -> list[QueuedMessage]: + """Remove and return all messages.""" + queue = self._get_queue(task_id) + messages = list(queue) + queue.clear() + return messages + + async def wait_for_message(self, task_id: str) -> None: + """Wait until a message is available.""" + # Check if there are already messages + if not await self.is_empty(task_id): + return + + # Create a fresh event for waiting (anyio.Event can't be cleared) + self._events[task_id] = anyio.Event() + event = self._events[task_id] + + # Double-check after creating event (avoid race condition) + if not await self.is_empty(task_id): + return + + # Wait for a new message + await event.wait() + + async def notify_message_available(self, task_id: str) -> None: + """Signal that a message is available.""" + if task_id in self._events: + self._events[task_id].set() + + def cleanup(self, task_id: str | None = None) -> None: + """ + Clean up queues and events. + + Args: + task_id: If provided, clean up only this task. Otherwise clean up all. + """ + if task_id is not None: + self._queues.pop(task_id, None) + self._events.pop(task_id, None) + else: + self._queues.clear() + self._events.clear() diff --git a/src/mcp/shared/experimental/tasks/polling.py b/src/mcp/shared/experimental/tasks/polling.py new file mode 100644 index 0000000000..39db2e6b68 --- /dev/null +++ b/src/mcp/shared/experimental/tasks/polling.py @@ -0,0 +1,45 @@ +""" +Shared polling utilities for task operations. + +This module provides generic polling logic that works for both client→server +and server→client task polling. + +WARNING: These APIs are experimental and may change without notice. +""" + +from collections.abc import AsyncIterator, Awaitable, Callable + +import anyio + +from mcp.shared.experimental.tasks.helpers import is_terminal +from mcp.types import GetTaskResult + + +async def poll_until_terminal( + get_task: Callable[[str], Awaitable[GetTaskResult]], + task_id: str, + default_interval_ms: int = 500, +) -> AsyncIterator[GetTaskResult]: + """ + Poll a task until it reaches terminal status. + + This is a generic utility that works for both client→server and server→client + polling. The caller provides the get_task function appropriate for their direction. + + Args: + get_task: Async function that takes task_id and returns GetTaskResult + task_id: The task to poll + default_interval_ms: Fallback poll interval if server doesn't specify + + Yields: + GetTaskResult for each poll + """ + while True: + status = await get_task(task_id) + yield status + + if is_terminal(status.status): + break + + interval_ms = status.pollInterval if status.pollInterval is not None else default_interval_ms + await anyio.sleep(interval_ms / 1000) diff --git a/src/mcp/shared/experimental/tasks/resolver.py b/src/mcp/shared/experimental/tasks/resolver.py new file mode 100644 index 0000000000..f27425b2c6 --- /dev/null +++ b/src/mcp/shared/experimental/tasks/resolver.py @@ -0,0 +1,60 @@ +""" +Resolver - An anyio-compatible future-like object for async result passing. + +This provides a simple way to pass a result (or exception) from one coroutine +to another without depending on asyncio.Future. +""" + +from typing import Generic, TypeVar, cast + +import anyio + +T = TypeVar("T") + + +class Resolver(Generic[T]): + """ + A simple resolver for passing results between coroutines. + + Unlike asyncio.Future, this works with any anyio-compatible async backend. + + Usage: + resolver: Resolver[str] = Resolver() + + # In one coroutine: + resolver.set_result("hello") + + # In another coroutine: + result = await resolver.wait() # returns "hello" + """ + + def __init__(self) -> None: + self._event = anyio.Event() + self._value: T | None = None + self._exception: BaseException | None = None + + def set_result(self, value: T) -> None: + """Set the result value and wake up waiters.""" + if self._event.is_set(): + raise RuntimeError("Resolver already completed") + self._value = value + self._event.set() + + def set_exception(self, exc: BaseException) -> None: + """Set an exception and wake up waiters.""" + if self._event.is_set(): + raise RuntimeError("Resolver already completed") + self._exception = exc + self._event.set() + + async def wait(self) -> T: + """Wait for the result and return it, or raise the exception.""" + await self._event.wait() + if self._exception is not None: + raise self._exception + # If we reach here, set_result() was called, so _value is set + return cast(T, self._value) + + def done(self) -> bool: + """Return True if the resolver has been completed.""" + return self._event.is_set() diff --git a/src/mcp/shared/experimental/tasks/store.py b/src/mcp/shared/experimental/tasks/store.py new file mode 100644 index 0000000000..71fb4511b8 --- /dev/null +++ b/src/mcp/shared/experimental/tasks/store.py @@ -0,0 +1,156 @@ +""" +TaskStore - Abstract interface for task state storage. +""" + +from abc import ABC, abstractmethod + +from mcp.types import Result, Task, TaskMetadata, TaskStatus + + +class TaskStore(ABC): + """ + Abstract interface for task state storage. + + This is a pure storage interface - it doesn't manage execution. + Implementations can use in-memory storage, databases, Redis, etc. + + All methods are async to support various backends. + """ + + @abstractmethod + async def create_task( + self, + metadata: TaskMetadata, + task_id: str | None = None, + ) -> Task: + """ + Create a new task. + + Args: + metadata: Task metadata (ttl, etc.) + task_id: Optional task ID. If None, implementation should generate one. + + Returns: + The created Task with status="working" + + Raises: + ValueError: If task_id already exists + """ + + @abstractmethod + async def get_task(self, task_id: str) -> Task | None: + """ + Get a task by ID. + + Args: + task_id: The task identifier + + Returns: + The Task, or None if not found + """ + + @abstractmethod + async def update_task( + self, + task_id: str, + status: TaskStatus | None = None, + status_message: str | None = None, + ) -> Task: + """ + Update a task's status and/or message. + + Args: + task_id: The task identifier + status: New status (if changing) + status_message: New status message (if changing) + + Returns: + The updated Task + + Raises: + ValueError: If task not found + ValueError: If attempting to transition from a terminal status + (completed, failed, cancelled). Per spec, terminal states + MUST NOT transition to any other status. + """ + + @abstractmethod + async def store_result(self, task_id: str, result: Result) -> None: + """ + Store the result for a task. + + Args: + task_id: The task identifier + result: The result to store + + Raises: + ValueError: If task not found + """ + + @abstractmethod + async def get_result(self, task_id: str) -> Result | None: + """ + Get the stored result for a task. + + Args: + task_id: The task identifier + + Returns: + The stored Result, or None if not available + """ + + @abstractmethod + async def list_tasks( + self, + cursor: str | None = None, + ) -> tuple[list[Task], str | None]: + """ + List tasks with pagination. + + Args: + cursor: Optional cursor for pagination + + Returns: + Tuple of (tasks, next_cursor). next_cursor is None if no more pages. + """ + + @abstractmethod + async def delete_task(self, task_id: str) -> bool: + """ + Delete a task. + + Args: + task_id: The task identifier + + Returns: + True if deleted, False if not found + """ + + @abstractmethod + async def wait_for_update(self, task_id: str) -> None: + """ + Wait until the task status changes. + + This blocks until either: + 1. The task status changes + 2. The wait is cancelled + + Used by tasks/result to wait for task completion or status changes. + + Args: + task_id: The task identifier + + Raises: + ValueError: If task not found + """ + + @abstractmethod + async def notify_update(self, task_id: str) -> None: + """ + Signal that a task has been updated. + + This wakes up any coroutines waiting in wait_for_update(). + + Args: + task_id: The task identifier + """ diff --git a/src/mcp/shared/response_router.py b/src/mcp/shared/response_router.py new file mode 100644 index 0000000000..31796157fe --- /dev/null +++ b/src/mcp/shared/response_router.py @@ -0,0 +1,63 @@ +""" +ResponseRouter - Protocol for pluggable response routing. + +This module defines a protocol for routing JSON-RPC responses to alternative +handlers before falling back to the default response stream mechanism. + +The primary use case is task-augmented requests: when a TaskSession enqueues +a request (like elicitation), the response needs to be routed back to the +waiting resolver instead of the normal response stream. + +Design: +- Protocol-based for testability and flexibility +- Returns bool to indicate if response was handled +- Supports both success responses and errors +""" + +from typing import Any, Protocol + +from mcp.types import ErrorData, RequestId + + +class ResponseRouter(Protocol): + """ + Protocol for routing responses to alternative handlers. + + Implementations check if they have a pending request for the given ID + and deliver the response/error to the appropriate handler. + + Example: + class TaskResultHandler(ResponseRouter): + def route_response(self, request_id, response): + resolver = self._pending_requests.pop(request_id, None) + if resolver: + resolver.set_result(response) + return True + return False + """ + + def route_response(self, request_id: RequestId, response: dict[str, Any]) -> bool: + """ + Try to route a response to a pending request handler. + + Args: + request_id: The JSON-RPC request ID from the response + response: The response result data + + Returns: + True if the response was handled, False otherwise + """ + ... # pragma: no cover + + def route_error(self, request_id: RequestId, error: ErrorData) -> bool: + """ + Try to route an error to a pending request handler. + + Args: + request_id: The JSON-RPC request ID from the error response + error: The error data + + Returns: + True if the error was handled, False otherwise + """ + ... # pragma: no cover diff --git a/src/mcp/shared/session.py b/src/mcp/shared/session.py index 3b2cd3ecb1..cceefccce6 100644 --- a/src/mcp/shared/session.py +++ b/src/mcp/shared/session.py @@ -13,6 +13,7 @@ from mcp.shared.exceptions import McpError from mcp.shared.message import MessageMetadata, ServerMessageMetadata, SessionMessage +from mcp.shared.response_router import ResponseRouter from mcp.types import ( CONNECTION_CLOSED, INVALID_PARAMS, @@ -179,6 +180,7 @@ class BaseSession( _request_id: int _in_flight: dict[RequestId, RequestResponder[ReceiveRequestT, SendResultT]] _progress_callbacks: dict[RequestId, ProgressFnT] + _response_routers: list["ResponseRouter"] def __init__( self, @@ -198,8 +200,24 @@ def __init__( self._session_read_timeout_seconds = read_timeout_seconds self._in_flight = {} self._progress_callbacks = {} + self._response_routers = [] self._exit_stack = AsyncExitStack() + def add_response_router(self, router: ResponseRouter) -> None: + """ + Register a response router to handle responses for non-standard requests. + + Response routers are checked in order before falling back to the default + response stream mechanism. This is used by TaskResultHandler to route + responses for queued task requests back to their resolvers. + + WARNING: This is an experimental API that may change without notice. + + Args: + router: A ResponseRouter implementation + """ + self._response_routers.append(router) + async def __aenter__(self) -> Self: self._task_group = anyio.create_task_group() await self._task_group.__aenter__() @@ -413,13 +431,7 @@ async def _receive_loop(self) -> None: f"Failed to validate notification: {e}. Message was: {message.message.root}" ) else: # Response or error - stream = self._response_streams.pop(message.message.root.id, None) - if stream: # pragma: no cover - await stream.send(message.message.root) - else: # pragma: no cover - await self._handle_incoming( - RuntimeError(f"Received response with an unknown request ID: {message}") - ) + await self._handle_response(message) except anyio.ClosedResourceError: # This is expected when the client disconnects abruptly. @@ -443,6 +455,44 @@ async def _receive_loop(self) -> None: pass self._response_streams.clear() + async def _handle_response(self, message: SessionMessage) -> None: + """ + Handle an incoming response or error message. + + Checks response routers first (e.g., for task-related responses), + then falls back to the normal response stream mechanism. + """ + root = message.message.root + + # This check is always true at runtime: the caller (_receive_loop) only invokes + # this method in the else branch after checking for JSONRPCRequest and + # JSONRPCNotification. However, the type checker can't infer this from the + # method signature, so we need this guard for type narrowing. + if not isinstance(root, JSONRPCResponse | JSONRPCError): + return # pragma: no cover + + response_id: RequestId = root.id + + # First, check response routers (e.g., TaskResultHandler) + if isinstance(root, JSONRPCError): + # Route error to routers + for router in self._response_routers: + if router.route_error(response_id, root.error): + return # Handled + else: + # Route success response to routers + response_data: dict[str, Any] = root.result or {} + for router in self._response_routers: + if router.route_response(response_id, response_data): + return # Handled + + # Fall back to normal response streams + stream = self._response_streams.pop(response_id, None) + if stream: # pragma: no cover + await stream.send(root) + else: # pragma: no cover + await self._handle_incoming(RuntimeError(f"Received response with an unknown request ID: {message}")) + async def _received_request(self, responder: RequestResponder[ReceiveRequestT, SendResultT]) -> None: """ Can be overridden by subclasses to handle a request without needing to diff --git a/src/mcp/types.py b/src/mcp/types.py index dd9775f8c8..1246219a45 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -1,5 +1,6 @@ from collections.abc import Callable -from typing import Annotated, Any, Generic, Literal, TypeAlias, TypeVar +from datetime import datetime +from typing import Annotated, Any, Final, Generic, Literal, TypeAlias, TypeVar from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel from pydantic.networks import AnyUrl, UrlConstraints @@ -39,6 +40,23 @@ RequestId = Annotated[int, Field(strict=True)] | str AnyFunction: TypeAlias = Callable[..., Any] +TaskExecutionMode = Literal["forbidden", "optional", "required"] +TASK_FORBIDDEN: Final[Literal["forbidden"]] = "forbidden" +TASK_OPTIONAL: Final[Literal["optional"]] = "optional" +TASK_REQUIRED: Final[Literal["required"]] = "required" + + +class TaskMetadata(BaseModel): + """ + Metadata for augmenting a request with task execution. + Include this in the `task` field of the request parameters. + """ + + model_config = ConfigDict(extra="allow") + + ttl: Annotated[int, Field(strict=True)] | None = None + """Requested duration in milliseconds to retain task from creation.""" + class RequestParams(BaseModel): class Meta(BaseModel): @@ -52,6 +70,16 @@ class Meta(BaseModel): model_config = ConfigDict(extra="allow") + task: TaskMetadata | None = None + """ + If specified, the caller is requesting task-augmented execution for this request. + The request will return a CreateTaskResult immediately, and the actual result can be + retrieved later via tasks/result. + + Task augmentation is subject to capability negotiation - receivers MUST declare support + for task augmentation of specific request types in their capabilities. + """ + meta: Meta | None = Field(alias="_meta", default=None) @@ -321,6 +349,71 @@ class SamplingCapability(BaseModel): model_config = ConfigDict(extra="allow") +class TasksListCapability(BaseModel): + """Capability for tasks listing operations.""" + + model_config = ConfigDict(extra="allow") + + +class TasksCancelCapability(BaseModel): + """Capability for tasks cancel operations.""" + + model_config = ConfigDict(extra="allow") + + +class TasksCreateMessageCapability(BaseModel): + """Capability for tasks create messages.""" + + model_config = ConfigDict(extra="allow") + + +class TasksSamplingCapability(BaseModel): + """Capability for tasks sampling operations.""" + + model_config = ConfigDict(extra="allow") + + createMessage: TasksCreateMessageCapability | None = None + + +class TasksCreateElicitationCapability(BaseModel): + """Capability for tasks create elicitation operations.""" + + model_config = ConfigDict(extra="allow") + + +class TasksElicitationCapability(BaseModel): + """Capability for tasks elicitation operations.""" + + model_config = ConfigDict(extra="allow") + + create: TasksCreateElicitationCapability | None = None + + +class ClientTasksRequestsCapability(BaseModel): + """Capability for tasks requests operations.""" + + model_config = ConfigDict(extra="allow") + + sampling: TasksSamplingCapability | None = None + + elicitation: TasksElicitationCapability | None = None + + +class ClientTasksCapability(BaseModel): + """Capability for client tasks operations.""" + + model_config = ConfigDict(extra="allow") + + list: TasksListCapability | None = None + """Whether this client supports tasks/list.""" + + cancel: TasksCancelCapability | None = None + """Whether this client supports tasks/cancel.""" + + requests: ClientTasksRequestsCapability | None = None + """Specifies which request types can be augmented with tasks.""" + + class ClientCapabilities(BaseModel): """Capabilities a client may support.""" @@ -335,6 +428,9 @@ class ClientCapabilities(BaseModel): """Present if the client supports elicitation from the user.""" roots: RootsCapability | None = None """Present if the client supports listing roots.""" + tasks: ClientTasksCapability | None = None + """Present if the client supports task-augmented requests.""" + model_config = ConfigDict(extra="allow") @@ -376,6 +472,37 @@ class CompletionsCapability(BaseModel): model_config = ConfigDict(extra="allow") +class TasksCallCapability(BaseModel): + """Capability for tasks call operations.""" + + model_config = ConfigDict(extra="allow") + + +class TasksToolsCapability(BaseModel): + """Capability for tasks tools operations.""" + + model_config = ConfigDict(extra="allow") + call: TasksCallCapability | None = None + + +class ServerTasksRequestsCapability(BaseModel): + """Capability for tasks requests operations.""" + + model_config = ConfigDict(extra="allow") + + tools: TasksToolsCapability | None = None + + +class ServerTasksCapability(BaseModel): + """Capability for server tasks operations.""" + + model_config = ConfigDict(extra="allow") + + list: TasksListCapability | None = None + cancel: TasksCancelCapability | None = None + requests: ServerTasksRequestsCapability | None = None + + class ServerCapabilities(BaseModel): """Capabilities that a server may support.""" @@ -391,7 +518,154 @@ class ServerCapabilities(BaseModel): """Present if the server offers any tools to call.""" completions: CompletionsCapability | None = None """Present if the server offers autocompletion suggestions for prompts and resources.""" + tasks: ServerTasksCapability | None = None + """Present if the server supports task-augmented requests.""" + model_config = ConfigDict(extra="allow") + + +TaskStatus = Literal["working", "input_required", "completed", "failed", "cancelled"] + +# Task status constants +TASK_STATUS_WORKING: Final[Literal["working"]] = "working" +TASK_STATUS_INPUT_REQUIRED: Final[Literal["input_required"]] = "input_required" +TASK_STATUS_COMPLETED: Final[Literal["completed"]] = "completed" +TASK_STATUS_FAILED: Final[Literal["failed"]] = "failed" +TASK_STATUS_CANCELLED: Final[Literal["cancelled"]] = "cancelled" + + +class RelatedTaskMetadata(BaseModel): + """ + Metadata for associating messages with a task. + + Include this in the `_meta` field under the key `io.modelcontextprotocol/related-task`. + """ + model_config = ConfigDict(extra="allow") + taskId: str + """The task identifier this message is associated with.""" + + +class Task(BaseModel): + """Data associated with a task.""" + + model_config = ConfigDict(extra="allow") + + taskId: str + """The task identifier.""" + + status: TaskStatus + """Current task state.""" + + statusMessage: str | None = None + """ + Optional human-readable message describing the current task state. + This can provide context for any status, including: + - Reasons for "cancelled" status + - Summaries for "completed" status + - Diagnostic information for "failed" status (e.g., error details, what went wrong) + """ + + createdAt: datetime # Pydantic will enforce ISO 8601 and re-serialize as a string later + """ISO 8601 timestamp when the task was created.""" + + lastUpdatedAt: datetime + """ISO 8601 timestamp when the task was last updated.""" + + ttl: Annotated[int, Field(strict=True)] | None + """Actual retention duration from creation in milliseconds, null for unlimited.""" + + pollInterval: Annotated[int, Field(strict=True)] | None = None + """Suggested polling interval in milliseconds.""" + + +class CreateTaskResult(Result): + """A response to a task-augmented request.""" + + task: Task + + +class GetTaskRequestParams(RequestParams): + model_config = ConfigDict(extra="allow") + taskId: str + """The task identifier to query.""" + + +class GetTaskRequest(Request[GetTaskRequestParams, Literal["tasks/get"]]): + """A request to retrieve the state of a task.""" + + method: Literal["tasks/get"] = "tasks/get" + + params: GetTaskRequestParams + + +class GetTaskResult(Result, Task): + """The response to a tasks/get request.""" + + +class GetTaskPayloadRequestParams(RequestParams): + model_config = ConfigDict(extra="allow") + + taskId: str + """The task identifier to retrieve results for.""" + + +class GetTaskPayloadRequest(Request[GetTaskPayloadRequestParams, Literal["tasks/result"]]): + """A request to retrieve the result of a completed task.""" + + method: Literal["tasks/result"] = "tasks/result" + params: GetTaskPayloadRequestParams + + +class GetTaskPayloadResult(Result): + """ + The response to a tasks/result request. + The structure matches the result type of the original request. + For example, a tools/call task would return the CallToolResult structure. + """ + + +class CancelTaskRequestParams(RequestParams): + model_config = ConfigDict(extra="allow") + + taskId: str + """The task identifier to cancel.""" + + +class CancelTaskRequest(Request[CancelTaskRequestParams, Literal["tasks/cancel"]]): + """A request to cancel a task.""" + + method: Literal["tasks/cancel"] = "tasks/cancel" + params: CancelTaskRequestParams + + +class CancelTaskResult(Result, Task): + """The response to a tasks/cancel request.""" + + +class ListTasksRequest(PaginatedRequest[Literal["tasks/list"]]): + """A request to retrieve a list of tasks.""" + + method: Literal["tasks/list"] = "tasks/list" + + +class ListTasksResult(PaginatedResult): + """The response to a tasks/list request.""" + + tasks: list[Task] + + +class TaskStatusNotificationParams(NotificationParams, Task): + """Parameters for a `notifications/tasks/status` notification.""" + + +class TaskStatusNotification(Notification[TaskStatusNotificationParams, Literal["notifications/tasks/status"]]): + """ + An optional notification from the receiver to the requestor, informing them that a task's status has changed. + Receivers are not required to send these notifications + """ + + method: Literal["notifications/tasks/status"] = "notifications/tasks/status" + params: TaskStatusNotificationParams class InitializeRequestParams(RequestParams): @@ -1011,8 +1285,28 @@ class ToolAnnotations(BaseModel): of a memory tool is not. Default: true """ + + model_config = ConfigDict(extra="allow") + + +class ToolExecution(BaseModel): + """Execution-related properties for a tool.""" + model_config = ConfigDict(extra="allow") + taskSupport: TaskExecutionMode | None = None + """ + Indicates whether this tool supports task-augmented execution. + This allows clients to handle long-running operations through polling + the task system. + + - "forbidden": Tool does not support task-augmented execution (default when absent) + - "optional": Tool may support task-augmented execution + - "required": Tool requires task-augmented execution + + Default: "forbidden" + """ + class Tool(BaseMetadata): """Definition for a tool the client can call.""" @@ -1035,6 +1329,9 @@ class Tool(BaseMetadata): See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) for notes on _meta usage. """ + + execution: ToolExecution | None = None + model_config = ConfigDict(extra="allow") @@ -1419,10 +1716,17 @@ class RootsListChangedNotification( class CancelledNotificationParams(NotificationParams): """Parameters for cancellation notifications.""" - requestId: RequestId - """The ID of the request to cancel.""" + requestId: RequestId | None = None + """ + The ID of the request to cancel. + + This MUST correspond to the ID of a request previously issued in the same direction. + This MUST be provided for cancelling non-task requests. + This MUST NOT be used for cancelling tasks (use the `tasks/cancel` request instead). + """ reason: str | None = None """An optional string describing the reason for the cancellation.""" + model_config = ConfigDict(extra="allow") @@ -1462,29 +1766,41 @@ class ElicitCompleteNotification( params: ElicitCompleteNotificationParams -class ClientRequest( - RootModel[ - PingRequest - | InitializeRequest - | CompleteRequest - | SetLevelRequest - | GetPromptRequest - | ListPromptsRequest - | ListResourcesRequest - | ListResourceTemplatesRequest - | ReadResourceRequest - | SubscribeRequest - | UnsubscribeRequest - | CallToolRequest - | ListToolsRequest - ] -): +ClientRequestType: TypeAlias = ( + PingRequest + | InitializeRequest + | CompleteRequest + | SetLevelRequest + | GetPromptRequest + | ListPromptsRequest + | ListResourcesRequest + | ListResourceTemplatesRequest + | ReadResourceRequest + | SubscribeRequest + | UnsubscribeRequest + | CallToolRequest + | ListToolsRequest + | GetTaskRequest + | GetTaskPayloadRequest + | ListTasksRequest + | CancelTaskRequest +) + + +class ClientRequest(RootModel[ClientRequestType]): pass -class ClientNotification( - RootModel[CancelledNotification | ProgressNotification | InitializedNotification | RootsListChangedNotification] -): +ClientNotificationType: TypeAlias = ( + CancelledNotification + | ProgressNotification + | InitializedNotification + | RootsListChangedNotification + | TaskStatusNotification +) + + +class ClientNotification(RootModel[ClientNotificationType]): pass @@ -1585,41 +1901,74 @@ class ElicitationRequiredErrorData(BaseModel): model_config = ConfigDict(extra="allow") -class ClientResult(RootModel[EmptyResult | CreateMessageResult | ListRootsResult | ElicitResult]): +ClientResultType: TypeAlias = ( + EmptyResult + | CreateMessageResult + | ListRootsResult + | ElicitResult + | GetTaskResult + | GetTaskPayloadResult + | ListTasksResult + | CancelTaskResult + | CreateTaskResult +) + + +class ClientResult(RootModel[ClientResultType]): pass -class ServerRequest(RootModel[PingRequest | CreateMessageRequest | ListRootsRequest | ElicitRequest]): +ServerRequestType: TypeAlias = ( + PingRequest + | CreateMessageRequest + | ListRootsRequest + | ElicitRequest + | GetTaskRequest + | GetTaskPayloadRequest + | ListTasksRequest + | CancelTaskRequest +) + + +class ServerRequest(RootModel[ServerRequestType]): pass -class ServerNotification( - RootModel[ - CancelledNotification - | ProgressNotification - | LoggingMessageNotification - | ResourceUpdatedNotification - | ResourceListChangedNotification - | ToolListChangedNotification - | PromptListChangedNotification - | ElicitCompleteNotification - ] -): +ServerNotificationType: TypeAlias = ( + CancelledNotification + | ProgressNotification + | LoggingMessageNotification + | ResourceUpdatedNotification + | ResourceListChangedNotification + | ToolListChangedNotification + | PromptListChangedNotification + | ElicitCompleteNotification + | TaskStatusNotification +) + + +class ServerNotification(RootModel[ServerNotificationType]): pass -class ServerResult( - RootModel[ - EmptyResult - | InitializeResult - | CompleteResult - | GetPromptResult - | ListPromptsResult - | ListResourcesResult - | ListResourceTemplatesResult - | ReadResourceResult - | CallToolResult - | ListToolsResult - ] -): +ServerResultType: TypeAlias = ( + EmptyResult + | InitializeResult + | CompleteResult + | GetPromptResult + | ListPromptsResult + | ListResourcesResult + | ListResourceTemplatesResult + | ReadResourceResult + | CallToolResult + | ListToolsResult + | GetTaskResult + | GetTaskPayloadResult + | ListTasksResult + | CancelTaskResult + | CreateTaskResult +) + + +class ServerResult(RootModel[ServerResultType]): pass diff --git a/tests/experimental/__init__.py b/tests/experimental/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/experimental/tasks/__init__.py b/tests/experimental/tasks/__init__.py new file mode 100644 index 0000000000..6e8649d283 --- /dev/null +++ b/tests/experimental/tasks/__init__.py @@ -0,0 +1 @@ +"""Tests for MCP task support.""" diff --git a/tests/experimental/tasks/client/__init__.py b/tests/experimental/tasks/client/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/experimental/tasks/client/test_capabilities.py b/tests/experimental/tasks/client/test_capabilities.py new file mode 100644 index 0000000000..f2def4e3a6 --- /dev/null +++ b/tests/experimental/tasks/client/test_capabilities.py @@ -0,0 +1,331 @@ +"""Tests for client task capabilities declaration during initialization.""" + +import anyio +import pytest + +import mcp.types as types +from mcp import ClientCapabilities +from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers +from mcp.client.session import ClientSession +from mcp.shared.context import RequestContext +from mcp.shared.message import SessionMessage +from mcp.types import ( + LATEST_PROTOCOL_VERSION, + ClientRequest, + Implementation, + InitializeRequest, + InitializeResult, + JSONRPCMessage, + JSONRPCRequest, + JSONRPCResponse, + ServerCapabilities, + ServerResult, +) + + +@pytest.mark.anyio +async def test_client_capabilities_without_tasks(): + """Test that tasks capability is None when not provided.""" + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1) + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1) + + received_capabilities = None + + async def mock_server(): + nonlocal received_capabilities + + session_message = await client_to_server_receive.receive() + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + received_capabilities = request.root.params.capabilities + + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + async with server_to_client_send: + await server_to_client_send.send( + SessionMessage( + JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + await client_to_server_receive.receive() + + async with ( + ClientSession( + server_to_client_receive, + client_to_server_send, + ) as session, + anyio.create_task_group() as tg, + client_to_server_send, + client_to_server_receive, + server_to_client_send, + server_to_client_receive, + ): + tg.start_soon(mock_server) + await session.initialize() + + # Assert that tasks capability is None when not provided + assert received_capabilities is not None + assert received_capabilities.tasks is None + + +@pytest.mark.anyio +async def test_client_capabilities_with_tasks(): + """Test that tasks capability is properly set when handlers are provided.""" + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1) + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1) + + received_capabilities: ClientCapabilities | None = None + + # Define custom handlers to trigger capability building (never actually called) + async def my_list_tasks_handler( + context: RequestContext[ClientSession, None], + params: types.PaginatedRequestParams | None, + ) -> types.ListTasksResult | types.ErrorData: + raise NotImplementedError + + async def my_cancel_task_handler( + context: RequestContext[ClientSession, None], + params: types.CancelTaskRequestParams, + ) -> types.CancelTaskResult | types.ErrorData: + raise NotImplementedError + + async def mock_server(): + nonlocal received_capabilities + + session_message = await client_to_server_receive.receive() + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + received_capabilities = request.root.params.capabilities + + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + async with server_to_client_send: + await server_to_client_send.send( + SessionMessage( + JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + await client_to_server_receive.receive() + + # Create handlers container + task_handlers = ExperimentalTaskHandlers( + list_tasks=my_list_tasks_handler, + cancel_task=my_cancel_task_handler, + ) + + async with ( + ClientSession( + server_to_client_receive, + client_to_server_send, + experimental_task_handlers=task_handlers, + ) as session, + anyio.create_task_group() as tg, + client_to_server_send, + client_to_server_receive, + server_to_client_send, + server_to_client_receive, + ): + tg.start_soon(mock_server) + await session.initialize() + + # Assert that tasks capability is properly set from handlers + assert received_capabilities is not None + assert received_capabilities.tasks is not None + assert isinstance(received_capabilities.tasks, types.ClientTasksCapability) + assert received_capabilities.tasks.list is not None + assert received_capabilities.tasks.cancel is not None + + +@pytest.mark.anyio +async def test_client_capabilities_auto_built_from_handlers(): + """Test that tasks capability is automatically built from provided handlers.""" + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1) + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1) + + received_capabilities: ClientCapabilities | None = None + + # Define custom handlers (not defaults) + async def my_list_tasks_handler( + context: RequestContext[ClientSession, None], + params: types.PaginatedRequestParams | None, + ) -> types.ListTasksResult | types.ErrorData: + raise NotImplementedError + + async def my_cancel_task_handler( + context: RequestContext[ClientSession, None], + params: types.CancelTaskRequestParams, + ) -> types.CancelTaskResult | types.ErrorData: + raise NotImplementedError + + async def mock_server(): + nonlocal received_capabilities + + session_message = await client_to_server_receive.receive() + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + received_capabilities = request.root.params.capabilities + + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + async with server_to_client_send: + await server_to_client_send.send( + SessionMessage( + JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + await client_to_server_receive.receive() + + # Provide handlers via ExperimentalTaskHandlers + task_handlers = ExperimentalTaskHandlers( + list_tasks=my_list_tasks_handler, + cancel_task=my_cancel_task_handler, + ) + + async with ( + ClientSession( + server_to_client_receive, + client_to_server_send, + experimental_task_handlers=task_handlers, + ) as session, + anyio.create_task_group() as tg, + client_to_server_send, + client_to_server_receive, + server_to_client_send, + server_to_client_receive, + ): + tg.start_soon(mock_server) + await session.initialize() + + # Assert that tasks capability was auto-built from handlers + assert received_capabilities is not None + assert received_capabilities.tasks is not None + assert received_capabilities.tasks.list is not None + assert received_capabilities.tasks.cancel is not None + # requests should be None since we didn't provide task-augmented handlers + assert received_capabilities.tasks.requests is None + + +@pytest.mark.anyio +async def test_client_capabilities_with_task_augmented_handlers(): + """Test that requests capability is built when augmented handlers are provided.""" + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1) + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1) + + received_capabilities: ClientCapabilities | None = None + + # Define task-augmented handler + async def my_augmented_sampling_handler( + context: RequestContext[ClientSession, None], + params: types.CreateMessageRequestParams, + task_metadata: types.TaskMetadata, + ) -> types.CreateTaskResult | types.ErrorData: + raise NotImplementedError + + async def mock_server(): + nonlocal received_capabilities + + session_message = await client_to_server_receive.receive() + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + received_capabilities = request.root.params.capabilities + + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + async with server_to_client_send: + await server_to_client_send.send( + SessionMessage( + JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + await client_to_server_receive.receive() + + # Provide task-augmented sampling handler + task_handlers = ExperimentalTaskHandlers( + augmented_sampling=my_augmented_sampling_handler, + ) + + async with ( + ClientSession( + server_to_client_receive, + client_to_server_send, + experimental_task_handlers=task_handlers, + ) as session, + anyio.create_task_group() as tg, + client_to_server_send, + client_to_server_receive, + server_to_client_send, + server_to_client_receive, + ): + tg.start_soon(mock_server) + await session.initialize() + + # Assert that tasks capability includes requests.sampling + assert received_capabilities is not None + assert received_capabilities.tasks is not None + assert received_capabilities.tasks.requests is not None + assert received_capabilities.tasks.requests.sampling is not None + assert received_capabilities.tasks.requests.elicitation is None # Not provided diff --git a/tests/experimental/tasks/client/test_handlers.py b/tests/experimental/tasks/client/test_handlers.py new file mode 100644 index 0000000000..86cea42ae1 --- /dev/null +++ b/tests/experimental/tasks/client/test_handlers.py @@ -0,0 +1,878 @@ +"""Tests for client-side task management handlers (server -> client requests). + +These tests verify that clients can handle task-related requests from servers: +- GetTaskRequest - server polling client's task status +- GetTaskPayloadRequest - server getting result from client's task +- ListTasksRequest - server listing client's tasks +- CancelTaskRequest - server cancelling client's task + +This is the inverse of the existing tests in test_tasks.py, which test +client -> server task requests. +""" + +from collections.abc import AsyncIterator +from dataclasses import dataclass + +import anyio +import pytest +from anyio import Event +from anyio.abc import TaskGroup +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream + +import mcp.types as types +from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers +from mcp.client.session import ClientSession +from mcp.shared.context import RequestContext +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 ( + CancelTaskRequest, + CancelTaskRequestParams, + CancelTaskResult, + ClientResult, + CreateMessageRequest, + CreateMessageRequestParams, + CreateMessageResult, + CreateTaskResult, + ElicitRequest, + ElicitRequestFormParams, + ElicitRequestParams, + ElicitResult, + ErrorData, + GetTaskPayloadRequest, + GetTaskPayloadRequestParams, + GetTaskPayloadResult, + GetTaskRequest, + GetTaskRequestParams, + GetTaskResult, + ListTasksRequest, + ListTasksResult, + SamplingMessage, + ServerNotification, + ServerRequest, + TaskMetadata, + TextContent, +) + +# Buffer size for test streams +STREAM_BUFFER_SIZE = 10 + + +@dataclass +class ClientTestStreams: + """Bidirectional message streams for client/server communication in tests.""" + + server_send: MemoryObjectSendStream[SessionMessage] + server_receive: MemoryObjectReceiveStream[SessionMessage] + client_send: MemoryObjectSendStream[SessionMessage] + client_receive: MemoryObjectReceiveStream[SessionMessage] + + +@pytest.fixture +async def client_streams() -> AsyncIterator[ClientTestStreams]: + """Create bidirectional message streams for client tests. + + Automatically closes all streams after the test completes. + """ + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage]( + STREAM_BUFFER_SIZE + ) + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage]( + STREAM_BUFFER_SIZE + ) + + streams = ClientTestStreams( + server_send=server_to_client_send, + server_receive=client_to_server_receive, + client_send=client_to_server_send, + client_receive=server_to_client_receive, + ) + + yield streams + + # Cleanup + await server_to_client_send.aclose() + await server_to_client_receive.aclose() + await client_to_server_send.aclose() + await client_to_server_receive.aclose() + + +async def _default_message_handler( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, +) -> None: + """Default message handler that ignores messages (tests handle them explicitly).""" + ... + + +@pytest.mark.anyio +async def test_client_handles_get_task_request(client_streams: ClientTestStreams) -> None: + """Test that client can respond to GetTaskRequest from server.""" + with anyio.fail_after(10): + store = InMemoryTaskStore() + received_task_id: str | None = None + + async def get_task_handler( + context: RequestContext[ClientSession, None], + params: GetTaskRequestParams, + ) -> GetTaskResult | ErrorData: + nonlocal received_task_id + received_task_id = params.taskId + task = await store.get_task(params.taskId) + assert task is not None, f"Test setup error: task {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, + ) + + await store.create_task(TaskMetadata(ttl=60000), task_id="test-task-123") + + task_handlers = ExperimentalTaskHandlers(get_task=get_task_handler) + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + experimental_task_handlers=task_handlers, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + typed_request = GetTaskRequest(params=GetTaskRequestParams(taskId="test-task-123")) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-1", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCResponse) + assert response.id == "req-1" + + result = GetTaskResult.model_validate(response.result) + assert result.taskId == "test-task-123" + assert result.status == "working" + assert received_task_id == "test-task-123" + + tg.cancel_scope.cancel() + + store.cleanup() + + +@pytest.mark.anyio +async def test_client_handles_get_task_result_request(client_streams: ClientTestStreams) -> None: + """Test that client can respond to GetTaskPayloadRequest from server.""" + with anyio.fail_after(10): + store = InMemoryTaskStore() + + async def get_task_result_handler( + context: RequestContext[ClientSession, None], + params: GetTaskPayloadRequestParams, + ) -> GetTaskPayloadResult | ErrorData: + result = await store.get_result(params.taskId) + assert result is not None, f"Test setup error: result for {params.taskId} should exist" + assert isinstance(result, types.CallToolResult) + return GetTaskPayloadResult(**result.model_dump()) + + await store.create_task(TaskMetadata(ttl=60000), task_id="test-task-456") + await store.store_result( + "test-task-456", + types.CallToolResult(content=[TextContent(type="text", text="Task completed successfully!")]), + ) + await store.update_task("test-task-456", status="completed") + + task_handlers = ExperimentalTaskHandlers(get_task_result=get_task_result_handler) + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + experimental_task_handlers=task_handlers, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + typed_request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId="test-task-456")) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-2", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCResponse) + + assert isinstance(response.result, dict) + result_dict = response.result + assert "content" in result_dict + assert len(result_dict["content"]) == 1 + assert result_dict["content"][0]["text"] == "Task completed successfully!" + + tg.cancel_scope.cancel() + + store.cleanup() + + +@pytest.mark.anyio +async def test_client_handles_list_tasks_request(client_streams: ClientTestStreams) -> None: + """Test that client can respond to ListTasksRequest from server.""" + with anyio.fail_after(10): + store = InMemoryTaskStore() + + async def list_tasks_handler( + context: RequestContext[ClientSession, None], + params: types.PaginatedRequestParams | None, + ) -> ListTasksResult | ErrorData: + cursor = params.cursor if params else None + tasks_list, next_cursor = await store.list_tasks(cursor=cursor) + return ListTasksResult(tasks=tasks_list, nextCursor=next_cursor) + + await store.create_task(TaskMetadata(ttl=60000), task_id="task-1") + await store.create_task(TaskMetadata(ttl=60000), task_id="task-2") + + task_handlers = ExperimentalTaskHandlers(list_tasks=list_tasks_handler) + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + experimental_task_handlers=task_handlers, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + typed_request = ListTasksRequest() + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-3", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCResponse) + + result = ListTasksResult.model_validate(response.result) + assert len(result.tasks) == 2 + + tg.cancel_scope.cancel() + + store.cleanup() + + +@pytest.mark.anyio +async def test_client_handles_cancel_task_request(client_streams: ClientTestStreams) -> None: + """Test that client can respond to CancelTaskRequest from server.""" + with anyio.fail_after(10): + store = InMemoryTaskStore() + + async def cancel_task_handler( + context: RequestContext[ClientSession, None], + params: CancelTaskRequestParams, + ) -> CancelTaskResult | ErrorData: + task = await store.get_task(params.taskId) + assert task is not None, f"Test setup error: task {params.taskId} should exist" + await store.update_task(params.taskId, status="cancelled") + updated = await store.get_task(params.taskId) + assert updated is not None + return CancelTaskResult( + taskId=updated.taskId, + status=updated.status, + createdAt=updated.createdAt, + lastUpdatedAt=updated.lastUpdatedAt, + ttl=updated.ttl, + ) + + await store.create_task(TaskMetadata(ttl=60000), task_id="task-to-cancel") + + task_handlers = ExperimentalTaskHandlers(cancel_task=cancel_task_handler) + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + experimental_task_handlers=task_handlers, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + typed_request = CancelTaskRequest(params=CancelTaskRequestParams(taskId="task-to-cancel")) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-4", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCResponse) + + result = CancelTaskResult.model_validate(response.result) + assert result.taskId == "task-to-cancel" + assert result.status == "cancelled" + + tg.cancel_scope.cancel() + + store.cleanup() + + +@pytest.mark.anyio +async def test_client_task_augmented_sampling(client_streams: ClientTestStreams) -> None: + """Test that client can handle task-augmented sampling request from server.""" + with anyio.fail_after(10): + store = InMemoryTaskStore() + sampling_completed = Event() + created_task_id: list[str | None] = [None] + background_tg: list[TaskGroup | None] = [None] + + async def task_augmented_sampling_callback( + context: RequestContext[ClientSession, None], + params: CreateMessageRequestParams, + task_metadata: TaskMetadata, + ) -> CreateTaskResult: + task = await store.create_task(task_metadata) + created_task_id[0] = task.taskId + + async def do_sampling() -> None: + result = CreateMessageResult( + role="assistant", + content=TextContent(type="text", text="Sampled response"), + model="test-model", + stopReason="endTurn", + ) + await store.store_result(task.taskId, result) + await store.update_task(task.taskId, status="completed") + sampling_completed.set() + + assert background_tg[0] is not None + background_tg[0].start_soon(do_sampling) + return CreateTaskResult(task=task) + + async def get_task_handler( + context: RequestContext[ClientSession, None], + params: GetTaskRequestParams, + ) -> GetTaskResult | ErrorData: + task = await store.get_task(params.taskId) + assert task is not None, f"Test setup error: task {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, + ) + + async def get_task_result_handler( + context: RequestContext[ClientSession, None], + params: GetTaskPayloadRequestParams, + ) -> GetTaskPayloadResult | ErrorData: + result = await store.get_result(params.taskId) + assert result is not None, f"Test setup error: result for {params.taskId} should exist" + assert isinstance(result, CreateMessageResult) + return GetTaskPayloadResult(**result.model_dump()) + + task_handlers = ExperimentalTaskHandlers( + augmented_sampling=task_augmented_sampling_callback, + get_task=get_task_handler, + get_task_result=get_task_result_handler, + ) + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + background_tg[0] = tg + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + experimental_task_handlers=task_handlers, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + # Step 1: Server sends task-augmented CreateMessageRequest + typed_request = CreateMessageRequest( + params=CreateMessageRequestParams( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))], + maxTokens=100, + task=TaskMetadata(ttl=60000), + ) + ) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-sampling", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + # Step 2: Client responds with CreateTaskResult + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCResponse) + + task_result = CreateTaskResult.model_validate(response.result) + task_id = task_result.task.taskId + assert task_id == created_task_id[0] + + # Step 3: Wait for background sampling + await sampling_completed.wait() + + # Step 4: Server polls task status + typed_poll = GetTaskRequest(params=GetTaskRequestParams(taskId=task_id)) + poll_request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-poll", + **typed_poll.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(poll_request))) + + poll_response_msg = await client_streams.server_receive.receive() + poll_response = poll_response_msg.message.root + assert isinstance(poll_response, types.JSONRPCResponse) + + status = GetTaskResult.model_validate(poll_response.result) + assert status.status == "completed" + + # Step 5: Server gets result + typed_result_req = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task_id)) + result_request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-result", + **typed_result_req.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(result_request))) + + result_response_msg = await client_streams.server_receive.receive() + result_response = result_response_msg.message.root + assert isinstance(result_response, types.JSONRPCResponse) + + assert isinstance(result_response.result, dict) + assert result_response.result["role"] == "assistant" + + tg.cancel_scope.cancel() + + store.cleanup() + + +@pytest.mark.anyio +async def test_client_task_augmented_elicitation(client_streams: ClientTestStreams) -> None: + """Test that client can handle task-augmented elicitation request from server.""" + with anyio.fail_after(10): + store = InMemoryTaskStore() + elicitation_completed = Event() + created_task_id: list[str | None] = [None] + background_tg: list[TaskGroup | None] = [None] + + async def task_augmented_elicitation_callback( + context: RequestContext[ClientSession, None], + params: ElicitRequestParams, + task_metadata: TaskMetadata, + ) -> CreateTaskResult | ErrorData: + task = await store.create_task(task_metadata) + created_task_id[0] = task.taskId + + async def do_elicitation() -> None: + # Simulate user providing elicitation response + result = ElicitResult(action="accept", content={"name": "Test User"}) + await store.store_result(task.taskId, result) + await store.update_task(task.taskId, status="completed") + elicitation_completed.set() + + assert background_tg[0] is not None + background_tg[0].start_soon(do_elicitation) + return CreateTaskResult(task=task) + + async def get_task_handler( + context: RequestContext[ClientSession, None], + params: GetTaskRequestParams, + ) -> GetTaskResult | ErrorData: + task = await store.get_task(params.taskId) + assert task is not None, f"Test setup error: task {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, + ) + + async def get_task_result_handler( + context: RequestContext[ClientSession, None], + params: GetTaskPayloadRequestParams, + ) -> GetTaskPayloadResult | ErrorData: + result = await store.get_result(params.taskId) + assert result is not None, f"Test setup error: result for {params.taskId} should exist" + assert isinstance(result, ElicitResult) + return GetTaskPayloadResult(**result.model_dump()) + + task_handlers = ExperimentalTaskHandlers( + augmented_elicitation=task_augmented_elicitation_callback, + get_task=get_task_handler, + get_task_result=get_task_result_handler, + ) + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + background_tg[0] = tg + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + experimental_task_handlers=task_handlers, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + # Step 1: Server sends task-augmented ElicitRequest + typed_request = ElicitRequest( + params=ElicitRequestFormParams( + message="What is your name?", + requestedSchema={"type": "object", "properties": {"name": {"type": "string"}}}, + task=TaskMetadata(ttl=60000), + ) + ) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-elicit", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + # Step 2: Client responds with CreateTaskResult + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCResponse) + + task_result = CreateTaskResult.model_validate(response.result) + task_id = task_result.task.taskId + assert task_id == created_task_id[0] + + # Step 3: Wait for background elicitation + await elicitation_completed.wait() + + # Step 4: Server polls task status + typed_poll = GetTaskRequest(params=GetTaskRequestParams(taskId=task_id)) + poll_request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-poll", + **typed_poll.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(poll_request))) + + poll_response_msg = await client_streams.server_receive.receive() + poll_response = poll_response_msg.message.root + assert isinstance(poll_response, types.JSONRPCResponse) + + status = GetTaskResult.model_validate(poll_response.result) + assert status.status == "completed" + + # Step 5: Server gets result + typed_result_req = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task_id)) + result_request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-result", + **typed_result_req.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(result_request))) + + result_response_msg = await client_streams.server_receive.receive() + result_response = result_response_msg.message.root + assert isinstance(result_response, types.JSONRPCResponse) + + # Verify the elicitation result + assert isinstance(result_response.result, dict) + assert result_response.result["action"] == "accept" + assert result_response.result["content"] == {"name": "Test User"} + + tg.cancel_scope.cancel() + + store.cleanup() + + +@pytest.mark.anyio +async def test_client_returns_error_for_unhandled_task_request(client_streams: ClientTestStreams) -> None: + """Test that client returns error when no handler is registered for task request.""" + with anyio.fail_after(10): + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + typed_request = GetTaskRequest(params=GetTaskRequestParams(taskId="nonexistent")) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-unhandled", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCError) + assert ( + "not supported" in response.error.message.lower() + or "method not found" in response.error.message.lower() + ) + + tg.cancel_scope.cancel() + + +@pytest.mark.anyio +async def test_client_returns_error_for_unhandled_task_result_request(client_streams: ClientTestStreams) -> None: + """Test that client returns error for unhandled tasks/result request.""" + with anyio.fail_after(10): + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + typed_request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId="nonexistent")) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-result", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCError) + assert "not supported" in response.error.message.lower() + + tg.cancel_scope.cancel() + + +@pytest.mark.anyio +async def test_client_returns_error_for_unhandled_list_tasks_request(client_streams: ClientTestStreams) -> None: + """Test that client returns error for unhandled tasks/list request.""" + with anyio.fail_after(10): + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + typed_request = ListTasksRequest() + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-list", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCError) + assert "not supported" in response.error.message.lower() + + tg.cancel_scope.cancel() + + +@pytest.mark.anyio +async def test_client_returns_error_for_unhandled_cancel_task_request(client_streams: ClientTestStreams) -> None: + """Test that client returns error for unhandled tasks/cancel request.""" + with anyio.fail_after(10): + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + typed_request = CancelTaskRequest(params=CancelTaskRequestParams(taskId="nonexistent")) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-cancel", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCError) + assert "not supported" in response.error.message.lower() + + tg.cancel_scope.cancel() + + +@pytest.mark.anyio +async def test_client_returns_error_for_unhandled_task_augmented_sampling(client_streams: ClientTestStreams) -> None: + """Test that client returns error for task-augmented sampling without handler.""" + with anyio.fail_after(10): + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + # No task handlers provided - uses defaults + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + # Send task-augmented sampling request + typed_request = CreateMessageRequest( + params=CreateMessageRequestParams( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))], + maxTokens=100, + task=TaskMetadata(ttl=60000), + ) + ) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-sampling", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCError) + assert "not supported" in response.error.message.lower() + + tg.cancel_scope.cancel() + + +@pytest.mark.anyio +async def test_client_returns_error_for_unhandled_task_augmented_elicitation( + client_streams: ClientTestStreams, +) -> None: + """Test that client returns error for task-augmented elicitation without handler.""" + with anyio.fail_after(10): + client_ready = anyio.Event() + + async with anyio.create_task_group() as tg: + + async def run_client() -> None: + # No task handlers provided - uses defaults + async with ClientSession( + client_streams.client_receive, + client_streams.client_send, + message_handler=_default_message_handler, + ): + client_ready.set() + await anyio.sleep_forever() + + tg.start_soon(run_client) + await client_ready.wait() + + # Send task-augmented elicitation request + typed_request = ElicitRequest( + params=ElicitRequestFormParams( + message="What is your name?", + requestedSchema={"type": "object", "properties": {"name": {"type": "string"}}}, + task=TaskMetadata(ttl=60000), + ) + ) + request = types.JSONRPCRequest( + jsonrpc="2.0", + id="req-elicit", + **typed_request.model_dump(by_alias=True), + ) + await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) + + response_msg = await client_streams.server_receive.receive() + response = response_msg.message.root + assert isinstance(response, types.JSONRPCError) + assert "not supported" in response.error.message.lower() + + tg.cancel_scope.cancel() diff --git a/tests/experimental/tasks/client/test_poll_task.py b/tests/experimental/tasks/client/test_poll_task.py new file mode 100644 index 0000000000..8275dc668e --- /dev/null +++ b/tests/experimental/tasks/client/test_poll_task.py @@ -0,0 +1,121 @@ +"""Tests for poll_task async iterator.""" + +from collections.abc import Callable, Coroutine +from datetime import datetime, timezone +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from mcp.client.experimental.tasks import ExperimentalClientFeatures +from mcp.types import GetTaskResult, TaskStatus + + +def make_task_result( + status: TaskStatus = "working", + poll_interval: int = 0, + task_id: str = "test-task", + status_message: str | None = None, +) -> GetTaskResult: + """Create GetTaskResult with sensible defaults.""" + now = datetime.now(timezone.utc) + return GetTaskResult( + taskId=task_id, + status=status, + statusMessage=status_message, + createdAt=now, + lastUpdatedAt=now, + ttl=60000, + pollInterval=poll_interval, + ) + + +def make_status_sequence( + *statuses: TaskStatus, + task_id: str = "test-task", +) -> Callable[[str], Coroutine[Any, Any, GetTaskResult]]: + """Create mock get_task that returns statuses in sequence.""" + status_iter = iter(statuses) + + async def mock_get_task(tid: str) -> GetTaskResult: + return make_task_result(status=next(status_iter), task_id=tid) + + return mock_get_task + + +@pytest.fixture +def mock_session() -> AsyncMock: + return AsyncMock() + + +@pytest.fixture +def features(mock_session: AsyncMock) -> ExperimentalClientFeatures: + return ExperimentalClientFeatures(mock_session) + + +@pytest.mark.anyio +async def test_poll_task_yields_until_completed(features: ExperimentalClientFeatures) -> None: + """poll_task yields each status until terminal.""" + features.get_task = make_status_sequence("working", "working", "completed") # type: ignore[method-assign] + + statuses = [s.status async for s in features.poll_task("test-task")] + + assert statuses == ["working", "working", "completed"] + + +@pytest.mark.anyio +@pytest.mark.parametrize("terminal_status", ["completed", "failed", "cancelled"]) +async def test_poll_task_exits_on_terminal(features: ExperimentalClientFeatures, terminal_status: TaskStatus) -> None: + """poll_task exits immediately when task is already terminal.""" + features.get_task = make_status_sequence(terminal_status) # type: ignore[method-assign] + + statuses = [s.status async for s in features.poll_task("test-task")] + + assert statuses == [terminal_status] + + +@pytest.mark.anyio +async def test_poll_task_continues_through_input_required(features: ExperimentalClientFeatures) -> None: + """poll_task yields input_required and continues (non-terminal).""" + features.get_task = make_status_sequence("working", "input_required", "working", "completed") # type: ignore[method-assign] + + statuses = [s.status async for s in features.poll_task("test-task")] + + assert statuses == ["working", "input_required", "working", "completed"] + + +@pytest.mark.anyio +async def test_poll_task_passes_task_id(features: ExperimentalClientFeatures) -> None: + """poll_task passes correct task_id to get_task.""" + received_ids: list[str] = [] + + async def mock_get_task(task_id: str) -> GetTaskResult: + received_ids.append(task_id) + return make_task_result(status="completed", task_id=task_id) + + features.get_task = mock_get_task # type: ignore[method-assign] + + _ = [s async for s in features.poll_task("my-task-123")] + + assert received_ids == ["my-task-123"] + + +@pytest.mark.anyio +async def test_poll_task_yields_full_result(features: ExperimentalClientFeatures) -> None: + """poll_task yields complete GetTaskResult objects.""" + + async def mock_get_task(task_id: str) -> GetTaskResult: + return make_task_result( + status="completed", + task_id=task_id, + status_message="All done!", + ) + + features.get_task = mock_get_task # type: ignore[method-assign] + + results = [r async for r in features.poll_task("test-task")] + + assert len(results) == 1 + assert results[0].status == "completed" + assert results[0].statusMessage == "All done!" + assert results[0].taskId == "test-task" diff --git a/tests/experimental/tasks/client/test_tasks.py b/tests/experimental/tasks/client/test_tasks.py new file mode 100644 index 0000000000..24c8891def --- /dev/null +++ b/tests/experimental/tasks/client/test_tasks.py @@ -0,0 +1,483 @@ +"""Tests for the experimental client task methods (session.experimental).""" + +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 ( + CallToolRequest, + CallToolRequestParams, + CallToolResult, + CancelTaskRequest, + CancelTaskResult, + ClientRequest, + ClientResult, + CreateTaskResult, + GetTaskPayloadRequest, + GetTaskPayloadResult, + GetTaskRequest, + GetTaskResult, + ListTasksRequest, + ListTasksResult, + ServerNotification, + ServerRequest, + TaskMetadata, + TextContent, + Tool, +) + + +@dataclass +class AppContext: + """Application context passed via lifespan_context.""" + + task_group: TaskGroup + store: InMemoryTaskStore + task_done_events: dict[str, Event] = field(default_factory=lambda: {}) + + +@pytest.mark.anyio +async def test_session_experimental_get_task() -> None: + """Test session.experimental.get_task() method.""" + # Note: We bypass the normal lifespan mechanism + server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment] + store = InMemoryTaskStore() + + @server.list_tools() + async def list_tools(): + return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult: + ctx = server.request_context + app = ctx.lifespan_context + if ctx.experimental.is_task: + task_metadata = ctx.experimental.task_metadata + assert task_metadata is not None + task = await app.store.create_task(task_metadata) + + done_event = Event() + app.task_done_events[task.taskId] = done_event + + async def do_work(): + async with task_execution(task.taskId, app.store) as task_ctx: + await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")])) + done_event.set() + + app.task_group.start_soon(do_work) + return CreateTaskResult(task=task) + + raise NotImplementedError + + @server.experimental.get_task() + async def handle_get_task(request: GetTaskRequest) -> GetTaskResult: + app = server.request_context.lifespan_context + task = await app.store.get_task(request.params.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 branch + + async def run_server(app_context: AppContext): + async with ServerSession( + client_to_server_receive, + server_to_client_send, + InitializationOptions( + server_name="test-server", + server_version="1.0.0", + capabilities=server.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ), + ) as server_session: + async for message in server_session.incoming_messages: + await server._handle_message(message, server_session, app_context, raise_exceptions=False) + + async with anyio.create_task_group() as tg: + app_context = AppContext(task_group=tg, store=store) + tg.start_soon(run_server, app_context) + + async with ClientSession( + server_to_client_receive, + client_to_server_send, + message_handler=message_handler, + ) as client_session: + await client_session.initialize() + + # Create a task + create_result = await client_session.send_request( + ClientRequest( + CallToolRequest( + params=CallToolRequestParams( + name="test_tool", + arguments={}, + task=TaskMetadata(ttl=60000), + ) + ) + ), + CreateTaskResult, + ) + task_id = create_result.task.taskId + + # Wait for task to complete + await app_context.task_done_events[task_id].wait() + + # Use session.experimental to get task status + task_status = await client_session.experimental.get_task(task_id) + + assert task_status.taskId == task_id + assert task_status.status == "completed" + + tg.cancel_scope.cancel() + + +@pytest.mark.anyio +async def test_session_experimental_get_task_result() -> None: + """Test session.experimental.get_task_result() method.""" + server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment] + store = InMemoryTaskStore() + + @server.list_tools() + async def list_tools(): + return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult: + ctx = server.request_context + app = ctx.lifespan_context + if ctx.experimental.is_task: + task_metadata = ctx.experimental.task_metadata + assert task_metadata is not None + task = await app.store.create_task(task_metadata) + + done_event = Event() + app.task_done_events[task.taskId] = done_event + + async def do_work(): + async with task_execution(task.taskId, app.store) as task_ctx: + await task_ctx.complete( + CallToolResult(content=[TextContent(type="text", text="Task result content")]) + ) + done_event.set() + + app.task_group.start_soon(do_work) + return CreateTaskResult(task=task) + + raise NotImplementedError + + @server.experimental.get_task_result() + async def handle_get_task_result( + request: GetTaskPayloadRequest, + ) -> 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 GetTaskPayloadResult(**result.model_dump()) + + # Set up streams + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) + + async def message_handler( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + ) -> None: ... # pragma: no branch + + async def run_server(app_context: AppContext): + async with ServerSession( + client_to_server_receive, + server_to_client_send, + InitializationOptions( + server_name="test-server", + server_version="1.0.0", + capabilities=server.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ), + ) as server_session: + async for message in server_session.incoming_messages: + await server._handle_message(message, server_session, app_context, raise_exceptions=False) + + async with anyio.create_task_group() as tg: + app_context = AppContext(task_group=tg, store=store) + tg.start_soon(run_server, app_context) + + async with ClientSession( + server_to_client_receive, + client_to_server_send, + message_handler=message_handler, + ) as client_session: + await client_session.initialize() + + # Create a task + create_result = await client_session.send_request( + ClientRequest( + CallToolRequest( + params=CallToolRequestParams( + name="test_tool", + arguments={}, + task=TaskMetadata(ttl=60000), + ) + ) + ), + CreateTaskResult, + ) + task_id = create_result.task.taskId + + # Wait for task to complete + await app_context.task_done_events[task_id].wait() + + # Use TaskClient to get task result + task_result = await client_session.experimental.get_task_result(task_id, CallToolResult) + + assert len(task_result.content) == 1 + content = task_result.content[0] + assert isinstance(content, TextContent) + assert content.text == "Task result content" + + tg.cancel_scope.cancel() + + +@pytest.mark.anyio +async def test_session_experimental_list_tasks() -> None: + """Test TaskClient.list_tasks() method.""" + server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment] + store = InMemoryTaskStore() + + @server.list_tools() + async def list_tools(): + return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult: + ctx = server.request_context + app = ctx.lifespan_context + if ctx.experimental.is_task: + task_metadata = ctx.experimental.task_metadata + assert task_metadata is not None + task = await app.store.create_task(task_metadata) + + done_event = Event() + app.task_done_events[task.taskId] = done_event + + async def do_work(): + async with task_execution(task.taskId, app.store) as task_ctx: + await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")])) + done_event.set() + + app.task_group.start_soon(do_work) + return CreateTaskResult(task=task) + + raise NotImplementedError + + @server.experimental.list_tasks() + async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult: + app = server.request_context.lifespan_context + tasks_list, next_cursor = await app.store.list_tasks(cursor=request.params.cursor if request.params else None) + return ListTasksResult(tasks=tasks_list, nextCursor=next_cursor) + + # Set up streams + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) + + async def message_handler( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + ) -> None: ... # pragma: no branch + + async def run_server(app_context: AppContext): + async with ServerSession( + client_to_server_receive, + server_to_client_send, + InitializationOptions( + server_name="test-server", + server_version="1.0.0", + capabilities=server.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ), + ) as server_session: + async for message in server_session.incoming_messages: + await server._handle_message(message, server_session, app_context, raise_exceptions=False) + + async with anyio.create_task_group() as tg: + app_context = AppContext(task_group=tg, store=store) + tg.start_soon(run_server, app_context) + + async with ClientSession( + server_to_client_receive, + client_to_server_send, + message_handler=message_handler, + ) as client_session: + await client_session.initialize() + + # Create two tasks + for _ in range(2): + create_result = await client_session.send_request( + ClientRequest( + CallToolRequest( + params=CallToolRequestParams( + name="test_tool", + arguments={}, + task=TaskMetadata(ttl=60000), + ) + ) + ), + CreateTaskResult, + ) + await app_context.task_done_events[create_result.task.taskId].wait() + + # Use TaskClient to list tasks + list_result = await client_session.experimental.list_tasks() + + assert len(list_result.tasks) == 2 + + tg.cancel_scope.cancel() + + +@pytest.mark.anyio +async def test_session_experimental_cancel_task() -> None: + """Test TaskClient.cancel_task() method.""" + server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment] + store = InMemoryTaskStore() + + @server.list_tools() + async def list_tools(): + return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult: + ctx = server.request_context + app = ctx.lifespan_context + if ctx.experimental.is_task: + task_metadata = ctx.experimental.task_metadata + assert task_metadata is not None + task = await app.store.create_task(task_metadata) + # Don't start any work - task stays in "working" status + return CreateTaskResult(task=task) + + raise NotImplementedError + + @server.experimental.get_task() + async def handle_get_task(request: GetTaskRequest) -> GetTaskResult: + app = server.request_context.lifespan_context + task = await app.store.get_task(request.params.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.cancel_task() + async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult: + 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" + await app.store.update_task(request.params.taskId, status="cancelled") + # CancelTaskResult extends Task, so we need to return the updated task info + updated_task = await app.store.get_task(request.params.taskId) + assert updated_task is not None + return CancelTaskResult( + taskId=updated_task.taskId, + status=updated_task.status, + createdAt=updated_task.createdAt, + lastUpdatedAt=updated_task.lastUpdatedAt, + ttl=updated_task.ttl, + ) + + # Set up streams + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) + + async def message_handler( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + ) -> None: ... # pragma: no branch + + async def run_server(app_context: AppContext): + async with ServerSession( + client_to_server_receive, + server_to_client_send, + InitializationOptions( + server_name="test-server", + server_version="1.0.0", + capabilities=server.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ), + ) as server_session: + async for message in server_session.incoming_messages: + await server._handle_message(message, server_session, app_context, raise_exceptions=False) + + async with anyio.create_task_group() as tg: + app_context = AppContext(task_group=tg, store=store) + tg.start_soon(run_server, app_context) + + async with ClientSession( + server_to_client_receive, + client_to_server_send, + message_handler=message_handler, + ) as client_session: + await client_session.initialize() + + # Create a task (but don't complete it) + create_result = await client_session.send_request( + ClientRequest( + CallToolRequest( + params=CallToolRequestParams( + name="test_tool", + arguments={}, + task=TaskMetadata(ttl=60000), + ) + ) + ), + CreateTaskResult, + ) + task_id = create_result.task.taskId + + # Verify task is working + status_before = await client_session.experimental.get_task(task_id) + assert status_before.status == "working" + + # Cancel the task + await client_session.experimental.cancel_task(task_id) + + # Verify task is cancelled + status_after = await client_session.experimental.get_task(task_id) + assert status_after.status == "cancelled" + + tg.cancel_scope.cancel() diff --git a/tests/experimental/tasks/server/__init__.py b/tests/experimental/tasks/server/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/experimental/tasks/server/test_context.py b/tests/experimental/tasks/server/test_context.py new file mode 100644 index 0000000000..2f09ff1540 --- /dev/null +++ b/tests/experimental/tasks/server/test_context.py @@ -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): + ... diff --git a/tests/experimental/tasks/server/test_integration.py b/tests/experimental/tasks/server/test_integration.py new file mode 100644 index 0000000000..ba61dfcead --- /dev/null +++ b/tests/experimental/tasks/server/test_integration.py @@ -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() diff --git a/tests/experimental/tasks/server/test_run_task_flow.py b/tests/experimental/tasks/server/test_run_task_flow.py new file mode 100644 index 0000000000..7f680beb66 --- /dev/null +++ b/tests/experimental/tasks/server/test_run_task_flow.py @@ -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) diff --git a/tests/experimental/tasks/server/test_server.py b/tests/experimental/tasks/server/test_server.py new file mode 100644 index 0000000000..7209ed412a --- /dev/null +++ b/tests/experimental/tasks/server/test_server.py @@ -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() diff --git a/tests/experimental/tasks/server/test_server_task_context.py b/tests/experimental/tasks/server/test_server_task_context.py new file mode 100644 index 0000000000..3d6b16f482 --- /dev/null +++ b/tests/experimental/tasks/server/test_server_task_context.py @@ -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() diff --git a/tests/experimental/tasks/server/test_store.py b/tests/experimental/tasks/server/test_store.py new file mode 100644 index 0000000000..2eac31dfe6 --- /dev/null +++ b/tests/experimental/tasks/server/test_store.py @@ -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" diff --git a/tests/experimental/tasks/server/test_task_result_handler.py b/tests/experimental/tasks/server/test_task_result_handler.py new file mode 100644 index 0000000000..db5b9edc70 --- /dev/null +++ b/tests/experimental/tasks/server/test_task_result_handler.py @@ -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) diff --git a/tests/experimental/tasks/test_capabilities.py b/tests/experimental/tasks/test_capabilities.py new file mode 100644 index 0000000000..e78f16fe3f --- /dev/null +++ b/tests/experimental/tasks/test_capabilities.py @@ -0,0 +1,283 @@ +"""Tests for tasks capability checking utilities.""" + +import pytest + +from mcp.shared.exceptions import McpError +from mcp.shared.experimental.tasks.capabilities import ( + check_tasks_capability, + has_task_augmented_elicitation, + has_task_augmented_sampling, + require_task_augmented_elicitation, + require_task_augmented_sampling, +) +from mcp.types import ( + ClientCapabilities, + ClientTasksCapability, + ClientTasksRequestsCapability, + TasksCreateElicitationCapability, + TasksCreateMessageCapability, + TasksElicitationCapability, + TasksSamplingCapability, +) + + +class TestCheckTasksCapability: + """Tests for check_tasks_capability function.""" + + def test_required_requests_none_returns_true(self) -> None: + """When required.requests is None, should return True.""" + required = ClientTasksCapability() + client = ClientTasksCapability() + assert check_tasks_capability(required, client) is True + + def test_client_requests_none_returns_false(self) -> None: + """When client.requests is None but required.requests is set, should return False.""" + required = ClientTasksCapability(requests=ClientTasksRequestsCapability()) + client = ClientTasksCapability() + assert check_tasks_capability(required, client) is False + + def test_elicitation_required_but_client_missing(self) -> None: + """When elicitation is required but client doesn't have it.""" + required = ClientTasksCapability( + requests=ClientTasksRequestsCapability(elicitation=TasksElicitationCapability()) + ) + client = ClientTasksCapability(requests=ClientTasksRequestsCapability()) + assert check_tasks_capability(required, client) is False + + def test_elicitation_create_required_but_client_missing(self) -> None: + """When elicitation.create is required but client doesn't have it.""" + required = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()) + ) + ) + client = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability() # No create + ) + ) + assert check_tasks_capability(required, client) is False + + def test_elicitation_create_present(self) -> None: + """When elicitation.create is required and client has it.""" + required = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()) + ) + ) + client = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()) + ) + ) + assert check_tasks_capability(required, client) is True + + def test_sampling_required_but_client_missing(self) -> None: + """When sampling is required but client doesn't have it.""" + required = ClientTasksCapability(requests=ClientTasksRequestsCapability(sampling=TasksSamplingCapability())) + client = ClientTasksCapability(requests=ClientTasksRequestsCapability()) + assert check_tasks_capability(required, client) is False + + def test_sampling_create_message_required_but_client_missing(self) -> None: + """When sampling.createMessage is required but client doesn't have it.""" + required = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()) + ) + ) + client = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + sampling=TasksSamplingCapability() # No createMessage + ) + ) + assert check_tasks_capability(required, client) is False + + def test_sampling_create_message_present(self) -> None: + """When sampling.createMessage is required and client has it.""" + required = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()) + ) + ) + client = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()) + ) + ) + assert check_tasks_capability(required, client) is True + + def test_both_elicitation_and_sampling_present(self) -> None: + """When both elicitation.create and sampling.createMessage are required and client has both.""" + required = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()), + sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()), + ) + ) + client = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()), + sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()), + ) + ) + assert check_tasks_capability(required, client) is True + + def test_elicitation_without_create_required(self) -> None: + """When elicitation is required but not create specifically.""" + required = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability() # No create + ) + ) + client = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()) + ) + ) + assert check_tasks_capability(required, client) is True + + def test_sampling_without_create_message_required(self) -> None: + """When sampling is required but not createMessage specifically.""" + required = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + sampling=TasksSamplingCapability() # No createMessage + ) + ) + client = ClientTasksCapability( + requests=ClientTasksRequestsCapability( + sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()) + ) + ) + assert check_tasks_capability(required, client) is True + + +class TestHasTaskAugmentedElicitation: + """Tests for has_task_augmented_elicitation function.""" + + def test_tasks_none(self) -> None: + """Returns False when caps.tasks is None.""" + caps = ClientCapabilities() + assert has_task_augmented_elicitation(caps) is False + + def test_requests_none(self) -> None: + """Returns False when caps.tasks.requests is None.""" + caps = ClientCapabilities(tasks=ClientTasksCapability()) + assert has_task_augmented_elicitation(caps) is False + + def test_elicitation_none(self) -> None: + """Returns False when caps.tasks.requests.elicitation is None.""" + caps = ClientCapabilities(tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability())) + assert has_task_augmented_elicitation(caps) is False + + def test_create_none(self) -> None: + """Returns False when caps.tasks.requests.elicitation.create is None.""" + caps = ClientCapabilities( + tasks=ClientTasksCapability( + requests=ClientTasksRequestsCapability(elicitation=TasksElicitationCapability()) + ) + ) + assert has_task_augmented_elicitation(caps) is False + + def test_create_present(self) -> None: + """Returns True when full capability path is present.""" + caps = ClientCapabilities( + tasks=ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()) + ) + ) + ) + assert has_task_augmented_elicitation(caps) is True + + +class TestHasTaskAugmentedSampling: + """Tests for has_task_augmented_sampling function.""" + + def test_tasks_none(self) -> None: + """Returns False when caps.tasks is None.""" + caps = ClientCapabilities() + assert has_task_augmented_sampling(caps) is False + + def test_requests_none(self) -> None: + """Returns False when caps.tasks.requests is None.""" + caps = ClientCapabilities(tasks=ClientTasksCapability()) + assert has_task_augmented_sampling(caps) is False + + def test_sampling_none(self) -> None: + """Returns False when caps.tasks.requests.sampling is None.""" + caps = ClientCapabilities(tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability())) + assert has_task_augmented_sampling(caps) is False + + def test_create_message_none(self) -> None: + """Returns False when caps.tasks.requests.sampling.createMessage is None.""" + caps = ClientCapabilities( + tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability(sampling=TasksSamplingCapability())) + ) + assert has_task_augmented_sampling(caps) is False + + def test_create_message_present(self) -> None: + """Returns True when full capability path is present.""" + caps = ClientCapabilities( + tasks=ClientTasksCapability( + requests=ClientTasksRequestsCapability( + sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()) + ) + ) + ) + assert has_task_augmented_sampling(caps) is True + + +class TestRequireTaskAugmentedElicitation: + """Tests for require_task_augmented_elicitation function.""" + + def test_raises_when_none(self) -> None: + """Raises McpError when client_caps is None.""" + with pytest.raises(McpError) as exc_info: + require_task_augmented_elicitation(None) + assert "task-augmented elicitation" in str(exc_info.value) + + def test_raises_when_missing(self) -> None: + """Raises McpError when capability is missing.""" + caps = ClientCapabilities() + with pytest.raises(McpError) as exc_info: + require_task_augmented_elicitation(caps) + assert "task-augmented elicitation" in str(exc_info.value) + + def test_passes_when_present(self) -> None: + """Does not raise when capability is present.""" + caps = ClientCapabilities( + tasks=ClientTasksCapability( + requests=ClientTasksRequestsCapability( + elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()) + ) + ) + ) + require_task_augmented_elicitation(caps) + + +class TestRequireTaskAugmentedSampling: + """Tests for require_task_augmented_sampling function.""" + + def test_raises_when_none(self) -> None: + """Raises McpError when client_caps is None.""" + with pytest.raises(McpError) as exc_info: + require_task_augmented_sampling(None) + assert "task-augmented sampling" in str(exc_info.value) + + def test_raises_when_missing(self) -> None: + """Raises McpError when capability is missing.""" + caps = ClientCapabilities() + with pytest.raises(McpError) as exc_info: + require_task_augmented_sampling(caps) + assert "task-augmented sampling" in str(exc_info.value) + + def test_passes_when_present(self) -> None: + """Does not raise when capability is present.""" + caps = ClientCapabilities( + tasks=ClientTasksCapability( + requests=ClientTasksRequestsCapability( + sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()) + ) + ) + ) + require_task_augmented_sampling(caps) diff --git a/tests/experimental/tasks/test_elicitation_scenarios.py b/tests/experimental/tasks/test_elicitation_scenarios.py new file mode 100644 index 0000000000..be2b616018 --- /dev/null +++ b/tests/experimental/tasks/test_elicitation_scenarios.py @@ -0,0 +1,737 @@ +""" +Tests for the four elicitation scenarios with tasks. + +This tests all combinations of tool call types and elicitation types: +1. Normal tool call + Normal elicitation (session.elicit) +2. Normal tool call + Task-augmented elicitation (session.experimental.elicit_as_task) +3. Task-augmented tool call + Normal elicitation (task.elicit) +4. Task-augmented tool call + Task-augmented elicitation (task.elicit_as_task) + +And the same for sampling (create_message). +""" + +from typing import Any + +import anyio +import pytest +from anyio import Event + +from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers +from mcp.client.session import ClientSession +from mcp.server import Server +from mcp.server.experimental.task_context import ServerTaskContext +from mcp.server.lowlevel import NotificationOptions +from mcp.shared.context import RequestContext +from mcp.shared.experimental.tasks.helpers import is_terminal +from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore +from mcp.shared.message import SessionMessage +from mcp.types import ( + TASK_REQUIRED, + CallToolResult, + CreateMessageRequestParams, + CreateMessageResult, + CreateTaskResult, + ElicitRequestParams, + ElicitResult, + ErrorData, + GetTaskPayloadResult, + GetTaskResult, + SamplingMessage, + TaskMetadata, + TextContent, + Tool, + ToolExecution, +) + + +def create_client_task_handlers( + client_task_store: InMemoryTaskStore, + elicit_received: Event, +) -> ExperimentalTaskHandlers: + """Create task handlers for client to handle task-augmented elicitation from server.""" + + elicit_response = ElicitResult(action="accept", content={"confirm": True}) + task_complete_events: dict[str, Event] = {} + + async def handle_augmented_elicitation( + context: RequestContext[ClientSession, Any], + params: ElicitRequestParams, + task_metadata: TaskMetadata, + ) -> CreateTaskResult: + """Handle task-augmented elicitation by creating a client-side task.""" + elicit_received.set() + task = await client_task_store.create_task(task_metadata) + task_complete_events[task.taskId] = Event() + + async def complete_task() -> None: + # Store result before updating status to avoid race condition + await client_task_store.store_result(task.taskId, elicit_response) + await client_task_store.update_task(task.taskId, status="completed") + task_complete_events[task.taskId].set() + + context.session._task_group.start_soon(complete_task) # pyright: ignore[reportPrivateUsage] + return CreateTaskResult(task=task) + + async def handle_get_task( + context: RequestContext[ClientSession, Any], + params: Any, + ) -> GetTaskResult: + """Handle tasks/get from server.""" + task = await client_task_store.get_task(params.taskId) + assert task is not None, f"Task not found: {params.taskId}" + return GetTaskResult( + taskId=task.taskId, + status=task.status, + statusMessage=task.statusMessage, + createdAt=task.createdAt, + lastUpdatedAt=task.lastUpdatedAt, + ttl=task.ttl, + pollInterval=100, + ) + + async def handle_get_task_result( + context: RequestContext[ClientSession, Any], + params: Any, + ) -> GetTaskPayloadResult | ErrorData: + """Handle tasks/result from server.""" + event = task_complete_events.get(params.taskId) + assert event is not None, f"No completion event for task: {params.taskId}" + await event.wait() + result = await client_task_store.get_result(params.taskId) + assert result is not None, f"Result not found for task: {params.taskId}" + return GetTaskPayloadResult.model_validate(result.model_dump(by_alias=True)) + + return ExperimentalTaskHandlers( + augmented_elicitation=handle_augmented_elicitation, + get_task=handle_get_task, + get_task_result=handle_get_task_result, + ) + + +def create_sampling_task_handlers( + client_task_store: InMemoryTaskStore, + sampling_received: Event, +) -> ExperimentalTaskHandlers: + """Create task handlers for client to handle task-augmented sampling from server.""" + + sampling_response = CreateMessageResult( + role="assistant", + content=TextContent(type="text", text="Hello from the model!"), + model="test-model", + ) + task_complete_events: dict[str, Event] = {} + + async def handle_augmented_sampling( + context: RequestContext[ClientSession, Any], + params: CreateMessageRequestParams, + task_metadata: TaskMetadata, + ) -> CreateTaskResult: + """Handle task-augmented sampling by creating a client-side task.""" + sampling_received.set() + task = await client_task_store.create_task(task_metadata) + task_complete_events[task.taskId] = Event() + + async def complete_task() -> None: + # Store result before updating status to avoid race condition + await client_task_store.store_result(task.taskId, sampling_response) + await client_task_store.update_task(task.taskId, status="completed") + task_complete_events[task.taskId].set() + + context.session._task_group.start_soon(complete_task) # pyright: ignore[reportPrivateUsage] + return CreateTaskResult(task=task) + + async def handle_get_task( + context: RequestContext[ClientSession, Any], + params: Any, + ) -> GetTaskResult: + """Handle tasks/get from server.""" + task = await client_task_store.get_task(params.taskId) + assert task is not None, f"Task not found: {params.taskId}" + return GetTaskResult( + taskId=task.taskId, + status=task.status, + statusMessage=task.statusMessage, + createdAt=task.createdAt, + lastUpdatedAt=task.lastUpdatedAt, + ttl=task.ttl, + pollInterval=100, + ) + + async def handle_get_task_result( + context: RequestContext[ClientSession, Any], + params: Any, + ) -> GetTaskPayloadResult | ErrorData: + """Handle tasks/result from server.""" + event = task_complete_events.get(params.taskId) + assert event is not None, f"No completion event for task: {params.taskId}" + await event.wait() + result = await client_task_store.get_result(params.taskId) + assert result is not None, f"Result not found for task: {params.taskId}" + return GetTaskPayloadResult.model_validate(result.model_dump(by_alias=True)) + + return ExperimentalTaskHandlers( + augmented_sampling=handle_augmented_sampling, + get_task=handle_get_task, + get_task_result=handle_get_task_result, + ) + + +@pytest.mark.anyio +async def test_scenario1_normal_tool_normal_elicitation() -> None: + """ + Scenario 1: Normal tool call with normal elicitation. + + Server calls session.elicit() directly, client responds immediately. + """ + server = Server("test-scenario1") + elicit_received = Event() + tool_result: list[str] = [] + + @server.list_tools() + async def list_tools() -> list[Tool]: + return [ + Tool( + name="confirm_action", + description="Confirm an action", + inputSchema={"type": "object"}, + ) + ] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult: + ctx = server.request_context + + # Normal elicitation - expects immediate response + result = await ctx.session.elicit( + message="Please confirm the action", + requestedSchema={"type": "object", "properties": {"confirm": {"type": "boolean"}}}, + ) + + confirmed = result.content.get("confirm", False) if result.content else False + tool_result.append("confirmed" if confirmed else "cancelled") + return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")]) + + # Elicitation callback for client + async def elicitation_callback( + context: RequestContext[ClientSession, Any], + params: ElicitRequestParams, + ) -> ElicitResult: + elicit_received.set() + return ElicitResult(action="accept", content={"confirm": True}) + + # 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, + elicitation_callback=elicitation_callback, + ) as client_session: + await client_session.initialize() + + # Call tool normally (not as task) + result = await client_session.call_tool("confirm_action", {}) + + # Verify elicitation was received and tool completed + assert elicit_received.is_set() + assert len(result.content) > 0 + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == "confirmed" + + async with anyio.create_task_group() as tg: + tg.start_soon(run_server) + tg.start_soon(run_client) + + assert tool_result[0] == "confirmed" + + +@pytest.mark.anyio +async def test_scenario2_normal_tool_task_augmented_elicitation() -> None: + """ + Scenario 2: Normal tool call with task-augmented elicitation. + + Server calls session.experimental.elicit_as_task(), client creates a task + for the elicitation and returns CreateTaskResult. Server polls client. + """ + server = Server("test-scenario2") + elicit_received = Event() + tool_result: list[str] = [] + + # Client-side task store for handling task-augmented elicitation + client_task_store = InMemoryTaskStore() + + @server.list_tools() + async def list_tools() -> list[Tool]: + return [ + Tool( + name="confirm_action", + description="Confirm an action", + inputSchema={"type": "object"}, + ) + ] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult: + ctx = server.request_context + + # Task-augmented elicitation - server polls client + result = await ctx.session.experimental.elicit_as_task( + message="Please confirm the action", + requestedSchema={"type": "object", "properties": {"confirm": {"type": "boolean"}}}, + ttl=60000, + ) + + confirmed = result.content.get("confirm", False) if result.content else False + tool_result.append("confirmed" if confirmed else "cancelled") + return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")]) + + task_handlers = create_client_task_handlers(client_task_store, elicit_received) + + # 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, + experimental_task_handlers=task_handlers, + ) as client_session: + await client_session.initialize() + + # Call tool normally (not as task) + result = await client_session.call_tool("confirm_action", {}) + + # Verify elicitation was received and tool completed + assert elicit_received.is_set() + assert len(result.content) > 0 + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == "confirmed" + + async with anyio.create_task_group() as tg: + tg.start_soon(run_server) + tg.start_soon(run_client) + + assert tool_result[0] == "confirmed" + client_task_store.cleanup() + + +@pytest.mark.anyio +async def test_scenario3_task_augmented_tool_normal_elicitation() -> None: + """ + Scenario 3: Task-augmented tool call with normal elicitation. + + Client calls tool as task. Inside the task, server uses task.elicit() + which queues the request and delivers via tasks/result. + """ + server = Server("test-scenario3") + server.experimental.enable_tasks() + + elicit_received = Event() + work_completed = Event() + + @server.list_tools() + async def list_tools() -> list[Tool]: + return [ + Tool( + name="confirm_action", + description="Confirm an action", + inputSchema={"type": "object"}, + execution=ToolExecution(taskSupport=TASK_REQUIRED), + ) + ] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult: + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + async def work(task: ServerTaskContext) -> CallToolResult: + # Normal elicitation within task - queued and delivered via tasks/result + result = await task.elicit( + message="Please confirm the action", + requestedSchema={"type": "object", "properties": {"confirm": {"type": "boolean"}}}, + ) + + confirmed = result.content.get("confirm", False) if result.content else False + work_completed.set() + return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")]) + + return await ctx.experimental.run_task(work) + + # Elicitation callback for client + async def elicitation_callback( + context: RequestContext[ClientSession, Any], + params: ElicitRequestParams, + ) -> ElicitResult: + elicit_received.set() + return ElicitResult(action="accept", content={"confirm": True}) + + # 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, + elicitation_callback=elicitation_callback, + ) as client_session: + await client_session.initialize() + + # Call tool as task + create_result = await client_session.experimental.call_tool_as_task("confirm_action", {}) + task_id = create_result.task.taskId + assert create_result.task.status == "working" + + # Poll until input_required, then call tasks/result + found_input_required = False + async for status in client_session.experimental.poll_task(task_id): # pragma: no branch + if status.status == "input_required": # pragma: no branch + found_input_required = True + break + assert found_input_required, "Expected to see input_required status" + + # This will deliver the elicitation and get the response + final_result = await client_session.experimental.get_task_result(task_id, CallToolResult) + + # Verify + assert elicit_received.is_set() + assert len(final_result.content) > 0 + assert isinstance(final_result.content[0], TextContent) + assert final_result.content[0].text == "confirmed" + + async with anyio.create_task_group() as tg: + tg.start_soon(run_server) + tg.start_soon(run_client) + + assert work_completed.is_set() + + +@pytest.mark.anyio +async def test_scenario4_task_augmented_tool_task_augmented_elicitation() -> None: + """ + Scenario 4: Task-augmented tool call with task-augmented elicitation. + + Client calls tool as task. Inside the task, server uses task.elicit_as_task() + which sends task-augmented elicitation. Client creates its own task for the + elicitation, and server polls the client. + + This tests the full bidirectional flow where: + 1. Client calls tasks/result on server (for tool task) + 2. Server delivers task-augmented elicitation through that stream + 3. Client creates its own task and returns CreateTaskResult + 4. Server polls the client's task while the client's tasks/result is still open + 5. Server gets the ElicitResult and completes the tool task + 6. Client's tasks/result returns with the CallToolResult + """ + server = Server("test-scenario4") + server.experimental.enable_tasks() + + elicit_received = Event() + work_completed = Event() + + # Client-side task store for handling task-augmented elicitation + client_task_store = InMemoryTaskStore() + + @server.list_tools() + async def list_tools() -> list[Tool]: + return [ + Tool( + name="confirm_action", + description="Confirm an action", + inputSchema={"type": "object"}, + execution=ToolExecution(taskSupport=TASK_REQUIRED), + ) + ] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult: + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + async def work(task: ServerTaskContext) -> CallToolResult: + # Task-augmented elicitation within task - server polls client + result = await task.elicit_as_task( + message="Please confirm the action", + requestedSchema={"type": "object", "properties": {"confirm": {"type": "boolean"}}}, + ttl=60000, + ) + + confirmed = result.content.get("confirm", False) if result.content else False + work_completed.set() + return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")]) + + return await ctx.experimental.run_task(work) + + task_handlers = create_client_task_handlers(client_task_store, elicit_received) + + # 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, + experimental_task_handlers=task_handlers, + ) as client_session: + await client_session.initialize() + + # Call tool as task + create_result = await client_session.experimental.call_tool_as_task("confirm_action", {}) + task_id = create_result.task.taskId + assert create_result.task.status == "working" + + # Poll until input_required or terminal, then call tasks/result + found_expected_status = False + async for status in client_session.experimental.poll_task(task_id): # pragma: no branch + if status.status == "input_required" or is_terminal(status.status): # pragma: no branch + found_expected_status = True + break + assert found_expected_status, "Expected to see input_required or terminal status" + + # This will deliver the task-augmented elicitation, + # server will poll client, and eventually return the tool result + final_result = await client_session.experimental.get_task_result(task_id, CallToolResult) + + # Verify + assert elicit_received.is_set() + assert len(final_result.content) > 0 + assert isinstance(final_result.content[0], TextContent) + assert final_result.content[0].text == "confirmed" + + async with anyio.create_task_group() as tg: + tg.start_soon(run_server) + tg.start_soon(run_client) + + assert work_completed.is_set() + client_task_store.cleanup() + + +@pytest.mark.anyio +async def test_scenario2_sampling_normal_tool_task_augmented_sampling() -> None: + """ + Scenario 2 for sampling: Normal tool call with task-augmented sampling. + + Server calls session.experimental.create_message_as_task(), client creates + a task for the sampling and returns CreateTaskResult. Server polls client. + """ + server = Server("test-scenario2-sampling") + sampling_received = Event() + tool_result: list[str] = [] + + # Client-side task store for handling task-augmented sampling + client_task_store = InMemoryTaskStore() + + @server.list_tools() + async def list_tools() -> list[Tool]: + return [ + Tool( + name="generate_text", + description="Generate text using sampling", + inputSchema={"type": "object"}, + ) + ] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult: + ctx = server.request_context + + # Task-augmented sampling - server polls client + result = await ctx.session.experimental.create_message_as_task( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))], + max_tokens=100, + ttl=60000, + ) + + assert isinstance(result.content, TextContent), "Expected TextContent response" + response_text = result.content.text + + tool_result.append(response_text) + return CallToolResult(content=[TextContent(type="text", text=response_text)]) + + task_handlers = create_sampling_task_handlers(client_task_store, sampling_received) + + # 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, + experimental_task_handlers=task_handlers, + ) as client_session: + await client_session.initialize() + + # Call tool normally (not as task) + result = await client_session.call_tool("generate_text", {}) + + # Verify sampling was received and tool completed + assert sampling_received.is_set() + assert len(result.content) > 0 + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == "Hello from the model!" + + async with anyio.create_task_group() as tg: + tg.start_soon(run_server) + tg.start_soon(run_client) + + assert tool_result[0] == "Hello from the model!" + client_task_store.cleanup() + + +@pytest.mark.anyio +async def test_scenario4_sampling_task_augmented_tool_task_augmented_sampling() -> None: + """ + Scenario 4 for sampling: Task-augmented tool call with task-augmented sampling. + + Client calls tool as task. Inside the task, server uses task.create_message_as_task() + which sends task-augmented sampling. Client creates its own task for the sampling, + and server polls the client. + """ + server = Server("test-scenario4-sampling") + server.experimental.enable_tasks() + + sampling_received = Event() + work_completed = Event() + + # Client-side task store for handling task-augmented sampling + client_task_store = InMemoryTaskStore() + + @server.list_tools() + async def list_tools() -> list[Tool]: + return [ + Tool( + name="generate_text", + description="Generate text using sampling", + inputSchema={"type": "object"}, + execution=ToolExecution(taskSupport=TASK_REQUIRED), + ) + ] + + @server.call_tool() + async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult: + ctx = server.request_context + ctx.experimental.validate_task_mode(TASK_REQUIRED) + + async def work(task: ServerTaskContext) -> CallToolResult: + # Task-augmented sampling within task - server polls client + result = await task.create_message_as_task( + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))], + max_tokens=100, + ttl=60000, + ) + + assert isinstance(result.content, TextContent), "Expected TextContent response" + response_text = result.content.text + + work_completed.set() + return CallToolResult(content=[TextContent(type="text", text=response_text)]) + + return await ctx.experimental.run_task(work) + + task_handlers = create_sampling_task_handlers(client_task_store, sampling_received) + + # 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, + experimental_task_handlers=task_handlers, + ) as client_session: + await client_session.initialize() + + # Call tool as task + create_result = await client_session.experimental.call_tool_as_task("generate_text", {}) + task_id = create_result.task.taskId + assert create_result.task.status == "working" + + # Poll until input_required or terminal + found_expected_status = False + async for status in client_session.experimental.poll_task(task_id): # pragma: no branch + if status.status == "input_required" or is_terminal(status.status): # pragma: no branch + found_expected_status = True + break + assert found_expected_status, "Expected to see input_required or terminal status" + + final_result = await client_session.experimental.get_task_result(task_id, CallToolResult) + + # Verify + assert sampling_received.is_set() + assert len(final_result.content) > 0 + assert isinstance(final_result.content[0], TextContent) + assert final_result.content[0].text == "Hello from the model!" + + async with anyio.create_task_group() as tg: + tg.start_soon(run_server) + tg.start_soon(run_client) + + assert work_completed.is_set() + client_task_store.cleanup() diff --git a/tests/experimental/tasks/test_message_queue.py b/tests/experimental/tasks/test_message_queue.py new file mode 100644 index 0000000000..86d6875cc4 --- /dev/null +++ b/tests/experimental/tasks/test_message_queue.py @@ -0,0 +1,331 @@ +""" +Tests for TaskMessageQueue and InMemoryTaskMessageQueue. +""" + +from datetime import datetime, timezone + +import anyio +import pytest + +from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, QueuedMessage +from mcp.shared.experimental.tasks.resolver import Resolver +from mcp.types import JSONRPCNotification, JSONRPCRequest + + +@pytest.fixture +def queue() -> InMemoryTaskMessageQueue: + return InMemoryTaskMessageQueue() + + +def make_request(id: int = 1, method: str = "test/method") -> JSONRPCRequest: + return JSONRPCRequest(jsonrpc="2.0", id=id, method=method) + + +def make_notification(method: str = "test/notify") -> JSONRPCNotification: + return JSONRPCNotification(jsonrpc="2.0", method=method) + + +class TestInMemoryTaskMessageQueue: + @pytest.mark.anyio + async def test_enqueue_and_dequeue(self, queue: InMemoryTaskMessageQueue) -> None: + """Test basic enqueue and dequeue operations.""" + task_id = "task-1" + msg = QueuedMessage(type="request", message=make_request()) + + await queue.enqueue(task_id, msg) + result = await queue.dequeue(task_id) + + assert result is not None + assert result.type == "request" + assert result.message.method == "test/method" + + @pytest.mark.anyio + async def test_dequeue_empty_returns_none(self, queue: InMemoryTaskMessageQueue) -> None: + """Dequeue from empty queue returns None.""" + result = await queue.dequeue("nonexistent-task") + assert result is None + + @pytest.mark.anyio + async def test_fifo_ordering(self, queue: InMemoryTaskMessageQueue) -> None: + """Messages are dequeued in FIFO order.""" + task_id = "task-1" + + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(1, "first"))) + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(2, "second"))) + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(3, "third"))) + + msg1 = await queue.dequeue(task_id) + msg2 = await queue.dequeue(task_id) + msg3 = await queue.dequeue(task_id) + + assert msg1 is not None and msg1.message.method == "first" + assert msg2 is not None and msg2.message.method == "second" + assert msg3 is not None and msg3.message.method == "third" + + @pytest.mark.anyio + async def test_separate_queues_per_task(self, queue: InMemoryTaskMessageQueue) -> None: + """Each task has its own queue.""" + await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1, "task1-msg"))) + await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2, "task2-msg"))) + + msg1 = await queue.dequeue("task-1") + msg2 = await queue.dequeue("task-2") + + assert msg1 is not None and msg1.message.method == "task1-msg" + assert msg2 is not None and msg2.message.method == "task2-msg" + + @pytest.mark.anyio + async def test_peek_does_not_remove(self, queue: InMemoryTaskMessageQueue) -> None: + """Peek returns message without removing it.""" + task_id = "task-1" + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request())) + + peeked = await queue.peek(task_id) + dequeued = await queue.dequeue(task_id) + + assert peeked is not None + assert dequeued is not None + assert isinstance(peeked.message, JSONRPCRequest) + assert isinstance(dequeued.message, JSONRPCRequest) + assert peeked.message.id == dequeued.message.id + + @pytest.mark.anyio + async def test_is_empty(self, queue: InMemoryTaskMessageQueue) -> None: + """Test is_empty method.""" + task_id = "task-1" + + assert await queue.is_empty(task_id) is True + + await queue.enqueue(task_id, QueuedMessage(type="notification", message=make_notification())) + assert await queue.is_empty(task_id) is False + + await queue.dequeue(task_id) + assert await queue.is_empty(task_id) is True + + @pytest.mark.anyio + async def test_clear_returns_all_messages(self, queue: InMemoryTaskMessageQueue) -> None: + """Clear removes and returns all messages.""" + task_id = "task-1" + + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(1))) + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(2))) + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(3))) + + messages = await queue.clear(task_id) + + assert len(messages) == 3 + assert await queue.is_empty(task_id) is True + + @pytest.mark.anyio + async def test_clear_empty_queue(self, queue: InMemoryTaskMessageQueue) -> None: + """Clear on empty queue returns empty list.""" + messages = await queue.clear("nonexistent") + assert messages == [] + + @pytest.mark.anyio + async def test_notification_messages(self, queue: InMemoryTaskMessageQueue) -> None: + """Test queuing notification messages.""" + task_id = "task-1" + msg = QueuedMessage(type="notification", message=make_notification("log/message")) + + await queue.enqueue(task_id, msg) + result = await queue.dequeue(task_id) + + assert result is not None + assert result.type == "notification" + assert result.message.method == "log/message" + + @pytest.mark.anyio + async def test_message_timestamp(self, queue: InMemoryTaskMessageQueue) -> None: + """Messages have timestamps.""" + before = datetime.now(timezone.utc) + msg = QueuedMessage(type="request", message=make_request()) + after = datetime.now(timezone.utc) + + assert before <= msg.timestamp <= after + + @pytest.mark.anyio + async def test_message_with_resolver(self, queue: InMemoryTaskMessageQueue) -> None: + """Messages can have resolvers.""" + task_id = "task-1" + resolver: Resolver[dict[str, str]] = Resolver() + + msg = QueuedMessage( + type="request", + message=make_request(), + resolver=resolver, + original_request_id=42, + ) + + await queue.enqueue(task_id, msg) + result = await queue.dequeue(task_id) + + assert result is not None + assert result.resolver is resolver + assert result.original_request_id == 42 + + @pytest.mark.anyio + async def test_cleanup_specific_task(self, queue: InMemoryTaskMessageQueue) -> None: + """Cleanup removes specific task's data.""" + await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1))) + await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2))) + + queue.cleanup("task-1") + + assert await queue.is_empty("task-1") is True + assert await queue.is_empty("task-2") is False + + @pytest.mark.anyio + async def test_cleanup_all(self, queue: InMemoryTaskMessageQueue) -> None: + """Cleanup without task_id removes all data.""" + await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1))) + await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2))) + + queue.cleanup() + + assert await queue.is_empty("task-1") is True + assert await queue.is_empty("task-2") is True + + @pytest.mark.anyio + async def test_wait_for_message_returns_immediately_if_message_exists( + self, queue: InMemoryTaskMessageQueue + ) -> None: + """wait_for_message returns immediately if queue not empty.""" + task_id = "task-1" + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request())) + + # Should return immediately, not block + with anyio.fail_after(1): + await queue.wait_for_message(task_id) + + @pytest.mark.anyio + async def test_wait_for_message_blocks_until_message(self, queue: InMemoryTaskMessageQueue) -> None: + """wait_for_message blocks until a message is enqueued.""" + task_id = "task-1" + received = False + waiter_started = anyio.Event() + + async def enqueue_when_ready() -> None: + # Wait until the waiter has started before enqueueing + await waiter_started.wait() + await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request())) + + async def wait_for_msg() -> None: + nonlocal received + # Signal that we're about to start waiting + waiter_started.set() + await queue.wait_for_message(task_id) + received = True + + async with anyio.create_task_group() as tg: + tg.start_soon(wait_for_msg) + tg.start_soon(enqueue_when_ready) + + assert received is True + + @pytest.mark.anyio + async def test_notify_message_available_wakes_waiter(self, queue: InMemoryTaskMessageQueue) -> None: + """notify_message_available wakes up waiting coroutines.""" + task_id = "task-1" + notified = False + waiter_started = anyio.Event() + + async def notify_when_ready() -> None: + # Wait until the waiter has started before notifying + await waiter_started.wait() + await queue.notify_message_available(task_id) + + async def wait_for_notification() -> None: + nonlocal notified + # Signal that we're about to start waiting + waiter_started.set() + await queue.wait_for_message(task_id) + notified = True + + async with anyio.create_task_group() as tg: + tg.start_soon(wait_for_notification) + tg.start_soon(notify_when_ready) + + assert notified is True + + @pytest.mark.anyio + async def test_peek_empty_queue_returns_none(self, queue: InMemoryTaskMessageQueue) -> None: + """Peek on empty queue returns None.""" + result = await queue.peek("nonexistent-task") + assert result is None + + @pytest.mark.anyio + async def test_wait_for_message_double_check_race_condition(self, queue: InMemoryTaskMessageQueue) -> None: + """wait_for_message returns early if message arrives after event creation but before wait.""" + task_id = "task-1" + + # To test the double-check path (lines 223-225), we need a message to arrive + # after the event is created (line 220) but before event.wait() (line 228). + # We simulate this by injecting a message before is_empty is called the second time. + + original_is_empty = queue.is_empty + call_count = 0 + + async def is_empty_with_injection(tid: str) -> bool: + nonlocal call_count + call_count += 1 + if call_count == 2 and tid == task_id: + # Before second check, inject a message - this simulates a message + # arriving between event creation and the double-check + queue._queues[task_id] = [QueuedMessage(type="request", message=make_request())] + return await original_is_empty(tid) + + queue.is_empty = is_empty_with_injection # type: ignore[method-assign] + + # Should return immediately due to double-check finding the message + with anyio.fail_after(1): + await queue.wait_for_message(task_id) + + +class TestResolver: + @pytest.mark.anyio + async def test_set_result_and_wait(self) -> None: + """Test basic set_result and wait flow.""" + resolver: Resolver[str] = Resolver() + + resolver.set_result("hello") + result = await resolver.wait() + + assert result == "hello" + assert resolver.done() + + @pytest.mark.anyio + async def test_set_exception_and_wait(self) -> None: + """Test set_exception raises on wait.""" + resolver: Resolver[str] = Resolver() + + resolver.set_exception(ValueError("test error")) + + with pytest.raises(ValueError, match="test error"): + await resolver.wait() + + assert resolver.done() + + @pytest.mark.anyio + async def test_set_result_when_already_completed_raises(self) -> None: + """Test that set_result raises if resolver already completed.""" + resolver: Resolver[str] = Resolver() + resolver.set_result("first") + + with pytest.raises(RuntimeError, match="already completed"): + resolver.set_result("second") + + @pytest.mark.anyio + async def test_set_exception_when_already_completed_raises(self) -> None: + """Test that set_exception raises if resolver already completed.""" + resolver: Resolver[str] = Resolver() + resolver.set_result("done") + + with pytest.raises(RuntimeError, match="already completed"): + resolver.set_exception(ValueError("too late")) + + @pytest.mark.anyio + async def test_done_returns_false_before_completion(self) -> None: + """Test done() returns False before any result is set.""" + resolver: Resolver[str] = Resolver() + assert resolver.done() is False diff --git a/tests/experimental/tasks/test_request_context.py b/tests/experimental/tasks/test_request_context.py new file mode 100644 index 0000000000..5fa5da81af --- /dev/null +++ b/tests/experimental/tasks/test_request_context.py @@ -0,0 +1,166 @@ +"""Tests for the RequestContext.experimental (Experimental class) task validation helpers.""" + +import pytest + +from mcp.server.experimental.request_context import Experimental +from mcp.shared.exceptions import McpError +from mcp.types import ( + METHOD_NOT_FOUND, + TASK_FORBIDDEN, + TASK_OPTIONAL, + TASK_REQUIRED, + ClientCapabilities, + ClientTasksCapability, + TaskMetadata, + Tool, + ToolExecution, +) + + +def test_is_task_true_when_metadata_present() -> None: + exp = Experimental(task_metadata=TaskMetadata(ttl=60000)) + assert exp.is_task is True + + +def test_is_task_false_when_no_metadata() -> None: + exp = Experimental(task_metadata=None) + assert exp.is_task is False + + +def test_client_supports_tasks_true() -> None: + exp = Experimental(_client_capabilities=ClientCapabilities(tasks=ClientTasksCapability())) + assert exp.client_supports_tasks is True + + +def test_client_supports_tasks_false_no_tasks() -> None: + exp = Experimental(_client_capabilities=ClientCapabilities()) + assert exp.client_supports_tasks is False + + +def test_client_supports_tasks_false_no_capabilities() -> None: + exp = Experimental(_client_capabilities=None) + assert exp.client_supports_tasks is False + + +def test_validate_task_mode_required_with_task_is_valid() -> None: + exp = Experimental(task_metadata=TaskMetadata(ttl=60000)) + error = exp.validate_task_mode(TASK_REQUIRED, raise_error=False) + assert error is None + + +def test_validate_task_mode_required_without_task_returns_error() -> None: + exp = Experimental(task_metadata=None) + error = exp.validate_task_mode(TASK_REQUIRED, raise_error=False) + assert error is not None + assert error.code == METHOD_NOT_FOUND + assert "requires task-augmented" in error.message + + +def test_validate_task_mode_required_without_task_raises_by_default() -> None: + exp = Experimental(task_metadata=None) + with pytest.raises(McpError) as exc_info: + exp.validate_task_mode(TASK_REQUIRED) + assert exc_info.value.error.code == METHOD_NOT_FOUND + + +def test_validate_task_mode_forbidden_without_task_is_valid() -> None: + exp = Experimental(task_metadata=None) + error = exp.validate_task_mode(TASK_FORBIDDEN, raise_error=False) + assert error is None + + +def test_validate_task_mode_forbidden_with_task_returns_error() -> None: + exp = Experimental(task_metadata=TaskMetadata(ttl=60000)) + error = exp.validate_task_mode(TASK_FORBIDDEN, raise_error=False) + assert error is not None + assert error.code == METHOD_NOT_FOUND + assert "does not support task-augmented" in error.message + + +def test_validate_task_mode_forbidden_with_task_raises_by_default() -> None: + exp = Experimental(task_metadata=TaskMetadata(ttl=60000)) + with pytest.raises(McpError) as exc_info: + exp.validate_task_mode(TASK_FORBIDDEN) + assert exc_info.value.error.code == METHOD_NOT_FOUND + + +def test_validate_task_mode_none_treated_as_forbidden() -> None: + exp = Experimental(task_metadata=TaskMetadata(ttl=60000)) + error = exp.validate_task_mode(None, raise_error=False) + assert error is not None + assert "does not support task-augmented" in error.message + + +def test_validate_task_mode_optional_with_task_is_valid() -> None: + exp = Experimental(task_metadata=TaskMetadata(ttl=60000)) + error = exp.validate_task_mode(TASK_OPTIONAL, raise_error=False) + assert error is None + + +def test_validate_task_mode_optional_without_task_is_valid() -> None: + exp = Experimental(task_metadata=None) + error = exp.validate_task_mode(TASK_OPTIONAL, raise_error=False) + assert error is None + + +def test_validate_for_tool_with_execution_required() -> None: + exp = Experimental(task_metadata=None) + tool = Tool( + name="test", + description="test", + inputSchema={"type": "object"}, + execution=ToolExecution(taskSupport=TASK_REQUIRED), + ) + error = exp.validate_for_tool(tool, raise_error=False) + assert error is not None + assert "requires task-augmented" in error.message + + +def test_validate_for_tool_without_execution() -> None: + exp = Experimental(task_metadata=TaskMetadata(ttl=60000)) + tool = Tool( + name="test", + description="test", + inputSchema={"type": "object"}, + execution=None, + ) + error = exp.validate_for_tool(tool, raise_error=False) + assert error is not None + assert "does not support task-augmented" in error.message + + +def test_validate_for_tool_optional_with_task() -> None: + exp = Experimental(task_metadata=TaskMetadata(ttl=60000)) + tool = Tool( + name="test", + description="test", + inputSchema={"type": "object"}, + execution=ToolExecution(taskSupport=TASK_OPTIONAL), + ) + error = exp.validate_for_tool(tool, raise_error=False) + assert error is None + + +def test_can_use_tool_required_with_task_support() -> None: + exp = Experimental(_client_capabilities=ClientCapabilities(tasks=ClientTasksCapability())) + assert exp.can_use_tool(TASK_REQUIRED) is True + + +def test_can_use_tool_required_without_task_support() -> None: + exp = Experimental(_client_capabilities=ClientCapabilities()) + assert exp.can_use_tool(TASK_REQUIRED) is False + + +def test_can_use_tool_optional_without_task_support() -> None: + exp = Experimental(_client_capabilities=ClientCapabilities()) + assert exp.can_use_tool(TASK_OPTIONAL) is True + + +def test_can_use_tool_forbidden_without_task_support() -> None: + exp = Experimental(_client_capabilities=ClientCapabilities()) + assert exp.can_use_tool(TASK_FORBIDDEN) is True + + +def test_can_use_tool_none_without_task_support() -> None: + exp = Experimental(_client_capabilities=ClientCapabilities()) + assert exp.can_use_tool(None) is True diff --git a/tests/experimental/tasks/test_spec_compliance.py b/tests/experimental/tasks/test_spec_compliance.py new file mode 100644 index 0000000000..842bfa7e1f --- /dev/null +++ b/tests/experimental/tasks/test_spec_compliance.py @@ -0,0 +1,753 @@ +""" +Tasks Spec Compliance Tests +=========================== + +Test structure mirrors: https://modelcontextprotocol.io/specification/draft/basic/utilities/tasks.md + +Each section contains tests for normative requirements (MUST/SHOULD/MAY). +""" + +from datetime import datetime, timezone + +import pytest + +from mcp.server import Server +from mcp.server.lowlevel import NotificationOptions +from mcp.shared.experimental.tasks.helpers import MODEL_IMMEDIATE_RESPONSE_KEY +from mcp.types import ( + CancelTaskRequest, + CancelTaskResult, + CreateTaskResult, + GetTaskRequest, + GetTaskResult, + ListTasksRequest, + ListTasksResult, + ServerCapabilities, + Task, +) + +# Shared test datetime +TEST_DATETIME = datetime(2025, 1, 1, tzinfo=timezone.utc) + + +def _get_capabilities(server: Server) -> ServerCapabilities: + """Helper to get capabilities from a server.""" + return server.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ) + + +def test_server_without_task_handlers_has_no_tasks_capability() -> None: + """Server without any task handlers has no tasks capability.""" + server: Server = Server("test") + caps = _get_capabilities(server) + assert caps.tasks is None + + +def test_server_with_list_tasks_handler_declares_list_capability() -> None: + """Server with list_tasks handler declares tasks.list capability.""" + server: Server = Server("test") + + @server.experimental.list_tasks() + async def handle_list(req: ListTasksRequest) -> ListTasksResult: + raise NotImplementedError + + caps = _get_capabilities(server) + assert caps.tasks is not None + assert caps.tasks.list is not None + + +def test_server_with_cancel_task_handler_declares_cancel_capability() -> None: + """Server with cancel_task handler declares tasks.cancel capability.""" + server: Server = Server("test") + + @server.experimental.cancel_task() + async def handle_cancel(req: CancelTaskRequest) -> CancelTaskResult: + raise NotImplementedError + + caps = _get_capabilities(server) + assert caps.tasks is not None + assert caps.tasks.cancel is not None + + +def test_server_with_get_task_handler_declares_requests_tools_call_capability() -> None: + """ + Server with get_task handler declares tasks.requests.tools.call capability. + (get_task is required for task-augmented tools/call support) + """ + server: Server = Server("test") + + @server.experimental.get_task() + async def handle_get(req: GetTaskRequest) -> GetTaskResult: + raise NotImplementedError + + caps = _get_capabilities(server) + assert caps.tasks is not None + assert caps.tasks.requests is not None + assert caps.tasks.requests.tools is not None + + +def test_server_without_list_handler_has_no_list_capability() -> None: + """Server without list_tasks handler has no tasks.list capability.""" + server: Server = Server("test") + + # Register only get_task (not list_tasks) + @server.experimental.get_task() + async def handle_get(req: GetTaskRequest) -> GetTaskResult: + raise NotImplementedError + + caps = _get_capabilities(server) + assert caps.tasks is not None + assert caps.tasks.list is None + + +def test_server_without_cancel_handler_has_no_cancel_capability() -> None: + """Server without cancel_task handler has no tasks.cancel capability.""" + server: Server = Server("test") + + # Register only get_task (not cancel_task) + @server.experimental.get_task() + async def handle_get(req: GetTaskRequest) -> GetTaskResult: + raise NotImplementedError + + caps = _get_capabilities(server) + assert caps.tasks is not None + assert caps.tasks.cancel is None + + +def test_server_with_all_task_handlers_has_full_capability() -> None: + """Server with all task handlers declares complete tasks capability.""" + server: Server = Server("test") + + @server.experimental.list_tasks() + async def handle_list(req: ListTasksRequest) -> ListTasksResult: + raise NotImplementedError + + @server.experimental.cancel_task() + async def handle_cancel(req: CancelTaskRequest) -> CancelTaskResult: + raise NotImplementedError + + @server.experimental.get_task() + async def handle_get(req: GetTaskRequest) -> GetTaskResult: + raise NotImplementedError + + caps = _get_capabilities(server) + assert caps.tasks is not None + assert caps.tasks.list is not None + assert caps.tasks.cancel is not None + assert caps.tasks.requests is not None + assert caps.tasks.requests.tools is not None + + +class TestClientCapabilities: + """ + Clients declare: + - tasks.list — supports listing operations + - tasks.cancel — supports cancellation + - tasks.requests.sampling.createMessage — task-augmented sampling + - tasks.requests.elicitation.create — task-augmented elicitation + """ + + def test_client_declares_tasks_capability(self) -> None: + """Client can declare tasks capability.""" + pytest.skip("TODO") + + +class TestToolLevelNegotiation: + """ + Tools in tools/list responses include execution.taskSupport with values: + - Not present or "forbidden": No task augmentation allowed + - "optional": Task augmentation allowed at requestor discretion + - "required": Task augmentation is mandatory + """ + + def test_tool_execution_task_forbidden_rejects_task_augmented_call(self) -> None: + """Tool with execution.taskSupport="forbidden" MUST reject task-augmented calls (-32601).""" + pytest.skip("TODO") + + def test_tool_execution_task_absent_rejects_task_augmented_call(self) -> None: + """Tool without execution.taskSupport MUST reject task-augmented calls (-32601).""" + pytest.skip("TODO") + + def test_tool_execution_task_optional_accepts_normal_call(self) -> None: + """Tool with execution.taskSupport="optional" accepts normal calls.""" + pytest.skip("TODO") + + def test_tool_execution_task_optional_accepts_task_augmented_call(self) -> None: + """Tool with execution.taskSupport="optional" accepts task-augmented calls.""" + pytest.skip("TODO") + + def test_tool_execution_task_required_rejects_normal_call(self) -> None: + """Tool with execution.taskSupport="required" MUST reject non-task calls (-32601).""" + pytest.skip("TODO") + + def test_tool_execution_task_required_accepts_task_augmented_call(self) -> None: + """Tool with execution.taskSupport="required" accepts task-augmented calls.""" + pytest.skip("TODO") + + +class TestCapabilityNegotiation: + """ + Requestors SHOULD only augment requests with a task if the corresponding + capability has been declared by the receiver. + + Receivers that do not declare the task capability for a request type + MUST process requests of that type normally, ignoring any task-augmentation + metadata if present. + """ + + def test_receiver_without_capability_ignores_task_metadata(self) -> None: + """ + Receiver without task capability MUST process request normally, + ignoring task-augmentation metadata. + """ + pytest.skip("TODO") + + def test_receiver_with_capability_may_require_task_augmentation(self) -> None: + """ + Receivers that declare task capability MAY return error (-32600) + for non-task-augmented requests, requiring task augmentation. + """ + pytest.skip("TODO") + + +class TestTaskStatusLifecycle: + """ + Tasks begin in working status and follow valid transitions: + working → input_required → working → terminal + working → terminal (directly) + input_required → terminal (directly) + + Terminal states (no further transitions allowed): + - completed + - failed + - cancelled + """ + + def test_task_begins_in_working_status(self) -> None: + """Tasks MUST begin in working status.""" + pytest.skip("TODO") + + def test_working_to_completed_transition(self) -> None: + """working → completed is valid.""" + pytest.skip("TODO") + + def test_working_to_failed_transition(self) -> None: + """working → failed is valid.""" + pytest.skip("TODO") + + def test_working_to_cancelled_transition(self) -> None: + """working → cancelled is valid.""" + pytest.skip("TODO") + + def test_working_to_input_required_transition(self) -> None: + """working → input_required is valid.""" + pytest.skip("TODO") + + def test_input_required_to_working_transition(self) -> None: + """input_required → working is valid.""" + pytest.skip("TODO") + + def test_input_required_to_terminal_transition(self) -> None: + """input_required → terminal is valid.""" + pytest.skip("TODO") + + def test_terminal_state_no_further_transitions(self) -> None: + """Terminal states allow no further transitions.""" + pytest.skip("TODO") + + def test_completed_is_terminal(self) -> None: + """completed is a terminal state.""" + pytest.skip("TODO") + + def test_failed_is_terminal(self) -> None: + """failed is a terminal state.""" + pytest.skip("TODO") + + def test_cancelled_is_terminal(self) -> None: + """cancelled is a terminal state.""" + pytest.skip("TODO") + + +class TestInputRequiredStatus: + """ + When a receiver needs information to proceed, it moves the task to input_required. + The requestor should call tasks/result to retrieve input requests. + The task must include io.modelcontextprotocol/related-task metadata in associated requests. + """ + + def test_input_required_status_retrievable_via_tasks_get(self) -> None: + """Task in input_required status is retrievable via tasks/get.""" + pytest.skip("TODO") + + def test_input_required_related_task_metadata_in_requests(self) -> None: + """ + Task MUST include io.modelcontextprotocol/related-task metadata + in associated requests. + """ + pytest.skip("TODO") + + +class TestCreatingTask: + """ + Request structure: + {"method": "tools/call", "params": {"name": "...", "arguments": {...}, "task": {"ttl": 60000}}} + + Response (CreateTaskResult): + {"result": {"task": {"taskId": "...", "status": "working", ...}}} + + Receivers may include io.modelcontextprotocol/model-immediate-response in _meta. + """ + + def test_task_augmented_request_returns_create_task_result(self) -> None: + """Task-augmented request MUST return CreateTaskResult immediately.""" + pytest.skip("TODO") + + def test_create_task_result_contains_task_id(self) -> None: + """CreateTaskResult MUST contain taskId.""" + pytest.skip("TODO") + + def test_create_task_result_contains_status_working(self) -> None: + """CreateTaskResult MUST have status=working initially.""" + pytest.skip("TODO") + + def test_create_task_result_contains_created_at(self) -> None: + """CreateTaskResult MUST contain createdAt timestamp.""" + pytest.skip("TODO") + + def test_create_task_result_created_at_is_iso8601(self) -> None: + """createdAt MUST be ISO 8601 formatted.""" + pytest.skip("TODO") + + def test_create_task_result_may_contain_ttl(self) -> None: + """CreateTaskResult MAY contain ttl.""" + pytest.skip("TODO") + + def test_create_task_result_may_contain_poll_interval(self) -> None: + """CreateTaskResult MAY contain pollInterval.""" + pytest.skip("TODO") + + def test_create_task_result_may_contain_status_message(self) -> None: + """CreateTaskResult MAY contain statusMessage.""" + pytest.skip("TODO") + + def test_receiver_may_override_requested_ttl(self) -> None: + """Receiver MAY override requested ttl but MUST return actual value.""" + pytest.skip("TODO") + + def test_model_immediate_response_in_meta(self) -> None: + """ + Receiver MAY include io.modelcontextprotocol/model-immediate-response + in _meta to provide immediate response while task executes. + """ + # Verify the constant has the correct value per spec + assert MODEL_IMMEDIATE_RESPONSE_KEY == "io.modelcontextprotocol/model-immediate-response" + + # CreateTaskResult can include model-immediate-response in _meta + task = Task( + taskId="test-123", + status="working", + createdAt=TEST_DATETIME, + lastUpdatedAt=TEST_DATETIME, + ttl=60000, + ) + immediate_msg = "Task started, processing your request..." + # Note: Must use _meta= (alias) not meta= due to Pydantic alias handling + result = CreateTaskResult( + task=task, + **{"_meta": {MODEL_IMMEDIATE_RESPONSE_KEY: immediate_msg}}, + ) + + # Verify the metadata is present and correct + assert result.meta is not None + assert MODEL_IMMEDIATE_RESPONSE_KEY in result.meta + assert result.meta[MODEL_IMMEDIATE_RESPONSE_KEY] == immediate_msg + + # Verify it serializes correctly with _meta alias + serialized = result.model_dump(by_alias=True) + assert "_meta" in serialized + assert MODEL_IMMEDIATE_RESPONSE_KEY in serialized["_meta"] + assert serialized["_meta"][MODEL_IMMEDIATE_RESPONSE_KEY] == immediate_msg + + +class TestGettingTaskStatus: + """ + Request: {"method": "tasks/get", "params": {"taskId": "..."}} + Response: Returns full Task object with current status and pollInterval. + """ + + def test_tasks_get_returns_task_object(self) -> None: + """tasks/get MUST return full Task object.""" + pytest.skip("TODO") + + def test_tasks_get_returns_current_status(self) -> None: + """tasks/get MUST return current status.""" + pytest.skip("TODO") + + def test_tasks_get_may_return_poll_interval(self) -> None: + """tasks/get MAY return pollInterval.""" + pytest.skip("TODO") + + def test_tasks_get_invalid_task_id_returns_error(self) -> None: + """tasks/get with invalid taskId MUST return -32602.""" + pytest.skip("TODO") + + def test_tasks_get_nonexistent_task_id_returns_error(self) -> None: + """tasks/get with nonexistent taskId MUST return -32602.""" + pytest.skip("TODO") + + +class TestRetrievingResults: + """ + Request: {"method": "tasks/result", "params": {"taskId": "..."}} + Response: The actual operation result structure (e.g., CallToolResult). + + This call blocks until terminal status. + """ + + def test_tasks_result_returns_underlying_result(self) -> None: + """tasks/result MUST return exactly what underlying request would return.""" + pytest.skip("TODO") + + def test_tasks_result_blocks_until_terminal(self) -> None: + """tasks/result MUST block for non-terminal tasks.""" + pytest.skip("TODO") + + def test_tasks_result_unblocks_on_terminal(self) -> None: + """tasks/result MUST unblock upon reaching terminal status.""" + pytest.skip("TODO") + + def test_tasks_result_includes_related_task_metadata(self) -> None: + """tasks/result MUST include io.modelcontextprotocol/related-task in _meta.""" + pytest.skip("TODO") + + def test_tasks_result_returns_error_for_failed_task(self) -> None: + """ + tasks/result returns the same error the underlying request + would have produced for failed tasks. + """ + pytest.skip("TODO") + + def test_tasks_result_invalid_task_id_returns_error(self) -> None: + """tasks/result with invalid taskId MUST return -32602.""" + pytest.skip("TODO") + + +class TestListingTasks: + """ + Request: {"method": "tasks/list", "params": {"cursor": "optional"}} + Response: Array of tasks with pagination support via nextCursor. + """ + + def test_tasks_list_returns_array_of_tasks(self) -> None: + """tasks/list MUST return array of tasks.""" + pytest.skip("TODO") + + def test_tasks_list_pagination_with_cursor(self) -> None: + """tasks/list supports pagination via cursor.""" + pytest.skip("TODO") + + def test_tasks_list_returns_next_cursor_when_more_results(self) -> None: + """tasks/list MUST return nextCursor when more results available.""" + pytest.skip("TODO") + + def test_tasks_list_cursors_are_opaque(self) -> None: + """Implementers MUST treat cursors as opaque tokens.""" + pytest.skip("TODO") + + def test_tasks_list_invalid_cursor_returns_error(self) -> None: + """tasks/list with invalid cursor MUST return -32602.""" + pytest.skip("TODO") + + +class TestCancellingTasks: + """ + Request: {"method": "tasks/cancel", "params": {"taskId": "..."}} + Response: Returns the task object with status: "cancelled". + """ + + def test_tasks_cancel_returns_cancelled_task(self) -> None: + """tasks/cancel MUST return task with status=cancelled.""" + pytest.skip("TODO") + + def test_tasks_cancel_terminal_task_returns_error(self) -> None: + """Cancelling already-terminal task MUST return -32602.""" + pytest.skip("TODO") + + def test_tasks_cancel_completed_task_returns_error(self) -> None: + """Cancelling completed task MUST return -32602.""" + pytest.skip("TODO") + + def test_tasks_cancel_failed_task_returns_error(self) -> None: + """Cancelling failed task MUST return -32602.""" + pytest.skip("TODO") + + def test_tasks_cancel_already_cancelled_task_returns_error(self) -> None: + """Cancelling already-cancelled task MUST return -32602.""" + pytest.skip("TODO") + + def test_tasks_cancel_invalid_task_id_returns_error(self) -> None: + """tasks/cancel with invalid taskId MUST return -32602.""" + pytest.skip("TODO") + + +class TestStatusNotifications: + """ + Receivers MAY send: {"method": "notifications/tasks/status", "params": {...}} + These are optional; requestors MUST NOT rely on them and SHOULD continue polling. + """ + + def test_receiver_may_send_status_notification(self) -> None: + """Receiver MAY send notifications/tasks/status.""" + pytest.skip("TODO") + + def test_status_notification_contains_task_id(self) -> None: + """Status notification MUST contain taskId.""" + pytest.skip("TODO") + + def test_status_notification_contains_status(self) -> None: + """Status notification MUST contain status.""" + pytest.skip("TODO") + + +class TestTaskManagement: + """ + - Receivers generate unique task IDs as strings + - Tasks must begin in working status + - createdAt timestamps must be ISO 8601 formatted + - Receivers may override requested ttl but must return actual value + - Receivers may delete tasks after TTL expires + - All task-related messages must include io.modelcontextprotocol/related-task + in _meta except for tasks/get, tasks/list, tasks/cancel operations + """ + + def test_task_ids_are_unique_strings(self) -> None: + """Receivers MUST generate unique task IDs as strings.""" + pytest.skip("TODO") + + def test_multiple_tasks_have_unique_ids(self) -> None: + """Multiple tasks MUST have unique IDs.""" + pytest.skip("TODO") + + def test_receiver_may_delete_tasks_after_ttl(self) -> None: + """Receivers MAY delete tasks after TTL expires.""" + pytest.skip("TODO") + + def test_related_task_metadata_in_task_messages(self) -> None: + """ + All task-related messages MUST include io.modelcontextprotocol/related-task + in _meta. + """ + pytest.skip("TODO") + + def test_tasks_get_does_not_require_related_task_metadata(self) -> None: + """tasks/get does not require related-task metadata.""" + pytest.skip("TODO") + + def test_tasks_list_does_not_require_related_task_metadata(self) -> None: + """tasks/list does not require related-task metadata.""" + pytest.skip("TODO") + + def test_tasks_cancel_does_not_require_related_task_metadata(self) -> None: + """tasks/cancel does not require related-task metadata.""" + pytest.skip("TODO") + + +class TestResultHandling: + """ + - Receivers must return CreateTaskResult immediately upon accepting task-augmented requests + - tasks/result must return exactly what the underlying request would return + - tasks/result blocks for non-terminal tasks; must unblock upon reaching terminal status + """ + + def test_create_task_result_returned_immediately(self) -> None: + """Receiver MUST return CreateTaskResult immediately (not after work completes).""" + pytest.skip("TODO") + + def test_tasks_result_matches_underlying_result_structure(self) -> None: + """tasks/result MUST return same structure as underlying request.""" + pytest.skip("TODO") + + def test_tasks_result_for_tool_call_returns_call_tool_result(self) -> None: + """tasks/result for tools/call returns CallToolResult.""" + pytest.skip("TODO") + + +class TestProgressTracking: + """ + Task-augmented requests support progress notifications using the progressToken + mechanism, which remains valid throughout the task lifetime. + """ + + def test_progress_token_valid_throughout_task_lifetime(self) -> None: + """progressToken remains valid throughout task lifetime.""" + pytest.skip("TODO") + + def test_progress_notifications_sent_during_task_execution(self) -> None: + """Progress notifications can be sent during task execution.""" + pytest.skip("TODO") + + +class TestProtocolErrors: + """ + Protocol Errors (JSON-RPC standard codes): + - -32600 (Invalid request): Non-task requests to endpoint requiring task augmentation + - -32602 (Invalid params): Invalid/nonexistent taskId, invalid cursor, cancel terminal task + - -32603 (Internal error): Server-side execution failures + """ + + def test_invalid_request_for_required_task_augmentation(self) -> None: + """Non-task request to task-required endpoint returns -32600.""" + pytest.skip("TODO") + + def test_invalid_params_for_invalid_task_id(self) -> None: + """Invalid taskId returns -32602.""" + pytest.skip("TODO") + + def test_invalid_params_for_nonexistent_task_id(self) -> None: + """Nonexistent taskId returns -32602.""" + pytest.skip("TODO") + + def test_invalid_params_for_invalid_cursor(self) -> None: + """Invalid cursor in tasks/list returns -32602.""" + pytest.skip("TODO") + + def test_invalid_params_for_cancel_terminal_task(self) -> None: + """Attempt to cancel terminal task returns -32602.""" + pytest.skip("TODO") + + def test_internal_error_for_server_failure(self) -> None: + """Server-side execution failure returns -32603.""" + pytest.skip("TODO") + + +class TestTaskExecutionErrors: + """ + When underlying requests fail, the task moves to failed status. + - tasks/get response should include statusMessage explaining failure + - tasks/result returns same error the underlying request would have produced + - For tool calls, isError: true moves task to failed status + """ + + def test_underlying_failure_moves_task_to_failed(self) -> None: + """Underlying request failure moves task to failed status.""" + pytest.skip("TODO") + + def test_failed_task_has_status_message(self) -> None: + """Failed task SHOULD include statusMessage explaining failure.""" + pytest.skip("TODO") + + def test_tasks_result_returns_underlying_error(self) -> None: + """tasks/result returns same error underlying request would produce.""" + pytest.skip("TODO") + + def test_tool_call_is_error_true_moves_to_failed(self) -> None: + """Tool call with isError: true moves task to failed status.""" + pytest.skip("TODO") + + +class TestTaskObject: + """ + Task Object fields: + - taskId: String identifier + - status: Current execution state + - statusMessage: Optional human-readable description + - createdAt: ISO 8601 timestamp of creation + - ttl: Milliseconds before potential deletion + - pollInterval: Suggested milliseconds between polls + """ + + def test_task_has_task_id_string(self) -> None: + """Task MUST have taskId as string.""" + pytest.skip("TODO") + + def test_task_has_status(self) -> None: + """Task MUST have status.""" + pytest.skip("TODO") + + def test_task_status_message_is_optional(self) -> None: + """Task statusMessage is optional.""" + pytest.skip("TODO") + + def test_task_has_created_at(self) -> None: + """Task MUST have createdAt.""" + pytest.skip("TODO") + + def test_task_ttl_is_optional(self) -> None: + """Task ttl is optional.""" + pytest.skip("TODO") + + def test_task_poll_interval_is_optional(self) -> None: + """Task pollInterval is optional.""" + pytest.skip("TODO") + + +class TestRelatedTaskMetadata: + """ + Related Task Metadata structure: + {"_meta": {"io.modelcontextprotocol/related-task": {"taskId": "..."}}} + """ + + def test_related_task_metadata_structure(self) -> None: + """Related task metadata has correct structure.""" + pytest.skip("TODO") + + def test_related_task_metadata_contains_task_id(self) -> None: + """Related task metadata contains taskId.""" + pytest.skip("TODO") + + +class TestAccessAndIsolation: + """ + - Task IDs enable access to sensitive results + - Authorization context binding is essential where available + - For non-authorized environments: strong entropy IDs, strict TTL limits + """ + + def test_task_bound_to_authorization_context(self) -> None: + """ + Receivers receiving authorization context MUST bind tasks to that context. + """ + pytest.skip("TODO") + + def test_reject_task_operations_outside_authorization_context(self) -> None: + """ + Receivers MUST reject task operations for tasks outside + requestor's authorization context. + """ + pytest.skip("TODO") + + def test_non_authorized_environments_use_secure_ids(self) -> None: + """ + For non-authorized environments, receivers SHOULD use + cryptographically secure IDs. + """ + pytest.skip("TODO") + + def test_non_authorized_environments_use_shorter_ttls(self) -> None: + """ + For non-authorized environments, receivers SHOULD use shorter TTLs. + """ + pytest.skip("TODO") + + +class TestResourceLimits: + """ + Receivers should: + - Enforce concurrent task limits per requestor + - Implement maximum TTL constraints + - Clean up expired tasks promptly + """ + + def test_concurrent_task_limit_enforced(self) -> None: + """Receiver SHOULD enforce concurrent task limits per requestor.""" + pytest.skip("TODO") + + def test_maximum_ttl_constraint_enforced(self) -> None: + """Receiver SHOULD implement maximum TTL constraints.""" + pytest.skip("TODO") + + def test_expired_tasks_cleaned_up(self) -> None: + """Receiver SHOULD clean up expired tasks promptly.""" + pytest.skip("TODO") diff --git a/tests/server/test_validation.py b/tests/server/test_validation.py new file mode 100644 index 0000000000..56044460df --- /dev/null +++ b/tests/server/test_validation.py @@ -0,0 +1,141 @@ +"""Tests for server validation functions.""" + +import pytest + +from mcp.server.validation import ( + check_sampling_tools_capability, + validate_sampling_tools, + validate_tool_use_result_messages, +) +from mcp.shared.exceptions import McpError +from mcp.types import ( + ClientCapabilities, + SamplingCapability, + SamplingMessage, + SamplingToolsCapability, + TextContent, + Tool, + ToolChoice, + ToolResultContent, + ToolUseContent, +) + + +class TestCheckSamplingToolsCapability: + """Tests for check_sampling_tools_capability function.""" + + def test_returns_false_when_caps_none(self) -> None: + """Returns False when client_caps is None.""" + assert check_sampling_tools_capability(None) is False + + def test_returns_false_when_sampling_none(self) -> None: + """Returns False when client_caps.sampling is None.""" + caps = ClientCapabilities() + assert check_sampling_tools_capability(caps) is False + + def test_returns_false_when_tools_none(self) -> None: + """Returns False when client_caps.sampling.tools is None.""" + caps = ClientCapabilities(sampling=SamplingCapability()) + assert check_sampling_tools_capability(caps) is False + + def test_returns_true_when_tools_present(self) -> None: + """Returns True when sampling.tools is present.""" + caps = ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability())) + assert check_sampling_tools_capability(caps) is True + + +class TestValidateSamplingTools: + """Tests for validate_sampling_tools function.""" + + def test_no_error_when_tools_none(self) -> None: + """No error when tools and tool_choice are None.""" + validate_sampling_tools(None, None, None) # Should not raise + + def test_raises_when_tools_provided_but_no_capability(self) -> None: + """Raises McpError when tools provided but client doesn't support.""" + tool = Tool(name="test", inputSchema={"type": "object"}) + with pytest.raises(McpError) as exc_info: + validate_sampling_tools(None, [tool], None) + assert "sampling tools capability" in str(exc_info.value) + + def test_raises_when_tool_choice_provided_but_no_capability(self) -> None: + """Raises McpError when tool_choice provided but client doesn't support.""" + with pytest.raises(McpError) as exc_info: + validate_sampling_tools(None, None, ToolChoice(mode="auto")) + assert "sampling tools capability" in str(exc_info.value) + + def test_no_error_when_capability_present(self) -> None: + """No error when client has sampling.tools capability.""" + caps = ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability())) + tool = Tool(name="test", inputSchema={"type": "object"}) + validate_sampling_tools(caps, [tool], ToolChoice(mode="auto")) # Should not raise + + +class TestValidateToolUseResultMessages: + """Tests for validate_tool_use_result_messages function.""" + + def test_no_error_for_empty_messages(self) -> None: + """No error when messages list is empty.""" + validate_tool_use_result_messages([]) # Should not raise + + def test_no_error_for_simple_text_messages(self) -> None: + """No error for simple text messages.""" + messages = [ + SamplingMessage(role="user", content=TextContent(type="text", text="Hello")), + SamplingMessage(role="assistant", content=TextContent(type="text", text="Hi")), + ] + validate_tool_use_result_messages(messages) # Should not raise + + def test_raises_when_tool_result_mixed_with_other_content(self) -> None: + """Raises when tool_result is mixed with other content types.""" + messages = [ + SamplingMessage( + role="user", + content=[ + ToolResultContent(type="tool_result", toolUseId="123"), + TextContent(type="text", text="also this"), + ], + ), + ] + with pytest.raises(ValueError, match="only tool_result content"): + validate_tool_use_result_messages(messages) + + def test_raises_when_tool_result_without_previous_tool_use(self) -> None: + """Raises when tool_result appears without preceding tool_use.""" + messages = [ + SamplingMessage( + role="user", + content=ToolResultContent(type="tool_result", toolUseId="123"), + ), + ] + with pytest.raises(ValueError, match="previous message containing tool_use"): + validate_tool_use_result_messages(messages) + + def test_raises_when_tool_result_ids_dont_match_tool_use(self) -> None: + """Raises when tool_result IDs don't match tool_use IDs.""" + messages = [ + SamplingMessage( + role="assistant", + content=ToolUseContent(type="tool_use", id="tool-1", name="test", input={}), + ), + SamplingMessage( + role="user", + content=ToolResultContent(type="tool_result", toolUseId="tool-2"), + ), + ] + with pytest.raises(ValueError, match="do not match"): + validate_tool_use_result_messages(messages) + + def test_no_error_when_tool_result_matches_tool_use(self) -> None: + """No error when tool_result IDs match tool_use IDs.""" + messages = [ + SamplingMessage( + role="assistant", + content=ToolUseContent(type="tool_use", id="tool-1", name="test", input={}), + ), + SamplingMessage( + role="user", + content=ToolResultContent(type="tool_result", toolUseId="tool-1"), + ), + ] + validate_tool_use_result_messages(messages) # Should not raise diff --git a/uv.lock b/uv.lock index d1363aef41..2aec51e51c 100644 --- a/uv.lock +++ b/uv.lock @@ -15,6 +15,10 @@ members = [ "mcp-simple-resource", "mcp-simple-streamablehttp", "mcp-simple-streamablehttp-stateless", + "mcp-simple-task", + "mcp-simple-task-client", + "mcp-simple-task-interactive", + "mcp-simple-task-interactive-client", "mcp-simple-tool", "mcp-snippets", "mcp-structured-output-lowlevel", @@ -1196,6 +1200,126 @@ dev = [ { name = "ruff", specifier = ">=0.6.9" }, ] +[[package]] +name = "mcp-simple-task" +version = "0.1.0" +source = { editable = "examples/servers/simple-task" } +dependencies = [ + { name = "anyio" }, + { name = "click" }, + { name = "mcp" }, + { name = "starlette" }, + { name = "uvicorn" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pyright" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "anyio", specifier = ">=4.5" }, + { name = "click", specifier = ">=8.0" }, + { name = "mcp", editable = "." }, + { name = "starlette" }, + { name = "uvicorn" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "pyright", specifier = ">=1.1.378" }, + { name = "ruff", specifier = ">=0.6.9" }, +] + +[[package]] +name = "mcp-simple-task-client" +version = "0.1.0" +source = { editable = "examples/clients/simple-task-client" } +dependencies = [ + { name = "click" }, + { name = "mcp" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pyright" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "click", specifier = ">=8.0" }, + { name = "mcp", editable = "." }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "pyright", specifier = ">=1.1.378" }, + { name = "ruff", specifier = ">=0.6.9" }, +] + +[[package]] +name = "mcp-simple-task-interactive" +version = "0.1.0" +source = { editable = "examples/servers/simple-task-interactive" } +dependencies = [ + { name = "anyio" }, + { name = "click" }, + { name = "mcp" }, + { name = "starlette" }, + { name = "uvicorn" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pyright" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "anyio", specifier = ">=4.5" }, + { name = "click", specifier = ">=8.0" }, + { name = "mcp", editable = "." }, + { name = "starlette" }, + { name = "uvicorn" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "pyright", specifier = ">=1.1.378" }, + { name = "ruff", specifier = ">=0.6.9" }, +] + +[[package]] +name = "mcp-simple-task-interactive-client" +version = "0.1.0" +source = { editable = "examples/clients/simple-task-interactive-client" } +dependencies = [ + { name = "click" }, + { name = "mcp" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pyright" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "click", specifier = ">=8.0" }, + { name = "mcp", editable = "." }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "pyright", specifier = ">=1.1.378" }, + { name = "ruff", specifier = ">=0.6.9" }, +] + [[package]] name = "mcp-simple-tool" version = "0.1.0"