Fix #201: Move incoming message stream from BaseSession to ServerSession (#325)

Co-authored-by: Claude <noreply@anthropic.com>

tests/client/test_logging_callback.py

@@ -1,11 +1,12 @@
 from typing import Literal
 
-import anyio
 import pytest
 
+import mcp.types as types
 from mcp.shared.memory import (
     create_connected_server_and_client_session as create_session,
 )
+from mcp.shared.session import RequestResponder
 from mcp.types import (
     LoggingMessageNotificationParams,
     TextContent,
@@ -46,40 +47,37 @@ async def test_logging_callback():
         )
         return True
 
-    async with anyio.create_task_group() as tg:
-        async with create_session(
-            server._mcp_server, logging_callback=logging_collector
-        ) as client_session:
+    # Create a message handler to catch exceptions
+    async def message_handler(
+        message: RequestResponder[types.ServerRequest, types.ClientResult]
+        | types.ServerNotification
+        | Exception,
+    ) -> None:
+        if isinstance(message, Exception):
+            raise message
 
-            async def listen_session():
-                try:
-                    async for message in client_session.incoming_messages:
-                        if isinstance(message, Exception):
-                            raise message
-                except anyio.EndOfStream:
-                    pass
+    async with create_session(
+        server._mcp_server,
+        logging_callback=logging_collector,
+        message_handler=message_handler,
+    ) as client_session:
+        # First verify our test tool works
+        result = await client_session.call_tool("test_tool", {})
+        assert result.isError is False
+        assert isinstance(result.content[0], TextContent)
+        assert result.content[0].text == "true"
 
-            tg.start_soon(listen_session)
-
-            # First verify our test tool works
-            result = await client_session.call_tool("test_tool", {})
-            assert result.isError is False
-            assert isinstance(result.content[0], TextContent)
-            assert result.content[0].text == "true"
-
-            # Now send a log message via our tool
-            log_result = await client_session.call_tool(
-                "test_tool_with_log",
-                {
-                    "message": "Test log message",
-                    "level": "info",
-                    "logger": "test_logger",
-                },
-            )
-            assert log_result.isError is False
-            assert len(logging_collector.log_messages) == 1
-            assert logging_collector.log_messages[
-                0
-            ] == LoggingMessageNotificationParams(
-                level="info", logger="test_logger", data="Test log message"
-            )
+        # Now send a log message via our tool
+        log_result = await client_session.call_tool(
+            "test_tool_with_log",
+            {
+                "message": "Test log message",
+                "level": "info",
+                "logger": "test_logger",
+            },
+        )
+        assert log_result.isError is False
+        assert len(logging_collector.log_messages) == 1
+        assert logging_collector.log_messages[0] == LoggingMessageNotificationParams(
+            level="info", logger="test_logger", data="Test log message"
+        )

tests/client/test_session.py

@@ -1,7 +1,9 @@
 import anyio
 import pytest
 
+import mcp.types as types
 from mcp.client.session import ClientSession
+from mcp.shared.session import RequestResponder
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
     ClientNotification,
@@ -75,13 +77,21 @@ async def test_client_session_initialize():
                 )
             )
 
-    async def listen_session():
-        async for message in session.incoming_messages:
-            if isinstance(message, Exception):
-                raise message
+    # Create a message handler to catch exceptions
+    async def message_handler(
+        message: RequestResponder[types.ServerRequest, types.ClientResult]
+        | types.ServerNotification
+        | Exception,
+    ) -> None:
+        if isinstance(message, Exception):
+            raise message
 
     async with (
-        ClientSession(server_to_client_receive, client_to_server_send) as session,
+        ClientSession(
+            server_to_client_receive,
+            client_to_server_send,
+            message_handler=message_handler,
+        ) as session,
         anyio.create_task_group() as tg,
         client_to_server_send,
         client_to_server_receive,
@@ -89,7 +99,6 @@ async def test_client_session_initialize():
         server_to_client_receive,
     ):
         tg.start_soon(mock_server)
-        tg.start_soon(listen_session)
         result = await session.initialize()
 
     # Assert the result