Merge pull request #167 from modelcontextprotocol/davidsp/88-v2
feat: add request cancellation and cleanup
This commit is contained in:
David Soria Parra
GitHub
commit
c58adfe3f5
2 changed files with 364 additions and 28 deletions
111
tests/issues/test_88_random_error.py
Normal file
111
tests/issues/test_88_random_error.py
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
"""Test to reproduce issue #88: Random error thrown on response."""
|
||||
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Sequence
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.server.lowlevel import Server
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.types import (
|
||||
EmbeddedResource,
|
||||
ImageContent,
|
||||
TextContent,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_notification_validation_error(tmp_path: Path):
|
||||
"""Test that timeouts are handled gracefully and don't break the server.
|
||||
|
||||
This test verifies that when a client request times out:
|
||||
1. The server task stays alive
|
||||
2. The server can still handle new requests
|
||||
3. The client can make new requests
|
||||
4. No resources are leaked
|
||||
"""
|
||||
|
||||
server = Server(name="test")
|
||||
request_count = 0
|
||||
slow_request_started = anyio.Event()
|
||||
slow_request_complete = anyio.Event()
|
||||
|
||||
@server.call_tool()
|
||||
async def slow_tool(
|
||||
name: str, arg
|
||||
) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
|
||||
nonlocal request_count
|
||||
request_count += 1
|
||||
|
||||
if name == "slow":
|
||||
# Signal that slow request has started
|
||||
slow_request_started.set()
|
||||
# Long enough to ensure timeout
|
||||
await anyio.sleep(0.2)
|
||||
# Signal completion
|
||||
slow_request_complete.set()
|
||||
return [TextContent(type="text", text=f"slow {request_count}")]
|
||||
elif name == "fast":
|
||||
# Fast enough to complete before timeout
|
||||
await anyio.sleep(0.01)
|
||||
return [TextContent(type="text", text=f"fast {request_count}")]
|
||||
return [TextContent(type="text", text=f"unknown {request_count}")]
|
||||
|
||||
async def server_handler(read_stream, write_stream):
|
||||
await server.run(
|
||||
read_stream,
|
||||
write_stream,
|
||||
server.create_initialization_options(),
|
||||
raise_exceptions=True,
|
||||
)
|
||||
|
||||
async def client(read_stream, write_stream):
|
||||
# Use a timeout that's:
|
||||
# - Long enough for fast operations (>10ms)
|
||||
# - Short enough for slow operations (<200ms)
|
||||
# - Not too short to avoid flakiness
|
||||
async with ClientSession(
|
||||
read_stream, write_stream, read_timeout_seconds=timedelta(milliseconds=50)
|
||||
) as session:
|
||||
await session.initialize()
|
||||
|
||||
# First call should work (fast operation)
|
||||
result = await session.call_tool("fast")
|
||||
assert result.content == [TextContent(type="text", text="fast 1")]
|
||||
assert not slow_request_complete.is_set()
|
||||
|
||||
# Second call should timeout (slow operation)
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await session.call_tool("slow")
|
||||
assert "Timed out while waiting" in str(exc_info.value)
|
||||
|
||||
# Wait for slow request to complete in the background
|
||||
with anyio.fail_after(1): # Timeout after 1 second
|
||||
await slow_request_complete.wait()
|
||||
|
||||
# Third call should work (fast operation),
|
||||
# proving server is still responsive
|
||||
result = await session.call_tool("fast")
|
||||
assert result.content == [TextContent(type="text", text="fast 3")]
|
||||
|
||||
# Run server and client in separate task groups to avoid cancellation
|
||||
server_writer, server_reader = anyio.create_memory_object_stream(1)
|
||||
client_writer, client_reader = anyio.create_memory_object_stream(1)
|
||||
|
||||
server_ready = anyio.Event()
|
||||
|
||||
async def wrapped_server_handler(read_stream, write_stream):
|
||||
server_ready.set()
|
||||
await server_handler(read_stream, write_stream)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(wrapped_server_handler, server_reader, client_writer)
|
||||
# Wait for server to start and initialize
|
||||
with anyio.fail_after(1): # Timeout after 1 second
|
||||
await server_ready.wait()
|
||||
# Run client in a separate task to avoid cancellation
|
||||
async with anyio.create_task_group() as client_tg:
|
||||
client_tg.start_soon(client, client_reader, server_writer)
|
||||
126
tests/shared/test_session.py
Normal file
126
tests/shared/test_session.py
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
from typing import AsyncGenerator
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
import mcp.types as types
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.server.lowlevel.server import Server
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.memory import create_connected_server_and_client_session
|
||||
from mcp.types import (
|
||||
CancelledNotification,
|
||||
CancelledNotificationParams,
|
||||
ClientNotification,
|
||||
ClientRequest,
|
||||
EmptyResult,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mcp_server() -> Server:
|
||||
return Server(name="test server")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def client_connected_to_server(
|
||||
mcp_server: Server,
|
||||
) -> AsyncGenerator[ClientSession, None]:
|
||||
async with create_connected_server_and_client_session(mcp_server) as client_session:
|
||||
yield client_session
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_in_flight_requests_cleared_after_completion(
|
||||
client_connected_to_server: ClientSession,
|
||||
):
|
||||
"""Verify that _in_flight is empty after all requests complete."""
|
||||
# Send a request and wait for response
|
||||
response = await client_connected_to_server.send_ping()
|
||||
assert isinstance(response, EmptyResult)
|
||||
|
||||
# Verify _in_flight is empty
|
||||
assert len(client_connected_to_server._in_flight) == 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_request_cancellation():
|
||||
"""Test that requests can be cancelled while in-flight."""
|
||||
# The tool is already registered in the fixture
|
||||
|
||||
ev_tool_called = anyio.Event()
|
||||
ev_cancelled = anyio.Event()
|
||||
request_id = None
|
||||
|
||||
# Start the request in a separate task so we can cancel it
|
||||
def make_server() -> Server:
|
||||
server = Server(name="TestSessionServer")
|
||||
|
||||
# Register the tool handler
|
||||
@server.call_tool()
|
||||
async def handle_call_tool(name: str, arguments: dict | None) -> list:
|
||||
nonlocal request_id, ev_tool_called
|
||||
if name == "slow_tool":
|
||||
request_id = server.request_context.request_id
|
||||
ev_tool_called.set()
|
||||
await anyio.sleep(10) # Long enough to ensure we can cancel
|
||||
return []
|
||||
raise ValueError(f"Unknown tool: {name}")
|
||||
|
||||
# Register the tool so it shows up in list_tools
|
||||
@server.list_tools()
|
||||
async def handle_list_tools() -> list[types.Tool]:
|
||||
return [
|
||||
types.Tool(
|
||||
name="slow_tool",
|
||||
description="A slow tool that takes 10 seconds to complete",
|
||||
inputSchema={},
|
||||
)
|
||||
]
|
||||
|
||||
return server
|
||||
|
||||
async def make_request(client_session):
|
||||
nonlocal ev_cancelled
|
||||
try:
|
||||
await client_session.send_request(
|
||||
ClientRequest(
|
||||
types.CallToolRequest(
|
||||
method="tools/call",
|
||||
params=types.CallToolRequestParams(
|
||||
name="slow_tool", arguments={}
|
||||
),
|
||||
)
|
||||
),
|
||||
types.CallToolResult,
|
||||
)
|
||||
pytest.fail("Request should have been cancelled")
|
||||
except McpError as e:
|
||||
# Expected - request was cancelled
|
||||
assert "Request cancelled" in str(e)
|
||||
ev_cancelled.set()
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
make_server()
|
||||
) as client_session:
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(make_request, client_session)
|
||||
|
||||
# Wait for the request to be in-flight
|
||||
with anyio.fail_after(1): # Timeout after 1 second
|
||||
await ev_tool_called.wait()
|
||||
|
||||
# Send cancellation notification
|
||||
assert request_id is not None
|
||||
await client_session.send_notification(
|
||||
ClientNotification(
|
||||
CancelledNotification(
|
||||
method="notifications/cancelled",
|
||||
params=CancelledNotificationParams(requestId=request_id),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Give cancellation time to process
|
||||
with anyio.fail_after(1):
|
||||
await ev_cancelled.wait()
|
||||
Loading…
Reference in a new issue