874 lines
34 KiB
Python
874 lines
34 KiB
Python
"""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
|
|
|
|
from mcp import 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],
|
|
params: GetTaskRequestParams,
|
|
) -> GetTaskResult | ErrorData:
|
|
nonlocal received_task_id
|
|
received_task_id = params.task_id
|
|
task = await store.get_task(params.task_id)
|
|
assert task is not None, f"Test setup error: task {params.task_id} should exist"
|
|
return GetTaskResult(
|
|
task_id=task.task_id,
|
|
status=task.status,
|
|
status_message=task.status_message,
|
|
created_at=task.created_at,
|
|
last_updated_at=task.last_updated_at,
|
|
ttl=task.ttl,
|
|
poll_interval=task.poll_interval,
|
|
)
|
|
|
|
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(task_id="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(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
assert isinstance(response, types.JSONRPCResponse)
|
|
assert response.id == "req-1"
|
|
|
|
result = GetTaskResult.model_validate(response.result)
|
|
assert result.task_id == "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],
|
|
params: GetTaskPayloadRequestParams,
|
|
) -> GetTaskPayloadResult | ErrorData:
|
|
result = await store.get_result(params.task_id)
|
|
assert result is not None, f"Test setup error: result for {params.task_id} 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(task_id="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(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
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],
|
|
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, next_cursor=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(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
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],
|
|
params: CancelTaskRequestParams,
|
|
) -> CancelTaskResult | ErrorData:
|
|
task = await store.get_task(params.task_id)
|
|
assert task is not None, f"Test setup error: task {params.task_id} should exist"
|
|
await store.update_task(params.task_id, status="cancelled")
|
|
updated = await store.get_task(params.task_id)
|
|
assert updated is not None
|
|
return CancelTaskResult(
|
|
task_id=updated.task_id,
|
|
status=updated.status,
|
|
created_at=updated.created_at,
|
|
last_updated_at=updated.last_updated_at,
|
|
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(task_id="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(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
assert isinstance(response, types.JSONRPCResponse)
|
|
|
|
result = CancelTaskResult.model_validate(response.result)
|
|
assert result.task_id == "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],
|
|
params: CreateMessageRequestParams,
|
|
task_metadata: TaskMetadata,
|
|
) -> CreateTaskResult:
|
|
task = await store.create_task(task_metadata)
|
|
created_task_id[0] = task.task_id
|
|
|
|
async def do_sampling() -> None:
|
|
result = CreateMessageResult(
|
|
role="assistant",
|
|
content=TextContent(type="text", text="Sampled response"),
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
await store.store_result(task.task_id, result)
|
|
await store.update_task(task.task_id, 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],
|
|
params: GetTaskRequestParams,
|
|
) -> GetTaskResult | ErrorData:
|
|
task = await store.get_task(params.task_id)
|
|
assert task is not None, f"Test setup error: task {params.task_id} should exist"
|
|
return GetTaskResult(
|
|
task_id=task.task_id,
|
|
status=task.status,
|
|
status_message=task.status_message,
|
|
created_at=task.created_at,
|
|
last_updated_at=task.last_updated_at,
|
|
ttl=task.ttl,
|
|
poll_interval=task.poll_interval,
|
|
)
|
|
|
|
async def get_task_result_handler(
|
|
context: RequestContext[ClientSession],
|
|
params: GetTaskPayloadRequestParams,
|
|
) -> GetTaskPayloadResult | ErrorData:
|
|
result = await store.get_result(params.task_id)
|
|
assert result is not None, f"Test setup error: result for {params.task_id} 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"))],
|
|
max_tokens=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(request))
|
|
|
|
# Step 2: Client responds with CreateTaskResult
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
assert isinstance(response, types.JSONRPCResponse)
|
|
|
|
task_result = CreateTaskResult.model_validate(response.result)
|
|
task_id = task_result.task.task_id
|
|
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(task_id=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(poll_request))
|
|
|
|
poll_response_msg = await client_streams.server_receive.receive()
|
|
poll_response = poll_response_msg.message
|
|
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(task_id=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(result_request))
|
|
|
|
result_response_msg = await client_streams.server_receive.receive()
|
|
result_response = result_response_msg.message
|
|
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],
|
|
params: ElicitRequestParams,
|
|
task_metadata: TaskMetadata,
|
|
) -> CreateTaskResult | ErrorData:
|
|
task = await store.create_task(task_metadata)
|
|
created_task_id[0] = task.task_id
|
|
|
|
async def do_elicitation() -> None:
|
|
# Simulate user providing elicitation response
|
|
result = ElicitResult(action="accept", content={"name": "Test User"})
|
|
await store.store_result(task.task_id, result)
|
|
await store.update_task(task.task_id, 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],
|
|
params: GetTaskRequestParams,
|
|
) -> GetTaskResult | ErrorData:
|
|
task = await store.get_task(params.task_id)
|
|
assert task is not None, f"Test setup error: task {params.task_id} should exist"
|
|
return GetTaskResult(
|
|
task_id=task.task_id,
|
|
status=task.status,
|
|
status_message=task.status_message,
|
|
created_at=task.created_at,
|
|
last_updated_at=task.last_updated_at,
|
|
ttl=task.ttl,
|
|
poll_interval=task.poll_interval,
|
|
)
|
|
|
|
async def get_task_result_handler(
|
|
context: RequestContext[ClientSession],
|
|
params: GetTaskPayloadRequestParams,
|
|
) -> GetTaskPayloadResult | ErrorData:
|
|
result = await store.get_result(params.task_id)
|
|
assert result is not None, f"Test setup error: result for {params.task_id} 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?",
|
|
requested_schema={"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(request))
|
|
|
|
# Step 2: Client responds with CreateTaskResult
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
assert isinstance(response, types.JSONRPCResponse)
|
|
|
|
task_result = CreateTaskResult.model_validate(response.result)
|
|
task_id = task_result.task.task_id
|
|
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(task_id=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(poll_request))
|
|
|
|
poll_response_msg = await client_streams.server_receive.receive()
|
|
poll_response = poll_response_msg.message
|
|
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(task_id=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(result_request))
|
|
|
|
result_response_msg = await client_streams.server_receive.receive()
|
|
result_response = result_response_msg.message
|
|
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(task_id="nonexistent"))
|
|
request = types.JSONRPCRequest(
|
|
jsonrpc="2.0",
|
|
id="req-unhandled",
|
|
**typed_request.model_dump(by_alias=True),
|
|
)
|
|
await client_streams.server_send.send(SessionMessage(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
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(task_id="nonexistent"))
|
|
request = types.JSONRPCRequest(
|
|
jsonrpc="2.0",
|
|
id="req-result",
|
|
**typed_request.model_dump(by_alias=True),
|
|
)
|
|
await client_streams.server_send.send(SessionMessage(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
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(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
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(task_id="nonexistent"))
|
|
request = types.JSONRPCRequest(
|
|
jsonrpc="2.0",
|
|
id="req-cancel",
|
|
**typed_request.model_dump(by_alias=True),
|
|
)
|
|
await client_streams.server_send.send(SessionMessage(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
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"))],
|
|
max_tokens=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(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
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?",
|
|
requested_schema={"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(request))
|
|
|
|
response_msg = await client_streams.server_receive.receive()
|
|
response = response_msg.message
|
|
assert isinstance(response, types.JSONRPCError)
|
|
assert "not supported" in response.error.message.lower()
|
|
|
|
tg.cancel_scope.cancel()
|