115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
"""Test that cancelled requests don't cause double responses."""
|
|
|
|
from typing import Any
|
|
|
|
import anyio
|
|
import pytest
|
|
|
|
import mcp.types as types
|
|
from mcp.client._memory import InMemoryTransport
|
|
from mcp.client.session import ClientSession
|
|
from mcp.server.lowlevel.server import Server
|
|
from mcp.shared.exceptions import McpError
|
|
from mcp.types import (
|
|
CallToolRequest,
|
|
CallToolRequestParams,
|
|
CallToolResult,
|
|
CancelledNotification,
|
|
CancelledNotificationParams,
|
|
ClientNotification,
|
|
ClientRequest,
|
|
Tool,
|
|
)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_server_remains_functional_after_cancel():
|
|
"""Verify server can handle new requests after a cancellation."""
|
|
|
|
server = Server("test-server")
|
|
|
|
# Track tool calls
|
|
call_count = 0
|
|
ev_first_call = anyio.Event()
|
|
first_request_id = None
|
|
|
|
@server.list_tools()
|
|
async def handle_list_tools() -> list[Tool]:
|
|
return [
|
|
Tool(
|
|
name="test_tool",
|
|
description="Tool for testing",
|
|
input_schema={},
|
|
)
|
|
]
|
|
|
|
@server.call_tool()
|
|
async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[types.TextContent]:
|
|
nonlocal call_count, first_request_id
|
|
if name == "test_tool":
|
|
call_count += 1
|
|
if call_count == 1:
|
|
first_request_id = server.request_context.request_id
|
|
ev_first_call.set()
|
|
await anyio.sleep(5) # First call is slow
|
|
return [types.TextContent(type="text", text=f"Call number: {call_count}")]
|
|
raise ValueError(f"Unknown tool: {name}") # pragma: no cover
|
|
|
|
transport = InMemoryTransport(server)
|
|
async with transport.connect() as (read_stream, write_stream):
|
|
async with ClientSession(read_stream, write_stream) as client:
|
|
await client.initialize()
|
|
|
|
# First request (will be cancelled)
|
|
async def first_request():
|
|
try:
|
|
await client.send_request(
|
|
ClientRequest(
|
|
CallToolRequest(
|
|
params=CallToolRequestParams(name="test_tool", arguments={}),
|
|
)
|
|
),
|
|
CallToolResult,
|
|
)
|
|
pytest.fail("First request should have been cancelled") # pragma: no cover
|
|
except McpError:
|
|
pass # Expected
|
|
|
|
# Start first request
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(first_request)
|
|
|
|
# Wait for it to start
|
|
await ev_first_call.wait()
|
|
|
|
# Cancel it
|
|
assert first_request_id is not None
|
|
await client.send_notification(
|
|
ClientNotification(
|
|
CancelledNotification(
|
|
params=CancelledNotificationParams(
|
|
request_id=first_request_id,
|
|
reason="Testing server recovery",
|
|
),
|
|
)
|
|
)
|
|
)
|
|
|
|
# Second request (should work normally)
|
|
result = await client.send_request(
|
|
ClientRequest(
|
|
CallToolRequest(
|
|
params=CallToolRequestParams(name="test_tool", arguments={}),
|
|
)
|
|
),
|
|
CallToolResult,
|
|
)
|
|
|
|
# Verify second request completed successfully
|
|
assert len(result.content) == 1
|
|
# Type narrowing for pyright
|
|
content = result.content[0]
|
|
assert content.type == "text"
|
|
assert isinstance(content, types.TextContent)
|
|
assert content.text == "Call number: 2"
|
|
assert call_count == 2
|