SEP-1686: Tasks (#1645)

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

View file

@ -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"]

View file

@ -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

View file

@ -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

View file

@ -8,6 +8,8 @@ from pydantic import AnyUrl, TypeAdapter
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 @@ class ClientSession(
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 @@ class ClientSession(
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 @@ class ClientSession(
elicitation=elicitation,
experimental=None,
roots=roots,
tasks=self._task_handlers.build_capability(),
),
clientInfo=self._client_info,
),
@ -188,6 +197,20 @@ class ClientSession(
"""
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 @@ class ClientSession(
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 @@ class ClientSession(
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,

View file

@ -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
"""

View file

@ -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 {})

View file

@ -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)

View file

@ -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

View file

@ -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

View file

@ -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(),
)

View file

@ -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

View file

@ -67,6 +67,7 @@ messages from the client.
from __future__ import annotations as _annotations
import base64
import contextvars
import json
import logging
@ -82,6 +83,8 @@ from pydantic import AnyUrl
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 @@ class Server(Generic[LifespanResultT, RequestT]):
}
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 @@ class Server(Generic[LifespanResultT, RequestT]):
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 @@ class Server(Generic[LifespanResultT, RequestT]):
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 @@ class Server(Generic[LifespanResultT, RequestT]):
"""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 @@ class Server(Generic[LifespanResultT, RequestT]):
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 @@ class Server(Generic[LifespanResultT, RequestT]):
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 @@ class Server(Generic[LifespanResultT, RequestT]):
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 @@ class Server(Generic[LifespanResultT, RequestT]):
)
)
# 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 @@ class Server(Generic[LifespanResultT, RequestT]):
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 @@ class Server(Generic[LifespanResultT, RequestT]):
# 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,
)
)

View file

@ -46,8 +46,11 @@ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStre
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 @@ class ServerSession(
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 @@ class ServerSession(
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 @@ class ServerSession(
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 @@ class ServerSession(
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)

View file

@ -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")

View file

@ -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

View file

@ -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.
"""

View file

@ -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
"""

View file

@ -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",
)
)

View file

@ -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,
)

View file

@ -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

Some files were not shown because too many files have changed in this diff Show more