Skip to content

Commit c7e4fc3

Browse files
author
Samson Kolge
committed
Fix: Handle Exception case in lowlevel server _handle_message
- Add missing Exception case in _handle_message match statement - Use logger.exception() following CLAUDE.md guidelines - Respect raise_exceptions parameter for backward compatibility - Add comprehensive tests for exception handling scenarios - Remove TODO comment as issue is resolved This fixes a critical gap where Exception instances passed to _handle_message were not properly handled, potentially causing silent failures in production.
1 parent 959d4e3 commit c7e4fc3

File tree

2 files changed

+126
-2
lines changed

2 files changed

+126
-2
lines changed

src/mcp/server/lowlevel/server.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -603,13 +603,16 @@ async def _handle_message(
603603
raise_exceptions: bool = False,
604604
):
605605
with warnings.catch_warnings(record=True) as w:
606-
# TODO(Marcelo): We should be checking if message is Exception here.
607-
match message: # type: ignore[reportMatchNotExhaustive]
606+
match message:
608607
case RequestResponder(request=types.ClientRequest(root=req)) as responder:
609608
with responder:
610609
await self._handle_request(message, req, session, lifespan_context, raise_exceptions)
611610
case types.ClientNotification(root=notify):
612611
await self._handle_notification(notify)
612+
case Exception() as error:
613+
logger.exception("Error in message processing: %s", error)
614+
if raise_exceptions:
615+
raise error
613616

614617
for warning in w:
615618
logger.info("Warning: %s: %s", warning.category.__name__, warning.message)
Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
"""Test exception handling in lowlevel server message processing."""
2+
3+
import logging
4+
from unittest.mock import Mock
5+
6+
import anyio
7+
import pytest
8+
9+
from mcp.server import Server
10+
from mcp.server.lowlevel import NotificationOptions
11+
from mcp.server.models import InitializationOptions
12+
from mcp.server.session import ServerSession
13+
from mcp.shared.message import SessionMessage
14+
from mcp.types import ServerCapabilities
15+
16+
17+
@pytest.mark.anyio
18+
async def test_handle_message_with_exception_logging(caplog):
19+
"""Test that Exception instances passed to _handle_message are properly logged."""
20+
server = Server("test")
21+
22+
# Create in-memory streams for testing
23+
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
24+
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
25+
26+
# Create a server session
27+
session = ServerSession(
28+
read_stream=client_to_server_receive,
29+
write_stream=server_to_client_send,
30+
init_options=InitializationOptions(
31+
server_name="test",
32+
server_version="1.0.0",
33+
capabilities=ServerCapabilities(),
34+
),
35+
)
36+
37+
# Create a test exception
38+
test_exception = ValueError("Test exception for logging")
39+
40+
# Test the _handle_message method directly with an Exception
41+
with caplog.at_level(logging.ERROR):
42+
await server._handle_message(
43+
message=test_exception,
44+
session=session,
45+
lifespan_context=None,
46+
raise_exceptions=False,
47+
)
48+
49+
# Verify that the exception was logged
50+
assert len(caplog.records) == 1
51+
record = caplog.records[0]
52+
assert record.levelno == logging.ERROR
53+
assert "Error in message processing" in record.getMessage()
54+
assert "Test exception for logging" in record.getMessage()
55+
56+
57+
@pytest.mark.anyio
58+
async def test_handle_message_with_exception_raising():
59+
"""Test that Exception instances are re-raised when raise_exceptions=True."""
60+
server = Server("test")
61+
62+
# Create in-memory streams for testing
63+
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
64+
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
65+
66+
# Create a server session
67+
session = ServerSession(
68+
read_stream=client_to_server_receive,
69+
write_stream=server_to_client_send,
70+
init_options=InitializationOptions(
71+
server_name="test",
72+
server_version="1.0.0",
73+
capabilities=ServerCapabilities(),
74+
),
75+
)
76+
77+
# Create a test exception
78+
test_exception = ValueError("Test exception for raising")
79+
80+
# Test that the exception is re-raised when raise_exceptions=True
81+
with pytest.raises(ValueError, match="Test exception for raising"):
82+
await server._handle_message(
83+
message=test_exception,
84+
session=session,
85+
lifespan_context=None,
86+
raise_exceptions=True,
87+
)
88+
89+
90+
@pytest.mark.anyio
91+
async def test_handle_message_with_exception_no_raise():
92+
"""Test that Exception instances are not re-raised when raise_exceptions=False."""
93+
server = Server("test")
94+
95+
# Create in-memory streams for testing
96+
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
97+
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
98+
99+
# Create a server session
100+
session = ServerSession(
101+
read_stream=client_to_server_receive,
102+
write_stream=server_to_client_send,
103+
init_options=InitializationOptions(
104+
server_name="test",
105+
server_version="1.0.0",
106+
capabilities=ServerCapabilities(),
107+
),
108+
)
109+
110+
# Create a test exception
111+
test_exception = RuntimeError("Test exception for no raising")
112+
113+
# Test that the exception is not re-raised when raise_exceptions=False
114+
# This should not raise an exception
115+
await server._handle_message(
116+
message=test_exception,
117+
session=session,
118+
lifespan_context=None,
119+
raise_exceptions=False,
120+
)
121+
# If we reach this point, the test passed

0 commit comments

Comments
 (0)