refactor: replace lowlevel Server decorators with on_* constructor kwargs (#1985)
This commit is contained in:
Max Isbey
GitHub
parent
d6d3ad9a7c
commit
0a22a9dc33
1 changed files with 3563 additions and 5481 deletions
|
|
@ -1,43 +1,38 @@
|
|||
"""Tests for the experimental client task methods (session.experimental)."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
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 import Client
|
||||
from mcp.server import Server, ServerRequestContext
|
||||
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,
|
||||
CancelTaskRequestParams,
|
||||
CancelTaskResult,
|
||||
ClientResult,
|
||||
CreateTaskResult,
|
||||
GetTaskPayloadRequest,
|
||||
GetTaskPayloadRequestParams,
|
||||
GetTaskPayloadResult,
|
||||
GetTaskRequest,
|
||||
GetTaskRequestParams,
|
||||
GetTaskResult,
|
||||
ListTasksRequest,
|
||||
ListTasksResult,
|
||||
ServerNotification,
|
||||
ServerRequest,
|
||||
ListToolsResult,
|
||||
PaginatedRequestParams,
|
||||
TaskMetadata,
|
||||
TextContent,
|
||||
Tool,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@dataclass
|
||||
class AppContext:
|
||||
|
|
@ -48,44 +43,53 @@ class AppContext:
|
|||
task_done_events: dict[str, Event] = field(default_factory=lambda: {})
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def _handle_list_tools(
|
||||
ctx: ServerRequestContext[AppContext], params: PaginatedRequestParams | None
|
||||
) -> ListToolsResult:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
async def _handle_call_tool_with_done_event(
|
||||
ctx: ServerRequestContext[AppContext], params: CallToolRequestParams, *, result_text: str = "Done"
|
||||
) -> CallToolResult | CreateTaskResult:
|
||||
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() -> None:
|
||||
async with task_execution(task.task_id, app.store) as task_ctx:
|
||||
await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text=result_text)]))
|
||||
done_event.set()
|
||||
|
||||
app.task_group.start_soon(do_work)
|
||||
return CreateTaskResult(task=task)
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def _make_lifespan(store: InMemoryTaskStore, task_done_events: dict[str, Event]):
|
||||
@asynccontextmanager
|
||||
async def app_lifespan(server: Server[AppContext]) -> AsyncIterator[AppContext]:
|
||||
async with anyio.create_task_group() as tg:
|
||||
yield AppContext(task_group=tg, store=store, task_done_events=task_done_events)
|
||||
|
||||
return app_lifespan
|
||||
|
||||
|
||||
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()
|
||||
task_done_events: dict[str, Event] = {}
|
||||
|
||||
@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
|
||||
async def handle_get_task(ctx: ServerRequestContext[AppContext], params: GetTaskRequestParams) -> GetTaskResult:
|
||||
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"
|
||||
task = await app.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,
|
||||
|
|
@ -96,147 +100,116 @@ async def test_session_experimental_get_task() -> None:
|
|||
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)
|
||||
server: Server[AppContext] = Server(
|
||||
"test-server",
|
||||
lifespan=_make_lifespan(store, task_done_events),
|
||||
on_list_tools=_handle_list_tools,
|
||||
on_call_tool=_handle_call_tool_with_done_event,
|
||||
)
|
||||
server.experimental.enable_tasks(on_get_task=handle_get_task)
|
||||
|
||||
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={},
|
||||
),
|
||||
async with Client(server) as client:
|
||||
# Create a task
|
||||
create_result = await client.session.send_request(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
CreateTaskResult,
|
||||
)
|
||||
task_id = create_result.task.task_id
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
# Wait for task to complete
|
||||
await task_done_events[task_id].wait()
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
# Use session.experimental to get task status
|
||||
task_status = await client.session.experimental.get_task(task_id)
|
||||
|
||||
# Create a task
|
||||
create_result = await client_session.send_request(
|
||||
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()
|
||||
assert task_status.task_id == task_id
|
||||
assert task_status.status == "completed"
|
||||
|
||||
|
||||
@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()
|
||||
task_done_events: dict[str, Event] = {}
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
return [Tool(name="test_tool", description="Test", input_schema={"type": "object"})]
|
||||
async def handle_call_tool(
|
||||
ctx: ServerRequestContext[AppContext], params: CallToolRequestParams
|
||||
) -> CallToolResult | CreateTaskResult:
|
||||
return await _handle_call_tool_with_done_event(ctx, params, result_text="Task result content")
|
||||
|
||||
@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,
|
||||
ctx: ServerRequestContext[AppContext], params: GetTaskPayloadRequestParams
|
||||
) -> 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"
|
||||
app = ctx.lifespan_context
|
||||
result = await app.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, 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)
|
||||
server: Server[AppContext] = Server(
|
||||
"test-server",
|
||||
lifespan=_make_lifespan(store, task_done_events),
|
||||
on_list_tools=_handle_list_tools,
|
||||
on_call_tool=handle_call_tool,
|
||||
)
|
||||
server.experimental.enable_tasks(on_task_result=handle_get_task_result)
|
||||
|
||||
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={},
|
||||
),
|
||||
async with Client(server) as client:
|
||||
# Create a task
|
||||
create_result = await client.session.send_request(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
CreateTaskResult,
|
||||
)
|
||||
task_id = create_result.task.task_id
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
# Wait for task to complete
|
||||
await task_done_events[task_id].wait()
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
# Use TaskClient to get task result
|
||||
task_result = await client.session.experimental.get_task_result(task_id, CallToolResult)
|
||||
|
||||
# Create a task
|
||||
create_result = await client_session.send_request(
|
||||
assert len(task_result.content) == 1
|
||||
content = task_result.content[0]
|
||||
assert isinstance(content, TextContent)
|
||||
assert content.text == "Task result content"
|
||||
|
||||
|
||||
async def test_session_experimental_list_tasks() -> None:
|
||||
"""Test TaskClient.list_tasks() method."""
|
||||
store = InMemoryTaskStore()
|
||||
task_done_events: dict[str, Event] = {}
|
||||
|
||||
async def handle_list_tasks(
|
||||
ctx: ServerRequestContext[AppContext], params: PaginatedRequestParams | None
|
||||
) -> ListTasksResult:
|
||||
app = ctx.lifespan_context
|
||||
cursor = params.cursor if params else None
|
||||
tasks_list, next_cursor = await app.store.list_tasks(cursor=cursor)
|
||||
return ListTasksResult(tasks=tasks_list, next_cursor=next_cursor)
|
||||
|
||||
server: Server[AppContext] = Server(
|
||||
"test-server",
|
||||
lifespan=_make_lifespan(store, task_done_events),
|
||||
on_list_tools=_handle_list_tools,
|
||||
on_call_tool=_handle_call_tool_with_done_event,
|
||||
)
|
||||
server.experimental.enable_tasks(on_list_tasks=handle_list_tasks)
|
||||
|
||||
async with Client(server) as client:
|
||||
# Create two tasks
|
||||
for _ in range(2):
|
||||
create_result = await client.session.send_request(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
|
|
@ -246,130 +219,22 @@ async def test_session_experimental_get_task_result() -> None:
|
|||
),
|
||||
CreateTaskResult,
|
||||
)
|
||||
task_id = create_result.task.task_id
|
||||
await task_done_events[create_result.task.task_id].wait()
|
||||
|
||||
# Wait for task to complete
|
||||
await app_context.task_done_events[task_id].wait()
|
||||
# Use TaskClient to list tasks
|
||||
list_result = await client.session.experimental.list_tasks()
|
||||
|
||||
# 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()
|
||||
assert len(list_result.tasks) == 2
|
||||
|
||||
|
||||
@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(
|
||||
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()
|
||||
task_done_events: dict[str, Event] = {}
|
||||
|
||||
@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
|
||||
async def handle_call_tool_no_work(
|
||||
ctx: ServerRequestContext[AppContext], params: CallToolRequestParams
|
||||
) -> CallToolResult | CreateTaskResult:
|
||||
app = ctx.lifespan_context
|
||||
if ctx.experimental.is_task:
|
||||
task_metadata = ctx.experimental.task_metadata
|
||||
|
|
@ -377,14 +242,12 @@ async def test_session_experimental_cancel_task() -> 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"
|
||||
async def handle_get_task(ctx: ServerRequestContext[AppContext], params: GetTaskRequestParams) -> GetTaskResult:
|
||||
app = ctx.lifespan_context
|
||||
task = await app.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,
|
||||
|
|
@ -395,14 +258,14 @@ async def test_session_experimental_cancel_task() -> None:
|
|||
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)
|
||||
async def handle_cancel_task(
|
||||
ctx: ServerRequestContext[AppContext], params: CancelTaskRequestParams
|
||||
) -> CancelTaskResult:
|
||||
app = ctx.lifespan_context
|
||||
task = await app.store.get_task(params.task_id)
|
||||
assert task is not None, f"Test setup error: task {params.task_id} should exist"
|
||||
await app.store.update_task(params.task_id, status="cancelled")
|
||||
updated_task = await app.store.get_task(params.task_id)
|
||||
assert updated_task is not None
|
||||
return CancelTaskResult(
|
||||
task_id=updated_task.task_id,
|
||||
|
|
@ -412,63 +275,35 @@ async def test_session_experimental_cancel_task() -> None:
|
|||
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)
|
||||
server: Server[AppContext] = Server(
|
||||
"test-server",
|
||||
lifespan=_make_lifespan(store, task_done_events),
|
||||
on_list_tools=_handle_list_tools,
|
||||
on_call_tool=handle_call_tool_no_work,
|
||||
)
|
||||
server.experimental.enable_tasks(on_get_task=handle_get_task, on_cancel_task=handle_cancel_task)
|
||||
|
||||
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={},
|
||||
),
|
||||
async with Client(server) as client:
|
||||
# Create a task (but don't complete it)
|
||||
create_result = await client.session.send_request(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(
|
||||
name="test_tool",
|
||||
arguments={},
|
||||
task=TaskMetadata(ttl=60000),
|
||||
)
|
||||
),
|
||||
) as server_session:
|
||||
async for message in server_session.incoming_messages:
|
||||
await server._handle_message(message, server_session, app_context, raise_exceptions=False)
|
||||
CreateTaskResult,
|
||||
)
|
||||
task_id = create_result.task.task_id
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
app_context = AppContext(task_group=tg, store=store)
|
||||
tg.start_soon(run_server, app_context)
|
||||
# Verify task is working
|
||||
status_before = await client.session.experimental.get_task(task_id)
|
||||
assert status_before.status == "working"
|
||||
|
||||
async with ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
message_handler=message_handler,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
# Cancel the task
|
||||
await client.session.experimental.cancel_task(task_id)
|
||||
|
||||
# Create a task (but don't complete it)
|
||||
create_result = await client_session.send_request(
|
||||
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()
|
||||
# Verify task is cancelled
|
||||
status_after = await client.session.experimental.get_task(task_id)
|
||||
assert status_after.status == "cancelled"
|
||||
|
|
|
|||
Loading…
Reference in a new issue