483 lines
19 KiB
Python
483 lines
19 KiB
Python
"""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", input_schema={"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.task_id] = done_event
|
|
|
|
async def do_work():
|
|
async with task_execution(task.task_id, 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.task_id)
|
|
assert task is not None, f"Test setup error: task {request.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,
|
|
)
|
|
|
|
# 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.task_id
|
|
|
|
# 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.task_id == 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", input_schema={"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.task_id] = done_event
|
|
|
|
async def do_work():
|
|
async with task_execution(task.task_id, 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.task_id)
|
|
assert result is not None, f"Test setup error: result for {request.params.task_id} 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.task_id
|
|
|
|
# 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", input_schema={"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.task_id] = done_event
|
|
|
|
async def do_work():
|
|
async with task_execution(task.task_id, 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, next_cursor=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.task_id].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", input_schema={"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.task_id)
|
|
assert task is not None, f"Test setup error: task {request.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,
|
|
)
|
|
|
|
@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.task_id)
|
|
assert task is not None, f"Test setup error: task {request.params.task_id} should exist"
|
|
await app.store.update_task(request.params.task_id, status="cancelled")
|
|
# CancelTaskResult extends Task, so we need to return the updated task info
|
|
updated_task = await app.store.get_task(request.params.task_id)
|
|
assert updated_task is not None
|
|
return CancelTaskResult(
|
|
task_id=updated_task.task_id,
|
|
status=updated_task.status,
|
|
created_at=updated_task.created_at,
|
|
last_updated_at=updated_task.last_updated_at,
|
|
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.task_id
|
|
|
|
# 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()
|