modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

SEP-1686: Tasks (#1645)

Browse source
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
1 changed files with 14230 additions and 120 deletions
Download patch file Download diff file Expand all files Collapse all files

737
tests/experimental/tasks/test_elicitation_scenarios.py Normal file
View file

@ -0,0 +1,737 @@
"""
Tests for the four elicitation scenarios with tasks.
This tests all combinations of tool call types and elicitation types:
1. Normal tool call + Normal elicitation (session.elicit)
2. Normal tool call + Task-augmented elicitation (session.experimental.elicit_as_task)
3. Task-augmented tool call + Normal elicitation (task.elicit)
4. Task-augmented tool call + Task-augmented elicitation (task.elicit_as_task)
And the same for sampling (create_message).
"""
from typing import Any
import anyio
import pytest
from anyio import Event
from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.experimental.task_context import ServerTaskContext
from mcp.server.lowlevel import NotificationOptions
from mcp.shared.context import RequestContext
from mcp.shared.experimental.tasks.helpers import is_terminal
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.message import SessionMessage
from mcp.types import (
TASK_REQUIRED,
CallToolResult,
CreateMessageRequestParams,
CreateMessageResult,
CreateTaskResult,
ElicitRequestParams,
ElicitResult,
ErrorData,
GetTaskPayloadResult,
GetTaskResult,
SamplingMessage,
TaskMetadata,
TextContent,
Tool,
ToolExecution,
)
def create_client_task_handlers(
client_task_store: InMemoryTaskStore,
elicit_received: Event,
) -> ExperimentalTaskHandlers:
"""Create task handlers for client to handle task-augmented elicitation from server."""
elicit_response = ElicitResult(action="accept", content={"confirm": True})
task_complete_events: dict[str, Event] = {}
async def handle_augmented_elicitation(
context: RequestContext[ClientSession, Any],
params: ElicitRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult:
"""Handle task-augmented elicitation by creating a client-side task."""
elicit_received.set()
task = await client_task_store.create_task(task_metadata)
task_complete_events[task.taskId] = Event()
async def complete_task() -> None:
# Store result before updating status to avoid race condition
await client_task_store.store_result(task.taskId, elicit_response)
await client_task_store.update_task(task.taskId, status="completed")
task_complete_events[task.taskId].set()
context.session._task_group.start_soon(complete_task) # pyright: ignore[reportPrivateUsage]
return CreateTaskResult(task=task)
async def handle_get_task(
context: RequestContext[ClientSession, Any],
params: Any,
) -> GetTaskResult:
"""Handle tasks/get from server."""
task = await client_task_store.get_task(params.taskId)
assert task is not None, f"Task not found: {params.taskId}"
return GetTaskResult(
taskId=task.taskId,
status=task.status,
statusMessage=task.statusMessage,
createdAt=task.createdAt,
lastUpdatedAt=task.lastUpdatedAt,
ttl=task.ttl,
pollInterval=100,
)
async def handle_get_task_result(
context: RequestContext[ClientSession, Any],
params: Any,
) -> GetTaskPayloadResult | ErrorData:
"""Handle tasks/result from server."""
event = task_complete_events.get(params.taskId)
assert event is not None, f"No completion event for task: {params.taskId}"
await event.wait()
result = await client_task_store.get_result(params.taskId)
assert result is not None, f"Result not found for task: {params.taskId}"
return GetTaskPayloadResult.model_validate(result.model_dump(by_alias=True))
return ExperimentalTaskHandlers(
augmented_elicitation=handle_augmented_elicitation,
get_task=handle_get_task,
get_task_result=handle_get_task_result,
)
def create_sampling_task_handlers(
client_task_store: InMemoryTaskStore,
sampling_received: Event,
) -> ExperimentalTaskHandlers:
"""Create task handlers for client to handle task-augmented sampling from server."""
sampling_response = CreateMessageResult(
role="assistant",
content=TextContent(type="text", text="Hello from the model!"),
model="test-model",
)
task_complete_events: dict[str, Event] = {}
async def handle_augmented_sampling(
context: RequestContext[ClientSession, Any],
params: CreateMessageRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult:
"""Handle task-augmented sampling by creating a client-side task."""
sampling_received.set()
task = await client_task_store.create_task(task_metadata)
task_complete_events[task.taskId] = Event()
async def complete_task() -> None:
# Store result before updating status to avoid race condition
await client_task_store.store_result(task.taskId, sampling_response)
await client_task_store.update_task(task.taskId, status="completed")
task_complete_events[task.taskId].set()
context.session._task_group.start_soon(complete_task) # pyright: ignore[reportPrivateUsage]
return CreateTaskResult(task=task)
async def handle_get_task(
context: RequestContext[ClientSession, Any],
params: Any,
) -> GetTaskResult:
"""Handle tasks/get from server."""
task = await client_task_store.get_task(params.taskId)
assert task is not None, f"Task not found: {params.taskId}"
return GetTaskResult(
taskId=task.taskId,
status=task.status,
statusMessage=task.statusMessage,
createdAt=task.createdAt,
lastUpdatedAt=task.lastUpdatedAt,
ttl=task.ttl,
pollInterval=100,
)
async def handle_get_task_result(
context: RequestContext[ClientSession, Any],
params: Any,
) -> GetTaskPayloadResult | ErrorData:
"""Handle tasks/result from server."""
event = task_complete_events.get(params.taskId)
assert event is not None, f"No completion event for task: {params.taskId}"
await event.wait()
result = await client_task_store.get_result(params.taskId)
assert result is not None, f"Result not found for task: {params.taskId}"
return GetTaskPayloadResult.model_validate(result.model_dump(by_alias=True))
return ExperimentalTaskHandlers(
augmented_sampling=handle_augmented_sampling,
get_task=handle_get_task,
get_task_result=handle_get_task_result,
)
@pytest.mark.anyio
async def test_scenario1_normal_tool_normal_elicitation() -> None:
"""
Scenario 1: Normal tool call with normal elicitation.
Server calls session.elicit() directly, client responds immediately.
"""
server = Server("test-scenario1")
elicit_received = Event()
tool_result: list[str] = []
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="confirm_action",
description="Confirm an action",
inputSchema={"type": "object"},
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult:
ctx = server.request_context
# Normal elicitation - expects immediate response
result = await ctx.session.elicit(
message="Please confirm the action",
requestedSchema={"type": "object", "properties": {"confirm": {"type": "boolean"}}},
)
confirmed = result.content.get("confirm", False) if result.content else False
tool_result.append("confirmed" if confirmed else "cancelled")
return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")])
# Elicitation callback for client
async def elicitation_callback(
context: RequestContext[ClientSession, Any],
params: ElicitRequestParams,
) -> ElicitResult:
elicit_received.set()
return ElicitResult(action="accept", content={"confirm": True})
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
)
async def run_client() -> None:
async with ClientSession(
server_to_client_receive,
client_to_server_send,
elicitation_callback=elicitation_callback,
) as client_session:
await client_session.initialize()
# Call tool normally (not as task)
result = await client_session.call_tool("confirm_action", {})
# Verify elicitation was received and tool completed
assert elicit_received.is_set()
assert len(result.content) > 0
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "confirmed"
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
assert tool_result[0] == "confirmed"
@pytest.mark.anyio
async def test_scenario2_normal_tool_task_augmented_elicitation() -> None:
"""
Scenario 2: Normal tool call with task-augmented elicitation.
Server calls session.experimental.elicit_as_task(), client creates a task
for the elicitation and returns CreateTaskResult. Server polls client.
"""
server = Server("test-scenario2")
elicit_received = Event()
tool_result: list[str] = []
# Client-side task store for handling task-augmented elicitation
client_task_store = InMemoryTaskStore()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="confirm_action",
description="Confirm an action",
inputSchema={"type": "object"},
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult:
ctx = server.request_context
# Task-augmented elicitation - server polls client
result = await ctx.session.experimental.elicit_as_task(
message="Please confirm the action",
requestedSchema={"type": "object", "properties": {"confirm": {"type": "boolean"}}},
ttl=60000,
)
confirmed = result.content.get("confirm", False) if result.content else False
tool_result.append("confirmed" if confirmed else "cancelled")
return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")])
task_handlers = create_client_task_handlers(client_task_store, elicit_received)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
)
async def run_client() -> None:
async with ClientSession(
server_to_client_receive,
client_to_server_send,
experimental_task_handlers=task_handlers,
) as client_session:
await client_session.initialize()
# Call tool normally (not as task)
result = await client_session.call_tool("confirm_action", {})
# Verify elicitation was received and tool completed
assert elicit_received.is_set()
assert len(result.content) > 0
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "confirmed"
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
assert tool_result[0] == "confirmed"
client_task_store.cleanup()
@pytest.mark.anyio
async def test_scenario3_task_augmented_tool_normal_elicitation() -> None:
"""
Scenario 3: Task-augmented tool call with normal elicitation.
Client calls tool as task. Inside the task, server uses task.elicit()
which queues the request and delivers via tasks/result.
"""
server = Server("test-scenario3")
server.experimental.enable_tasks()
elicit_received = Event()
work_completed = Event()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="confirm_action",
description="Confirm an action",
inputSchema={"type": "object"},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult:
ctx = server.request_context
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
# Normal elicitation within task - queued and delivered via tasks/result
result = await task.elicit(
message="Please confirm the action",
requestedSchema={"type": "object", "properties": {"confirm": {"type": "boolean"}}},
)
confirmed = result.content.get("confirm", False) if result.content else False
work_completed.set()
return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")])
return await ctx.experimental.run_task(work)
# Elicitation callback for client
async def elicitation_callback(
context: RequestContext[ClientSession, Any],
params: ElicitRequestParams,
) -> ElicitResult:
elicit_received.set()
return ElicitResult(action="accept", content={"confirm": True})
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
)
async def run_client() -> None:
async with ClientSession(
server_to_client_receive,
client_to_server_send,
elicitation_callback=elicitation_callback,
) as client_session:
await client_session.initialize()
# Call tool as task
create_result = await client_session.experimental.call_tool_as_task("confirm_action", {})
task_id = create_result.task.taskId
assert create_result.task.status == "working"
# Poll until input_required, then call tasks/result
found_input_required = False
async for status in client_session.experimental.poll_task(task_id): # pragma: no branch
if status.status == "input_required": # pragma: no branch
found_input_required = True
break
assert found_input_required, "Expected to see input_required status"
# This will deliver the elicitation and get the response
final_result = await client_session.experimental.get_task_result(task_id, CallToolResult)
# Verify
assert elicit_received.is_set()
assert len(final_result.content) > 0
assert isinstance(final_result.content[0], TextContent)
assert final_result.content[0].text == "confirmed"
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
assert work_completed.is_set()
@pytest.mark.anyio
async def test_scenario4_task_augmented_tool_task_augmented_elicitation() -> None:
"""
Scenario 4: Task-augmented tool call with task-augmented elicitation.
Client calls tool as task. Inside the task, server uses task.elicit_as_task()
which sends task-augmented elicitation. Client creates its own task for the
elicitation, and server polls the client.
This tests the full bidirectional flow where:
1. Client calls tasks/result on server (for tool task)
2. Server delivers task-augmented elicitation through that stream
3. Client creates its own task and returns CreateTaskResult
4. Server polls the client's task while the client's tasks/result is still open
5. Server gets the ElicitResult and completes the tool task
6. Client's tasks/result returns with the CallToolResult
"""
server = Server("test-scenario4")
server.experimental.enable_tasks()
elicit_received = Event()
work_completed = Event()
# Client-side task store for handling task-augmented elicitation
client_task_store = InMemoryTaskStore()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="confirm_action",
description="Confirm an action",
inputSchema={"type": "object"},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult:
ctx = server.request_context
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
# Task-augmented elicitation within task - server polls client
result = await task.elicit_as_task(
message="Please confirm the action",
requestedSchema={"type": "object", "properties": {"confirm": {"type": "boolean"}}},
ttl=60000,
)
confirmed = result.content.get("confirm", False) if result.content else False
work_completed.set()
return CallToolResult(content=[TextContent(type="text", text="confirmed" if confirmed else "cancelled")])
return await ctx.experimental.run_task(work)
task_handlers = create_client_task_handlers(client_task_store, elicit_received)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
)
async def run_client() -> None:
async with ClientSession(
server_to_client_receive,
client_to_server_send,
experimental_task_handlers=task_handlers,
) as client_session:
await client_session.initialize()
# Call tool as task
create_result = await client_session.experimental.call_tool_as_task("confirm_action", {})
task_id = create_result.task.taskId
assert create_result.task.status == "working"
# Poll until input_required or terminal, then call tasks/result
found_expected_status = False
async for status in client_session.experimental.poll_task(task_id): # pragma: no branch
if status.status == "input_required" or is_terminal(status.status): # pragma: no branch
found_expected_status = True
break
assert found_expected_status, "Expected to see input_required or terminal status"
# This will deliver the task-augmented elicitation,
# server will poll client, and eventually return the tool result
final_result = await client_session.experimental.get_task_result(task_id, CallToolResult)
# Verify
assert elicit_received.is_set()
assert len(final_result.content) > 0
assert isinstance(final_result.content[0], TextContent)
assert final_result.content[0].text == "confirmed"
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
assert work_completed.is_set()
client_task_store.cleanup()
@pytest.mark.anyio
async def test_scenario2_sampling_normal_tool_task_augmented_sampling() -> None:
"""
Scenario 2 for sampling: Normal tool call with task-augmented sampling.
Server calls session.experimental.create_message_as_task(), client creates
a task for the sampling and returns CreateTaskResult. Server polls client.
"""
server = Server("test-scenario2-sampling")
sampling_received = Event()
tool_result: list[str] = []
# Client-side task store for handling task-augmented sampling
client_task_store = InMemoryTaskStore()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="generate_text",
description="Generate text using sampling",
inputSchema={"type": "object"},
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult:
ctx = server.request_context
# Task-augmented sampling - server polls client
result = await ctx.session.experimental.create_message_as_task(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
max_tokens=100,
ttl=60000,
)
assert isinstance(result.content, TextContent), "Expected TextContent response"
response_text = result.content.text
tool_result.append(response_text)
return CallToolResult(content=[TextContent(type="text", text=response_text)])
task_handlers = create_sampling_task_handlers(client_task_store, sampling_received)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
)
async def run_client() -> None:
async with ClientSession(
server_to_client_receive,
client_to_server_send,
experimental_task_handlers=task_handlers,
) as client_session:
await client_session.initialize()
# Call tool normally (not as task)
result = await client_session.call_tool("generate_text", {})
# Verify sampling was received and tool completed
assert sampling_received.is_set()
assert len(result.content) > 0
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "Hello from the model!"
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
assert tool_result[0] == "Hello from the model!"
client_task_store.cleanup()
@pytest.mark.anyio
async def test_scenario4_sampling_task_augmented_tool_task_augmented_sampling() -> None:
"""
Scenario 4 for sampling: Task-augmented tool call with task-augmented sampling.
Client calls tool as task. Inside the task, server uses task.create_message_as_task()
which sends task-augmented sampling. Client creates its own task for the sampling,
and server polls the client.
"""
server = Server("test-scenario4-sampling")
server.experimental.enable_tasks()
sampling_received = Event()
work_completed = Event()
# Client-side task store for handling task-augmented sampling
client_task_store = InMemoryTaskStore()
@server.list_tools()
async def list_tools() -> list[Tool]:
return [
Tool(
name="generate_text",
description="Generate text using sampling",
inputSchema={"type": "object"},
execution=ToolExecution(taskSupport=TASK_REQUIRED),
)
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CreateTaskResult:
ctx = server.request_context
ctx.experimental.validate_task_mode(TASK_REQUIRED)
async def work(task: ServerTaskContext) -> CallToolResult:
# Task-augmented sampling within task - server polls client
result = await task.create_message_as_task(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
max_tokens=100,
ttl=60000,
)
assert isinstance(result.content, TextContent), "Expected TextContent response"
response_text = result.content.text
work_completed.set()
return CallToolResult(content=[TextContent(type="text", text=response_text)])
return await ctx.experimental.run_task(work)
task_handlers = create_sampling_task_handlers(client_task_store, sampling_received)
# Set up streams
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
async def run_server() -> None:
await server.run(
client_to_server_receive,
server_to_client_send,
server.create_initialization_options(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
)
async def run_client() -> None:
async with ClientSession(
server_to_client_receive,
client_to_server_send,
experimental_task_handlers=task_handlers,
) as client_session:
await client_session.initialize()
# Call tool as task
create_result = await client_session.experimental.call_tool_as_task("generate_text", {})
task_id = create_result.task.taskId
assert create_result.task.status == "working"
# Poll until input_required or terminal
found_expected_status = False
async for status in client_session.experimental.poll_task(task_id): # pragma: no branch
if status.status == "input_required" or is_terminal(status.status): # pragma: no branch
found_expected_status = True
break
assert found_expected_status, "Expected to see input_required or terminal status"
final_result = await client_session.experimental.get_task_result(task_id, CallToolResult)
# Verify
assert sampling_received.is_set()
assert len(final_result.content) > 0
assert isinstance(final_result.content[0], TextContent)
assert final_result.content[0].text == "Hello from the model!"
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
tg.start_soon(run_client)
assert work_completed.is_set()
client_task_store.cleanup()