SEP-1686: Tasks (#1645)
This commit is contained in:
parent
5983a650cc
commit
c92bb2f7ff
19 changed files with 14230 additions and 120 deletions
1
tests/experimental/tasks/__init__.py
Normal file
1
tests/experimental/tasks/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""Tests for MCP task support."""
|
||||
0
tests/experimental/tasks/client/__init__.py
Normal file
0
tests/experimental/tasks/client/__init__.py
Normal file
331
tests/experimental/tasks/client/test_capabilities.py
Normal file
331
tests/experimental/tasks/client/test_capabilities.py
Normal file
|
|
@ -0,0 +1,331 @@
|
|||
"""Tests for client task capabilities declaration during initialization."""
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
import mcp.types as types
|
||||
from mcp import ClientCapabilities
|
||||
from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.shared.context import RequestContext
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import (
|
||||
LATEST_PROTOCOL_VERSION,
|
||||
ClientRequest,
|
||||
Implementation,
|
||||
InitializeRequest,
|
||||
InitializeResult,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCResponse,
|
||||
ServerCapabilities,
|
||||
ServerResult,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_capabilities_without_tasks():
|
||||
"""Test that tasks capability is None when not provided."""
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
|
||||
received_capabilities = None
|
||||
|
||||
async def mock_server():
|
||||
nonlocal received_capabilities
|
||||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocolVersion=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
serverInfo=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await client_to_server_receive.receive()
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server_to_client_receive,
|
||||
):
|
||||
tg.start_soon(mock_server)
|
||||
await session.initialize()
|
||||
|
||||
# Assert that tasks capability is None when not provided
|
||||
assert received_capabilities is not None
|
||||
assert received_capabilities.tasks is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_capabilities_with_tasks():
|
||||
"""Test that tasks capability is properly set when handlers are provided."""
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
|
||||
received_capabilities: ClientCapabilities | None = None
|
||||
|
||||
# Define custom handlers to trigger capability building (never actually called)
|
||||
async def my_list_tasks_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: types.PaginatedRequestParams | None,
|
||||
) -> types.ListTasksResult | types.ErrorData:
|
||||
raise NotImplementedError
|
||||
|
||||
async def my_cancel_task_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: types.CancelTaskRequestParams,
|
||||
) -> types.CancelTaskResult | types.ErrorData:
|
||||
raise NotImplementedError
|
||||
|
||||
async def mock_server():
|
||||
nonlocal received_capabilities
|
||||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocolVersion=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
serverInfo=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await client_to_server_receive.receive()
|
||||
|
||||
# Create handlers container
|
||||
task_handlers = ExperimentalTaskHandlers(
|
||||
list_tasks=my_list_tasks_handler,
|
||||
cancel_task=my_cancel_task_handler,
|
||||
)
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
experimental_task_handlers=task_handlers,
|
||||
) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server_to_client_receive,
|
||||
):
|
||||
tg.start_soon(mock_server)
|
||||
await session.initialize()
|
||||
|
||||
# Assert that tasks capability is properly set from handlers
|
||||
assert received_capabilities is not None
|
||||
assert received_capabilities.tasks is not None
|
||||
assert isinstance(received_capabilities.tasks, types.ClientTasksCapability)
|
||||
assert received_capabilities.tasks.list is not None
|
||||
assert received_capabilities.tasks.cancel is not None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_capabilities_auto_built_from_handlers():
|
||||
"""Test that tasks capability is automatically built from provided handlers."""
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
|
||||
received_capabilities: ClientCapabilities | None = None
|
||||
|
||||
# Define custom handlers (not defaults)
|
||||
async def my_list_tasks_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: types.PaginatedRequestParams | None,
|
||||
) -> types.ListTasksResult | types.ErrorData:
|
||||
raise NotImplementedError
|
||||
|
||||
async def my_cancel_task_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: types.CancelTaskRequestParams,
|
||||
) -> types.CancelTaskResult | types.ErrorData:
|
||||
raise NotImplementedError
|
||||
|
||||
async def mock_server():
|
||||
nonlocal received_capabilities
|
||||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocolVersion=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
serverInfo=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await client_to_server_receive.receive()
|
||||
|
||||
# Provide handlers via ExperimentalTaskHandlers
|
||||
task_handlers = ExperimentalTaskHandlers(
|
||||
list_tasks=my_list_tasks_handler,
|
||||
cancel_task=my_cancel_task_handler,
|
||||
)
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
experimental_task_handlers=task_handlers,
|
||||
) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server_to_client_receive,
|
||||
):
|
||||
tg.start_soon(mock_server)
|
||||
await session.initialize()
|
||||
|
||||
# Assert that tasks capability was auto-built from handlers
|
||||
assert received_capabilities is not None
|
||||
assert received_capabilities.tasks is not None
|
||||
assert received_capabilities.tasks.list is not None
|
||||
assert received_capabilities.tasks.cancel is not None
|
||||
# requests should be None since we didn't provide task-augmented handlers
|
||||
assert received_capabilities.tasks.requests is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_capabilities_with_task_augmented_handlers():
|
||||
"""Test that requests capability is built when augmented handlers are provided."""
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
|
||||
received_capabilities: ClientCapabilities | None = None
|
||||
|
||||
# Define task-augmented handler
|
||||
async def my_augmented_sampling_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: types.CreateMessageRequestParams,
|
||||
task_metadata: types.TaskMetadata,
|
||||
) -> types.CreateTaskResult | types.ErrorData:
|
||||
raise NotImplementedError
|
||||
|
||||
async def mock_server():
|
||||
nonlocal received_capabilities
|
||||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocolVersion=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
serverInfo=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await client_to_server_receive.receive()
|
||||
|
||||
# Provide task-augmented sampling handler
|
||||
task_handlers = ExperimentalTaskHandlers(
|
||||
augmented_sampling=my_augmented_sampling_handler,
|
||||
)
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
experimental_task_handlers=task_handlers,
|
||||
) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server_to_client_receive,
|
||||
):
|
||||
tg.start_soon(mock_server)
|
||||
await session.initialize()
|
||||
|
||||
# Assert that tasks capability includes requests.sampling
|
||||
assert received_capabilities is not None
|
||||
assert received_capabilities.tasks is not None
|
||||
assert received_capabilities.tasks.requests is not None
|
||||
assert received_capabilities.tasks.requests.sampling is not None
|
||||
assert received_capabilities.tasks.requests.elicitation is None # Not provided
|
||||
878
tests/experimental/tasks/client/test_handlers.py
Normal file
878
tests/experimental/tasks/client/test_handlers.py
Normal file
|
|
@ -0,0 +1,878 @@
|
|||
"""Tests for client-side task management handlers (server -> client requests).
|
||||
|
||||
These tests verify that clients can handle task-related requests from servers:
|
||||
- GetTaskRequest - server polling client's task status
|
||||
- GetTaskPayloadRequest - server getting result from client's task
|
||||
- ListTasksRequest - server listing client's tasks
|
||||
- CancelTaskRequest - server cancelling client's task
|
||||
|
||||
This is the inverse of the existing tests in test_tasks.py, which test
|
||||
client -> server task requests.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from anyio import Event
|
||||
from anyio.abc import TaskGroup
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||
|
||||
import mcp.types as types
|
||||
from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.shared.context import RequestContext
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.types import (
|
||||
CancelTaskRequest,
|
||||
CancelTaskRequestParams,
|
||||
CancelTaskResult,
|
||||
ClientResult,
|
||||
CreateMessageRequest,
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
CreateTaskResult,
|
||||
ElicitRequest,
|
||||
ElicitRequestFormParams,
|
||||
ElicitRequestParams,
|
||||
ElicitResult,
|
||||
ErrorData,
|
||||
GetTaskPayloadRequest,
|
||||
GetTaskPayloadRequestParams,
|
||||
GetTaskPayloadResult,
|
||||
GetTaskRequest,
|
||||
GetTaskRequestParams,
|
||||
GetTaskResult,
|
||||
ListTasksRequest,
|
||||
ListTasksResult,
|
||||
SamplingMessage,
|
||||
ServerNotification,
|
||||
ServerRequest,
|
||||
TaskMetadata,
|
||||
TextContent,
|
||||
)
|
||||
|
||||
# Buffer size for test streams
|
||||
STREAM_BUFFER_SIZE = 10
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClientTestStreams:
|
||||
"""Bidirectional message streams for client/server communication in tests."""
|
||||
|
||||
server_send: MemoryObjectSendStream[SessionMessage]
|
||||
server_receive: MemoryObjectReceiveStream[SessionMessage]
|
||||
client_send: MemoryObjectSendStream[SessionMessage]
|
||||
client_receive: MemoryObjectReceiveStream[SessionMessage]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def client_streams() -> AsyncIterator[ClientTestStreams]:
|
||||
"""Create bidirectional message streams for client tests.
|
||||
|
||||
Automatically closes all streams after the test completes.
|
||||
"""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](
|
||||
STREAM_BUFFER_SIZE
|
||||
)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](
|
||||
STREAM_BUFFER_SIZE
|
||||
)
|
||||
|
||||
streams = ClientTestStreams(
|
||||
server_send=server_to_client_send,
|
||||
server_receive=client_to_server_receive,
|
||||
client_send=client_to_server_send,
|
||||
client_receive=server_to_client_receive,
|
||||
)
|
||||
|
||||
yield streams
|
||||
|
||||
# Cleanup
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
|
||||
|
||||
async def _default_message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None:
|
||||
"""Default message handler that ignores messages (tests handle them explicitly)."""
|
||||
...
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_handles_get_task_request(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client can respond to GetTaskRequest from server."""
|
||||
with anyio.fail_after(10):
|
||||
store = InMemoryTaskStore()
|
||||
received_task_id: str | None = None
|
||||
|
||||
async def get_task_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: GetTaskRequestParams,
|
||||
) -> GetTaskResult | ErrorData:
|
||||
nonlocal received_task_id
|
||||
received_task_id = params.taskId
|
||||
task = await store.get_task(params.taskId)
|
||||
assert task is not None, f"Test setup error: task {params.taskId} should exist"
|
||||
return GetTaskResult(
|
||||
taskId=task.taskId,
|
||||
status=task.status,
|
||||
statusMessage=task.statusMessage,
|
||||
createdAt=task.createdAt,
|
||||
lastUpdatedAt=task.lastUpdatedAt,
|
||||
ttl=task.ttl,
|
||||
pollInterval=task.pollInterval,
|
||||
)
|
||||
|
||||
await store.create_task(TaskMetadata(ttl=60000), task_id="test-task-123")
|
||||
|
||||
task_handlers = ExperimentalTaskHandlers(get_task=get_task_handler)
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
experimental_task_handlers=task_handlers,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
typed_request = GetTaskRequest(params=GetTaskRequestParams(taskId="test-task-123"))
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-1",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCResponse)
|
||||
assert response.id == "req-1"
|
||||
|
||||
result = GetTaskResult.model_validate(response.result)
|
||||
assert result.taskId == "test-task-123"
|
||||
assert result.status == "working"
|
||||
assert received_task_id == "test-task-123"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_handles_get_task_result_request(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client can respond to GetTaskPayloadRequest from server."""
|
||||
with anyio.fail_after(10):
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
async def get_task_result_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: GetTaskPayloadRequestParams,
|
||||
) -> GetTaskPayloadResult | ErrorData:
|
||||
result = await store.get_result(params.taskId)
|
||||
assert result is not None, f"Test setup error: result for {params.taskId} should exist"
|
||||
assert isinstance(result, types.CallToolResult)
|
||||
return GetTaskPayloadResult(**result.model_dump())
|
||||
|
||||
await store.create_task(TaskMetadata(ttl=60000), task_id="test-task-456")
|
||||
await store.store_result(
|
||||
"test-task-456",
|
||||
types.CallToolResult(content=[TextContent(type="text", text="Task completed successfully!")]),
|
||||
)
|
||||
await store.update_task("test-task-456", status="completed")
|
||||
|
||||
task_handlers = ExperimentalTaskHandlers(get_task_result=get_task_result_handler)
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
experimental_task_handlers=task_handlers,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
typed_request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId="test-task-456"))
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-2",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCResponse)
|
||||
|
||||
assert isinstance(response.result, dict)
|
||||
result_dict = response.result
|
||||
assert "content" in result_dict
|
||||
assert len(result_dict["content"]) == 1
|
||||
assert result_dict["content"][0]["text"] == "Task completed successfully!"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_handles_list_tasks_request(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client can respond to ListTasksRequest from server."""
|
||||
with anyio.fail_after(10):
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
async def list_tasks_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: types.PaginatedRequestParams | None,
|
||||
) -> ListTasksResult | ErrorData:
|
||||
cursor = params.cursor if params else None
|
||||
tasks_list, next_cursor = await store.list_tasks(cursor=cursor)
|
||||
return ListTasksResult(tasks=tasks_list, nextCursor=next_cursor)
|
||||
|
||||
await store.create_task(TaskMetadata(ttl=60000), task_id="task-1")
|
||||
await store.create_task(TaskMetadata(ttl=60000), task_id="task-2")
|
||||
|
||||
task_handlers = ExperimentalTaskHandlers(list_tasks=list_tasks_handler)
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
experimental_task_handlers=task_handlers,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
typed_request = ListTasksRequest()
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-3",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCResponse)
|
||||
|
||||
result = ListTasksResult.model_validate(response.result)
|
||||
assert len(result.tasks) == 2
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_handles_cancel_task_request(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client can respond to CancelTaskRequest from server."""
|
||||
with anyio.fail_after(10):
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
async def cancel_task_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: CancelTaskRequestParams,
|
||||
) -> CancelTaskResult | ErrorData:
|
||||
task = await store.get_task(params.taskId)
|
||||
assert task is not None, f"Test setup error: task {params.taskId} should exist"
|
||||
await store.update_task(params.taskId, status="cancelled")
|
||||
updated = await store.get_task(params.taskId)
|
||||
assert updated is not None
|
||||
return CancelTaskResult(
|
||||
taskId=updated.taskId,
|
||||
status=updated.status,
|
||||
createdAt=updated.createdAt,
|
||||
lastUpdatedAt=updated.lastUpdatedAt,
|
||||
ttl=updated.ttl,
|
||||
)
|
||||
|
||||
await store.create_task(TaskMetadata(ttl=60000), task_id="task-to-cancel")
|
||||
|
||||
task_handlers = ExperimentalTaskHandlers(cancel_task=cancel_task_handler)
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
experimental_task_handlers=task_handlers,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
typed_request = CancelTaskRequest(params=CancelTaskRequestParams(taskId="task-to-cancel"))
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-4",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCResponse)
|
||||
|
||||
result = CancelTaskResult.model_validate(response.result)
|
||||
assert result.taskId == "task-to-cancel"
|
||||
assert result.status == "cancelled"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_task_augmented_sampling(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client can handle task-augmented sampling request from server."""
|
||||
with anyio.fail_after(10):
|
||||
store = InMemoryTaskStore()
|
||||
sampling_completed = Event()
|
||||
created_task_id: list[str | None] = [None]
|
||||
background_tg: list[TaskGroup | None] = [None]
|
||||
|
||||
async def task_augmented_sampling_callback(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: CreateMessageRequestParams,
|
||||
task_metadata: TaskMetadata,
|
||||
) -> CreateTaskResult:
|
||||
task = await store.create_task(task_metadata)
|
||||
created_task_id[0] = task.taskId
|
||||
|
||||
async def do_sampling() -> None:
|
||||
result = CreateMessageResult(
|
||||
role="assistant",
|
||||
content=TextContent(type="text", text="Sampled response"),
|
||||
model="test-model",
|
||||
stopReason="endTurn",
|
||||
)
|
||||
await store.store_result(task.taskId, result)
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
sampling_completed.set()
|
||||
|
||||
assert background_tg[0] is not None
|
||||
background_tg[0].start_soon(do_sampling)
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
async def get_task_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: GetTaskRequestParams,
|
||||
) -> GetTaskResult | ErrorData:
|
||||
task = await store.get_task(params.taskId)
|
||||
assert task is not None, f"Test setup error: task {params.taskId} should exist"
|
||||
return GetTaskResult(
|
||||
taskId=task.taskId,
|
||||
status=task.status,
|
||||
statusMessage=task.statusMessage,
|
||||
createdAt=task.createdAt,
|
||||
lastUpdatedAt=task.lastUpdatedAt,
|
||||
ttl=task.ttl,
|
||||
pollInterval=task.pollInterval,
|
||||
)
|
||||
|
||||
async def get_task_result_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: GetTaskPayloadRequestParams,
|
||||
) -> GetTaskPayloadResult | ErrorData:
|
||||
result = await store.get_result(params.taskId)
|
||||
assert result is not None, f"Test setup error: result for {params.taskId} should exist"
|
||||
assert isinstance(result, CreateMessageResult)
|
||||
return GetTaskPayloadResult(**result.model_dump())
|
||||
|
||||
task_handlers = ExperimentalTaskHandlers(
|
||||
augmented_sampling=task_augmented_sampling_callback,
|
||||
get_task=get_task_handler,
|
||||
get_task_result=get_task_result_handler,
|
||||
)
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
background_tg[0] = tg
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
experimental_task_handlers=task_handlers,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
# Step 1: Server sends task-augmented CreateMessageRequest
|
||||
typed_request = CreateMessageRequest(
|
||||
params=CreateMessageRequestParams(
|
||||
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
|
||||
maxTokens=100,
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
)
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-sampling",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
# Step 2: Client responds with CreateTaskResult
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCResponse)
|
||||
|
||||
task_result = CreateTaskResult.model_validate(response.result)
|
||||
task_id = task_result.task.taskId
|
||||
assert task_id == created_task_id[0]
|
||||
|
||||
# Step 3: Wait for background sampling
|
||||
await sampling_completed.wait()
|
||||
|
||||
# Step 4: Server polls task status
|
||||
typed_poll = GetTaskRequest(params=GetTaskRequestParams(taskId=task_id))
|
||||
poll_request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-poll",
|
||||
**typed_poll.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(poll_request)))
|
||||
|
||||
poll_response_msg = await client_streams.server_receive.receive()
|
||||
poll_response = poll_response_msg.message.root
|
||||
assert isinstance(poll_response, types.JSONRPCResponse)
|
||||
|
||||
status = GetTaskResult.model_validate(poll_response.result)
|
||||
assert status.status == "completed"
|
||||
|
||||
# Step 5: Server gets result
|
||||
typed_result_req = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task_id))
|
||||
result_request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-result",
|
||||
**typed_result_req.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(result_request)))
|
||||
|
||||
result_response_msg = await client_streams.server_receive.receive()
|
||||
result_response = result_response_msg.message.root
|
||||
assert isinstance(result_response, types.JSONRPCResponse)
|
||||
|
||||
assert isinstance(result_response.result, dict)
|
||||
assert result_response.result["role"] == "assistant"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_task_augmented_elicitation(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client can handle task-augmented elicitation request from server."""
|
||||
with anyio.fail_after(10):
|
||||
store = InMemoryTaskStore()
|
||||
elicitation_completed = Event()
|
||||
created_task_id: list[str | None] = [None]
|
||||
background_tg: list[TaskGroup | None] = [None]
|
||||
|
||||
async def task_augmented_elicitation_callback(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: ElicitRequestParams,
|
||||
task_metadata: TaskMetadata,
|
||||
) -> CreateTaskResult | ErrorData:
|
||||
task = await store.create_task(task_metadata)
|
||||
created_task_id[0] = task.taskId
|
||||
|
||||
async def do_elicitation() -> None:
|
||||
# Simulate user providing elicitation response
|
||||
result = ElicitResult(action="accept", content={"name": "Test User"})
|
||||
await store.store_result(task.taskId, result)
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
elicitation_completed.set()
|
||||
|
||||
assert background_tg[0] is not None
|
||||
background_tg[0].start_soon(do_elicitation)
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
async def get_task_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: GetTaskRequestParams,
|
||||
) -> GetTaskResult | ErrorData:
|
||||
task = await store.get_task(params.taskId)
|
||||
assert task is not None, f"Test setup error: task {params.taskId} should exist"
|
||||
return GetTaskResult(
|
||||
taskId=task.taskId,
|
||||
status=task.status,
|
||||
statusMessage=task.statusMessage,
|
||||
createdAt=task.createdAt,
|
||||
lastUpdatedAt=task.lastUpdatedAt,
|
||||
ttl=task.ttl,
|
||||
pollInterval=task.pollInterval,
|
||||
)
|
||||
|
||||
async def get_task_result_handler(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: GetTaskPayloadRequestParams,
|
||||
) -> GetTaskPayloadResult | ErrorData:
|
||||
result = await store.get_result(params.taskId)
|
||||
assert result is not None, f"Test setup error: result for {params.taskId} should exist"
|
||||
assert isinstance(result, ElicitResult)
|
||||
return GetTaskPayloadResult(**result.model_dump())
|
||||
|
||||
task_handlers = ExperimentalTaskHandlers(
|
||||
augmented_elicitation=task_augmented_elicitation_callback,
|
||||
get_task=get_task_handler,
|
||||
get_task_result=get_task_result_handler,
|
||||
)
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
background_tg[0] = tg
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
experimental_task_handlers=task_handlers,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
# Step 1: Server sends task-augmented ElicitRequest
|
||||
typed_request = ElicitRequest(
|
||||
params=ElicitRequestFormParams(
|
||||
message="What is your name?",
|
||||
requestedSchema={"type": "object", "properties": {"name": {"type": "string"}}},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
)
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-elicit",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
# Step 2: Client responds with CreateTaskResult
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCResponse)
|
||||
|
||||
task_result = CreateTaskResult.model_validate(response.result)
|
||||
task_id = task_result.task.taskId
|
||||
assert task_id == created_task_id[0]
|
||||
|
||||
# Step 3: Wait for background elicitation
|
||||
await elicitation_completed.wait()
|
||||
|
||||
# Step 4: Server polls task status
|
||||
typed_poll = GetTaskRequest(params=GetTaskRequestParams(taskId=task_id))
|
||||
poll_request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-poll",
|
||||
**typed_poll.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(poll_request)))
|
||||
|
||||
poll_response_msg = await client_streams.server_receive.receive()
|
||||
poll_response = poll_response_msg.message.root
|
||||
assert isinstance(poll_response, types.JSONRPCResponse)
|
||||
|
||||
status = GetTaskResult.model_validate(poll_response.result)
|
||||
assert status.status == "completed"
|
||||
|
||||
# Step 5: Server gets result
|
||||
typed_result_req = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task_id))
|
||||
result_request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-result",
|
||||
**typed_result_req.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(result_request)))
|
||||
|
||||
result_response_msg = await client_streams.server_receive.receive()
|
||||
result_response = result_response_msg.message.root
|
||||
assert isinstance(result_response, types.JSONRPCResponse)
|
||||
|
||||
# Verify the elicitation result
|
||||
assert isinstance(result_response.result, dict)
|
||||
assert result_response.result["action"] == "accept"
|
||||
assert result_response.result["content"] == {"name": "Test User"}
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_returns_error_for_unhandled_task_request(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client returns error when no handler is registered for task request."""
|
||||
with anyio.fail_after(10):
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
typed_request = GetTaskRequest(params=GetTaskRequestParams(taskId="nonexistent"))
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-unhandled",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCError)
|
||||
assert (
|
||||
"not supported" in response.error.message.lower()
|
||||
or "method not found" in response.error.message.lower()
|
||||
)
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_returns_error_for_unhandled_task_result_request(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client returns error for unhandled tasks/result request."""
|
||||
with anyio.fail_after(10):
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
typed_request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId="nonexistent"))
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-result",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCError)
|
||||
assert "not supported" in response.error.message.lower()
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_returns_error_for_unhandled_list_tasks_request(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client returns error for unhandled tasks/list request."""
|
||||
with anyio.fail_after(10):
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
typed_request = ListTasksRequest()
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-list",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCError)
|
||||
assert "not supported" in response.error.message.lower()
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_returns_error_for_unhandled_cancel_task_request(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client returns error for unhandled tasks/cancel request."""
|
||||
with anyio.fail_after(10):
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
typed_request = CancelTaskRequest(params=CancelTaskRequestParams(taskId="nonexistent"))
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-cancel",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCError)
|
||||
assert "not supported" in response.error.message.lower()
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_returns_error_for_unhandled_task_augmented_sampling(client_streams: ClientTestStreams) -> None:
|
||||
"""Test that client returns error for task-augmented sampling without handler."""
|
||||
with anyio.fail_after(10):
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
# No task handlers provided - uses defaults
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
# Send task-augmented sampling request
|
||||
typed_request = CreateMessageRequest(
|
||||
params=CreateMessageRequestParams(
|
||||
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
|
||||
maxTokens=100,
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
)
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-sampling",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCError)
|
||||
assert "not supported" in response.error.message.lower()
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_returns_error_for_unhandled_task_augmented_elicitation(
|
||||
client_streams: ClientTestStreams,
|
||||
) -> None:
|
||||
"""Test that client returns error for task-augmented elicitation without handler."""
|
||||
with anyio.fail_after(10):
|
||||
client_ready = anyio.Event()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def run_client() -> None:
|
||||
# No task handlers provided - uses defaults
|
||||
async with ClientSession(
|
||||
client_streams.client_receive,
|
||||
client_streams.client_send,
|
||||
message_handler=_default_message_handler,
|
||||
):
|
||||
client_ready.set()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
tg.start_soon(run_client)
|
||||
await client_ready.wait()
|
||||
|
||||
# Send task-augmented elicitation request
|
||||
typed_request = ElicitRequest(
|
||||
params=ElicitRequestFormParams(
|
||||
message="What is your name?",
|
||||
requestedSchema={"type": "object", "properties": {"name": {"type": "string"}}},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
)
|
||||
request = types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="req-elicit",
|
||||
**typed_request.model_dump(by_alias=True),
|
||||
)
|
||||
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
|
||||
|
||||
response_msg = await client_streams.server_receive.receive()
|
||||
response = response_msg.message.root
|
||||
assert isinstance(response, types.JSONRPCError)
|
||||
assert "not supported" in response.error.message.lower()
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
121
tests/experimental/tasks/client/test_poll_task.py
Normal file
121
tests/experimental/tasks/client/test_poll_task.py
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
"""Tests for poll_task async iterator."""
|
||||
|
||||
from collections.abc import Callable, Coroutine
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp.client.experimental.tasks import ExperimentalClientFeatures
|
||||
from mcp.types import GetTaskResult, TaskStatus
|
||||
|
||||
|
||||
def make_task_result(
|
||||
status: TaskStatus = "working",
|
||||
poll_interval: int = 0,
|
||||
task_id: str = "test-task",
|
||||
status_message: str | None = None,
|
||||
) -> GetTaskResult:
|
||||
"""Create GetTaskResult with sensible defaults."""
|
||||
now = datetime.now(timezone.utc)
|
||||
return GetTaskResult(
|
||||
taskId=task_id,
|
||||
status=status,
|
||||
statusMessage=status_message,
|
||||
createdAt=now,
|
||||
lastUpdatedAt=now,
|
||||
ttl=60000,
|
||||
pollInterval=poll_interval,
|
||||
)
|
||||
|
||||
|
||||
def make_status_sequence(
|
||||
*statuses: TaskStatus,
|
||||
task_id: str = "test-task",
|
||||
) -> Callable[[str], Coroutine[Any, Any, GetTaskResult]]:
|
||||
"""Create mock get_task that returns statuses in sequence."""
|
||||
status_iter = iter(statuses)
|
||||
|
||||
async def mock_get_task(tid: str) -> GetTaskResult:
|
||||
return make_task_result(status=next(status_iter), task_id=tid)
|
||||
|
||||
return mock_get_task
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_session() -> AsyncMock:
|
||||
return AsyncMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def features(mock_session: AsyncMock) -> ExperimentalClientFeatures:
|
||||
return ExperimentalClientFeatures(mock_session)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_poll_task_yields_until_completed(features: ExperimentalClientFeatures) -> None:
|
||||
"""poll_task yields each status until terminal."""
|
||||
features.get_task = make_status_sequence("working", "working", "completed") # type: ignore[method-assign]
|
||||
|
||||
statuses = [s.status async for s in features.poll_task("test-task")]
|
||||
|
||||
assert statuses == ["working", "working", "completed"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("terminal_status", ["completed", "failed", "cancelled"])
|
||||
async def test_poll_task_exits_on_terminal(features: ExperimentalClientFeatures, terminal_status: TaskStatus) -> None:
|
||||
"""poll_task exits immediately when task is already terminal."""
|
||||
features.get_task = make_status_sequence(terminal_status) # type: ignore[method-assign]
|
||||
|
||||
statuses = [s.status async for s in features.poll_task("test-task")]
|
||||
|
||||
assert statuses == [terminal_status]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_poll_task_continues_through_input_required(features: ExperimentalClientFeatures) -> None:
|
||||
"""poll_task yields input_required and continues (non-terminal)."""
|
||||
features.get_task = make_status_sequence("working", "input_required", "working", "completed") # type: ignore[method-assign]
|
||||
|
||||
statuses = [s.status async for s in features.poll_task("test-task")]
|
||||
|
||||
assert statuses == ["working", "input_required", "working", "completed"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_poll_task_passes_task_id(features: ExperimentalClientFeatures) -> None:
|
||||
"""poll_task passes correct task_id to get_task."""
|
||||
received_ids: list[str] = []
|
||||
|
||||
async def mock_get_task(task_id: str) -> GetTaskResult:
|
||||
received_ids.append(task_id)
|
||||
return make_task_result(status="completed", task_id=task_id)
|
||||
|
||||
features.get_task = mock_get_task # type: ignore[method-assign]
|
||||
|
||||
_ = [s async for s in features.poll_task("my-task-123")]
|
||||
|
||||
assert received_ids == ["my-task-123"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_poll_task_yields_full_result(features: ExperimentalClientFeatures) -> None:
|
||||
"""poll_task yields complete GetTaskResult objects."""
|
||||
|
||||
async def mock_get_task(task_id: str) -> GetTaskResult:
|
||||
return make_task_result(
|
||||
status="completed",
|
||||
task_id=task_id,
|
||||
status_message="All done!",
|
||||
)
|
||||
|
||||
features.get_task = mock_get_task # type: ignore[method-assign]
|
||||
|
||||
results = [r async for r in features.poll_task("test-task")]
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].status == "completed"
|
||||
assert results[0].statusMessage == "All done!"
|
||||
assert results[0].taskId == "test-task"
|
||||
483
tests/experimental/tasks/client/test_tasks.py
Normal file
483
tests/experimental/tasks/client/test_tasks.py
Normal file
|
|
@ -0,0 +1,483 @@
|
|||
"""Tests for the experimental client task methods (session.experimental)."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from anyio import Event
|
||||
from anyio.abc import TaskGroup
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.server import Server
|
||||
from mcp.server.lowlevel import NotificationOptions
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared.experimental.tasks.helpers import task_execution
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.types import (
|
||||
CallToolRequest,
|
||||
CallToolRequestParams,
|
||||
CallToolResult,
|
||||
CancelTaskRequest,
|
||||
CancelTaskResult,
|
||||
ClientRequest,
|
||||
ClientResult,
|
||||
CreateTaskResult,
|
||||
GetTaskPayloadRequest,
|
||||
GetTaskPayloadResult,
|
||||
GetTaskRequest,
|
||||
GetTaskResult,
|
||||
ListTasksRequest,
|
||||
ListTasksResult,
|
||||
ServerNotification,
|
||||
ServerRequest,
|
||||
TaskMetadata,
|
||||
TextContent,
|
||||
Tool,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AppContext:
|
||||
"""Application context passed via lifespan_context."""
|
||||
|
||||
task_group: TaskGroup
|
||||
store: InMemoryTaskStore
|
||||
task_done_events: dict[str, Event] = field(default_factory=lambda: {})
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_experimental_get_task() -> None:
|
||||
"""Test session.experimental.get_task() method."""
|
||||
# Note: We bypass the normal lifespan mechanism
|
||||
server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment]
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
app = ctx.lifespan_context
|
||||
if ctx.experimental.is_task:
|
||||
task_metadata = ctx.experimental.task_metadata
|
||||
assert task_metadata is not None
|
||||
task = await app.store.create_task(task_metadata)
|
||||
|
||||
done_event = Event()
|
||||
app.task_done_events[task.taskId] = done_event
|
||||
|
||||
async def do_work():
|
||||
async with task_execution(task.taskId, app.store) as task_ctx:
|
||||
await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")]))
|
||||
done_event.set()
|
||||
|
||||
app.task_group.start_soon(do_work)
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.get_task()
|
||||
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
|
||||
app = server.request_context.lifespan_context
|
||||
task = await app.store.get_task(request.params.taskId)
|
||||
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
|
||||
return GetTaskResult(
|
||||
taskId=task.taskId,
|
||||
status=task.status,
|
||||
statusMessage=task.statusMessage,
|
||||
createdAt=task.createdAt,
|
||||
lastUpdatedAt=task.lastUpdatedAt,
|
||||
ttl=task.ttl,
|
||||
pollInterval=task.pollInterval,
|
||||
)
|
||||
|
||||
# Set up streams
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no branch
|
||||
|
||||
async def run_server(app_context: AppContext):
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Create a task
|
||||
create_result = await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
)
|
||||
),
|
||||
CreateTaskResult,
|
||||
)
|
||||
task_id = create_result.task.taskId
|
||||
|
||||
# Wait for task to complete
|
||||
await app_context.task_done_events[task_id].wait()
|
||||
|
||||
# Use session.experimental to get task status
|
||||
task_status = await client_session.experimental.get_task(task_id)
|
||||
|
||||
assert task_status.taskId == task_id
|
||||
assert task_status.status == "completed"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_experimental_get_task_result() -> None:
|
||||
"""Test session.experimental.get_task_result() method."""
|
||||
server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment]
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
app = ctx.lifespan_context
|
||||
if ctx.experimental.is_task:
|
||||
task_metadata = ctx.experimental.task_metadata
|
||||
assert task_metadata is not None
|
||||
task = await app.store.create_task(task_metadata)
|
||||
|
||||
done_event = Event()
|
||||
app.task_done_events[task.taskId] = done_event
|
||||
|
||||
async def do_work():
|
||||
async with task_execution(task.taskId, app.store) as task_ctx:
|
||||
await task_ctx.complete(
|
||||
CallToolResult(content=[TextContent(type="text", text="Task result content")])
|
||||
)
|
||||
done_event.set()
|
||||
|
||||
app.task_group.start_soon(do_work)
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.get_task_result()
|
||||
async def handle_get_task_result(
|
||||
request: GetTaskPayloadRequest,
|
||||
) -> GetTaskPayloadResult:
|
||||
app = server.request_context.lifespan_context
|
||||
result = await app.store.get_result(request.params.taskId)
|
||||
assert result is not None, f"Test setup error: result for {request.params.taskId} should exist"
|
||||
assert isinstance(result, CallToolResult)
|
||||
return GetTaskPayloadResult(**result.model_dump())
|
||||
|
||||
# Set up streams
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no branch
|
||||
|
||||
async def run_server(app_context: AppContext):
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Create a task
|
||||
create_result = await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
)
|
||||
),
|
||||
CreateTaskResult,
|
||||
)
|
||||
task_id = create_result.task.taskId
|
||||
|
||||
# Wait for task to complete
|
||||
await app_context.task_done_events[task_id].wait()
|
||||
|
||||
# Use TaskClient to get task result
|
||||
task_result = await client_session.experimental.get_task_result(task_id, CallToolResult)
|
||||
|
||||
assert len(task_result.content) == 1
|
||||
content = task_result.content[0]
|
||||
assert isinstance(content, TextContent)
|
||||
assert content.text == "Task result content"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_experimental_list_tasks() -> None:
|
||||
"""Test TaskClient.list_tasks() method."""
|
||||
server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment]
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
app = ctx.lifespan_context
|
||||
if ctx.experimental.is_task:
|
||||
task_metadata = ctx.experimental.task_metadata
|
||||
assert task_metadata is not None
|
||||
task = await app.store.create_task(task_metadata)
|
||||
|
||||
done_event = Event()
|
||||
app.task_done_events[task.taskId] = done_event
|
||||
|
||||
async def do_work():
|
||||
async with task_execution(task.taskId, app.store) as task_ctx:
|
||||
await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")]))
|
||||
done_event.set()
|
||||
|
||||
app.task_group.start_soon(do_work)
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.list_tasks()
|
||||
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
|
||||
app = server.request_context.lifespan_context
|
||||
tasks_list, next_cursor = await app.store.list_tasks(cursor=request.params.cursor if request.params else None)
|
||||
return ListTasksResult(tasks=tasks_list, nextCursor=next_cursor)
|
||||
|
||||
# Set up streams
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no branch
|
||||
|
||||
async def run_server(app_context: AppContext):
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Create two tasks
|
||||
for _ in range(2):
|
||||
create_result = await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
)
|
||||
),
|
||||
CreateTaskResult,
|
||||
)
|
||||
await app_context.task_done_events[create_result.task.taskId].wait()
|
||||
|
||||
# Use TaskClient to list tasks
|
||||
list_result = await client_session.experimental.list_tasks()
|
||||
|
||||
assert len(list_result.tasks) == 2
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_experimental_cancel_task() -> None:
|
||||
"""Test TaskClient.cancel_task() method."""
|
||||
server: Server[AppContext, Any] = Server("test-server") # type: ignore[assignment]
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
app = ctx.lifespan_context
|
||||
if ctx.experimental.is_task:
|
||||
task_metadata = ctx.experimental.task_metadata
|
||||
assert task_metadata is not None
|
||||
task = await app.store.create_task(task_metadata)
|
||||
# Don't start any work - task stays in "working" status
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.get_task()
|
||||
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
|
||||
app = server.request_context.lifespan_context
|
||||
task = await app.store.get_task(request.params.taskId)
|
||||
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
|
||||
return GetTaskResult(
|
||||
taskId=task.taskId,
|
||||
status=task.status,
|
||||
statusMessage=task.statusMessage,
|
||||
createdAt=task.createdAt,
|
||||
lastUpdatedAt=task.lastUpdatedAt,
|
||||
ttl=task.ttl,
|
||||
pollInterval=task.pollInterval,
|
||||
)
|
||||
|
||||
@server.experimental.cancel_task()
|
||||
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
|
||||
app = server.request_context.lifespan_context
|
||||
task = await app.store.get_task(request.params.taskId)
|
||||
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
|
||||
await app.store.update_task(request.params.taskId, status="cancelled")
|
||||
# CancelTaskResult extends Task, so we need to return the updated task info
|
||||
updated_task = await app.store.get_task(request.params.taskId)
|
||||
assert updated_task is not None
|
||||
return CancelTaskResult(
|
||||
taskId=updated_task.taskId,
|
||||
status=updated_task.status,
|
||||
createdAt=updated_task.createdAt,
|
||||
lastUpdatedAt=updated_task.lastUpdatedAt,
|
||||
ttl=updated_task.ttl,
|
||||
)
|
||||
|
||||
# Set up streams
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no branch
|
||||
|
||||
async def run_server(app_context: AppContext):
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Create a task (but don't complete it)
|
||||
create_result = await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
)
|
||||
),
|
||||
CreateTaskResult,
|
||||
)
|
||||
task_id = create_result.task.taskId
|
||||
|
||||
# Verify task is working
|
||||
status_before = await client_session.experimental.get_task(task_id)
|
||||
assert status_before.status == "working"
|
||||
|
||||
# Cancel the task
|
||||
await client_session.experimental.cancel_task(task_id)
|
||||
|
||||
# Verify task is cancelled
|
||||
status_after = await client_session.experimental.get_task(task_id)
|
||||
assert status_after.status == "cancelled"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
0
tests/experimental/tasks/server/__init__.py
Normal file
0
tests/experimental/tasks/server/__init__.py
Normal file
183
tests/experimental/tasks/server/test_context.py
Normal file
183
tests/experimental/tasks/server/test_context.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
"""Tests for TaskContext and helper functions."""
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp.shared.experimental.tasks.context import TaskContext
|
||||
from mcp.shared.experimental.tasks.helpers import create_task_state, task_execution
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.types import CallToolResult, TaskMetadata, TextContent
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_context_properties() -> None:
|
||||
"""Test TaskContext basic properties."""
|
||||
store = InMemoryTaskStore()
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
ctx = TaskContext(task, store)
|
||||
|
||||
assert ctx.task_id == task.taskId
|
||||
assert ctx.task.taskId == task.taskId
|
||||
assert ctx.task.status == "working"
|
||||
assert ctx.is_cancelled is False
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_context_update_status() -> None:
|
||||
"""Test TaskContext.update_status."""
|
||||
store = InMemoryTaskStore()
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
ctx = TaskContext(task, store)
|
||||
|
||||
await ctx.update_status("Processing step 1...")
|
||||
|
||||
# Check status message was updated
|
||||
updated = await store.get_task(task.taskId)
|
||||
assert updated is not None
|
||||
assert updated.statusMessage == "Processing step 1..."
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_context_complete() -> None:
|
||||
"""Test TaskContext.complete."""
|
||||
store = InMemoryTaskStore()
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
ctx = TaskContext(task, store)
|
||||
|
||||
result = CallToolResult(content=[TextContent(type="text", text="Done!")])
|
||||
await ctx.complete(result)
|
||||
|
||||
# Check task status
|
||||
updated = await store.get_task(task.taskId)
|
||||
assert updated is not None
|
||||
assert updated.status == "completed"
|
||||
|
||||
# Check result is stored
|
||||
stored_result = await store.get_result(task.taskId)
|
||||
assert stored_result is not None
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_context_fail() -> None:
|
||||
"""Test TaskContext.fail."""
|
||||
store = InMemoryTaskStore()
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
ctx = TaskContext(task, store)
|
||||
|
||||
await ctx.fail("Something went wrong!")
|
||||
|
||||
# Check task status
|
||||
updated = await store.get_task(task.taskId)
|
||||
assert updated is not None
|
||||
assert updated.status == "failed"
|
||||
assert updated.statusMessage == "Something went wrong!"
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_context_cancellation() -> None:
|
||||
"""Test TaskContext cancellation request."""
|
||||
store = InMemoryTaskStore()
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
ctx = TaskContext(task, store)
|
||||
|
||||
assert ctx.is_cancelled is False
|
||||
|
||||
ctx.request_cancellation()
|
||||
|
||||
assert ctx.is_cancelled is True
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
def test_create_task_state_generates_id() -> None:
|
||||
"""create_task_state generates a unique task ID when none provided."""
|
||||
task1 = create_task_state(TaskMetadata(ttl=60000))
|
||||
task2 = create_task_state(TaskMetadata(ttl=60000))
|
||||
|
||||
assert task1.taskId != task2.taskId
|
||||
|
||||
|
||||
def test_create_task_state_uses_provided_id() -> None:
|
||||
"""create_task_state uses the provided task ID."""
|
||||
task = create_task_state(TaskMetadata(ttl=60000), task_id="my-task-123")
|
||||
assert task.taskId == "my-task-123"
|
||||
|
||||
|
||||
def test_create_task_state_null_ttl() -> None:
|
||||
"""create_task_state handles null TTL."""
|
||||
task = create_task_state(TaskMetadata(ttl=None))
|
||||
assert task.ttl is None
|
||||
|
||||
|
||||
def test_create_task_state_has_created_at() -> None:
|
||||
"""create_task_state sets createdAt timestamp."""
|
||||
task = create_task_state(TaskMetadata(ttl=60000))
|
||||
assert task.createdAt is not None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_execution_provides_context() -> None:
|
||||
"""task_execution provides a TaskContext for the task."""
|
||||
store = InMemoryTaskStore()
|
||||
await store.create_task(TaskMetadata(ttl=60000), task_id="exec-test-1")
|
||||
|
||||
async with task_execution("exec-test-1", store) as ctx:
|
||||
assert ctx.task_id == "exec-test-1"
|
||||
assert ctx.task.status == "working"
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_execution_auto_fails_on_exception() -> None:
|
||||
"""task_execution automatically fails task on unhandled exception."""
|
||||
store = InMemoryTaskStore()
|
||||
await store.create_task(TaskMetadata(ttl=60000), task_id="exec-fail-1")
|
||||
|
||||
async with task_execution("exec-fail-1", store):
|
||||
raise RuntimeError("Oops!")
|
||||
|
||||
# Task should be failed
|
||||
failed_task = await store.get_task("exec-fail-1")
|
||||
assert failed_task is not None
|
||||
assert failed_task.status == "failed"
|
||||
assert "Oops!" in (failed_task.statusMessage or "")
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_execution_doesnt_fail_if_already_terminal() -> None:
|
||||
"""task_execution doesn't re-fail if task already terminal."""
|
||||
store = InMemoryTaskStore()
|
||||
await store.create_task(TaskMetadata(ttl=60000), task_id="exec-term-1")
|
||||
|
||||
async with task_execution("exec-term-1", store) as ctx:
|
||||
# Complete the task first
|
||||
await ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")]))
|
||||
# Then raise - shouldn't change status
|
||||
raise RuntimeError("This shouldn't matter")
|
||||
|
||||
# Task should remain completed
|
||||
final_task = await store.get_task("exec-term-1")
|
||||
assert final_task is not None
|
||||
assert final_task.status == "completed"
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_execution_not_found() -> None:
|
||||
"""task_execution raises ValueError for non-existent task."""
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
async with task_execution("nonexistent", store):
|
||||
...
|
||||
357
tests/experimental/tasks/server/test_integration.py
Normal file
357
tests/experimental/tasks/server/test_integration.py
Normal file
|
|
@ -0,0 +1,357 @@
|
|||
"""End-to-end integration tests for tasks functionality.
|
||||
|
||||
These tests demonstrate the full task lifecycle:
|
||||
1. Client sends task-augmented request (tools/call with task metadata)
|
||||
2. Server creates task and returns CreateTaskResult immediately
|
||||
3. Background work executes (using task_execution context manager)
|
||||
4. Client polls with tasks/get
|
||||
5. Client retrieves result with tasks/result
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from anyio import Event
|
||||
from anyio.abc import TaskGroup
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.server import Server
|
||||
from mcp.server.lowlevel import NotificationOptions
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared.experimental.tasks.helpers import task_execution
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.types import (
|
||||
TASK_REQUIRED,
|
||||
CallToolRequest,
|
||||
CallToolRequestParams,
|
||||
CallToolResult,
|
||||
ClientRequest,
|
||||
ClientResult,
|
||||
CreateTaskResult,
|
||||
GetTaskPayloadRequest,
|
||||
GetTaskPayloadRequestParams,
|
||||
GetTaskPayloadResult,
|
||||
GetTaskRequest,
|
||||
GetTaskRequestParams,
|
||||
GetTaskResult,
|
||||
ListTasksRequest,
|
||||
ListTasksResult,
|
||||
ServerNotification,
|
||||
ServerRequest,
|
||||
TaskMetadata,
|
||||
TextContent,
|
||||
Tool,
|
||||
ToolExecution,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AppContext:
|
||||
"""Application context passed via lifespan_context."""
|
||||
|
||||
task_group: TaskGroup
|
||||
store: InMemoryTaskStore
|
||||
# Events to signal when tasks complete (for testing without sleeps)
|
||||
task_done_events: dict[str, Event] = field(default_factory=lambda: {})
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_lifecycle_with_task_execution() -> None:
|
||||
"""
|
||||
Test the complete task lifecycle using the task_execution pattern.
|
||||
|
||||
This demonstrates the recommended way to implement task-augmented tools:
|
||||
1. Create task in store
|
||||
2. Spawn work using task_execution() context manager
|
||||
3. Return CreateTaskResult immediately
|
||||
4. Work executes in background, auto-fails on exception
|
||||
"""
|
||||
# Note: We bypass the normal lifespan mechanism and pass context directly to _handle_message
|
||||
server: Server[AppContext, Any] = Server("test-tasks") # type: ignore[assignment]
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [
|
||||
Tool(
|
||||
name="process_data",
|
||||
description="Process data asynchronously",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {"input": {"type": "string"}},
|
||||
},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
app = ctx.lifespan_context
|
||||
if name == "process_data" and ctx.experimental.is_task:
|
||||
# 1. Create task in store
|
||||
task_metadata = ctx.experimental.task_metadata
|
||||
assert task_metadata is not None
|
||||
task = await app.store.create_task(task_metadata)
|
||||
|
||||
# 2. Create event to signal completion (for testing)
|
||||
done_event = Event()
|
||||
app.task_done_events[task.taskId] = done_event
|
||||
|
||||
# 3. Define work function using task_execution for safety
|
||||
async def do_work():
|
||||
async with task_execution(task.taskId, app.store) as task_ctx:
|
||||
await task_ctx.update_status("Processing input...")
|
||||
# Simulate work
|
||||
input_value = arguments.get("input", "")
|
||||
result_text = f"Processed: {input_value.upper()}"
|
||||
await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text=result_text)]))
|
||||
# Signal completion
|
||||
done_event.set()
|
||||
|
||||
# 4. Spawn work in task group (from lifespan_context)
|
||||
app.task_group.start_soon(do_work)
|
||||
|
||||
# 5. Return CreateTaskResult immediately
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
# Register task query handlers (delegate to store)
|
||||
@server.experimental.get_task()
|
||||
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
|
||||
app = server.request_context.lifespan_context
|
||||
task = await app.store.get_task(request.params.taskId)
|
||||
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
|
||||
return GetTaskResult(
|
||||
taskId=task.taskId,
|
||||
status=task.status,
|
||||
statusMessage=task.statusMessage,
|
||||
createdAt=task.createdAt,
|
||||
lastUpdatedAt=task.lastUpdatedAt,
|
||||
ttl=task.ttl,
|
||||
pollInterval=task.pollInterval,
|
||||
)
|
||||
|
||||
@server.experimental.get_task_result()
|
||||
async def handle_get_task_result(
|
||||
request: GetTaskPayloadRequest,
|
||||
) -> GetTaskPayloadResult:
|
||||
app = server.request_context.lifespan_context
|
||||
result = await app.store.get_result(request.params.taskId)
|
||||
assert result is not None, f"Test setup error: result for {request.params.taskId} should exist"
|
||||
assert isinstance(result, CallToolResult)
|
||||
# Return as GetTaskPayloadResult (which accepts extra fields)
|
||||
return GetTaskPayloadResult(**result.model_dump())
|
||||
|
||||
@server.experimental.list_tasks()
|
||||
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
|
||||
raise NotImplementedError
|
||||
|
||||
# Set up client-server communication
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no cover
|
||||
|
||||
async def run_server(app_context: AppContext):
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
# Create app context with task group and store
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# === Step 1: Send task-augmented tool call ===
|
||||
create_result = await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="process_data",
|
||||
arguments={"input": "hello world"},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
),
|
||||
)
|
||||
),
|
||||
CreateTaskResult,
|
||||
)
|
||||
|
||||
assert isinstance(create_result, CreateTaskResult)
|
||||
assert create_result.task.status == "working"
|
||||
task_id = create_result.task.taskId
|
||||
|
||||
# === Step 2: Wait for task to complete ===
|
||||
await app_context.task_done_events[task_id].wait()
|
||||
|
||||
task_status = await client_session.send_request(
|
||||
ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId=task_id))),
|
||||
GetTaskResult,
|
||||
)
|
||||
|
||||
assert task_status.taskId == task_id
|
||||
assert task_status.status == "completed"
|
||||
|
||||
# === Step 3: Retrieve the actual result ===
|
||||
task_result = await client_session.send_request(
|
||||
ClientRequest(GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task_id))),
|
||||
CallToolResult,
|
||||
)
|
||||
|
||||
assert len(task_result.content) == 1
|
||||
content = task_result.content[0]
|
||||
assert isinstance(content, TextContent)
|
||||
assert content.text == "Processed: HELLO WORLD"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_auto_fails_on_exception() -> None:
|
||||
"""Test that task_execution automatically fails the task on unhandled exception."""
|
||||
# Note: We bypass the normal lifespan mechanism and pass context directly to _handle_message
|
||||
server: Server[AppContext, Any] = Server("test-tasks-failure") # type: ignore[assignment]
|
||||
store = InMemoryTaskStore()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [
|
||||
Tool(
|
||||
name="failing_task",
|
||||
description="A task that fails",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
app = ctx.lifespan_context
|
||||
if name == "failing_task" and ctx.experimental.is_task:
|
||||
task_metadata = ctx.experimental.task_metadata
|
||||
assert task_metadata is not None
|
||||
task = await app.store.create_task(task_metadata)
|
||||
|
||||
# Create event to signal completion (for testing)
|
||||
done_event = Event()
|
||||
app.task_done_events[task.taskId] = done_event
|
||||
|
||||
async def do_failing_work():
|
||||
async with task_execution(task.taskId, app.store) as task_ctx:
|
||||
await task_ctx.update_status("About to fail...")
|
||||
raise RuntimeError("Something went wrong!")
|
||||
# Note: complete() is never called, but task_execution
|
||||
# will automatically call fail() due to the exception
|
||||
# This line is reached because task_execution suppresses the exception
|
||||
done_event.set()
|
||||
|
||||
app.task_group.start_soon(do_failing_work)
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.get_task()
|
||||
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
|
||||
app = server.request_context.lifespan_context
|
||||
task = await app.store.get_task(request.params.taskId)
|
||||
assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
|
||||
return GetTaskResult(
|
||||
taskId=task.taskId,
|
||||
status=task.status,
|
||||
statusMessage=task.statusMessage,
|
||||
createdAt=task.createdAt,
|
||||
lastUpdatedAt=task.lastUpdatedAt,
|
||||
ttl=task.ttl,
|
||||
pollInterval=task.pollInterval,
|
||||
)
|
||||
|
||||
# Set up streams
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no cover
|
||||
|
||||
async def run_server(app_context: AppContext):
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Send task request
|
||||
create_result = await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="failing_task",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
),
|
||||
)
|
||||
),
|
||||
CreateTaskResult,
|
||||
)
|
||||
|
||||
task_id = create_result.task.taskId
|
||||
|
||||
# Wait for task to complete (even though it fails)
|
||||
await app_context.task_done_events[task_id].wait()
|
||||
|
||||
# Check that task was auto-failed
|
||||
task_status = await client_session.send_request(
|
||||
ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId=task_id))),
|
||||
GetTaskResult,
|
||||
)
|
||||
|
||||
assert task_status.status == "failed"
|
||||
assert task_status.statusMessage == "Something went wrong!"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
538
tests/experimental/tasks/server/test_run_task_flow.py
Normal file
538
tests/experimental/tasks/server/test_run_task_flow.py
Normal file
|
|
@ -0,0 +1,538 @@
|
|||
"""
|
||||
Tests for the simplified task API: enable_tasks() + run_task()
|
||||
|
||||
This tests the recommended user flow:
|
||||
1. server.experimental.enable_tasks() - one-line setup
|
||||
2. ctx.experimental.run_task(work) - spawns work, returns CreateTaskResult
|
||||
3. work function uses ServerTaskContext for elicit/create_message
|
||||
|
||||
These are integration tests that verify the complete flow works end-to-end.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from anyio import Event
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.server import Server
|
||||
from mcp.server.experimental.request_context import Experimental
|
||||
from mcp.server.experimental.task_context import ServerTaskContext
|
||||
from mcp.server.experimental.task_support import TaskSupport
|
||||
from mcp.server.lowlevel import NotificationOptions
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import (
|
||||
TASK_REQUIRED,
|
||||
CallToolResult,
|
||||
CancelTaskRequest,
|
||||
CancelTaskResult,
|
||||
CreateTaskResult,
|
||||
GetTaskPayloadRequest,
|
||||
GetTaskPayloadResult,
|
||||
GetTaskRequest,
|
||||
GetTaskResult,
|
||||
ListTasksRequest,
|
||||
ListTasksResult,
|
||||
TextContent,
|
||||
Tool,
|
||||
ToolExecution,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_basic_flow() -> None:
|
||||
"""
|
||||
Test the basic run_task flow without elicitation.
|
||||
|
||||
1. enable_tasks() sets up handlers
|
||||
2. Client calls tool with task field
|
||||
3. run_task() spawns work, returns CreateTaskResult
|
||||
4. Work completes in background
|
||||
5. Client polls and sees completed status
|
||||
"""
|
||||
server = Server("test-run-task")
|
||||
|
||||
# One-line setup
|
||||
server.experimental.enable_tasks()
|
||||
|
||||
# Track when work completes and capture received meta
|
||||
work_completed = Event()
|
||||
received_meta: list[str | None] = [None]
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools() -> list[Tool]:
|
||||
return [
|
||||
Tool(
|
||||
name="simple_task",
|
||||
description="A simple task",
|
||||
inputSchema={"type": "object", "properties": {"input": {"type": "string"}}},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
ctx.experimental.validate_task_mode(TASK_REQUIRED)
|
||||
|
||||
# Capture the meta from the request (if present)
|
||||
if ctx.meta is not None and ctx.meta.model_extra: # pragma: no branch
|
||||
received_meta[0] = ctx.meta.model_extra.get("custom_field")
|
||||
|
||||
async def work(task: ServerTaskContext) -> CallToolResult:
|
||||
await task.update_status("Working...")
|
||||
input_val = arguments.get("input", "default")
|
||||
result = CallToolResult(content=[TextContent(type="text", text=f"Processed: {input_val}")])
|
||||
work_completed.set()
|
||||
return result
|
||||
|
||||
return await ctx.experimental.run_task(work)
|
||||
|
||||
# Set up streams
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def run_server() -> None:
|
||||
await server.run(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server.create_initialization_options(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
)
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
|
||||
# Initialize
|
||||
await client_session.initialize()
|
||||
|
||||
# Call tool as task (with meta to test that code path)
|
||||
result = await client_session.experimental.call_tool_as_task(
|
||||
"simple_task",
|
||||
{"input": "hello"},
|
||||
meta={"custom_field": "test_value"},
|
||||
)
|
||||
|
||||
# Should get CreateTaskResult
|
||||
task_id = result.task.taskId
|
||||
assert result.task.status == "working"
|
||||
|
||||
# Wait for work to complete
|
||||
with anyio.fail_after(5):
|
||||
await work_completed.wait()
|
||||
|
||||
# Poll until task status is completed
|
||||
with anyio.fail_after(5):
|
||||
while True:
|
||||
task_status = await client_session.experimental.get_task(task_id)
|
||||
if task_status.status == "completed": # pragma: no branch
|
||||
break
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_server)
|
||||
tg.start_soon(run_client)
|
||||
|
||||
# Verify the meta was passed through correctly
|
||||
assert received_meta[0] == "test_value"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_auto_fails_on_exception() -> None:
|
||||
"""
|
||||
Test that run_task automatically fails the task when work raises.
|
||||
"""
|
||||
server = Server("test-run-task-fail")
|
||||
server.experimental.enable_tasks()
|
||||
|
||||
work_failed = Event()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools() -> list[Tool]:
|
||||
return [
|
||||
Tool(
|
||||
name="failing_task",
|
||||
description="A task that fails",
|
||||
inputSchema={"type": "object"},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
ctx.experimental.validate_task_mode(TASK_REQUIRED)
|
||||
|
||||
async def work(task: ServerTaskContext) -> CallToolResult:
|
||||
work_failed.set()
|
||||
raise RuntimeError("Something went wrong!")
|
||||
|
||||
return await ctx.experimental.run_task(work)
|
||||
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def run_server() -> None:
|
||||
await server.run(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server.create_initialization_options(),
|
||||
)
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.experimental.call_tool_as_task("failing_task", {})
|
||||
task_id = result.task.taskId
|
||||
|
||||
# Wait for work to fail
|
||||
with anyio.fail_after(5):
|
||||
await work_failed.wait()
|
||||
|
||||
# Poll until task status is failed
|
||||
with anyio.fail_after(5):
|
||||
while True:
|
||||
task_status = await client_session.experimental.get_task(task_id)
|
||||
if task_status.status == "failed": # pragma: no branch
|
||||
break
|
||||
|
||||
assert "Something went wrong" in (task_status.statusMessage or "")
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_server)
|
||||
tg.start_soon(run_client)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_enable_tasks_auto_registers_handlers() -> None:
|
||||
"""
|
||||
Test that enable_tasks() auto-registers get_task, list_tasks, cancel_task handlers.
|
||||
"""
|
||||
server = Server("test-enable-tasks")
|
||||
|
||||
# Before enable_tasks, no task capabilities
|
||||
caps_before = server.get_capabilities(NotificationOptions(), {})
|
||||
assert caps_before.tasks is None
|
||||
|
||||
# Enable tasks
|
||||
server.experimental.enable_tasks()
|
||||
|
||||
# After enable_tasks, should have task capabilities
|
||||
caps_after = server.get_capabilities(NotificationOptions(), {})
|
||||
assert caps_after.tasks is not None
|
||||
assert caps_after.tasks.list is not None
|
||||
assert caps_after.tasks.cancel is not None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_enable_tasks_with_custom_store_and_queue() -> None:
|
||||
"""Test that enable_tasks() uses provided store and queue instead of defaults."""
|
||||
server = Server("test-custom-store-queue")
|
||||
|
||||
# Create custom store and queue
|
||||
custom_store = InMemoryTaskStore()
|
||||
custom_queue = InMemoryTaskMessageQueue()
|
||||
|
||||
# Enable tasks with custom implementations
|
||||
task_support = server.experimental.enable_tasks(store=custom_store, queue=custom_queue)
|
||||
|
||||
# Verify our custom implementations are used
|
||||
assert task_support.store is custom_store
|
||||
assert task_support.queue is custom_queue
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_enable_tasks_skips_default_handlers_when_custom_registered() -> None:
|
||||
"""Test that enable_tasks() doesn't override already-registered handlers."""
|
||||
server = Server("test-custom-handlers")
|
||||
|
||||
# Register custom handlers BEFORE enable_tasks (never called, just for registration)
|
||||
@server.experimental.get_task()
|
||||
async def custom_get_task(req: GetTaskRequest) -> GetTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.get_task_result()
|
||||
async def custom_get_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.list_tasks()
|
||||
async def custom_list_tasks(req: ListTasksRequest) -> ListTasksResult:
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.cancel_task()
|
||||
async def custom_cancel_task(req: CancelTaskRequest) -> CancelTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
# Now enable tasks - should NOT override our custom handlers
|
||||
server.experimental.enable_tasks()
|
||||
|
||||
# Verify our custom handlers are still registered (not replaced by defaults)
|
||||
# The handlers dict should contain our custom handlers
|
||||
assert GetTaskRequest in server.request_handlers
|
||||
assert GetTaskPayloadRequest in server.request_handlers
|
||||
assert ListTasksRequest in server.request_handlers
|
||||
assert CancelTaskRequest in server.request_handlers
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_without_enable_tasks_raises() -> None:
|
||||
"""Test that run_task raises when enable_tasks() wasn't called."""
|
||||
experimental = Experimental(
|
||||
task_metadata=None,
|
||||
_client_capabilities=None,
|
||||
_session=None,
|
||||
_task_support=None, # Not enabled
|
||||
)
|
||||
|
||||
async def work(task: ServerTaskContext) -> CallToolResult:
|
||||
raise NotImplementedError
|
||||
|
||||
with pytest.raises(RuntimeError, match="Task support not enabled"):
|
||||
await experimental.run_task(work)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_support_task_group_before_run_raises() -> None:
|
||||
"""Test that accessing task_group before run() raises RuntimeError."""
|
||||
task_support = TaskSupport.in_memory()
|
||||
|
||||
with pytest.raises(RuntimeError, match="TaskSupport not running"):
|
||||
_ = task_support.task_group
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_without_session_raises() -> None:
|
||||
"""Test that run_task raises when session is not available."""
|
||||
task_support = TaskSupport.in_memory()
|
||||
|
||||
experimental = Experimental(
|
||||
task_metadata=None,
|
||||
_client_capabilities=None,
|
||||
_session=None, # No session
|
||||
_task_support=task_support,
|
||||
)
|
||||
|
||||
async def work(task: ServerTaskContext) -> CallToolResult:
|
||||
raise NotImplementedError
|
||||
|
||||
with pytest.raises(RuntimeError, match="Session not available"):
|
||||
await experimental.run_task(work)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_without_task_metadata_raises() -> None:
|
||||
"""Test that run_task raises when request is not task-augmented."""
|
||||
task_support = TaskSupport.in_memory()
|
||||
mock_session = Mock()
|
||||
|
||||
experimental = Experimental(
|
||||
task_metadata=None, # Not a task-augmented request
|
||||
_client_capabilities=None,
|
||||
_session=mock_session,
|
||||
_task_support=task_support,
|
||||
)
|
||||
|
||||
async def work(task: ServerTaskContext) -> CallToolResult:
|
||||
raise NotImplementedError
|
||||
|
||||
with pytest.raises(RuntimeError, match="Request is not task-augmented"):
|
||||
await experimental.run_task(work)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_with_model_immediate_response() -> None:
|
||||
"""Test that run_task includes model_immediate_response in CreateTaskResult._meta."""
|
||||
server = Server("test-run-task-immediate")
|
||||
server.experimental.enable_tasks()
|
||||
|
||||
work_completed = Event()
|
||||
immediate_response_text = "Processing your request..."
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools() -> list[Tool]:
|
||||
return [
|
||||
Tool(
|
||||
name="task_with_immediate",
|
||||
description="A task with immediate response",
|
||||
inputSchema={"type": "object"},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
ctx.experimental.validate_task_mode(TASK_REQUIRED)
|
||||
|
||||
async def work(task: ServerTaskContext) -> CallToolResult:
|
||||
work_completed.set()
|
||||
return CallToolResult(content=[TextContent(type="text", text="Done")])
|
||||
|
||||
return await ctx.experimental.run_task(work, model_immediate_response=immediate_response_text)
|
||||
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def run_server() -> None:
|
||||
await server.run(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server.create_initialization_options(),
|
||||
)
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.experimental.call_tool_as_task("task_with_immediate", {})
|
||||
|
||||
# Verify the immediate response is in _meta
|
||||
assert result.meta is not None
|
||||
assert "io.modelcontextprotocol/model-immediate-response" in result.meta
|
||||
assert result.meta["io.modelcontextprotocol/model-immediate-response"] == immediate_response_text
|
||||
|
||||
with anyio.fail_after(5):
|
||||
await work_completed.wait()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_server)
|
||||
tg.start_soon(run_client)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_doesnt_complete_if_already_terminal() -> None:
|
||||
"""Test that run_task doesn't auto-complete if work manually completed the task."""
|
||||
server = Server("test-already-complete")
|
||||
server.experimental.enable_tasks()
|
||||
|
||||
work_completed = Event()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools() -> list[Tool]:
|
||||
return [
|
||||
Tool(
|
||||
name="manual_complete_task",
|
||||
description="A task that manually completes",
|
||||
inputSchema={"type": "object"},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
ctx.experimental.validate_task_mode(TASK_REQUIRED)
|
||||
|
||||
async def work(task: ServerTaskContext) -> CallToolResult:
|
||||
# Manually complete the task before returning
|
||||
manual_result = CallToolResult(content=[TextContent(type="text", text="Manually completed")])
|
||||
await task.complete(manual_result, notify=False)
|
||||
work_completed.set()
|
||||
# Return a different result - but it should be ignored since task is already terminal
|
||||
return CallToolResult(content=[TextContent(type="text", text="This should be ignored")])
|
||||
|
||||
return await ctx.experimental.run_task(work)
|
||||
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def run_server() -> None:
|
||||
await server.run(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server.create_initialization_options(),
|
||||
)
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.experimental.call_tool_as_task("manual_complete_task", {})
|
||||
task_id = result.task.taskId
|
||||
|
||||
with anyio.fail_after(5):
|
||||
await work_completed.wait()
|
||||
|
||||
# Poll until task status is completed
|
||||
with anyio.fail_after(5):
|
||||
while True:
|
||||
status = await client_session.experimental.get_task(task_id)
|
||||
if status.status == "completed": # pragma: no branch
|
||||
break
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_server)
|
||||
tg.start_soon(run_client)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_doesnt_fail_if_already_terminal() -> None:
|
||||
"""Test that run_task doesn't auto-fail if work manually failed/cancelled the task."""
|
||||
server = Server("test-already-failed")
|
||||
server.experimental.enable_tasks()
|
||||
|
||||
work_completed = Event()
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools() -> list[Tool]:
|
||||
return [
|
||||
Tool(
|
||||
name="manual_cancel_task",
|
||||
description="A task that manually cancels then raises",
|
||||
inputSchema={"type": "object"},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
|
||||
ctx = server.request_context
|
||||
ctx.experimental.validate_task_mode(TASK_REQUIRED)
|
||||
|
||||
async def work(task: ServerTaskContext) -> CallToolResult:
|
||||
# Manually fail the task first
|
||||
await task.fail("Manually failed", notify=False)
|
||||
work_completed.set()
|
||||
# Then raise - but the auto-fail should be skipped since task is already terminal
|
||||
raise RuntimeError("This error should not change status")
|
||||
|
||||
return await ctx.experimental.run_task(work)
|
||||
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def run_server() -> None:
|
||||
await server.run(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server.create_initialization_options(),
|
||||
)
|
||||
|
||||
async def run_client() -> None:
|
||||
async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.experimental.call_tool_as_task("manual_cancel_task", {})
|
||||
task_id = result.task.taskId
|
||||
|
||||
with anyio.fail_after(5):
|
||||
await work_completed.wait()
|
||||
|
||||
# Poll until task status is failed
|
||||
with anyio.fail_after(5):
|
||||
while True:
|
||||
status = await client_session.experimental.get_task(task_id)
|
||||
if status.status == "failed": # pragma: no branch
|
||||
break
|
||||
|
||||
# Task should still be failed (from manual fail, not auto-fail from exception)
|
||||
assert status.statusMessage == "Manually failed" # Not "This error should not change status"
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_server)
|
||||
tg.start_soon(run_client)
|
||||
965
tests/experimental/tasks/server/test_server.py
Normal file
965
tests/experimental/tasks/server/test_server.py
Normal file
|
|
@ -0,0 +1,965 @@
|
|||
"""Tests for server-side task support (handlers, capabilities, integration)."""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.server import Server
|
||||
from mcp.server.lowlevel import NotificationOptions
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.message import ServerMessageMetadata, SessionMessage
|
||||
from mcp.shared.response_router import ResponseRouter
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.types import (
|
||||
INVALID_REQUEST,
|
||||
TASK_FORBIDDEN,
|
||||
TASK_OPTIONAL,
|
||||
TASK_REQUIRED,
|
||||
CallToolRequest,
|
||||
CallToolRequestParams,
|
||||
CallToolResult,
|
||||
CancelTaskRequest,
|
||||
CancelTaskRequestParams,
|
||||
CancelTaskResult,
|
||||
ClientRequest,
|
||||
ClientResult,
|
||||
ErrorData,
|
||||
GetTaskPayloadRequest,
|
||||
GetTaskPayloadRequestParams,
|
||||
GetTaskPayloadResult,
|
||||
GetTaskRequest,
|
||||
GetTaskRequestParams,
|
||||
GetTaskResult,
|
||||
JSONRPCError,
|
||||
JSONRPCMessage,
|
||||
JSONRPCNotification,
|
||||
JSONRPCResponse,
|
||||
ListTasksRequest,
|
||||
ListTasksResult,
|
||||
ListToolsRequest,
|
||||
ListToolsResult,
|
||||
SamplingMessage,
|
||||
ServerCapabilities,
|
||||
ServerNotification,
|
||||
ServerRequest,
|
||||
ServerResult,
|
||||
Task,
|
||||
TaskMetadata,
|
||||
TextContent,
|
||||
Tool,
|
||||
ToolExecution,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_tasks_handler() -> None:
|
||||
"""Test that experimental list_tasks handler works."""
|
||||
server = Server("test")
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
test_tasks = [
|
||||
Task(
|
||||
taskId="task-1",
|
||||
status="working",
|
||||
createdAt=now,
|
||||
lastUpdatedAt=now,
|
||||
ttl=60000,
|
||||
pollInterval=1000,
|
||||
),
|
||||
Task(
|
||||
taskId="task-2",
|
||||
status="completed",
|
||||
createdAt=now,
|
||||
lastUpdatedAt=now,
|
||||
ttl=60000,
|
||||
pollInterval=1000,
|
||||
),
|
||||
]
|
||||
|
||||
@server.experimental.list_tasks()
|
||||
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
|
||||
return ListTasksResult(tasks=test_tasks)
|
||||
|
||||
handler = server.request_handlers[ListTasksRequest]
|
||||
request = ListTasksRequest(method="tasks/list")
|
||||
result = await handler(request)
|
||||
|
||||
assert isinstance(result, ServerResult)
|
||||
assert isinstance(result.root, ListTasksResult)
|
||||
assert len(result.root.tasks) == 2
|
||||
assert result.root.tasks[0].taskId == "task-1"
|
||||
assert result.root.tasks[1].taskId == "task-2"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_get_task_handler() -> None:
|
||||
"""Test that experimental get_task handler works."""
|
||||
server = Server("test")
|
||||
|
||||
@server.experimental.get_task()
|
||||
async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
|
||||
now = datetime.now(timezone.utc)
|
||||
return GetTaskResult(
|
||||
taskId=request.params.taskId,
|
||||
status="working",
|
||||
createdAt=now,
|
||||
lastUpdatedAt=now,
|
||||
ttl=60000,
|
||||
pollInterval=1000,
|
||||
)
|
||||
|
||||
handler = server.request_handlers[GetTaskRequest]
|
||||
request = GetTaskRequest(
|
||||
method="tasks/get",
|
||||
params=GetTaskRequestParams(taskId="test-task-123"),
|
||||
)
|
||||
result = await handler(request)
|
||||
|
||||
assert isinstance(result, ServerResult)
|
||||
assert isinstance(result.root, GetTaskResult)
|
||||
assert result.root.taskId == "test-task-123"
|
||||
assert result.root.status == "working"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_get_task_result_handler() -> None:
|
||||
"""Test that experimental get_task_result handler works."""
|
||||
server = Server("test")
|
||||
|
||||
@server.experimental.get_task_result()
|
||||
async def handle_get_task_result(request: GetTaskPayloadRequest) -> GetTaskPayloadResult:
|
||||
return GetTaskPayloadResult()
|
||||
|
||||
handler = server.request_handlers[GetTaskPayloadRequest]
|
||||
request = GetTaskPayloadRequest(
|
||||
method="tasks/result",
|
||||
params=GetTaskPayloadRequestParams(taskId="test-task-123"),
|
||||
)
|
||||
result = await handler(request)
|
||||
|
||||
assert isinstance(result, ServerResult)
|
||||
assert isinstance(result.root, GetTaskPayloadResult)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancel_task_handler() -> None:
|
||||
"""Test that experimental cancel_task handler works."""
|
||||
server = Server("test")
|
||||
|
||||
@server.experimental.cancel_task()
|
||||
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
|
||||
now = datetime.now(timezone.utc)
|
||||
return CancelTaskResult(
|
||||
taskId=request.params.taskId,
|
||||
status="cancelled",
|
||||
createdAt=now,
|
||||
lastUpdatedAt=now,
|
||||
ttl=60000,
|
||||
)
|
||||
|
||||
handler = server.request_handlers[CancelTaskRequest]
|
||||
request = CancelTaskRequest(
|
||||
method="tasks/cancel",
|
||||
params=CancelTaskRequestParams(taskId="test-task-123"),
|
||||
)
|
||||
result = await handler(request)
|
||||
|
||||
assert isinstance(result, ServerResult)
|
||||
assert isinstance(result.root, CancelTaskResult)
|
||||
assert result.root.taskId == "test-task-123"
|
||||
assert result.root.status == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_server_capabilities_include_tasks() -> None:
|
||||
"""Test that server capabilities include tasks when handlers are registered."""
|
||||
server = Server("test")
|
||||
|
||||
@server.experimental.list_tasks()
|
||||
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.cancel_task()
|
||||
async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
capabilities = server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
)
|
||||
|
||||
assert capabilities.tasks is not None
|
||||
assert capabilities.tasks.list is not None
|
||||
assert capabilities.tasks.cancel is not None
|
||||
assert capabilities.tasks.requests is not None
|
||||
assert capabilities.tasks.requests.tools is not None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_server_capabilities_partial_tasks() -> None:
|
||||
"""Test capabilities with only some task handlers registered."""
|
||||
server = Server("test")
|
||||
|
||||
@server.experimental.list_tasks()
|
||||
async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
|
||||
raise NotImplementedError
|
||||
|
||||
# Only list_tasks registered, not cancel_task
|
||||
|
||||
capabilities = server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
)
|
||||
|
||||
assert capabilities.tasks is not None
|
||||
assert capabilities.tasks.list is not None
|
||||
assert capabilities.tasks.cancel is None # Not registered
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_tool_with_task_execution_metadata() -> None:
|
||||
"""Test that tools can declare task execution mode."""
|
||||
server = Server("test")
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [
|
||||
Tool(
|
||||
name="quick_tool",
|
||||
description="Fast tool",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
execution=ToolExecution(taskSupport=TASK_FORBIDDEN),
|
||||
),
|
||||
Tool(
|
||||
name="long_tool",
|
||||
description="Long running tool",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
),
|
||||
Tool(
|
||||
name="flexible_tool",
|
||||
description="Can be either",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
execution=ToolExecution(taskSupport=TASK_OPTIONAL),
|
||||
),
|
||||
]
|
||||
|
||||
tools_handler = server.request_handlers[ListToolsRequest]
|
||||
request = ListToolsRequest(method="tools/list")
|
||||
result = await tools_handler(request)
|
||||
|
||||
assert isinstance(result, ServerResult)
|
||||
assert isinstance(result.root, ListToolsResult)
|
||||
tools = result.root.tools
|
||||
|
||||
assert tools[0].execution is not None
|
||||
assert tools[0].execution.taskSupport == TASK_FORBIDDEN
|
||||
assert tools[1].execution is not None
|
||||
assert tools[1].execution.taskSupport == TASK_REQUIRED
|
||||
assert tools[2].execution is not None
|
||||
assert tools[2].execution.taskSupport == TASK_OPTIONAL
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_metadata_in_call_tool_request() -> None:
|
||||
"""Test that task metadata is accessible via RequestContext when calling a tool."""
|
||||
server = Server("test")
|
||||
captured_task_metadata: TaskMetadata | None = None
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [
|
||||
Tool(
|
||||
name="long_task",
|
||||
description="A long running task",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
execution=ToolExecution(taskSupport="optional"),
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
|
||||
nonlocal captured_task_metadata
|
||||
ctx = server.request_context
|
||||
captured_task_metadata = ctx.experimental.task_metadata
|
||||
return [TextContent(type="text", text="done")]
|
||||
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no branch
|
||||
|
||||
async def run_server():
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def handle_messages():
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, {}, False)
|
||||
|
||||
tg.start_soon(handle_messages)
|
||||
await anyio.sleep_forever()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_server)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Call tool with task metadata
|
||||
await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="long_task",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
),
|
||||
)
|
||||
),
|
||||
CallToolResult,
|
||||
)
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
assert captured_task_metadata is not None
|
||||
assert captured_task_metadata.ttl == 60000
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_metadata_is_task_property() -> None:
|
||||
"""Test that RequestContext.experimental.is_task works correctly."""
|
||||
server = Server("test")
|
||||
is_task_values: list[bool] = []
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [
|
||||
Tool(
|
||||
name="test_tool",
|
||||
description="Test tool",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
)
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
|
||||
ctx = server.request_context
|
||||
is_task_values.append(ctx.experimental.is_task)
|
||||
return [TextContent(type="text", text="done")]
|
||||
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no branch
|
||||
|
||||
async def run_server():
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def handle_messages():
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, {}, False)
|
||||
|
||||
tg.start_soon(handle_messages)
|
||||
await anyio.sleep_forever()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_server)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Call without task metadata
|
||||
await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(name="test_tool", arguments={}),
|
||||
)
|
||||
),
|
||||
CallToolResult,
|
||||
)
|
||||
|
||||
# Call with task metadata
|
||||
await client_session.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
),
|
||||
)
|
||||
),
|
||||
CallToolResult,
|
||||
)
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
assert len(is_task_values) == 2
|
||||
assert is_task_values[0] is False # First call without task
|
||||
assert is_task_values[1] is True # Second call with task
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_update_capabilities_no_handlers() -> None:
|
||||
"""Test that update_capabilities returns early when no task handlers are registered."""
|
||||
server = Server("test-no-handlers")
|
||||
# Access experimental to initialize it, but don't register any task handlers
|
||||
_ = server.experimental
|
||||
|
||||
caps = server.get_capabilities(NotificationOptions(), {})
|
||||
|
||||
# Without any task handlers registered, tasks capability should be None
|
||||
assert caps.tasks is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_default_task_handlers_via_enable_tasks() -> None:
|
||||
"""Test that enable_tasks() auto-registers working default handlers.
|
||||
|
||||
This exercises the default handlers in lowlevel/experimental.py:
|
||||
- _default_get_task (task not found)
|
||||
- _default_get_task_result
|
||||
- _default_list_tasks
|
||||
- _default_cancel_task
|
||||
"""
|
||||
server = Server("test-default-handlers")
|
||||
# Enable tasks with default handlers (no custom handlers registered)
|
||||
task_support = server.experimental.enable_tasks()
|
||||
store = task_support.store
|
||||
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
|
||||
) -> None: ... # pragma: no branch
|
||||
|
||||
async def run_server() -> None:
|
||||
async with task_support.run():
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
),
|
||||
),
|
||||
) as server_session:
|
||||
task_support.configure_session(server_session)
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, {}, False)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_server)
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Create a task directly in the store for testing
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
# Test list_tasks (default handler)
|
||||
list_result = await client_session.send_request(
|
||||
ClientRequest(ListTasksRequest()),
|
||||
ListTasksResult,
|
||||
)
|
||||
assert len(list_result.tasks) == 1
|
||||
assert list_result.tasks[0].taskId == task.taskId
|
||||
|
||||
# Test get_task (default handler - found)
|
||||
get_result = await client_session.send_request(
|
||||
ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId=task.taskId))),
|
||||
GetTaskResult,
|
||||
)
|
||||
assert get_result.taskId == task.taskId
|
||||
assert get_result.status == "working"
|
||||
|
||||
# Test get_task (default handler - not found path)
|
||||
with pytest.raises(McpError, match="not found"):
|
||||
await client_session.send_request(
|
||||
ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId="nonexistent-task"))),
|
||||
GetTaskResult,
|
||||
)
|
||||
|
||||
# Create a completed task to test get_task_result
|
||||
completed_task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
await store.store_result(
|
||||
completed_task.taskId, CallToolResult(content=[TextContent(type="text", text="Test result")])
|
||||
)
|
||||
await store.update_task(completed_task.taskId, status="completed")
|
||||
|
||||
# Test get_task_result (default handler)
|
||||
payload_result = await client_session.send_request(
|
||||
ClientRequest(GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=completed_task.taskId))),
|
||||
GetTaskPayloadResult,
|
||||
)
|
||||
# The result should have the related-task metadata
|
||||
assert payload_result.meta is not None
|
||||
assert "io.modelcontextprotocol/related-task" in payload_result.meta
|
||||
|
||||
# Test cancel_task (default handler)
|
||||
cancel_result = await client_session.send_request(
|
||||
ClientRequest(CancelTaskRequest(params=CancelTaskRequestParams(taskId=task.taskId))),
|
||||
CancelTaskResult,
|
||||
)
|
||||
assert cancel_result.taskId == task.taskId
|
||||
assert cancel_result.status == "cancelled"
|
||||
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_build_elicit_form_request() -> None:
|
||||
"""Test that _build_elicit_form_request builds a proper elicitation request."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
try:
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as server_session:
|
||||
# Test without task_id
|
||||
request = server_session._build_elicit_form_request(
|
||||
message="Test message",
|
||||
requestedSchema={"type": "object", "properties": {"answer": {"type": "string"}}},
|
||||
)
|
||||
assert request.method == "elicitation/create"
|
||||
assert request.params is not None
|
||||
assert request.params["message"] == "Test message"
|
||||
|
||||
# Test with related_task_id (adds related-task metadata)
|
||||
request_with_task = server_session._build_elicit_form_request(
|
||||
message="Task message",
|
||||
requestedSchema={"type": "object"},
|
||||
related_task_id="test-task-123",
|
||||
)
|
||||
assert request_with_task.method == "elicitation/create"
|
||||
assert request_with_task.params is not None
|
||||
assert "_meta" in request_with_task.params
|
||||
assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"]
|
||||
assert (
|
||||
request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"] == "test-task-123"
|
||||
)
|
||||
finally: # pragma: no cover
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_build_elicit_url_request() -> None:
|
||||
"""Test that _build_elicit_url_request builds a proper URL mode elicitation request."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
try:
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as server_session:
|
||||
# Test without related_task_id
|
||||
request = server_session._build_elicit_url_request(
|
||||
message="Please authorize with GitHub",
|
||||
url="https://github.com/login/oauth/authorize",
|
||||
elicitation_id="oauth-123",
|
||||
)
|
||||
assert request.method == "elicitation/create"
|
||||
assert request.params is not None
|
||||
assert request.params["message"] == "Please authorize with GitHub"
|
||||
assert request.params["url"] == "https://github.com/login/oauth/authorize"
|
||||
assert request.params["elicitationId"] == "oauth-123"
|
||||
assert request.params["mode"] == "url"
|
||||
|
||||
# Test with related_task_id (adds related-task metadata)
|
||||
request_with_task = server_session._build_elicit_url_request(
|
||||
message="OAuth required",
|
||||
url="https://example.com/oauth",
|
||||
elicitation_id="oauth-456",
|
||||
related_task_id="test-task-789",
|
||||
)
|
||||
assert request_with_task.method == "elicitation/create"
|
||||
assert request_with_task.params is not None
|
||||
assert "_meta" in request_with_task.params
|
||||
assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"]
|
||||
assert (
|
||||
request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"] == "test-task-789"
|
||||
)
|
||||
finally: # pragma: no cover
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_build_create_message_request() -> None:
|
||||
"""Test that _build_create_message_request builds a proper sampling request."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
try:
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as server_session:
|
||||
messages = [
|
||||
SamplingMessage(role="user", content=TextContent(type="text", text="Hello")),
|
||||
]
|
||||
|
||||
# Test without task_id
|
||||
request = server_session._build_create_message_request(
|
||||
messages=messages,
|
||||
max_tokens=100,
|
||||
system_prompt="You are helpful",
|
||||
)
|
||||
assert request.method == "sampling/createMessage"
|
||||
assert request.params is not None
|
||||
assert request.params["maxTokens"] == 100
|
||||
|
||||
# Test with related_task_id (adds related-task metadata)
|
||||
request_with_task = server_session._build_create_message_request(
|
||||
messages=messages,
|
||||
max_tokens=50,
|
||||
related_task_id="sampling-task-456",
|
||||
)
|
||||
assert request_with_task.method == "sampling/createMessage"
|
||||
assert request_with_task.params is not None
|
||||
assert "_meta" in request_with_task.params
|
||||
assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"]
|
||||
assert (
|
||||
request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"]
|
||||
== "sampling-task-456"
|
||||
)
|
||||
finally: # pragma: no cover
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_send_message() -> None:
|
||||
"""Test that send_message sends a raw session message."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
try:
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as server_session:
|
||||
# Create a test message
|
||||
notification = JSONRPCNotification(jsonrpc="2.0", method="test/notification")
|
||||
message = SessionMessage(
|
||||
message=JSONRPCMessage(notification),
|
||||
metadata=ServerMessageMetadata(related_request_id="test-req-1"),
|
||||
)
|
||||
|
||||
# Send the message
|
||||
await server_session.send_message(message)
|
||||
|
||||
# Verify it was sent to the stream
|
||||
received = await server_to_client_receive.receive()
|
||||
assert isinstance(received.message.root, JSONRPCNotification)
|
||||
assert received.message.root.method == "test/notification"
|
||||
finally: # pragma: no cover
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_response_routing_success() -> None:
|
||||
"""Test that response routing works for success responses."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
# Track routed responses with event for synchronization
|
||||
routed_responses: list[dict[str, Any]] = []
|
||||
response_received = anyio.Event()
|
||||
|
||||
class TestRouter(ResponseRouter):
|
||||
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
|
||||
routed_responses.append({"id": request_id, "response": response})
|
||||
response_received.set()
|
||||
return True # Handled
|
||||
|
||||
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
try:
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as server_session:
|
||||
router = TestRouter()
|
||||
server_session.add_response_router(router)
|
||||
|
||||
# Simulate receiving a response from client
|
||||
response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"})
|
||||
message = SessionMessage(message=JSONRPCMessage(response))
|
||||
|
||||
# Send from "client" side
|
||||
await client_to_server_send.send(message)
|
||||
|
||||
# Wait for response to be routed
|
||||
with anyio.fail_after(5):
|
||||
await response_received.wait()
|
||||
|
||||
# Verify response was routed
|
||||
assert len(routed_responses) == 1
|
||||
assert routed_responses[0]["id"] == "test-req-1"
|
||||
assert routed_responses[0]["response"]["status"] == "ok"
|
||||
finally: # pragma: no cover
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_response_routing_error() -> None:
|
||||
"""Test that error routing works for error responses."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
# Track routed errors with event for synchronization
|
||||
routed_errors: list[dict[str, Any]] = []
|
||||
error_received = anyio.Event()
|
||||
|
||||
class TestRouter(ResponseRouter):
|
||||
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
|
||||
routed_errors.append({"id": request_id, "error": error})
|
||||
error_received.set()
|
||||
return True # Handled
|
||||
|
||||
try:
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as server_session:
|
||||
router = TestRouter()
|
||||
server_session.add_response_router(router)
|
||||
|
||||
# Simulate receiving an error response from client
|
||||
error_data = ErrorData(code=INVALID_REQUEST, message="Test error")
|
||||
error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data)
|
||||
message = SessionMessage(message=JSONRPCMessage(error_response))
|
||||
|
||||
# Send from "client" side
|
||||
await client_to_server_send.send(message)
|
||||
|
||||
# Wait for error to be routed
|
||||
with anyio.fail_after(5):
|
||||
await error_received.wait()
|
||||
|
||||
# Verify error was routed
|
||||
assert len(routed_errors) == 1
|
||||
assert routed_errors[0]["id"] == "test-req-2"
|
||||
assert routed_errors[0]["error"].message == "Test error"
|
||||
finally: # pragma: no cover
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_response_routing_skips_non_matching_routers() -> None:
|
||||
"""Test that routing continues to next router when first doesn't match."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
# Track which routers were called
|
||||
router_calls: list[str] = []
|
||||
response_received = anyio.Event()
|
||||
|
||||
class NonMatchingRouter(ResponseRouter):
|
||||
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
|
||||
router_calls.append("non_matching_response")
|
||||
return False # Doesn't handle it
|
||||
|
||||
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
class MatchingRouter(ResponseRouter):
|
||||
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
|
||||
router_calls.append("matching_response")
|
||||
response_received.set()
|
||||
return True # Handles it
|
||||
|
||||
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
try:
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as server_session:
|
||||
# Add non-matching router first, then matching router
|
||||
server_session.add_response_router(NonMatchingRouter())
|
||||
server_session.add_response_router(MatchingRouter())
|
||||
|
||||
# Send a response - should skip first router and be handled by second
|
||||
response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"})
|
||||
message = SessionMessage(message=JSONRPCMessage(response))
|
||||
await client_to_server_send.send(message)
|
||||
|
||||
with anyio.fail_after(5):
|
||||
await response_received.wait()
|
||||
|
||||
# Verify both routers were called (first returned False, second returned True)
|
||||
assert router_calls == ["non_matching_response", "matching_response"]
|
||||
finally: # pragma: no cover
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_error_routing_skips_non_matching_routers() -> None:
|
||||
"""Test that error routing continues to next router when first doesn't match."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
# Track which routers were called
|
||||
router_calls: list[str] = []
|
||||
error_received = anyio.Event()
|
||||
|
||||
class NonMatchingRouter(ResponseRouter):
|
||||
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
|
||||
router_calls.append("non_matching_error")
|
||||
return False # Doesn't handle it
|
||||
|
||||
class MatchingRouter(ResponseRouter):
|
||||
def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def route_error(self, request_id: str | int, error: ErrorData) -> bool:
|
||||
router_calls.append("matching_error")
|
||||
error_received.set()
|
||||
return True # Handles it
|
||||
|
||||
try:
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test-server",
|
||||
server_version="1.0.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as server_session:
|
||||
# Add non-matching router first, then matching router
|
||||
server_session.add_response_router(NonMatchingRouter())
|
||||
server_session.add_response_router(MatchingRouter())
|
||||
|
||||
# Send an error - should skip first router and be handled by second
|
||||
error_data = ErrorData(code=INVALID_REQUEST, message="Test error")
|
||||
error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data)
|
||||
message = SessionMessage(message=JSONRPCMessage(error_response))
|
||||
await client_to_server_send.send(message)
|
||||
|
||||
with anyio.fail_after(5):
|
||||
await error_received.wait()
|
||||
|
||||
# Verify both routers were called (first returned False, second returned True)
|
||||
assert router_calls == ["non_matching_error", "matching_error"]
|
||||
finally: # pragma: no cover
|
||||
await server_to_client_send.aclose()
|
||||
await server_to_client_receive.aclose()
|
||||
await client_to_server_send.aclose()
|
||||
await client_to_server_receive.aclose()
|
||||
709
tests/experimental/tasks/server/test_server_task_context.py
Normal file
709
tests/experimental/tasks/server/test_server_task_context.py
Normal file
|
|
@ -0,0 +1,709 @@
|
|||
"""Tests for ServerTaskContext."""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from mcp.server.experimental.task_context import ServerTaskContext
|
||||
from mcp.server.experimental.task_result_handler import TaskResultHandler
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
|
||||
from mcp.types import (
|
||||
CallToolResult,
|
||||
ClientCapabilities,
|
||||
ClientTasksCapability,
|
||||
ClientTasksRequestsCapability,
|
||||
Implementation,
|
||||
InitializeRequestParams,
|
||||
JSONRPCRequest,
|
||||
SamplingMessage,
|
||||
TaskMetadata,
|
||||
TasksCreateElicitationCapability,
|
||||
TasksCreateMessageCapability,
|
||||
TasksElicitationCapability,
|
||||
TasksSamplingCapability,
|
||||
TextContent,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_server_task_context_properties() -> None:
|
||||
"""Test ServerTaskContext property accessors."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-123")
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
)
|
||||
|
||||
assert ctx.task_id == "test-123"
|
||||
assert ctx.task.taskId == "test-123"
|
||||
assert ctx.is_cancelled is False
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_server_task_context_request_cancellation() -> None:
|
||||
"""Test ServerTaskContext.request_cancellation()."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
)
|
||||
|
||||
assert ctx.is_cancelled is False
|
||||
ctx.request_cancellation()
|
||||
assert ctx.is_cancelled is True
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_server_task_context_update_status_with_notify() -> None:
|
||||
"""Test update_status sends notification when notify=True."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.send_notification = AsyncMock()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
)
|
||||
|
||||
await ctx.update_status("Working...", notify=True)
|
||||
|
||||
mock_session.send_notification.assert_called_once()
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_server_task_context_update_status_without_notify() -> None:
|
||||
"""Test update_status skips notification when notify=False."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.send_notification = AsyncMock()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
)
|
||||
|
||||
await ctx.update_status("Working...", notify=False)
|
||||
|
||||
mock_session.send_notification.assert_not_called()
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_server_task_context_complete_with_notify() -> None:
|
||||
"""Test complete sends notification when notify=True."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.send_notification = AsyncMock()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
)
|
||||
|
||||
result = CallToolResult(content=[TextContent(type="text", text="Done")])
|
||||
await ctx.complete(result, notify=True)
|
||||
|
||||
mock_session.send_notification.assert_called_once()
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_server_task_context_fail_with_notify() -> None:
|
||||
"""Test fail sends notification when notify=True."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.send_notification = AsyncMock()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
)
|
||||
|
||||
await ctx.fail("Something went wrong", notify=True)
|
||||
|
||||
mock_session.send_notification.assert_called_once()
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicit_raises_when_client_lacks_capability() -> None:
|
||||
"""Test that elicit() raises McpError when client doesn't support elicitation."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=False)
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
handler = TaskResultHandler(store, queue)
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=handler,
|
||||
)
|
||||
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await ctx.elicit(message="Test?", requestedSchema={"type": "object"})
|
||||
|
||||
assert "elicitation capability" in exc_info.value.error.message
|
||||
mock_session.check_client_capability.assert_called_once()
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_raises_when_client_lacks_capability() -> None:
|
||||
"""Test that create_message() raises McpError when client doesn't support sampling."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=False)
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
handler = TaskResultHandler(store, queue)
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=handler,
|
||||
)
|
||||
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await ctx.create_message(messages=[], max_tokens=100)
|
||||
|
||||
assert "sampling capability" in exc_info.value.error.message
|
||||
mock_session.check_client_capability.assert_called_once()
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicit_raises_without_handler() -> None:
|
||||
"""Test that elicit() raises when handler is not provided."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=True)
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=None,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="handler is required"):
|
||||
await ctx.elicit(message="Test?", requestedSchema={"type": "object"})
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicit_url_raises_without_handler() -> None:
|
||||
"""Test that elicit_url() raises when handler is not provided."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=True)
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=None,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="handler is required for elicit_url"):
|
||||
await ctx.elicit_url(
|
||||
message="Please authorize",
|
||||
url="https://example.com/oauth",
|
||||
elicitation_id="oauth-123",
|
||||
)
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_raises_without_handler() -> None:
|
||||
"""Test that create_message() raises when handler is not provided."""
|
||||
store = InMemoryTaskStore()
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=True)
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=None,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="handler is required"):
|
||||
await ctx.create_message(messages=[], max_tokens=100)
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicit_queues_request_and_waits_for_response() -> None:
|
||||
"""Test that elicit() queues request and waits for response."""
|
||||
store = InMemoryTaskStore()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
handler = TaskResultHandler(store, queue)
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=True)
|
||||
mock_session._build_elicit_form_request = Mock(
|
||||
return_value=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="test-req-1",
|
||||
method="elicitation/create",
|
||||
params={"message": "Test?", "_meta": {}},
|
||||
)
|
||||
)
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=handler,
|
||||
)
|
||||
|
||||
elicit_result = None
|
||||
|
||||
async def run_elicit() -> None:
|
||||
nonlocal elicit_result
|
||||
elicit_result = await ctx.elicit(
|
||||
message="Test?",
|
||||
requestedSchema={"type": "object"},
|
||||
)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_elicit)
|
||||
|
||||
# Wait for request to be queued
|
||||
await queue.wait_for_message(task.taskId)
|
||||
|
||||
# Verify task is in input_required status
|
||||
updated_task = await store.get_task(task.taskId)
|
||||
assert updated_task is not None
|
||||
assert updated_task.status == "input_required"
|
||||
|
||||
# Dequeue and simulate response
|
||||
msg = await queue.dequeue(task.taskId)
|
||||
assert msg is not None
|
||||
assert msg.resolver is not None
|
||||
|
||||
# Resolve with mock elicitation response
|
||||
msg.resolver.set_result({"action": "accept", "content": {"name": "Alice"}})
|
||||
|
||||
# Verify result
|
||||
assert elicit_result is not None
|
||||
assert elicit_result.action == "accept"
|
||||
assert elicit_result.content == {"name": "Alice"}
|
||||
|
||||
# Verify task is back to working
|
||||
final_task = await store.get_task(task.taskId)
|
||||
assert final_task is not None
|
||||
assert final_task.status == "working"
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicit_url_queues_request_and_waits_for_response() -> None:
|
||||
"""Test that elicit_url() queues request and waits for response."""
|
||||
store = InMemoryTaskStore()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
handler = TaskResultHandler(store, queue)
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=True)
|
||||
mock_session._build_elicit_url_request = Mock(
|
||||
return_value=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="test-url-req-1",
|
||||
method="elicitation/create",
|
||||
params={"message": "Authorize", "url": "https://example.com", "elicitationId": "123", "mode": "url"},
|
||||
)
|
||||
)
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=handler,
|
||||
)
|
||||
|
||||
elicit_result = None
|
||||
|
||||
async def run_elicit_url() -> None:
|
||||
nonlocal elicit_result
|
||||
elicit_result = await ctx.elicit_url(
|
||||
message="Authorize",
|
||||
url="https://example.com/oauth",
|
||||
elicitation_id="oauth-123",
|
||||
)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_elicit_url)
|
||||
|
||||
# Wait for request to be queued
|
||||
await queue.wait_for_message(task.taskId)
|
||||
|
||||
# Verify task is in input_required status
|
||||
updated_task = await store.get_task(task.taskId)
|
||||
assert updated_task is not None
|
||||
assert updated_task.status == "input_required"
|
||||
|
||||
# Dequeue and simulate response
|
||||
msg = await queue.dequeue(task.taskId)
|
||||
assert msg is not None
|
||||
assert msg.resolver is not None
|
||||
|
||||
# Resolve with mock elicitation response (URL mode just returns action)
|
||||
msg.resolver.set_result({"action": "accept"})
|
||||
|
||||
# Verify result
|
||||
assert elicit_result is not None
|
||||
assert elicit_result.action == "accept"
|
||||
|
||||
# Verify task is back to working
|
||||
final_task = await store.get_task(task.taskId)
|
||||
assert final_task is not None
|
||||
assert final_task.status == "working"
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_queues_request_and_waits_for_response() -> None:
|
||||
"""Test that create_message() queues request and waits for response."""
|
||||
store = InMemoryTaskStore()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
handler = TaskResultHandler(store, queue)
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=True)
|
||||
mock_session._build_create_message_request = Mock(
|
||||
return_value=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="test-req-2",
|
||||
method="sampling/createMessage",
|
||||
params={"messages": [], "maxTokens": 100, "_meta": {}},
|
||||
)
|
||||
)
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=handler,
|
||||
)
|
||||
|
||||
sampling_result = None
|
||||
|
||||
async def run_sampling() -> None:
|
||||
nonlocal sampling_result
|
||||
sampling_result = await ctx.create_message(
|
||||
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_sampling)
|
||||
|
||||
# Wait for request to be queued
|
||||
await queue.wait_for_message(task.taskId)
|
||||
|
||||
# Verify task is in input_required status
|
||||
updated_task = await store.get_task(task.taskId)
|
||||
assert updated_task is not None
|
||||
assert updated_task.status == "input_required"
|
||||
|
||||
# Dequeue and simulate response
|
||||
msg = await queue.dequeue(task.taskId)
|
||||
assert msg is not None
|
||||
assert msg.resolver is not None
|
||||
|
||||
# Resolve with mock sampling response
|
||||
msg.resolver.set_result(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": {"type": "text", "text": "Hello back!"},
|
||||
"model": "test-model",
|
||||
"stopReason": "endTurn",
|
||||
}
|
||||
)
|
||||
|
||||
# Verify result
|
||||
assert sampling_result is not None
|
||||
assert sampling_result.role == "assistant"
|
||||
assert sampling_result.model == "test-model"
|
||||
|
||||
# Verify task is back to working
|
||||
final_task = await store.get_task(task.taskId)
|
||||
assert final_task is not None
|
||||
assert final_task.status == "working"
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicit_restores_status_on_cancellation() -> None:
|
||||
"""Test that elicit() restores task status to working when cancelled."""
|
||||
store = InMemoryTaskStore()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
handler = TaskResultHandler(store, queue)
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=True)
|
||||
mock_session._build_elicit_form_request = Mock(
|
||||
return_value=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="test-req-cancel",
|
||||
method="elicitation/create",
|
||||
params={"message": "Test?", "_meta": {}},
|
||||
)
|
||||
)
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=handler,
|
||||
)
|
||||
|
||||
cancelled_error_raised = False
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def do_elicit() -> None:
|
||||
nonlocal cancelled_error_raised
|
||||
try:
|
||||
await ctx.elicit(
|
||||
message="Test?",
|
||||
requestedSchema={"type": "object"},
|
||||
)
|
||||
except anyio.get_cancelled_exc_class():
|
||||
cancelled_error_raised = True
|
||||
# Don't re-raise - let the test continue
|
||||
|
||||
tg.start_soon(do_elicit)
|
||||
|
||||
# Wait for request to be queued
|
||||
await queue.wait_for_message(task.taskId)
|
||||
|
||||
# Verify task is in input_required status
|
||||
updated_task = await store.get_task(task.taskId)
|
||||
assert updated_task is not None
|
||||
assert updated_task.status == "input_required"
|
||||
|
||||
# Get the queued message and set cancellation exception on its resolver
|
||||
msg = await queue.dequeue(task.taskId)
|
||||
assert msg is not None
|
||||
assert msg.resolver is not None
|
||||
|
||||
# Trigger cancellation by setting exception (use asyncio.CancelledError directly)
|
||||
msg.resolver.set_exception(asyncio.CancelledError())
|
||||
|
||||
# Verify task is back to working after cancellation
|
||||
final_task = await store.get_task(task.taskId)
|
||||
assert final_task is not None
|
||||
assert final_task.status == "working"
|
||||
assert cancelled_error_raised
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_restores_status_on_cancellation() -> None:
|
||||
"""Test that create_message() restores task status to working when cancelled."""
|
||||
store = InMemoryTaskStore()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
handler = TaskResultHandler(store, queue)
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.check_client_capability = Mock(return_value=True)
|
||||
mock_session._build_create_message_request = Mock(
|
||||
return_value=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="test-req-cancel-2",
|
||||
method="sampling/createMessage",
|
||||
params={"messages": [], "maxTokens": 100, "_meta": {}},
|
||||
)
|
||||
)
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=handler,
|
||||
)
|
||||
|
||||
cancelled_error_raised = False
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def do_sampling() -> None:
|
||||
nonlocal cancelled_error_raised
|
||||
try:
|
||||
await ctx.create_message(
|
||||
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
|
||||
max_tokens=100,
|
||||
)
|
||||
except anyio.get_cancelled_exc_class():
|
||||
cancelled_error_raised = True
|
||||
# Don't re-raise
|
||||
|
||||
tg.start_soon(do_sampling)
|
||||
|
||||
# Wait for request to be queued
|
||||
await queue.wait_for_message(task.taskId)
|
||||
|
||||
# Verify task is in input_required status
|
||||
updated_task = await store.get_task(task.taskId)
|
||||
assert updated_task is not None
|
||||
assert updated_task.status == "input_required"
|
||||
|
||||
# Get the queued message and set cancellation exception on its resolver
|
||||
msg = await queue.dequeue(task.taskId)
|
||||
assert msg is not None
|
||||
assert msg.resolver is not None
|
||||
|
||||
# Trigger cancellation by setting exception (use asyncio.CancelledError directly)
|
||||
msg.resolver.set_exception(asyncio.CancelledError())
|
||||
|
||||
# Verify task is back to working after cancellation
|
||||
final_task = await store.get_task(task.taskId)
|
||||
assert final_task is not None
|
||||
assert final_task.status == "working"
|
||||
assert cancelled_error_raised
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicit_as_task_raises_without_handler() -> None:
|
||||
"""Test that elicit_as_task() raises when handler is not provided."""
|
||||
store = InMemoryTaskStore()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
# Create mock session with proper client capabilities
|
||||
mock_session = Mock()
|
||||
mock_session.client_params = InitializeRequestParams(
|
||||
protocolVersion="2025-01-01",
|
||||
capabilities=ClientCapabilities(
|
||||
tasks=ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
|
||||
)
|
||||
)
|
||||
),
|
||||
clientInfo=Implementation(name="test", version="1.0"),
|
||||
)
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=None,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="handler is required for elicit_as_task"):
|
||||
await ctx.elicit_as_task(message="Test?", requestedSchema={"type": "object"})
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_as_task_raises_without_handler() -> None:
|
||||
"""Test that create_message_as_task() raises when handler is not provided."""
|
||||
store = InMemoryTaskStore()
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
# Create mock session with proper client capabilities
|
||||
mock_session = Mock()
|
||||
mock_session.client_params = InitializeRequestParams(
|
||||
protocolVersion="2025-01-01",
|
||||
capabilities=ClientCapabilities(
|
||||
tasks=ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
|
||||
)
|
||||
)
|
||||
),
|
||||
clientInfo=Implementation(name="test", version="1.0"),
|
||||
)
|
||||
|
||||
ctx = ServerTaskContext(
|
||||
task=task,
|
||||
store=store,
|
||||
session=mock_session,
|
||||
queue=queue,
|
||||
handler=None,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="handler is required for create_message_as_task"):
|
||||
await ctx.create_message_as_task(
|
||||
messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
store.cleanup()
|
||||
406
tests/experimental/tasks/server/test_store.py
Normal file
406
tests/experimental/tasks/server/test_store.py
Normal file
|
|
@ -0,0 +1,406 @@
|
|||
"""Tests for InMemoryTaskStore."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.experimental.tasks.helpers import cancel_task
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.types import INVALID_PARAMS, CallToolResult, TaskMetadata, TextContent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def store() -> AsyncIterator[InMemoryTaskStore]:
|
||||
"""Provide a clean InMemoryTaskStore for each test with automatic cleanup."""
|
||||
store = InMemoryTaskStore()
|
||||
yield store
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_and_get(store: InMemoryTaskStore) -> None:
|
||||
"""Test InMemoryTaskStore create and get operations."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
assert task.taskId is not None
|
||||
assert task.status == "working"
|
||||
assert task.ttl == 60000
|
||||
|
||||
retrieved = await store.get_task(task.taskId)
|
||||
assert retrieved is not None
|
||||
assert retrieved.taskId == task.taskId
|
||||
assert retrieved.status == "working"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_with_custom_id(store: InMemoryTaskStore) -> None:
|
||||
"""Test InMemoryTaskStore create with custom task ID."""
|
||||
task = await store.create_task(
|
||||
metadata=TaskMetadata(ttl=60000),
|
||||
task_id="my-custom-id",
|
||||
)
|
||||
|
||||
assert task.taskId == "my-custom-id"
|
||||
assert task.status == "working"
|
||||
|
||||
retrieved = await store.get_task("my-custom-id")
|
||||
assert retrieved is not None
|
||||
assert retrieved.taskId == "my-custom-id"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_duplicate_id_raises(store: InMemoryTaskStore) -> None:
|
||||
"""Test that creating a task with duplicate ID raises."""
|
||||
await store.create_task(metadata=TaskMetadata(ttl=60000), task_id="duplicate")
|
||||
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
await store.create_task(metadata=TaskMetadata(ttl=60000), task_id="duplicate")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_get_nonexistent_returns_none(store: InMemoryTaskStore) -> None:
|
||||
"""Test that getting a nonexistent task returns None."""
|
||||
retrieved = await store.get_task("nonexistent")
|
||||
assert retrieved is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_update_status(store: InMemoryTaskStore) -> None:
|
||||
"""Test InMemoryTaskStore status updates."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
updated = await store.update_task(task.taskId, status="completed", status_message="All done!")
|
||||
|
||||
assert updated.status == "completed"
|
||||
assert updated.statusMessage == "All done!"
|
||||
|
||||
retrieved = await store.get_task(task.taskId)
|
||||
assert retrieved is not None
|
||||
assert retrieved.status == "completed"
|
||||
assert retrieved.statusMessage == "All done!"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_update_nonexistent_raises(store: InMemoryTaskStore) -> None:
|
||||
"""Test that updating a nonexistent task raises."""
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
await store.update_task("nonexistent", status="completed")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_store_and_get_result(store: InMemoryTaskStore) -> None:
|
||||
"""Test InMemoryTaskStore result storage and retrieval."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
# Store result
|
||||
result = CallToolResult(content=[TextContent(type="text", text="Result data")])
|
||||
await store.store_result(task.taskId, result)
|
||||
|
||||
# Retrieve result
|
||||
retrieved_result = await store.get_result(task.taskId)
|
||||
assert retrieved_result == result
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_get_result_nonexistent_returns_none(store: InMemoryTaskStore) -> None:
|
||||
"""Test that getting result for nonexistent task returns None."""
|
||||
result = await store.get_result("nonexistent")
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_get_result_no_result_returns_none(store: InMemoryTaskStore) -> None:
|
||||
"""Test that getting result when none stored returns None."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
result = await store.get_result(task.taskId)
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_tasks(store: InMemoryTaskStore) -> None:
|
||||
"""Test InMemoryTaskStore list operation."""
|
||||
# Create multiple tasks
|
||||
for _ in range(3):
|
||||
await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
tasks, next_cursor = await store.list_tasks()
|
||||
assert len(tasks) == 3
|
||||
assert next_cursor is None # Less than page size
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_tasks_pagination() -> None:
|
||||
"""Test InMemoryTaskStore pagination."""
|
||||
# Needs custom page_size, can't use fixture
|
||||
store = InMemoryTaskStore(page_size=2)
|
||||
|
||||
# Create 5 tasks
|
||||
for _ in range(5):
|
||||
await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
# First page
|
||||
tasks, next_cursor = await store.list_tasks()
|
||||
assert len(tasks) == 2
|
||||
assert next_cursor is not None
|
||||
|
||||
# Second page
|
||||
tasks, next_cursor = await store.list_tasks(cursor=next_cursor)
|
||||
assert len(tasks) == 2
|
||||
assert next_cursor is not None
|
||||
|
||||
# Third page (last)
|
||||
tasks, next_cursor = await store.list_tasks(cursor=next_cursor)
|
||||
assert len(tasks) == 1
|
||||
assert next_cursor is None
|
||||
|
||||
store.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_tasks_invalid_cursor(store: InMemoryTaskStore) -> None:
|
||||
"""Test that invalid cursor raises."""
|
||||
await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid cursor"):
|
||||
await store.list_tasks(cursor="invalid-cursor")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_delete_task(store: InMemoryTaskStore) -> None:
|
||||
"""Test InMemoryTaskStore delete operation."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
deleted = await store.delete_task(task.taskId)
|
||||
assert deleted is True
|
||||
|
||||
retrieved = await store.get_task(task.taskId)
|
||||
assert retrieved is None
|
||||
|
||||
# Delete non-existent
|
||||
deleted = await store.delete_task(task.taskId)
|
||||
assert deleted is False
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_get_all_tasks_helper(store: InMemoryTaskStore) -> None:
|
||||
"""Test the get_all_tasks debugging helper."""
|
||||
await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
all_tasks = store.get_all_tasks()
|
||||
assert len(all_tasks) == 2
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_store_result_nonexistent_raises(store: InMemoryTaskStore) -> None:
|
||||
"""Test that storing result for nonexistent task raises ValueError."""
|
||||
result = CallToolResult(content=[TextContent(type="text", text="Result")])
|
||||
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
await store.store_result("nonexistent-id", result)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_task_with_null_ttl(store: InMemoryTaskStore) -> None:
|
||||
"""Test creating task with null TTL (never expires)."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=None))
|
||||
|
||||
assert task.ttl is None
|
||||
|
||||
# Task should persist (not expire)
|
||||
retrieved = await store.get_task(task.taskId)
|
||||
assert retrieved is not None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_expiration_cleanup(store: InMemoryTaskStore) -> None:
|
||||
"""Test that expired tasks are cleaned up lazily."""
|
||||
# Create a task with very short TTL
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=1)) # 1ms TTL
|
||||
|
||||
# Manually force the expiry to be in the past
|
||||
stored = store._tasks.get(task.taskId)
|
||||
assert stored is not None
|
||||
stored.expires_at = datetime.now(timezone.utc) - timedelta(seconds=10)
|
||||
|
||||
# Task should still exist in internal dict but be expired
|
||||
assert task.taskId in store._tasks
|
||||
|
||||
# Any access operation should clean up expired tasks
|
||||
# list_tasks triggers cleanup
|
||||
tasks, _ = await store.list_tasks()
|
||||
|
||||
# Expired task should be cleaned up
|
||||
assert task.taskId not in store._tasks
|
||||
assert len(tasks) == 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_task_with_null_ttl_never_expires(store: InMemoryTaskStore) -> None:
|
||||
"""Test that tasks with null TTL never expire during cleanup."""
|
||||
# Create task with null TTL
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=None))
|
||||
|
||||
# Verify internal storage has no expiry
|
||||
stored = store._tasks.get(task.taskId)
|
||||
assert stored is not None
|
||||
assert stored.expires_at is None
|
||||
|
||||
# Access operations should NOT remove this task
|
||||
await store.list_tasks()
|
||||
await store.get_task(task.taskId)
|
||||
|
||||
# Task should still exist
|
||||
assert task.taskId in store._tasks
|
||||
retrieved = await store.get_task(task.taskId)
|
||||
assert retrieved is not None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_terminal_task_ttl_reset(store: InMemoryTaskStore) -> None:
|
||||
"""Test that TTL is reset when task enters terminal state."""
|
||||
# Create task with short TTL
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000)) # 60s
|
||||
|
||||
# Get the initial expiry
|
||||
stored = store._tasks.get(task.taskId)
|
||||
assert stored is not None
|
||||
initial_expiry = stored.expires_at
|
||||
assert initial_expiry is not None
|
||||
|
||||
# Update to terminal state (completed)
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
|
||||
# Expiry should be reset to a new time (from now + TTL)
|
||||
new_expiry = stored.expires_at
|
||||
assert new_expiry is not None
|
||||
assert new_expiry >= initial_expiry
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_terminal_status_transition_rejected(store: InMemoryTaskStore) -> None:
|
||||
"""Test that transitions from terminal states are rejected.
|
||||
|
||||
Per spec: Terminal states (completed, failed, cancelled) MUST NOT
|
||||
transition to any other status.
|
||||
"""
|
||||
# Test each terminal status
|
||||
for terminal_status in ("completed", "failed", "cancelled"):
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
# Move to terminal state
|
||||
await store.update_task(task.taskId, status=terminal_status)
|
||||
|
||||
# Attempting to transition to any other status should raise
|
||||
with pytest.raises(ValueError, match="Cannot transition from terminal status"):
|
||||
await store.update_task(task.taskId, status="working")
|
||||
|
||||
# Also test transitioning to another terminal state
|
||||
other_terminal = "failed" if terminal_status != "failed" else "completed"
|
||||
with pytest.raises(ValueError, match="Cannot transition from terminal status"):
|
||||
await store.update_task(task.taskId, status=other_terminal)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_terminal_status_allows_same_status(store: InMemoryTaskStore) -> None:
|
||||
"""Test that setting the same terminal status doesn't raise.
|
||||
|
||||
This is not a transition, so it should be allowed (no-op).
|
||||
"""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
|
||||
# Setting the same status should not raise
|
||||
updated = await store.update_task(task.taskId, status="completed")
|
||||
assert updated.status == "completed"
|
||||
|
||||
# Updating just the message should also work
|
||||
updated = await store.update_task(task.taskId, status_message="Updated message")
|
||||
assert updated.statusMessage == "Updated message"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_update_nonexistent_raises(store: InMemoryTaskStore) -> None:
|
||||
"""Test that wait_for_update raises for nonexistent task."""
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
await store.wait_for_update("nonexistent-task-id")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancel_task_succeeds_for_working_task(store: InMemoryTaskStore) -> None:
|
||||
"""Test cancel_task helper succeeds for a working task."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
assert task.status == "working"
|
||||
|
||||
result = await cancel_task(store, task.taskId)
|
||||
|
||||
assert result.taskId == task.taskId
|
||||
assert result.status == "cancelled"
|
||||
|
||||
# Verify store is updated
|
||||
retrieved = await store.get_task(task.taskId)
|
||||
assert retrieved is not None
|
||||
assert retrieved.status == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancel_task_rejects_nonexistent_task(store: InMemoryTaskStore) -> None:
|
||||
"""Test cancel_task raises McpError with INVALID_PARAMS for nonexistent task."""
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await cancel_task(store, "nonexistent-task-id")
|
||||
|
||||
assert exc_info.value.error.code == INVALID_PARAMS
|
||||
assert "not found" in exc_info.value.error.message
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancel_task_rejects_completed_task(store: InMemoryTaskStore) -> None:
|
||||
"""Test cancel_task raises McpError with INVALID_PARAMS for completed task."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await cancel_task(store, task.taskId)
|
||||
|
||||
assert exc_info.value.error.code == INVALID_PARAMS
|
||||
assert "terminal state 'completed'" in exc_info.value.error.message
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancel_task_rejects_failed_task(store: InMemoryTaskStore) -> None:
|
||||
"""Test cancel_task raises McpError with INVALID_PARAMS for failed task."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
await store.update_task(task.taskId, status="failed")
|
||||
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await cancel_task(store, task.taskId)
|
||||
|
||||
assert exc_info.value.error.code == INVALID_PARAMS
|
||||
assert "terminal state 'failed'" in exc_info.value.error.message
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancel_task_rejects_already_cancelled_task(store: InMemoryTaskStore) -> None:
|
||||
"""Test cancel_task raises McpError with INVALID_PARAMS for already cancelled task."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
await store.update_task(task.taskId, status="cancelled")
|
||||
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await cancel_task(store, task.taskId)
|
||||
|
||||
assert exc_info.value.error.code == INVALID_PARAMS
|
||||
assert "terminal state 'cancelled'" in exc_info.value.error.message
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancel_task_succeeds_for_input_required_task(store: InMemoryTaskStore) -> None:
|
||||
"""Test cancel_task helper succeeds for a task in input_required status."""
|
||||
task = await store.create_task(metadata=TaskMetadata(ttl=60000))
|
||||
await store.update_task(task.taskId, status="input_required")
|
||||
|
||||
result = await cancel_task(store, task.taskId)
|
||||
|
||||
assert result.taskId == task.taskId
|
||||
assert result.status == "cancelled"
|
||||
354
tests/experimental/tasks/server/test_task_result_handler.py
Normal file
354
tests/experimental/tasks/server/test_task_result_handler.py
Normal file
|
|
@ -0,0 +1,354 @@
|
|||
"""Tests for TaskResultHandler."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from mcp.server.experimental.task_result_handler import TaskResultHandler
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, QueuedMessage
|
||||
from mcp.shared.experimental.tasks.resolver import Resolver
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import (
|
||||
INVALID_REQUEST,
|
||||
CallToolResult,
|
||||
ErrorData,
|
||||
GetTaskPayloadRequest,
|
||||
GetTaskPayloadRequestParams,
|
||||
GetTaskPayloadResult,
|
||||
JSONRPCRequest,
|
||||
TaskMetadata,
|
||||
TextContent,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def store() -> AsyncIterator[InMemoryTaskStore]:
|
||||
"""Provide a clean store for each test."""
|
||||
s = InMemoryTaskStore()
|
||||
yield s
|
||||
s.cleanup()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def queue() -> InMemoryTaskMessageQueue:
|
||||
"""Provide a clean queue for each test."""
|
||||
return InMemoryTaskMessageQueue()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def handler(store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue) -> TaskResultHandler:
|
||||
"""Provide a handler for each test."""
|
||||
return TaskResultHandler(store, queue)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_handle_returns_result_for_completed_task(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that handle() returns the stored result for a completed task."""
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
|
||||
result = CallToolResult(content=[TextContent(type="text", text="Done!")])
|
||||
await store.store_result(task.taskId, result)
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.send_message = AsyncMock()
|
||||
|
||||
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
|
||||
response = await handler.handle(request, mock_session, "req-1")
|
||||
|
||||
assert response is not None
|
||||
assert response.meta is not None
|
||||
assert "io.modelcontextprotocol/related-task" in response.meta
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_handle_raises_for_nonexistent_task(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that handle() raises McpError for nonexistent task."""
|
||||
mock_session = Mock()
|
||||
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId="nonexistent"))
|
||||
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await handler.handle(request, mock_session, "req-1")
|
||||
|
||||
assert "not found" in exc_info.value.error.message
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_handle_returns_empty_result_when_no_result_stored(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that handle() returns minimal result when task completed without stored result."""
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.send_message = AsyncMock()
|
||||
|
||||
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
|
||||
response = await handler.handle(request, mock_session, "req-1")
|
||||
|
||||
assert response is not None
|
||||
assert response.meta is not None
|
||||
assert "io.modelcontextprotocol/related-task" in response.meta
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_handle_delivers_queued_messages(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that handle() delivers queued messages before returning."""
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
|
||||
|
||||
queued_msg = QueuedMessage(
|
||||
type="notification",
|
||||
message=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="notif-1",
|
||||
method="test/notification",
|
||||
params={},
|
||||
),
|
||||
)
|
||||
await queue.enqueue(task.taskId, queued_msg)
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
|
||||
sent_messages: list[SessionMessage] = []
|
||||
|
||||
async def track_send(msg: SessionMessage) -> None:
|
||||
sent_messages.append(msg)
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.send_message = track_send
|
||||
|
||||
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
|
||||
await handler.handle(request, mock_session, "req-1")
|
||||
|
||||
assert len(sent_messages) == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_handle_waits_for_task_completion(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that handle() waits for task to complete before returning."""
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.send_message = AsyncMock()
|
||||
|
||||
request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
|
||||
result_holder: list[GetTaskPayloadResult | None] = [None]
|
||||
|
||||
async def run_handle() -> None:
|
||||
result_holder[0] = await handler.handle(request, mock_session, "req-1")
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(run_handle)
|
||||
|
||||
# Wait for handler to start waiting (event gets created when wait starts)
|
||||
while task.taskId not in store._update_events:
|
||||
await anyio.sleep(0)
|
||||
|
||||
await store.store_result(task.taskId, CallToolResult(content=[TextContent(type="text", text="Done")]))
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
|
||||
assert result_holder[0] is not None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_route_response_resolves_pending_request(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that route_response() resolves a pending request."""
|
||||
resolver: Resolver[dict[str, Any]] = Resolver()
|
||||
handler._pending_requests["req-123"] = resolver
|
||||
|
||||
result = handler.route_response("req-123", {"status": "ok"})
|
||||
|
||||
assert result is True
|
||||
assert resolver.done()
|
||||
assert await resolver.wait() == {"status": "ok"}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_route_response_returns_false_for_unknown_request(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that route_response() returns False for unknown request ID."""
|
||||
result = handler.route_response("unknown-req", {"status": "ok"})
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_route_response_returns_false_for_already_done_resolver(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that route_response() returns False if resolver already completed."""
|
||||
resolver: Resolver[dict[str, Any]] = Resolver()
|
||||
resolver.set_result({"already": "done"})
|
||||
handler._pending_requests["req-123"] = resolver
|
||||
|
||||
result = handler.route_response("req-123", {"new": "data"})
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_route_error_resolves_pending_request_with_exception(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that route_error() sets exception on pending request."""
|
||||
resolver: Resolver[dict[str, Any]] = Resolver()
|
||||
handler._pending_requests["req-123"] = resolver
|
||||
|
||||
error = ErrorData(code=INVALID_REQUEST, message="Something went wrong")
|
||||
result = handler.route_error("req-123", error)
|
||||
|
||||
assert result is True
|
||||
assert resolver.done()
|
||||
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await resolver.wait()
|
||||
assert exc_info.value.error.message == "Something went wrong"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_route_error_returns_false_for_unknown_request(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that route_error() returns False for unknown request ID."""
|
||||
error = ErrorData(code=INVALID_REQUEST, message="Error")
|
||||
result = handler.route_error("unknown-req", error)
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_deliver_registers_resolver_for_request_messages(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that _deliver_queued_messages registers resolvers for request messages."""
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
|
||||
|
||||
resolver: Resolver[dict[str, Any]] = Resolver()
|
||||
queued_msg = QueuedMessage(
|
||||
type="request",
|
||||
message=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="inner-req-1",
|
||||
method="elicitation/create",
|
||||
params={},
|
||||
),
|
||||
resolver=resolver,
|
||||
original_request_id="inner-req-1",
|
||||
)
|
||||
await queue.enqueue(task.taskId, queued_msg)
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.send_message = AsyncMock()
|
||||
|
||||
await handler._deliver_queued_messages(task.taskId, mock_session, "outer-req-1")
|
||||
|
||||
assert "inner-req-1" in handler._pending_requests
|
||||
assert handler._pending_requests["inner-req-1"] is resolver
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_deliver_skips_resolver_registration_when_no_original_id(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that _deliver_queued_messages skips resolver registration when original_request_id is None."""
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
|
||||
|
||||
resolver: Resolver[dict[str, Any]] = Resolver()
|
||||
queued_msg = QueuedMessage(
|
||||
type="request",
|
||||
message=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="inner-req-1",
|
||||
method="elicitation/create",
|
||||
params={},
|
||||
),
|
||||
resolver=resolver,
|
||||
original_request_id=None, # No original request ID
|
||||
)
|
||||
await queue.enqueue(task.taskId, queued_msg)
|
||||
|
||||
mock_session = Mock()
|
||||
mock_session.send_message = AsyncMock()
|
||||
|
||||
await handler._deliver_queued_messages(task.taskId, mock_session, "outer-req-1")
|
||||
|
||||
# Resolver should NOT be registered since original_request_id is None
|
||||
assert len(handler._pending_requests) == 0
|
||||
# But the message should still be sent
|
||||
mock_session.send_message.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_task_update_handles_store_exception(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that _wait_for_task_update handles store exception gracefully."""
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
|
||||
|
||||
# Make wait_for_update raise an exception
|
||||
async def failing_wait(task_id: str) -> None:
|
||||
raise RuntimeError("Store error")
|
||||
|
||||
store.wait_for_update = failing_wait # type: ignore[method-assign]
|
||||
|
||||
# Queue a message to unblock the race via the queue path
|
||||
async def enqueue_later() -> None:
|
||||
# Wait for queue to start waiting (event gets created when wait starts)
|
||||
while task.taskId not in queue._events:
|
||||
await anyio.sleep(0)
|
||||
await queue.enqueue(
|
||||
task.taskId,
|
||||
QueuedMessage(
|
||||
type="notification",
|
||||
message=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id="notif-1",
|
||||
method="test/notification",
|
||||
params={},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(enqueue_later)
|
||||
# This should complete via the queue path even though store raises
|
||||
await handler._wait_for_task_update(task.taskId)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_task_update_handles_queue_exception(
|
||||
store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
|
||||
) -> None:
|
||||
"""Test that _wait_for_task_update handles queue exception gracefully."""
|
||||
task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
|
||||
|
||||
# Make wait_for_message raise an exception
|
||||
async def failing_wait(task_id: str) -> None:
|
||||
raise RuntimeError("Queue error")
|
||||
|
||||
queue.wait_for_message = failing_wait # type: ignore[method-assign]
|
||||
|
||||
# Update the store to unblock the race via the store path
|
||||
async def update_later() -> None:
|
||||
# Wait for store to start waiting (event gets created when wait starts)
|
||||
while task.taskId not in store._update_events:
|
||||
await anyio.sleep(0)
|
||||
await store.update_task(task.taskId, status="completed")
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(update_later)
|
||||
# This should complete via the store path even though queue raises
|
||||
await handler._wait_for_task_update(task.taskId)
|
||||
283
tests/experimental/tasks/test_capabilities.py
Normal file
283
tests/experimental/tasks/test_capabilities.py
Normal file
|
|
@ -0,0 +1,283 @@
|
|||
"""Tests for tasks capability checking utilities."""
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.experimental.tasks.capabilities import (
|
||||
check_tasks_capability,
|
||||
has_task_augmented_elicitation,
|
||||
has_task_augmented_sampling,
|
||||
require_task_augmented_elicitation,
|
||||
require_task_augmented_sampling,
|
||||
)
|
||||
from mcp.types import (
|
||||
ClientCapabilities,
|
||||
ClientTasksCapability,
|
||||
ClientTasksRequestsCapability,
|
||||
TasksCreateElicitationCapability,
|
||||
TasksCreateMessageCapability,
|
||||
TasksElicitationCapability,
|
||||
TasksSamplingCapability,
|
||||
)
|
||||
|
||||
|
||||
class TestCheckTasksCapability:
|
||||
"""Tests for check_tasks_capability function."""
|
||||
|
||||
def test_required_requests_none_returns_true(self) -> None:
|
||||
"""When required.requests is None, should return True."""
|
||||
required = ClientTasksCapability()
|
||||
client = ClientTasksCapability()
|
||||
assert check_tasks_capability(required, client) is True
|
||||
|
||||
def test_client_requests_none_returns_false(self) -> None:
|
||||
"""When client.requests is None but required.requests is set, should return False."""
|
||||
required = ClientTasksCapability(requests=ClientTasksRequestsCapability())
|
||||
client = ClientTasksCapability()
|
||||
assert check_tasks_capability(required, client) is False
|
||||
|
||||
def test_elicitation_required_but_client_missing(self) -> None:
|
||||
"""When elicitation is required but client doesn't have it."""
|
||||
required = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(elicitation=TasksElicitationCapability())
|
||||
)
|
||||
client = ClientTasksCapability(requests=ClientTasksRequestsCapability())
|
||||
assert check_tasks_capability(required, client) is False
|
||||
|
||||
def test_elicitation_create_required_but_client_missing(self) -> None:
|
||||
"""When elicitation.create is required but client doesn't have it."""
|
||||
required = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
|
||||
)
|
||||
)
|
||||
client = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability() # No create
|
||||
)
|
||||
)
|
||||
assert check_tasks_capability(required, client) is False
|
||||
|
||||
def test_elicitation_create_present(self) -> None:
|
||||
"""When elicitation.create is required and client has it."""
|
||||
required = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
|
||||
)
|
||||
)
|
||||
client = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
|
||||
)
|
||||
)
|
||||
assert check_tasks_capability(required, client) is True
|
||||
|
||||
def test_sampling_required_but_client_missing(self) -> None:
|
||||
"""When sampling is required but client doesn't have it."""
|
||||
required = ClientTasksCapability(requests=ClientTasksRequestsCapability(sampling=TasksSamplingCapability()))
|
||||
client = ClientTasksCapability(requests=ClientTasksRequestsCapability())
|
||||
assert check_tasks_capability(required, client) is False
|
||||
|
||||
def test_sampling_create_message_required_but_client_missing(self) -> None:
|
||||
"""When sampling.createMessage is required but client doesn't have it."""
|
||||
required = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
|
||||
)
|
||||
)
|
||||
client = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability() # No createMessage
|
||||
)
|
||||
)
|
||||
assert check_tasks_capability(required, client) is False
|
||||
|
||||
def test_sampling_create_message_present(self) -> None:
|
||||
"""When sampling.createMessage is required and client has it."""
|
||||
required = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
|
||||
)
|
||||
)
|
||||
client = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
|
||||
)
|
||||
)
|
||||
assert check_tasks_capability(required, client) is True
|
||||
|
||||
def test_both_elicitation_and_sampling_present(self) -> None:
|
||||
"""When both elicitation.create and sampling.createMessage are required and client has both."""
|
||||
required = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()),
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()),
|
||||
)
|
||||
)
|
||||
client = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()),
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()),
|
||||
)
|
||||
)
|
||||
assert check_tasks_capability(required, client) is True
|
||||
|
||||
def test_elicitation_without_create_required(self) -> None:
|
||||
"""When elicitation is required but not create specifically."""
|
||||
required = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability() # No create
|
||||
)
|
||||
)
|
||||
client = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
|
||||
)
|
||||
)
|
||||
assert check_tasks_capability(required, client) is True
|
||||
|
||||
def test_sampling_without_create_message_required(self) -> None:
|
||||
"""When sampling is required but not createMessage specifically."""
|
||||
required = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability() # No createMessage
|
||||
)
|
||||
)
|
||||
client = ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
|
||||
)
|
||||
)
|
||||
assert check_tasks_capability(required, client) is True
|
||||
|
||||
|
||||
class TestHasTaskAugmentedElicitation:
|
||||
"""Tests for has_task_augmented_elicitation function."""
|
||||
|
||||
def test_tasks_none(self) -> None:
|
||||
"""Returns False when caps.tasks is None."""
|
||||
caps = ClientCapabilities()
|
||||
assert has_task_augmented_elicitation(caps) is False
|
||||
|
||||
def test_requests_none(self) -> None:
|
||||
"""Returns False when caps.tasks.requests is None."""
|
||||
caps = ClientCapabilities(tasks=ClientTasksCapability())
|
||||
assert has_task_augmented_elicitation(caps) is False
|
||||
|
||||
def test_elicitation_none(self) -> None:
|
||||
"""Returns False when caps.tasks.requests.elicitation is None."""
|
||||
caps = ClientCapabilities(tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability()))
|
||||
assert has_task_augmented_elicitation(caps) is False
|
||||
|
||||
def test_create_none(self) -> None:
|
||||
"""Returns False when caps.tasks.requests.elicitation.create is None."""
|
||||
caps = ClientCapabilities(
|
||||
tasks=ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(elicitation=TasksElicitationCapability())
|
||||
)
|
||||
)
|
||||
assert has_task_augmented_elicitation(caps) is False
|
||||
|
||||
def test_create_present(self) -> None:
|
||||
"""Returns True when full capability path is present."""
|
||||
caps = ClientCapabilities(
|
||||
tasks=ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
|
||||
)
|
||||
)
|
||||
)
|
||||
assert has_task_augmented_elicitation(caps) is True
|
||||
|
||||
|
||||
class TestHasTaskAugmentedSampling:
|
||||
"""Tests for has_task_augmented_sampling function."""
|
||||
|
||||
def test_tasks_none(self) -> None:
|
||||
"""Returns False when caps.tasks is None."""
|
||||
caps = ClientCapabilities()
|
||||
assert has_task_augmented_sampling(caps) is False
|
||||
|
||||
def test_requests_none(self) -> None:
|
||||
"""Returns False when caps.tasks.requests is None."""
|
||||
caps = ClientCapabilities(tasks=ClientTasksCapability())
|
||||
assert has_task_augmented_sampling(caps) is False
|
||||
|
||||
def test_sampling_none(self) -> None:
|
||||
"""Returns False when caps.tasks.requests.sampling is None."""
|
||||
caps = ClientCapabilities(tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability()))
|
||||
assert has_task_augmented_sampling(caps) is False
|
||||
|
||||
def test_create_message_none(self) -> None:
|
||||
"""Returns False when caps.tasks.requests.sampling.createMessage is None."""
|
||||
caps = ClientCapabilities(
|
||||
tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability(sampling=TasksSamplingCapability()))
|
||||
)
|
||||
assert has_task_augmented_sampling(caps) is False
|
||||
|
||||
def test_create_message_present(self) -> None:
|
||||
"""Returns True when full capability path is present."""
|
||||
caps = ClientCapabilities(
|
||||
tasks=ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
|
||||
)
|
||||
)
|
||||
)
|
||||
assert has_task_augmented_sampling(caps) is True
|
||||
|
||||
|
||||
class TestRequireTaskAugmentedElicitation:
|
||||
"""Tests for require_task_augmented_elicitation function."""
|
||||
|
||||
def test_raises_when_none(self) -> None:
|
||||
"""Raises McpError when client_caps is None."""
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
require_task_augmented_elicitation(None)
|
||||
assert "task-augmented elicitation" in str(exc_info.value)
|
||||
|
||||
def test_raises_when_missing(self) -> None:
|
||||
"""Raises McpError when capability is missing."""
|
||||
caps = ClientCapabilities()
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
require_task_augmented_elicitation(caps)
|
||||
assert "task-augmented elicitation" in str(exc_info.value)
|
||||
|
||||
def test_passes_when_present(self) -> None:
|
||||
"""Does not raise when capability is present."""
|
||||
caps = ClientCapabilities(
|
||||
tasks=ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
|
||||
)
|
||||
)
|
||||
)
|
||||
require_task_augmented_elicitation(caps)
|
||||
|
||||
|
||||
class TestRequireTaskAugmentedSampling:
|
||||
"""Tests for require_task_augmented_sampling function."""
|
||||
|
||||
def test_raises_when_none(self) -> None:
|
||||
"""Raises McpError when client_caps is None."""
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
require_task_augmented_sampling(None)
|
||||
assert "task-augmented sampling" in str(exc_info.value)
|
||||
|
||||
def test_raises_when_missing(self) -> None:
|
||||
"""Raises McpError when capability is missing."""
|
||||
caps = ClientCapabilities()
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
require_task_augmented_sampling(caps)
|
||||
assert "task-augmented sampling" in str(exc_info.value)
|
||||
|
||||
def test_passes_when_present(self) -> None:
|
||||
"""Does not raise when capability is present."""
|
||||
caps = ClientCapabilities(
|
||||
tasks=ClientTasksCapability(
|
||||
requests=ClientTasksRequestsCapability(
|
||||
sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
|
||||
)
|
||||
)
|
||||
)
|
||||
require_task_augmented_sampling(caps)
|
||||
737
tests/experimental/tasks/test_elicitation_scenarios.py
Normal file
737
tests/experimental/tasks/test_elicitation_scenarios.py
Normal 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()
|
||||
331
tests/experimental/tasks/test_message_queue.py
Normal file
331
tests/experimental/tasks/test_message_queue.py
Normal file
|
|
@ -0,0 +1,331 @@
|
|||
"""
|
||||
Tests for TaskMessageQueue and InMemoryTaskMessageQueue.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, QueuedMessage
|
||||
from mcp.shared.experimental.tasks.resolver import Resolver
|
||||
from mcp.types import JSONRPCNotification, JSONRPCRequest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def queue() -> InMemoryTaskMessageQueue:
|
||||
return InMemoryTaskMessageQueue()
|
||||
|
||||
|
||||
def make_request(id: int = 1, method: str = "test/method") -> JSONRPCRequest:
|
||||
return JSONRPCRequest(jsonrpc="2.0", id=id, method=method)
|
||||
|
||||
|
||||
def make_notification(method: str = "test/notify") -> JSONRPCNotification:
|
||||
return JSONRPCNotification(jsonrpc="2.0", method=method)
|
||||
|
||||
|
||||
class TestInMemoryTaskMessageQueue:
|
||||
@pytest.mark.anyio
|
||||
async def test_enqueue_and_dequeue(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Test basic enqueue and dequeue operations."""
|
||||
task_id = "task-1"
|
||||
msg = QueuedMessage(type="request", message=make_request())
|
||||
|
||||
await queue.enqueue(task_id, msg)
|
||||
result = await queue.dequeue(task_id)
|
||||
|
||||
assert result is not None
|
||||
assert result.type == "request"
|
||||
assert result.message.method == "test/method"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_dequeue_empty_returns_none(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Dequeue from empty queue returns None."""
|
||||
result = await queue.dequeue("nonexistent-task")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_fifo_ordering(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Messages are dequeued in FIFO order."""
|
||||
task_id = "task-1"
|
||||
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(1, "first")))
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(2, "second")))
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(3, "third")))
|
||||
|
||||
msg1 = await queue.dequeue(task_id)
|
||||
msg2 = await queue.dequeue(task_id)
|
||||
msg3 = await queue.dequeue(task_id)
|
||||
|
||||
assert msg1 is not None and msg1.message.method == "first"
|
||||
assert msg2 is not None and msg2.message.method == "second"
|
||||
assert msg3 is not None and msg3.message.method == "third"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_separate_queues_per_task(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Each task has its own queue."""
|
||||
await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1, "task1-msg")))
|
||||
await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2, "task2-msg")))
|
||||
|
||||
msg1 = await queue.dequeue("task-1")
|
||||
msg2 = await queue.dequeue("task-2")
|
||||
|
||||
assert msg1 is not None and msg1.message.method == "task1-msg"
|
||||
assert msg2 is not None and msg2.message.method == "task2-msg"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_peek_does_not_remove(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Peek returns message without removing it."""
|
||||
task_id = "task-1"
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
|
||||
|
||||
peeked = await queue.peek(task_id)
|
||||
dequeued = await queue.dequeue(task_id)
|
||||
|
||||
assert peeked is not None
|
||||
assert dequeued is not None
|
||||
assert isinstance(peeked.message, JSONRPCRequest)
|
||||
assert isinstance(dequeued.message, JSONRPCRequest)
|
||||
assert peeked.message.id == dequeued.message.id
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_is_empty(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Test is_empty method."""
|
||||
task_id = "task-1"
|
||||
|
||||
assert await queue.is_empty(task_id) is True
|
||||
|
||||
await queue.enqueue(task_id, QueuedMessage(type="notification", message=make_notification()))
|
||||
assert await queue.is_empty(task_id) is False
|
||||
|
||||
await queue.dequeue(task_id)
|
||||
assert await queue.is_empty(task_id) is True
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_clear_returns_all_messages(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Clear removes and returns all messages."""
|
||||
task_id = "task-1"
|
||||
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(1)))
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(2)))
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(3)))
|
||||
|
||||
messages = await queue.clear(task_id)
|
||||
|
||||
assert len(messages) == 3
|
||||
assert await queue.is_empty(task_id) is True
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_clear_empty_queue(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Clear on empty queue returns empty list."""
|
||||
messages = await queue.clear("nonexistent")
|
||||
assert messages == []
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_notification_messages(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Test queuing notification messages."""
|
||||
task_id = "task-1"
|
||||
msg = QueuedMessage(type="notification", message=make_notification("log/message"))
|
||||
|
||||
await queue.enqueue(task_id, msg)
|
||||
result = await queue.dequeue(task_id)
|
||||
|
||||
assert result is not None
|
||||
assert result.type == "notification"
|
||||
assert result.message.method == "log/message"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_message_timestamp(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Messages have timestamps."""
|
||||
before = datetime.now(timezone.utc)
|
||||
msg = QueuedMessage(type="request", message=make_request())
|
||||
after = datetime.now(timezone.utc)
|
||||
|
||||
assert before <= msg.timestamp <= after
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_message_with_resolver(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Messages can have resolvers."""
|
||||
task_id = "task-1"
|
||||
resolver: Resolver[dict[str, str]] = Resolver()
|
||||
|
||||
msg = QueuedMessage(
|
||||
type="request",
|
||||
message=make_request(),
|
||||
resolver=resolver,
|
||||
original_request_id=42,
|
||||
)
|
||||
|
||||
await queue.enqueue(task_id, msg)
|
||||
result = await queue.dequeue(task_id)
|
||||
|
||||
assert result is not None
|
||||
assert result.resolver is resolver
|
||||
assert result.original_request_id == 42
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cleanup_specific_task(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Cleanup removes specific task's data."""
|
||||
await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1)))
|
||||
await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2)))
|
||||
|
||||
queue.cleanup("task-1")
|
||||
|
||||
assert await queue.is_empty("task-1") is True
|
||||
assert await queue.is_empty("task-2") is False
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cleanup_all(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Cleanup without task_id removes all data."""
|
||||
await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1)))
|
||||
await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2)))
|
||||
|
||||
queue.cleanup()
|
||||
|
||||
assert await queue.is_empty("task-1") is True
|
||||
assert await queue.is_empty("task-2") is True
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_message_returns_immediately_if_message_exists(
|
||||
self, queue: InMemoryTaskMessageQueue
|
||||
) -> None:
|
||||
"""wait_for_message returns immediately if queue not empty."""
|
||||
task_id = "task-1"
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
|
||||
|
||||
# Should return immediately, not block
|
||||
with anyio.fail_after(1):
|
||||
await queue.wait_for_message(task_id)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_message_blocks_until_message(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""wait_for_message blocks until a message is enqueued."""
|
||||
task_id = "task-1"
|
||||
received = False
|
||||
waiter_started = anyio.Event()
|
||||
|
||||
async def enqueue_when_ready() -> None:
|
||||
# Wait until the waiter has started before enqueueing
|
||||
await waiter_started.wait()
|
||||
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
|
||||
|
||||
async def wait_for_msg() -> None:
|
||||
nonlocal received
|
||||
# Signal that we're about to start waiting
|
||||
waiter_started.set()
|
||||
await queue.wait_for_message(task_id)
|
||||
received = True
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(wait_for_msg)
|
||||
tg.start_soon(enqueue_when_ready)
|
||||
|
||||
assert received is True
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_notify_message_available_wakes_waiter(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""notify_message_available wakes up waiting coroutines."""
|
||||
task_id = "task-1"
|
||||
notified = False
|
||||
waiter_started = anyio.Event()
|
||||
|
||||
async def notify_when_ready() -> None:
|
||||
# Wait until the waiter has started before notifying
|
||||
await waiter_started.wait()
|
||||
await queue.notify_message_available(task_id)
|
||||
|
||||
async def wait_for_notification() -> None:
|
||||
nonlocal notified
|
||||
# Signal that we're about to start waiting
|
||||
waiter_started.set()
|
||||
await queue.wait_for_message(task_id)
|
||||
notified = True
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(wait_for_notification)
|
||||
tg.start_soon(notify_when_ready)
|
||||
|
||||
assert notified is True
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_peek_empty_queue_returns_none(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""Peek on empty queue returns None."""
|
||||
result = await queue.peek("nonexistent-task")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_message_double_check_race_condition(self, queue: InMemoryTaskMessageQueue) -> None:
|
||||
"""wait_for_message returns early if message arrives after event creation but before wait."""
|
||||
task_id = "task-1"
|
||||
|
||||
# To test the double-check path (lines 223-225), we need a message to arrive
|
||||
# after the event is created (line 220) but before event.wait() (line 228).
|
||||
# We simulate this by injecting a message before is_empty is called the second time.
|
||||
|
||||
original_is_empty = queue.is_empty
|
||||
call_count = 0
|
||||
|
||||
async def is_empty_with_injection(tid: str) -> bool:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 2 and tid == task_id:
|
||||
# Before second check, inject a message - this simulates a message
|
||||
# arriving between event creation and the double-check
|
||||
queue._queues[task_id] = [QueuedMessage(type="request", message=make_request())]
|
||||
return await original_is_empty(tid)
|
||||
|
||||
queue.is_empty = is_empty_with_injection # type: ignore[method-assign]
|
||||
|
||||
# Should return immediately due to double-check finding the message
|
||||
with anyio.fail_after(1):
|
||||
await queue.wait_for_message(task_id)
|
||||
|
||||
|
||||
class TestResolver:
|
||||
@pytest.mark.anyio
|
||||
async def test_set_result_and_wait(self) -> None:
|
||||
"""Test basic set_result and wait flow."""
|
||||
resolver: Resolver[str] = Resolver()
|
||||
|
||||
resolver.set_result("hello")
|
||||
result = await resolver.wait()
|
||||
|
||||
assert result == "hello"
|
||||
assert resolver.done()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_set_exception_and_wait(self) -> None:
|
||||
"""Test set_exception raises on wait."""
|
||||
resolver: Resolver[str] = Resolver()
|
||||
|
||||
resolver.set_exception(ValueError("test error"))
|
||||
|
||||
with pytest.raises(ValueError, match="test error"):
|
||||
await resolver.wait()
|
||||
|
||||
assert resolver.done()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_set_result_when_already_completed_raises(self) -> None:
|
||||
"""Test that set_result raises if resolver already completed."""
|
||||
resolver: Resolver[str] = Resolver()
|
||||
resolver.set_result("first")
|
||||
|
||||
with pytest.raises(RuntimeError, match="already completed"):
|
||||
resolver.set_result("second")
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_set_exception_when_already_completed_raises(self) -> None:
|
||||
"""Test that set_exception raises if resolver already completed."""
|
||||
resolver: Resolver[str] = Resolver()
|
||||
resolver.set_result("done")
|
||||
|
||||
with pytest.raises(RuntimeError, match="already completed"):
|
||||
resolver.set_exception(ValueError("too late"))
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_done_returns_false_before_completion(self) -> None:
|
||||
"""Test done() returns False before any result is set."""
|
||||
resolver: Resolver[str] = Resolver()
|
||||
assert resolver.done() is False
|
||||
166
tests/experimental/tasks/test_request_context.py
Normal file
166
tests/experimental/tasks/test_request_context.py
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
"""Tests for the RequestContext.experimental (Experimental class) task validation helpers."""
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp.server.experimental.request_context import Experimental
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.types import (
|
||||
METHOD_NOT_FOUND,
|
||||
TASK_FORBIDDEN,
|
||||
TASK_OPTIONAL,
|
||||
TASK_REQUIRED,
|
||||
ClientCapabilities,
|
||||
ClientTasksCapability,
|
||||
TaskMetadata,
|
||||
Tool,
|
||||
ToolExecution,
|
||||
)
|
||||
|
||||
|
||||
def test_is_task_true_when_metadata_present() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
assert exp.is_task is True
|
||||
|
||||
|
||||
def test_is_task_false_when_no_metadata() -> None:
|
||||
exp = Experimental(task_metadata=None)
|
||||
assert exp.is_task is False
|
||||
|
||||
|
||||
def test_client_supports_tasks_true() -> None:
|
||||
exp = Experimental(_client_capabilities=ClientCapabilities(tasks=ClientTasksCapability()))
|
||||
assert exp.client_supports_tasks is True
|
||||
|
||||
|
||||
def test_client_supports_tasks_false_no_tasks() -> None:
|
||||
exp = Experimental(_client_capabilities=ClientCapabilities())
|
||||
assert exp.client_supports_tasks is False
|
||||
|
||||
|
||||
def test_client_supports_tasks_false_no_capabilities() -> None:
|
||||
exp = Experimental(_client_capabilities=None)
|
||||
assert exp.client_supports_tasks is False
|
||||
|
||||
|
||||
def test_validate_task_mode_required_with_task_is_valid() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
error = exp.validate_task_mode(TASK_REQUIRED, raise_error=False)
|
||||
assert error is None
|
||||
|
||||
|
||||
def test_validate_task_mode_required_without_task_returns_error() -> None:
|
||||
exp = Experimental(task_metadata=None)
|
||||
error = exp.validate_task_mode(TASK_REQUIRED, raise_error=False)
|
||||
assert error is not None
|
||||
assert error.code == METHOD_NOT_FOUND
|
||||
assert "requires task-augmented" in error.message
|
||||
|
||||
|
||||
def test_validate_task_mode_required_without_task_raises_by_default() -> None:
|
||||
exp = Experimental(task_metadata=None)
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
exp.validate_task_mode(TASK_REQUIRED)
|
||||
assert exc_info.value.error.code == METHOD_NOT_FOUND
|
||||
|
||||
|
||||
def test_validate_task_mode_forbidden_without_task_is_valid() -> None:
|
||||
exp = Experimental(task_metadata=None)
|
||||
error = exp.validate_task_mode(TASK_FORBIDDEN, raise_error=False)
|
||||
assert error is None
|
||||
|
||||
|
||||
def test_validate_task_mode_forbidden_with_task_returns_error() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
error = exp.validate_task_mode(TASK_FORBIDDEN, raise_error=False)
|
||||
assert error is not None
|
||||
assert error.code == METHOD_NOT_FOUND
|
||||
assert "does not support task-augmented" in error.message
|
||||
|
||||
|
||||
def test_validate_task_mode_forbidden_with_task_raises_by_default() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
exp.validate_task_mode(TASK_FORBIDDEN)
|
||||
assert exc_info.value.error.code == METHOD_NOT_FOUND
|
||||
|
||||
|
||||
def test_validate_task_mode_none_treated_as_forbidden() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
error = exp.validate_task_mode(None, raise_error=False)
|
||||
assert error is not None
|
||||
assert "does not support task-augmented" in error.message
|
||||
|
||||
|
||||
def test_validate_task_mode_optional_with_task_is_valid() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
error = exp.validate_task_mode(TASK_OPTIONAL, raise_error=False)
|
||||
assert error is None
|
||||
|
||||
|
||||
def test_validate_task_mode_optional_without_task_is_valid() -> None:
|
||||
exp = Experimental(task_metadata=None)
|
||||
error = exp.validate_task_mode(TASK_OPTIONAL, raise_error=False)
|
||||
assert error is None
|
||||
|
||||
|
||||
def test_validate_for_tool_with_execution_required() -> None:
|
||||
exp = Experimental(task_metadata=None)
|
||||
tool = Tool(
|
||||
name="test",
|
||||
description="test",
|
||||
inputSchema={"type": "object"},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
)
|
||||
error = exp.validate_for_tool(tool, raise_error=False)
|
||||
assert error is not None
|
||||
assert "requires task-augmented" in error.message
|
||||
|
||||
|
||||
def test_validate_for_tool_without_execution() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
tool = Tool(
|
||||
name="test",
|
||||
description="test",
|
||||
inputSchema={"type": "object"},
|
||||
execution=None,
|
||||
)
|
||||
error = exp.validate_for_tool(tool, raise_error=False)
|
||||
assert error is not None
|
||||
assert "does not support task-augmented" in error.message
|
||||
|
||||
|
||||
def test_validate_for_tool_optional_with_task() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
tool = Tool(
|
||||
name="test",
|
||||
description="test",
|
||||
inputSchema={"type": "object"},
|
||||
execution=ToolExecution(taskSupport=TASK_OPTIONAL),
|
||||
)
|
||||
error = exp.validate_for_tool(tool, raise_error=False)
|
||||
assert error is None
|
||||
|
||||
|
||||
def test_can_use_tool_required_with_task_support() -> None:
|
||||
exp = Experimental(_client_capabilities=ClientCapabilities(tasks=ClientTasksCapability()))
|
||||
assert exp.can_use_tool(TASK_REQUIRED) is True
|
||||
|
||||
|
||||
def test_can_use_tool_required_without_task_support() -> None:
|
||||
exp = Experimental(_client_capabilities=ClientCapabilities())
|
||||
assert exp.can_use_tool(TASK_REQUIRED) is False
|
||||
|
||||
|
||||
def test_can_use_tool_optional_without_task_support() -> None:
|
||||
exp = Experimental(_client_capabilities=ClientCapabilities())
|
||||
assert exp.can_use_tool(TASK_OPTIONAL) is True
|
||||
|
||||
|
||||
def test_can_use_tool_forbidden_without_task_support() -> None:
|
||||
exp = Experimental(_client_capabilities=ClientCapabilities())
|
||||
assert exp.can_use_tool(TASK_FORBIDDEN) is True
|
||||
|
||||
|
||||
def test_can_use_tool_none_without_task_support() -> None:
|
||||
exp = Experimental(_client_capabilities=ClientCapabilities())
|
||||
assert exp.can_use_tool(None) is True
|
||||
753
tests/experimental/tasks/test_spec_compliance.py
Normal file
753
tests/experimental/tasks/test_spec_compliance.py
Normal file
|
|
@ -0,0 +1,753 @@
|
|||
"""
|
||||
Tasks Spec Compliance Tests
|
||||
===========================
|
||||
|
||||
Test structure mirrors: https://modelcontextprotocol.io/specification/draft/basic/utilities/tasks.md
|
||||
|
||||
Each section contains tests for normative requirements (MUST/SHOULD/MAY).
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp.server import Server
|
||||
from mcp.server.lowlevel import NotificationOptions
|
||||
from mcp.shared.experimental.tasks.helpers import MODEL_IMMEDIATE_RESPONSE_KEY
|
||||
from mcp.types import (
|
||||
CancelTaskRequest,
|
||||
CancelTaskResult,
|
||||
CreateTaskResult,
|
||||
GetTaskRequest,
|
||||
GetTaskResult,
|
||||
ListTasksRequest,
|
||||
ListTasksResult,
|
||||
ServerCapabilities,
|
||||
Task,
|
||||
)
|
||||
|
||||
# Shared test datetime
|
||||
TEST_DATETIME = datetime(2025, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _get_capabilities(server: Server) -> ServerCapabilities:
|
||||
"""Helper to get capabilities from a server."""
|
||||
return server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
)
|
||||
|
||||
|
||||
def test_server_without_task_handlers_has_no_tasks_capability() -> None:
|
||||
"""Server without any task handlers has no tasks capability."""
|
||||
server: Server = Server("test")
|
||||
caps = _get_capabilities(server)
|
||||
assert caps.tasks is None
|
||||
|
||||
|
||||
def test_server_with_list_tasks_handler_declares_list_capability() -> None:
|
||||
"""Server with list_tasks handler declares tasks.list capability."""
|
||||
server: Server = Server("test")
|
||||
|
||||
@server.experimental.list_tasks()
|
||||
async def handle_list(req: ListTasksRequest) -> ListTasksResult:
|
||||
raise NotImplementedError
|
||||
|
||||
caps = _get_capabilities(server)
|
||||
assert caps.tasks is not None
|
||||
assert caps.tasks.list is not None
|
||||
|
||||
|
||||
def test_server_with_cancel_task_handler_declares_cancel_capability() -> None:
|
||||
"""Server with cancel_task handler declares tasks.cancel capability."""
|
||||
server: Server = Server("test")
|
||||
|
||||
@server.experimental.cancel_task()
|
||||
async def handle_cancel(req: CancelTaskRequest) -> CancelTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
caps = _get_capabilities(server)
|
||||
assert caps.tasks is not None
|
||||
assert caps.tasks.cancel is not None
|
||||
|
||||
|
||||
def test_server_with_get_task_handler_declares_requests_tools_call_capability() -> None:
|
||||
"""
|
||||
Server with get_task handler declares tasks.requests.tools.call capability.
|
||||
(get_task is required for task-augmented tools/call support)
|
||||
"""
|
||||
server: Server = Server("test")
|
||||
|
||||
@server.experimental.get_task()
|
||||
async def handle_get(req: GetTaskRequest) -> GetTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
caps = _get_capabilities(server)
|
||||
assert caps.tasks is not None
|
||||
assert caps.tasks.requests is not None
|
||||
assert caps.tasks.requests.tools is not None
|
||||
|
||||
|
||||
def test_server_without_list_handler_has_no_list_capability() -> None:
|
||||
"""Server without list_tasks handler has no tasks.list capability."""
|
||||
server: Server = Server("test")
|
||||
|
||||
# Register only get_task (not list_tasks)
|
||||
@server.experimental.get_task()
|
||||
async def handle_get(req: GetTaskRequest) -> GetTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
caps = _get_capabilities(server)
|
||||
assert caps.tasks is not None
|
||||
assert caps.tasks.list is None
|
||||
|
||||
|
||||
def test_server_without_cancel_handler_has_no_cancel_capability() -> None:
|
||||
"""Server without cancel_task handler has no tasks.cancel capability."""
|
||||
server: Server = Server("test")
|
||||
|
||||
# Register only get_task (not cancel_task)
|
||||
@server.experimental.get_task()
|
||||
async def handle_get(req: GetTaskRequest) -> GetTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
caps = _get_capabilities(server)
|
||||
assert caps.tasks is not None
|
||||
assert caps.tasks.cancel is None
|
||||
|
||||
|
||||
def test_server_with_all_task_handlers_has_full_capability() -> None:
|
||||
"""Server with all task handlers declares complete tasks capability."""
|
||||
server: Server = Server("test")
|
||||
|
||||
@server.experimental.list_tasks()
|
||||
async def handle_list(req: ListTasksRequest) -> ListTasksResult:
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.cancel_task()
|
||||
async def handle_cancel(req: CancelTaskRequest) -> CancelTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
@server.experimental.get_task()
|
||||
async def handle_get(req: GetTaskRequest) -> GetTaskResult:
|
||||
raise NotImplementedError
|
||||
|
||||
caps = _get_capabilities(server)
|
||||
assert caps.tasks is not None
|
||||
assert caps.tasks.list is not None
|
||||
assert caps.tasks.cancel is not None
|
||||
assert caps.tasks.requests is not None
|
||||
assert caps.tasks.requests.tools is not None
|
||||
|
||||
|
||||
class TestClientCapabilities:
|
||||
"""
|
||||
Clients declare:
|
||||
- tasks.list — supports listing operations
|
||||
- tasks.cancel — supports cancellation
|
||||
- tasks.requests.sampling.createMessage — task-augmented sampling
|
||||
- tasks.requests.elicitation.create — task-augmented elicitation
|
||||
"""
|
||||
|
||||
def test_client_declares_tasks_capability(self) -> None:
|
||||
"""Client can declare tasks capability."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestToolLevelNegotiation:
|
||||
"""
|
||||
Tools in tools/list responses include execution.taskSupport with values:
|
||||
- Not present or "forbidden": No task augmentation allowed
|
||||
- "optional": Task augmentation allowed at requestor discretion
|
||||
- "required": Task augmentation is mandatory
|
||||
"""
|
||||
|
||||
def test_tool_execution_task_forbidden_rejects_task_augmented_call(self) -> None:
|
||||
"""Tool with execution.taskSupport="forbidden" MUST reject task-augmented calls (-32601)."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tool_execution_task_absent_rejects_task_augmented_call(self) -> None:
|
||||
"""Tool without execution.taskSupport MUST reject task-augmented calls (-32601)."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tool_execution_task_optional_accepts_normal_call(self) -> None:
|
||||
"""Tool with execution.taskSupport="optional" accepts normal calls."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tool_execution_task_optional_accepts_task_augmented_call(self) -> None:
|
||||
"""Tool with execution.taskSupport="optional" accepts task-augmented calls."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tool_execution_task_required_rejects_normal_call(self) -> None:
|
||||
"""Tool with execution.taskSupport="required" MUST reject non-task calls (-32601)."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tool_execution_task_required_accepts_task_augmented_call(self) -> None:
|
||||
"""Tool with execution.taskSupport="required" accepts task-augmented calls."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestCapabilityNegotiation:
|
||||
"""
|
||||
Requestors SHOULD only augment requests with a task if the corresponding
|
||||
capability has been declared by the receiver.
|
||||
|
||||
Receivers that do not declare the task capability for a request type
|
||||
MUST process requests of that type normally, ignoring any task-augmentation
|
||||
metadata if present.
|
||||
"""
|
||||
|
||||
def test_receiver_without_capability_ignores_task_metadata(self) -> None:
|
||||
"""
|
||||
Receiver without task capability MUST process request normally,
|
||||
ignoring task-augmentation metadata.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_receiver_with_capability_may_require_task_augmentation(self) -> None:
|
||||
"""
|
||||
Receivers that declare task capability MAY return error (-32600)
|
||||
for non-task-augmented requests, requiring task augmentation.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestTaskStatusLifecycle:
|
||||
"""
|
||||
Tasks begin in working status and follow valid transitions:
|
||||
working → input_required → working → terminal
|
||||
working → terminal (directly)
|
||||
input_required → terminal (directly)
|
||||
|
||||
Terminal states (no further transitions allowed):
|
||||
- completed
|
||||
- failed
|
||||
- cancelled
|
||||
"""
|
||||
|
||||
def test_task_begins_in_working_status(self) -> None:
|
||||
"""Tasks MUST begin in working status."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_working_to_completed_transition(self) -> None:
|
||||
"""working → completed is valid."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_working_to_failed_transition(self) -> None:
|
||||
"""working → failed is valid."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_working_to_cancelled_transition(self) -> None:
|
||||
"""working → cancelled is valid."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_working_to_input_required_transition(self) -> None:
|
||||
"""working → input_required is valid."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_input_required_to_working_transition(self) -> None:
|
||||
"""input_required → working is valid."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_input_required_to_terminal_transition(self) -> None:
|
||||
"""input_required → terminal is valid."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_terminal_state_no_further_transitions(self) -> None:
|
||||
"""Terminal states allow no further transitions."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_completed_is_terminal(self) -> None:
|
||||
"""completed is a terminal state."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_failed_is_terminal(self) -> None:
|
||||
"""failed is a terminal state."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_cancelled_is_terminal(self) -> None:
|
||||
"""cancelled is a terminal state."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestInputRequiredStatus:
|
||||
"""
|
||||
When a receiver needs information to proceed, it moves the task to input_required.
|
||||
The requestor should call tasks/result to retrieve input requests.
|
||||
The task must include io.modelcontextprotocol/related-task metadata in associated requests.
|
||||
"""
|
||||
|
||||
def test_input_required_status_retrievable_via_tasks_get(self) -> None:
|
||||
"""Task in input_required status is retrievable via tasks/get."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_input_required_related_task_metadata_in_requests(self) -> None:
|
||||
"""
|
||||
Task MUST include io.modelcontextprotocol/related-task metadata
|
||||
in associated requests.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestCreatingTask:
|
||||
"""
|
||||
Request structure:
|
||||
{"method": "tools/call", "params": {"name": "...", "arguments": {...}, "task": {"ttl": 60000}}}
|
||||
|
||||
Response (CreateTaskResult):
|
||||
{"result": {"task": {"taskId": "...", "status": "working", ...}}}
|
||||
|
||||
Receivers may include io.modelcontextprotocol/model-immediate-response in _meta.
|
||||
"""
|
||||
|
||||
def test_task_augmented_request_returns_create_task_result(self) -> None:
|
||||
"""Task-augmented request MUST return CreateTaskResult immediately."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_create_task_result_contains_task_id(self) -> None:
|
||||
"""CreateTaskResult MUST contain taskId."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_create_task_result_contains_status_working(self) -> None:
|
||||
"""CreateTaskResult MUST have status=working initially."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_create_task_result_contains_created_at(self) -> None:
|
||||
"""CreateTaskResult MUST contain createdAt timestamp."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_create_task_result_created_at_is_iso8601(self) -> None:
|
||||
"""createdAt MUST be ISO 8601 formatted."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_create_task_result_may_contain_ttl(self) -> None:
|
||||
"""CreateTaskResult MAY contain ttl."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_create_task_result_may_contain_poll_interval(self) -> None:
|
||||
"""CreateTaskResult MAY contain pollInterval."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_create_task_result_may_contain_status_message(self) -> None:
|
||||
"""CreateTaskResult MAY contain statusMessage."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_receiver_may_override_requested_ttl(self) -> None:
|
||||
"""Receiver MAY override requested ttl but MUST return actual value."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_model_immediate_response_in_meta(self) -> None:
|
||||
"""
|
||||
Receiver MAY include io.modelcontextprotocol/model-immediate-response
|
||||
in _meta to provide immediate response while task executes.
|
||||
"""
|
||||
# Verify the constant has the correct value per spec
|
||||
assert MODEL_IMMEDIATE_RESPONSE_KEY == "io.modelcontextprotocol/model-immediate-response"
|
||||
|
||||
# CreateTaskResult can include model-immediate-response in _meta
|
||||
task = Task(
|
||||
taskId="test-123",
|
||||
status="working",
|
||||
createdAt=TEST_DATETIME,
|
||||
lastUpdatedAt=TEST_DATETIME,
|
||||
ttl=60000,
|
||||
)
|
||||
immediate_msg = "Task started, processing your request..."
|
||||
# Note: Must use _meta= (alias) not meta= due to Pydantic alias handling
|
||||
result = CreateTaskResult(
|
||||
task=task,
|
||||
**{"_meta": {MODEL_IMMEDIATE_RESPONSE_KEY: immediate_msg}},
|
||||
)
|
||||
|
||||
# Verify the metadata is present and correct
|
||||
assert result.meta is not None
|
||||
assert MODEL_IMMEDIATE_RESPONSE_KEY in result.meta
|
||||
assert result.meta[MODEL_IMMEDIATE_RESPONSE_KEY] == immediate_msg
|
||||
|
||||
# Verify it serializes correctly with _meta alias
|
||||
serialized = result.model_dump(by_alias=True)
|
||||
assert "_meta" in serialized
|
||||
assert MODEL_IMMEDIATE_RESPONSE_KEY in serialized["_meta"]
|
||||
assert serialized["_meta"][MODEL_IMMEDIATE_RESPONSE_KEY] == immediate_msg
|
||||
|
||||
|
||||
class TestGettingTaskStatus:
|
||||
"""
|
||||
Request: {"method": "tasks/get", "params": {"taskId": "..."}}
|
||||
Response: Returns full Task object with current status and pollInterval.
|
||||
"""
|
||||
|
||||
def test_tasks_get_returns_task_object(self) -> None:
|
||||
"""tasks/get MUST return full Task object."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_get_returns_current_status(self) -> None:
|
||||
"""tasks/get MUST return current status."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_get_may_return_poll_interval(self) -> None:
|
||||
"""tasks/get MAY return pollInterval."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_get_invalid_task_id_returns_error(self) -> None:
|
||||
"""tasks/get with invalid taskId MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_get_nonexistent_task_id_returns_error(self) -> None:
|
||||
"""tasks/get with nonexistent taskId MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestRetrievingResults:
|
||||
"""
|
||||
Request: {"method": "tasks/result", "params": {"taskId": "..."}}
|
||||
Response: The actual operation result structure (e.g., CallToolResult).
|
||||
|
||||
This call blocks until terminal status.
|
||||
"""
|
||||
|
||||
def test_tasks_result_returns_underlying_result(self) -> None:
|
||||
"""tasks/result MUST return exactly what underlying request would return."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_result_blocks_until_terminal(self) -> None:
|
||||
"""tasks/result MUST block for non-terminal tasks."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_result_unblocks_on_terminal(self) -> None:
|
||||
"""tasks/result MUST unblock upon reaching terminal status."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_result_includes_related_task_metadata(self) -> None:
|
||||
"""tasks/result MUST include io.modelcontextprotocol/related-task in _meta."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_result_returns_error_for_failed_task(self) -> None:
|
||||
"""
|
||||
tasks/result returns the same error the underlying request
|
||||
would have produced for failed tasks.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_result_invalid_task_id_returns_error(self) -> None:
|
||||
"""tasks/result with invalid taskId MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestListingTasks:
|
||||
"""
|
||||
Request: {"method": "tasks/list", "params": {"cursor": "optional"}}
|
||||
Response: Array of tasks with pagination support via nextCursor.
|
||||
"""
|
||||
|
||||
def test_tasks_list_returns_array_of_tasks(self) -> None:
|
||||
"""tasks/list MUST return array of tasks."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_list_pagination_with_cursor(self) -> None:
|
||||
"""tasks/list supports pagination via cursor."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_list_returns_next_cursor_when_more_results(self) -> None:
|
||||
"""tasks/list MUST return nextCursor when more results available."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_list_cursors_are_opaque(self) -> None:
|
||||
"""Implementers MUST treat cursors as opaque tokens."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_list_invalid_cursor_returns_error(self) -> None:
|
||||
"""tasks/list with invalid cursor MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestCancellingTasks:
|
||||
"""
|
||||
Request: {"method": "tasks/cancel", "params": {"taskId": "..."}}
|
||||
Response: Returns the task object with status: "cancelled".
|
||||
"""
|
||||
|
||||
def test_tasks_cancel_returns_cancelled_task(self) -> None:
|
||||
"""tasks/cancel MUST return task with status=cancelled."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_cancel_terminal_task_returns_error(self) -> None:
|
||||
"""Cancelling already-terminal task MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_cancel_completed_task_returns_error(self) -> None:
|
||||
"""Cancelling completed task MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_cancel_failed_task_returns_error(self) -> None:
|
||||
"""Cancelling failed task MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_cancel_already_cancelled_task_returns_error(self) -> None:
|
||||
"""Cancelling already-cancelled task MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_cancel_invalid_task_id_returns_error(self) -> None:
|
||||
"""tasks/cancel with invalid taskId MUST return -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestStatusNotifications:
|
||||
"""
|
||||
Receivers MAY send: {"method": "notifications/tasks/status", "params": {...}}
|
||||
These are optional; requestors MUST NOT rely on them and SHOULD continue polling.
|
||||
"""
|
||||
|
||||
def test_receiver_may_send_status_notification(self) -> None:
|
||||
"""Receiver MAY send notifications/tasks/status."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_status_notification_contains_task_id(self) -> None:
|
||||
"""Status notification MUST contain taskId."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_status_notification_contains_status(self) -> None:
|
||||
"""Status notification MUST contain status."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestTaskManagement:
|
||||
"""
|
||||
- Receivers generate unique task IDs as strings
|
||||
- Tasks must begin in working status
|
||||
- createdAt timestamps must be ISO 8601 formatted
|
||||
- Receivers may override requested ttl but must return actual value
|
||||
- Receivers may delete tasks after TTL expires
|
||||
- All task-related messages must include io.modelcontextprotocol/related-task
|
||||
in _meta except for tasks/get, tasks/list, tasks/cancel operations
|
||||
"""
|
||||
|
||||
def test_task_ids_are_unique_strings(self) -> None:
|
||||
"""Receivers MUST generate unique task IDs as strings."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_multiple_tasks_have_unique_ids(self) -> None:
|
||||
"""Multiple tasks MUST have unique IDs."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_receiver_may_delete_tasks_after_ttl(self) -> None:
|
||||
"""Receivers MAY delete tasks after TTL expires."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_related_task_metadata_in_task_messages(self) -> None:
|
||||
"""
|
||||
All task-related messages MUST include io.modelcontextprotocol/related-task
|
||||
in _meta.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_get_does_not_require_related_task_metadata(self) -> None:
|
||||
"""tasks/get does not require related-task metadata."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_list_does_not_require_related_task_metadata(self) -> None:
|
||||
"""tasks/list does not require related-task metadata."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_cancel_does_not_require_related_task_metadata(self) -> None:
|
||||
"""tasks/cancel does not require related-task metadata."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestResultHandling:
|
||||
"""
|
||||
- Receivers must return CreateTaskResult immediately upon accepting task-augmented requests
|
||||
- tasks/result must return exactly what the underlying request would return
|
||||
- tasks/result blocks for non-terminal tasks; must unblock upon reaching terminal status
|
||||
"""
|
||||
|
||||
def test_create_task_result_returned_immediately(self) -> None:
|
||||
"""Receiver MUST return CreateTaskResult immediately (not after work completes)."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_result_matches_underlying_result_structure(self) -> None:
|
||||
"""tasks/result MUST return same structure as underlying request."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_result_for_tool_call_returns_call_tool_result(self) -> None:
|
||||
"""tasks/result for tools/call returns CallToolResult."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestProgressTracking:
|
||||
"""
|
||||
Task-augmented requests support progress notifications using the progressToken
|
||||
mechanism, which remains valid throughout the task lifetime.
|
||||
"""
|
||||
|
||||
def test_progress_token_valid_throughout_task_lifetime(self) -> None:
|
||||
"""progressToken remains valid throughout task lifetime."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_progress_notifications_sent_during_task_execution(self) -> None:
|
||||
"""Progress notifications can be sent during task execution."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestProtocolErrors:
|
||||
"""
|
||||
Protocol Errors (JSON-RPC standard codes):
|
||||
- -32600 (Invalid request): Non-task requests to endpoint requiring task augmentation
|
||||
- -32602 (Invalid params): Invalid/nonexistent taskId, invalid cursor, cancel terminal task
|
||||
- -32603 (Internal error): Server-side execution failures
|
||||
"""
|
||||
|
||||
def test_invalid_request_for_required_task_augmentation(self) -> None:
|
||||
"""Non-task request to task-required endpoint returns -32600."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_invalid_params_for_invalid_task_id(self) -> None:
|
||||
"""Invalid taskId returns -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_invalid_params_for_nonexistent_task_id(self) -> None:
|
||||
"""Nonexistent taskId returns -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_invalid_params_for_invalid_cursor(self) -> None:
|
||||
"""Invalid cursor in tasks/list returns -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_invalid_params_for_cancel_terminal_task(self) -> None:
|
||||
"""Attempt to cancel terminal task returns -32602."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_internal_error_for_server_failure(self) -> None:
|
||||
"""Server-side execution failure returns -32603."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestTaskExecutionErrors:
|
||||
"""
|
||||
When underlying requests fail, the task moves to failed status.
|
||||
- tasks/get response should include statusMessage explaining failure
|
||||
- tasks/result returns same error the underlying request would have produced
|
||||
- For tool calls, isError: true moves task to failed status
|
||||
"""
|
||||
|
||||
def test_underlying_failure_moves_task_to_failed(self) -> None:
|
||||
"""Underlying request failure moves task to failed status."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_failed_task_has_status_message(self) -> None:
|
||||
"""Failed task SHOULD include statusMessage explaining failure."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tasks_result_returns_underlying_error(self) -> None:
|
||||
"""tasks/result returns same error underlying request would produce."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_tool_call_is_error_true_moves_to_failed(self) -> None:
|
||||
"""Tool call with isError: true moves task to failed status."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestTaskObject:
|
||||
"""
|
||||
Task Object fields:
|
||||
- taskId: String identifier
|
||||
- status: Current execution state
|
||||
- statusMessage: Optional human-readable description
|
||||
- createdAt: ISO 8601 timestamp of creation
|
||||
- ttl: Milliseconds before potential deletion
|
||||
- pollInterval: Suggested milliseconds between polls
|
||||
"""
|
||||
|
||||
def test_task_has_task_id_string(self) -> None:
|
||||
"""Task MUST have taskId as string."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_task_has_status(self) -> None:
|
||||
"""Task MUST have status."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_task_status_message_is_optional(self) -> None:
|
||||
"""Task statusMessage is optional."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_task_has_created_at(self) -> None:
|
||||
"""Task MUST have createdAt."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_task_ttl_is_optional(self) -> None:
|
||||
"""Task ttl is optional."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_task_poll_interval_is_optional(self) -> None:
|
||||
"""Task pollInterval is optional."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestRelatedTaskMetadata:
|
||||
"""
|
||||
Related Task Metadata structure:
|
||||
{"_meta": {"io.modelcontextprotocol/related-task": {"taskId": "..."}}}
|
||||
"""
|
||||
|
||||
def test_related_task_metadata_structure(self) -> None:
|
||||
"""Related task metadata has correct structure."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_related_task_metadata_contains_task_id(self) -> None:
|
||||
"""Related task metadata contains taskId."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestAccessAndIsolation:
|
||||
"""
|
||||
- Task IDs enable access to sensitive results
|
||||
- Authorization context binding is essential where available
|
||||
- For non-authorized environments: strong entropy IDs, strict TTL limits
|
||||
"""
|
||||
|
||||
def test_task_bound_to_authorization_context(self) -> None:
|
||||
"""
|
||||
Receivers receiving authorization context MUST bind tasks to that context.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_reject_task_operations_outside_authorization_context(self) -> None:
|
||||
"""
|
||||
Receivers MUST reject task operations for tasks outside
|
||||
requestor's authorization context.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_non_authorized_environments_use_secure_ids(self) -> None:
|
||||
"""
|
||||
For non-authorized environments, receivers SHOULD use
|
||||
cryptographically secure IDs.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_non_authorized_environments_use_shorter_ttls(self) -> None:
|
||||
"""
|
||||
For non-authorized environments, receivers SHOULD use shorter TTLs.
|
||||
"""
|
||||
pytest.skip("TODO")
|
||||
|
||||
|
||||
class TestResourceLimits:
|
||||
"""
|
||||
Receivers should:
|
||||
- Enforce concurrent task limits per requestor
|
||||
- Implement maximum TTL constraints
|
||||
- Clean up expired tasks promptly
|
||||
"""
|
||||
|
||||
def test_concurrent_task_limit_enforced(self) -> None:
|
||||
"""Receiver SHOULD enforce concurrent task limits per requestor."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_maximum_ttl_constraint_enforced(self) -> None:
|
||||
"""Receiver SHOULD implement maximum TTL constraints."""
|
||||
pytest.skip("TODO")
|
||||
|
||||
def test_expired_tasks_cleaned_up(self) -> None:
|
||||
"""Receiver SHOULD clean up expired tasks promptly."""
|
||||
pytest.skip("TODO")
|
||||
Loading…
Reference in a new issue