107 lines
4.3 KiB
Python
107 lines
4.3 KiB
Python
from unittest.mock import AsyncMock, Mock
|
|
|
|
import anyio
|
|
import pytest
|
|
|
|
from mcp import types
|
|
from mcp.server.lowlevel.server import Server
|
|
from mcp.server.session import ServerSession
|
|
from mcp.shared.message import SessionMessage
|
|
from mcp.shared.session import RequestResponder
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_exception_handling_with_raise_exceptions_true():
|
|
"""Transport exceptions are re-raised when raise_exceptions=True."""
|
|
server = Server("test-server")
|
|
session = Mock(spec=ServerSession)
|
|
|
|
test_exception = RuntimeError("Test error")
|
|
|
|
with pytest.raises(RuntimeError, match="Test error"):
|
|
await server._handle_message(test_exception, session, {}, raise_exceptions=True)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_exception_handling_with_raise_exceptions_false():
|
|
"""Transport exceptions are logged locally but not sent to the client.
|
|
|
|
The transport that reported the error is likely broken; writing back
|
|
through it races with stream closure (#1967, #2064). The TypeScript,
|
|
Go, and C# SDKs all log locally only.
|
|
"""
|
|
server = Server("test-server")
|
|
session = Mock(spec=ServerSession)
|
|
session.send_log_message = AsyncMock()
|
|
|
|
await server._handle_message(RuntimeError("Test error"), session, {}, raise_exceptions=False)
|
|
|
|
session.send_log_message.assert_not_called()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_normal_message_handling_not_affected():
|
|
"""Test that normal messages still work correctly"""
|
|
server = Server("test-server")
|
|
session = Mock(spec=ServerSession)
|
|
|
|
# Create a mock RequestResponder
|
|
responder = Mock(spec=RequestResponder)
|
|
responder.request = types.PingRequest(method="ping")
|
|
responder.__enter__ = Mock(return_value=responder)
|
|
responder.__exit__ = Mock(return_value=None)
|
|
|
|
# Mock the _handle_request method to avoid complex setup
|
|
server._handle_request = AsyncMock()
|
|
|
|
# Should handle normally without any exception handling
|
|
await server._handle_message(responder, session, {}, raise_exceptions=False)
|
|
|
|
# Verify _handle_request was called
|
|
server._handle_request.assert_called_once()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_server_run_exits_cleanly_when_transport_yields_exception_then_closes():
|
|
"""Regression test for #1967 / #2064.
|
|
|
|
Exercises the real Server.run() path with real memory streams, reproducing
|
|
what happens in stateless streamable HTTP when a POST handler throws:
|
|
|
|
1. Transport yields an Exception into the read stream
|
|
(streamable_http.py does this in its broad POST-handler except).
|
|
2. Transport closes the read stream (terminate() in stateless mode).
|
|
3. _receive_loop exits its `async with read_stream, write_stream:` block,
|
|
closing the write stream.
|
|
4. Meanwhile _handle_message(exc) was spawned via tg.start_soon and runs
|
|
after the write stream is closed.
|
|
|
|
Before the fix, _handle_message tried to send_log_message through the
|
|
closed write stream, raising ClosedResourceError inside the TaskGroup
|
|
and crashing server.run(). After the fix, it only logs locally.
|
|
"""
|
|
server = Server("test-server")
|
|
|
|
read_send, read_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
|
# Zero-buffer on the write stream forces send() to block until received.
|
|
# With no receiver, a send() sits blocked until _receive_loop exits its
|
|
# `async with self._read_stream, self._write_stream:` block and closes the
|
|
# stream, at which point the blocked send raises ClosedResourceError.
|
|
# This deterministically reproduces the race without sleeps.
|
|
write_send, write_recv = anyio.create_memory_object_stream[SessionMessage](0)
|
|
|
|
# What the streamable HTTP transport does: push the exception, then close.
|
|
read_send.send_nowait(RuntimeError("simulated transport error"))
|
|
read_send.close()
|
|
|
|
with anyio.fail_after(5):
|
|
# stateless=True so server.run doesn't wait for initialize handshake.
|
|
# Before this fix, this raised ExceptionGroup(ClosedResourceError).
|
|
await server.run(read_recv, write_send, server.create_initialization_options(), stateless=True)
|
|
|
|
# write_send was closed inside _receive_loop's `async with`; receive_nowait
|
|
# raises EndOfStream iff the buffer is empty (i.e., server wrote nothing).
|
|
with pytest.raises(anyio.EndOfStream):
|
|
write_recv.receive_nowait()
|
|
write_recv.close()
|