Handles message type Exception in lowlevel/server.py _handle_message function. Mentioned as TODO on line 528. (#786)
Co-authored-by: Felix Weinberger <fweinberger@anthropic.com> Co-authored-by: Felix Weinberger <3823880+felixweinberger@users.noreply.github.com>
This commit is contained in:
parent
0e29cc4130
commit
b4e50aa9f9
2 changed files with 84 additions and 2 deletions
|
|
@ -642,13 +642,21 @@ class Server(Generic[LifespanResultT, RequestT]):
|
||||||
raise_exceptions: bool = False,
|
raise_exceptions: bool = False,
|
||||||
):
|
):
|
||||||
with warnings.catch_warnings(record=True) as w:
|
with warnings.catch_warnings(record=True) as w:
|
||||||
# TODO(Marcelo): We should be checking if message is Exception here.
|
match message:
|
||||||
match message: # type: ignore[reportMatchNotExhaustive]
|
|
||||||
case RequestResponder(request=types.ClientRequest(root=req)) as responder:
|
case RequestResponder(request=types.ClientRequest(root=req)) as responder:
|
||||||
with responder:
|
with responder:
|
||||||
await self._handle_request(message, req, session, lifespan_context, raise_exceptions)
|
await self._handle_request(message, req, session, lifespan_context, raise_exceptions)
|
||||||
case types.ClientNotification(root=notify):
|
case types.ClientNotification(root=notify):
|
||||||
await self._handle_notification(notify)
|
await self._handle_notification(notify)
|
||||||
|
case Exception():
|
||||||
|
logger.error(f"Received exception from stream: {message}")
|
||||||
|
await session.send_log_message(
|
||||||
|
level="error",
|
||||||
|
data="Internal Server Error",
|
||||||
|
logger="mcp.server.exception_handler",
|
||||||
|
)
|
||||||
|
if raise_exceptions:
|
||||||
|
raise message
|
||||||
|
|
||||||
for warning in w:
|
for warning in w:
|
||||||
logger.info("Warning: %s: %s", warning.category.__name__, warning.message)
|
logger.info("Warning: %s: %s", warning.category.__name__, warning.message)
|
||||||
|
|
|
||||||
74
tests/server/test_lowlevel_exception_handling.py
Normal file
74
tests/server/test_lowlevel_exception_handling.py
Normal file
|
|
@ -0,0 +1,74 @@
|
||||||
|
from unittest.mock import AsyncMock, Mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
import mcp.types as types
|
||||||
|
from mcp.server.lowlevel.server import Server
|
||||||
|
from mcp.server.session import ServerSession
|
||||||
|
from mcp.shared.session import RequestResponder
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_exception_handling_with_raise_exceptions_true():
|
||||||
|
"""Test that exceptions are re-raised when raise_exceptions=True"""
|
||||||
|
server = Server("test-server")
|
||||||
|
session = Mock(spec=ServerSession)
|
||||||
|
session.send_log_message = AsyncMock()
|
||||||
|
|
||||||
|
test_exception = RuntimeError("Test error")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Test error"):
|
||||||
|
await server._handle_message(test_exception, session, {}, raise_exceptions=True)
|
||||||
|
|
||||||
|
session.send_log_message.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"exception_class,message",
|
||||||
|
[
|
||||||
|
(ValueError, "Test validation error"),
|
||||||
|
(RuntimeError, "Test runtime error"),
|
||||||
|
(KeyError, "Test key error"),
|
||||||
|
(Exception, "Basic error"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def test_exception_handling_with_raise_exceptions_false(exception_class: type[Exception], message: str):
|
||||||
|
"""Test that exceptions are logged when raise_exceptions=False"""
|
||||||
|
server = Server("test-server")
|
||||||
|
session = Mock(spec=ServerSession)
|
||||||
|
session.send_log_message = AsyncMock()
|
||||||
|
|
||||||
|
test_exception = exception_class(message)
|
||||||
|
|
||||||
|
await server._handle_message(test_exception, session, {}, raise_exceptions=False)
|
||||||
|
|
||||||
|
# Should send log message
|
||||||
|
session.send_log_message.assert_called_once()
|
||||||
|
call_args = session.send_log_message.call_args
|
||||||
|
|
||||||
|
assert call_args.kwargs["level"] == "error"
|
||||||
|
assert call_args.kwargs["data"] == "Internal Server Error"
|
||||||
|
assert call_args.kwargs["logger"] == "mcp.server.exception_handler"
|
||||||
|
|
||||||
|
|
||||||
|
@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.ClientRequest(root=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()
|
||||||
Loading…
Reference in a new issue