2025-03-20 09:13:08 +00:00
|
|
|
from collections.abc import AsyncGenerator
|
2025-08-11 19:56:37 +02:00
|
|
|
from typing import Any
|
2025-01-23 20:10:02 +00:00
|
|
|
|
|
|
|
|
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
|
2025-08-11 19:56:37 +02:00
|
|
|
from mcp.shared.memory import create_client_server_memory_streams, create_connected_server_and_client_session
|
2025-12-10 16:15:21 +00:00
|
|
|
from mcp.shared.message import SessionMessage
|
2025-01-23 20:10:02 +00:00
|
|
|
from mcp.types import (
|
|
|
|
|
CancelledNotification,
|
|
|
|
|
CancelledNotificationParams,
|
|
|
|
|
ClientNotification,
|
|
|
|
|
ClientRequest,
|
|
|
|
|
EmptyResult,
|
2025-12-10 16:15:21 +00:00
|
|
|
ErrorData,
|
|
|
|
|
JSONRPCError,
|
|
|
|
|
JSONRPCMessage,
|
|
|
|
|
JSONRPCRequest,
|
|
|
|
|
JSONRPCResponse,
|
2025-08-11 19:56:37 +02:00
|
|
|
TextContent,
|
2025-01-23 20:10:02 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@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()
|
2025-08-11 19:56:37 +02:00
|
|
|
async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[TextContent]:
|
2025-01-23 20:10:02 +00:00
|
|
|
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 [] # pragma: no cover
|
|
|
|
|
raise ValueError(f"Unknown tool: {name}") # pragma: no cover
|
|
|
|
|
|
|
|
|
|
# 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",
|
2026-01-16 15:51:27 +01:00
|
|
|
input_schema={},
|
2025-01-23 20:10:02 +00:00
|
|
|
)
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
return server
|
|
|
|
|
|
2025-08-11 19:56:37 +02:00
|
|
|
async def make_request(client_session: ClientSession):
|
2025-01-23 20:10:02 +00:00
|
|
|
nonlocal ev_cancelled
|
|
|
|
|
try:
|
|
|
|
|
await client_session.send_request(
|
|
|
|
|
ClientRequest(
|
|
|
|
|
types.CallToolRequest(
|
|
|
|
|
params=types.CallToolRequestParams(name="slow_tool", arguments={}),
|
|
|
|
|
)
|
|
|
|
|
),
|
|
|
|
|
types.CallToolResult,
|
|
|
|
|
)
|
|
|
|
|
pytest.fail("Request should have been cancelled") # pragma: no cover
|
|
|
|
|
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(
|
2026-01-16 15:51:27 +01:00
|
|
|
params=CancelledNotificationParams(request_id=request_id),
|
2025-01-23 20:10:02 +00:00
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# Give cancellation time to process
|
2026-01-07 17:28:23 +01:00
|
|
|
# TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed.
|
|
|
|
|
with anyio.fail_after(1): # pragma: no cover
|
2025-01-23 20:10:02 +00:00
|
|
|
await ev_cancelled.wait()
|
2025-05-27 17:55:27 -04:00
|
|
|
|
|
|
|
|
|
2025-12-10 16:15:21 +00:00
|
|
|
@pytest.mark.anyio
|
|
|
|
|
async def test_response_id_type_mismatch_string_to_int():
|
|
|
|
|
"""
|
|
|
|
|
Test that responses with string IDs are correctly matched to requests sent with
|
|
|
|
|
integer IDs.
|
|
|
|
|
|
|
|
|
|
This handles the case where a server returns "id": "0" (string) but the client
|
|
|
|
|
sent "id": 0 (integer). Without ID type normalization, this would cause a timeout.
|
|
|
|
|
"""
|
|
|
|
|
ev_response_received = anyio.Event()
|
|
|
|
|
result_holder: list[types.EmptyResult] = []
|
|
|
|
|
|
|
|
|
|
async with create_client_server_memory_streams() as (client_streams, server_streams):
|
|
|
|
|
client_read, client_write = client_streams
|
|
|
|
|
server_read, server_write = server_streams
|
|
|
|
|
|
|
|
|
|
async def mock_server():
|
|
|
|
|
"""Receive a request and respond with a string ID instead of integer."""
|
|
|
|
|
message = await server_read.receive()
|
|
|
|
|
assert isinstance(message, SessionMessage)
|
|
|
|
|
root = message.message.root
|
|
|
|
|
assert isinstance(root, JSONRPCRequest)
|
|
|
|
|
# Get the original request ID (which is an integer)
|
|
|
|
|
request_id = root.id
|
|
|
|
|
assert isinstance(request_id, int), f"Expected int, got {type(request_id)}"
|
|
|
|
|
|
|
|
|
|
# Respond with the ID as a string (simulating a buggy server)
|
|
|
|
|
response = JSONRPCResponse(
|
|
|
|
|
jsonrpc="2.0",
|
|
|
|
|
id=str(request_id), # Convert to string to simulate mismatch
|
|
|
|
|
result={},
|
|
|
|
|
)
|
|
|
|
|
await server_write.send(SessionMessage(message=JSONRPCMessage(response)))
|
|
|
|
|
|
|
|
|
|
async def make_request(client_session: ClientSession):
|
|
|
|
|
nonlocal result_holder
|
|
|
|
|
# Send a ping request (uses integer ID internally)
|
|
|
|
|
result = await client_session.send_ping()
|
|
|
|
|
result_holder.append(result)
|
|
|
|
|
ev_response_received.set()
|
|
|
|
|
|
|
|
|
|
async with (
|
|
|
|
|
anyio.create_task_group() as tg,
|
|
|
|
|
ClientSession(read_stream=client_read, write_stream=client_write) as client_session,
|
|
|
|
|
):
|
|
|
|
|
tg.start_soon(mock_server)
|
|
|
|
|
tg.start_soon(make_request, client_session)
|
|
|
|
|
|
2026-01-07 17:28:23 +01:00
|
|
|
# TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed.
|
|
|
|
|
with anyio.fail_after(2): # pragma: no cover
|
2025-12-10 16:15:21 +00:00
|
|
|
await ev_response_received.wait()
|
|
|
|
|
|
|
|
|
|
assert len(result_holder) == 1
|
|
|
|
|
assert isinstance(result_holder[0], EmptyResult)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.anyio
|
|
|
|
|
async def test_error_response_id_type_mismatch_string_to_int():
|
|
|
|
|
"""
|
|
|
|
|
Test that error responses with string IDs are correctly matched to requests
|
|
|
|
|
sent with integer IDs.
|
|
|
|
|
|
|
|
|
|
This handles the case where a server returns an error with "id": "0" (string)
|
|
|
|
|
but the client sent "id": 0 (integer).
|
|
|
|
|
"""
|
|
|
|
|
ev_error_received = anyio.Event()
|
|
|
|
|
error_holder: list[McpError] = []
|
|
|
|
|
|
|
|
|
|
async with create_client_server_memory_streams() as (client_streams, server_streams):
|
|
|
|
|
client_read, client_write = client_streams
|
|
|
|
|
server_read, server_write = server_streams
|
|
|
|
|
|
|
|
|
|
async def mock_server():
|
|
|
|
|
"""Receive a request and respond with an error using a string ID."""
|
|
|
|
|
message = await server_read.receive()
|
|
|
|
|
assert isinstance(message, SessionMessage)
|
|
|
|
|
root = message.message.root
|
|
|
|
|
assert isinstance(root, JSONRPCRequest)
|
|
|
|
|
request_id = root.id
|
|
|
|
|
assert isinstance(request_id, int)
|
|
|
|
|
|
|
|
|
|
# Respond with an error, using the ID as a string
|
|
|
|
|
error_response = JSONRPCError(
|
|
|
|
|
jsonrpc="2.0",
|
|
|
|
|
id=str(request_id), # Convert to string to simulate mismatch
|
|
|
|
|
error=ErrorData(code=-32600, message="Test error"),
|
|
|
|
|
)
|
|
|
|
|
await server_write.send(SessionMessage(message=JSONRPCMessage(error_response)))
|
|
|
|
|
|
|
|
|
|
async def make_request(client_session: ClientSession):
|
|
|
|
|
nonlocal error_holder
|
|
|
|
|
try:
|
|
|
|
|
await client_session.send_ping()
|
|
|
|
|
pytest.fail("Expected McpError to be raised") # pragma: no cover
|
|
|
|
|
except McpError as e:
|
|
|
|
|
error_holder.append(e)
|
|
|
|
|
ev_error_received.set()
|
|
|
|
|
|
|
|
|
|
async with (
|
|
|
|
|
anyio.create_task_group() as tg,
|
|
|
|
|
ClientSession(read_stream=client_read, write_stream=client_write) as client_session,
|
|
|
|
|
):
|
|
|
|
|
tg.start_soon(mock_server)
|
|
|
|
|
tg.start_soon(make_request, client_session)
|
|
|
|
|
|
2026-01-07 17:28:23 +01:00
|
|
|
# TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed.
|
|
|
|
|
with anyio.fail_after(2): # pragma: no cover
|
2025-12-10 16:15:21 +00:00
|
|
|
await ev_error_received.wait()
|
|
|
|
|
|
|
|
|
|
assert len(error_holder) == 1
|
|
|
|
|
assert "Test error" in str(error_holder[0])
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.anyio
|
|
|
|
|
async def test_response_id_non_numeric_string_no_match():
|
|
|
|
|
"""
|
|
|
|
|
Test that responses with non-numeric string IDs don't incorrectly match
|
|
|
|
|
integer request IDs.
|
|
|
|
|
|
|
|
|
|
If a server returns "id": "abc" (non-numeric string), it should not match
|
|
|
|
|
a request sent with "id": 0 (integer).
|
|
|
|
|
"""
|
|
|
|
|
ev_timeout = anyio.Event()
|
|
|
|
|
|
|
|
|
|
async with create_client_server_memory_streams() as (client_streams, server_streams):
|
|
|
|
|
client_read, client_write = client_streams
|
|
|
|
|
server_read, server_write = server_streams
|
|
|
|
|
|
|
|
|
|
async def mock_server():
|
|
|
|
|
"""Receive a request and respond with a non-numeric string ID."""
|
|
|
|
|
message = await server_read.receive()
|
|
|
|
|
assert isinstance(message, SessionMessage)
|
|
|
|
|
|
|
|
|
|
# Respond with a non-numeric string ID (should not match)
|
|
|
|
|
response = JSONRPCResponse(
|
|
|
|
|
jsonrpc="2.0",
|
|
|
|
|
id="not_a_number", # Non-numeric string
|
|
|
|
|
result={},
|
|
|
|
|
)
|
|
|
|
|
await server_write.send(SessionMessage(message=JSONRPCMessage(response)))
|
|
|
|
|
|
|
|
|
|
async def make_request(client_session: ClientSession):
|
|
|
|
|
try:
|
|
|
|
|
# Use a short timeout since we expect this to fail
|
|
|
|
|
await client_session.send_request(
|
|
|
|
|
ClientRequest(types.PingRequest()),
|
|
|
|
|
types.EmptyResult,
|
2025-12-19 17:52:56 +05:30
|
|
|
request_read_timeout_seconds=0.5,
|
2025-12-10 16:15:21 +00:00
|
|
|
)
|
|
|
|
|
pytest.fail("Expected timeout") # pragma: no cover
|
|
|
|
|
except McpError as e:
|
|
|
|
|
assert "Timed out" in str(e)
|
|
|
|
|
ev_timeout.set()
|
|
|
|
|
|
|
|
|
|
async with (
|
|
|
|
|
anyio.create_task_group() as tg,
|
|
|
|
|
ClientSession(read_stream=client_read, write_stream=client_write) as client_session,
|
|
|
|
|
):
|
|
|
|
|
tg.start_soon(mock_server)
|
|
|
|
|
tg.start_soon(make_request, client_session)
|
|
|
|
|
|
2026-01-07 17:28:23 +01:00
|
|
|
# TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed.
|
|
|
|
|
with anyio.fail_after(2): # pragma: no cover
|
2025-12-10 16:15:21 +00:00
|
|
|
await ev_timeout.wait()
|
|
|
|
|
|
|
|
|
|
|
2025-05-27 17:55:27 -04:00
|
|
|
@pytest.mark.anyio
|
|
|
|
|
async def test_connection_closed():
|
|
|
|
|
"""
|
|
|
|
|
Test that pending requests are cancelled when the connection is closed remotely.
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
ev_closed = anyio.Event()
|
|
|
|
|
ev_response = anyio.Event()
|
|
|
|
|
|
2025-08-11 19:56:37 +02:00
|
|
|
async with create_client_server_memory_streams() as (client_streams, server_streams):
|
2025-05-27 17:55:27 -04:00
|
|
|
client_read, client_write = client_streams
|
|
|
|
|
server_read, server_write = server_streams
|
|
|
|
|
|
2025-08-11 19:56:37 +02:00
|
|
|
async def make_request(client_session: ClientSession):
|
2025-05-27 17:55:27 -04:00
|
|
|
"""Send a request in a separate task"""
|
|
|
|
|
nonlocal ev_response
|
|
|
|
|
try:
|
|
|
|
|
# any request will do
|
|
|
|
|
await client_session.initialize()
|
|
|
|
|
pytest.fail("Request should have errored") # pragma: no cover
|
|
|
|
|
except McpError as e:
|
|
|
|
|
# Expected - request errored
|
|
|
|
|
assert "Connection closed" in str(e)
|
|
|
|
|
ev_response.set()
|
|
|
|
|
|
|
|
|
|
async def mock_server():
|
|
|
|
|
"""Wait for a request, then close the connection"""
|
|
|
|
|
nonlocal ev_closed
|
|
|
|
|
# Wait for a request
|
|
|
|
|
await server_read.receive()
|
|
|
|
|
# Close the connection, as if the server exited
|
|
|
|
|
server_write.close()
|
|
|
|
|
server_read.close()
|
|
|
|
|
ev_closed.set()
|
|
|
|
|
|
|
|
|
|
async with (
|
|
|
|
|
anyio.create_task_group() as tg,
|
2025-08-11 19:56:37 +02:00
|
|
|
ClientSession(read_stream=client_read, write_stream=client_write) as client_session,
|
2025-05-27 17:55:27 -04:00
|
|
|
):
|
|
|
|
|
tg.start_soon(make_request, client_session)
|
|
|
|
|
tg.start_soon(mock_server)
|
|
|
|
|
|
2026-01-07 17:28:23 +01:00
|
|
|
# TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed.
|
|
|
|
|
with anyio.fail_after(1): # pragma: no cover
|
2025-05-27 17:55:27 -04:00
|
|
|
await ev_closed.wait()
|
2026-01-07 17:28:23 +01:00
|
|
|
with anyio.fail_after(1): # pragma: no cover
|
2025-05-27 17:55:27 -04:00
|
|
|
await ev_response.wait()
|
