mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 14:54:24 +01:00
Fix uncaught exception in MCP server (#967)
This commit is contained in:
@@ -333,6 +333,7 @@ class BaseSession(
|
|||||||
self._read_stream,
|
self._read_stream,
|
||||||
self._write_stream,
|
self._write_stream,
|
||||||
):
|
):
|
||||||
|
try:
|
||||||
async for message in self._read_stream:
|
async for message in self._read_stream:
|
||||||
if isinstance(message, Exception):
|
if isinstance(message, Exception):
|
||||||
await self._handle_incoming(message)
|
await self._handle_incoming(message)
|
||||||
@@ -343,7 +344,9 @@ class BaseSession(
|
|||||||
)
|
)
|
||||||
responder = RequestResponder(
|
responder = RequestResponder(
|
||||||
request_id=message.message.root.id,
|
request_id=message.message.root.id,
|
||||||
request_meta=validated_request.root.params.meta if validated_request.root.params else None,
|
request_meta=validated_request.root.params.meta
|
||||||
|
if validated_request.root.params
|
||||||
|
else None,
|
||||||
request=validated_request,
|
request=validated_request,
|
||||||
session=self,
|
session=self,
|
||||||
on_complete=lambda r: self._in_flight.pop(r.request_id, None),
|
on_complete=lambda r: self._in_flight.pop(r.request_id, None),
|
||||||
@@ -410,12 +413,26 @@ class BaseSession(
|
|||||||
RuntimeError("Received response with an unknown " f"request ID: {message}")
|
RuntimeError("Received response with an unknown " f"request ID: {message}")
|
||||||
)
|
)
|
||||||
|
|
||||||
|
except anyio.ClosedResourceError:
|
||||||
|
# This is expected when the client disconnects abruptly.
|
||||||
|
# Without this handler, the exception would propagate up and
|
||||||
|
# crash the server's task group.
|
||||||
|
logging.debug("Read stream closed by client")
|
||||||
|
except Exception as e:
|
||||||
|
# Other exceptions are not expected and should be logged. We purposefully
|
||||||
|
# catch all exceptions here to avoid crashing the server.
|
||||||
|
logging.exception(f"Unhandled exception in receive loop: {e}")
|
||||||
|
finally:
|
||||||
# after the read stream is closed, we need to send errors
|
# after the read stream is closed, we need to send errors
|
||||||
# to any pending requests
|
# to any pending requests
|
||||||
for id, stream in self._response_streams.items():
|
for id, stream in self._response_streams.items():
|
||||||
error = ErrorData(code=CONNECTION_CLOSED, message="Connection closed")
|
error = ErrorData(code=CONNECTION_CLOSED, message="Connection closed")
|
||||||
|
try:
|
||||||
await stream.send(JSONRPCError(jsonrpc="2.0", id=id, error=error))
|
await stream.send(JSONRPCError(jsonrpc="2.0", id=id, error=error))
|
||||||
await stream.aclose()
|
await stream.aclose()
|
||||||
|
except Exception:
|
||||||
|
# Stream might already be closed
|
||||||
|
pass
|
||||||
self._response_streams.clear()
|
self._response_streams.clear()
|
||||||
|
|
||||||
async def _received_request(self, responder: RequestResponder[ReceiveRequestT, SendResultT]) -> None:
|
async def _received_request(self, responder: RequestResponder[ReceiveRequestT, SendResultT]) -> None:
|
||||||
|
|||||||
@@ -1521,3 +1521,40 @@ def test_server_backwards_compatibility_no_protocol_version(basic_server, basic_
|
|||||||
)
|
)
|
||||||
assert response.status_code == 200 # Should succeed for backwards compatibility
|
assert response.status_code == 200 # Should succeed for backwards compatibility
|
||||||
assert response.headers.get("Content-Type") == "text/event-stream"
|
assert response.headers.get("Content-Type") == "text/event-stream"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_client_crash_handled(basic_server, basic_server_url):
|
||||||
|
"""Test that cases where the client crashes are handled gracefully."""
|
||||||
|
|
||||||
|
# Simulate bad client that crashes after init
|
||||||
|
async def bad_client():
|
||||||
|
"""Client that triggers ClosedResourceError"""
|
||||||
|
async with streamablehttp_client(f"{basic_server_url}/mcp") as (
|
||||||
|
read_stream,
|
||||||
|
write_stream,
|
||||||
|
_,
|
||||||
|
):
|
||||||
|
async with ClientSession(read_stream, write_stream) as session:
|
||||||
|
await session.initialize()
|
||||||
|
raise Exception("client crash")
|
||||||
|
|
||||||
|
# Run bad client a few times to trigger the crash
|
||||||
|
for _ in range(3):
|
||||||
|
try:
|
||||||
|
await bad_client()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
await anyio.sleep(0.1)
|
||||||
|
|
||||||
|
# Try a good client, it should still be able to connect and list tools
|
||||||
|
async with streamablehttp_client(f"{basic_server_url}/mcp") as (
|
||||||
|
read_stream,
|
||||||
|
write_stream,
|
||||||
|
_,
|
||||||
|
):
|
||||||
|
async with ClientSession(read_stream, write_stream) as session:
|
||||||
|
result = await session.initialize()
|
||||||
|
assert isinstance(result, InitializeResult)
|
||||||
|
tools = await session.list_tools()
|
||||||
|
assert tools.tools
|
||||||
|
|||||||
Reference in New Issue
Block a user