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

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

src/mcp/client/__main__.py

@@ -7,9 +7,11 @@ from urllib.parse import urlparse
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
+import mcp.types as types
 from mcp.client.session import ClientSession
 from mcp.client.sse import sse_client
 from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.shared.session import RequestResponder
 from mcp.types import JSONRPCMessage
 
 if not sys.warnoptions:
@@ -21,26 +23,25 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("client")
 
 
-async def receive_loop(session: ClientSession):
-    logger.info("Starting receive loop")
-    async for message in session.incoming_messages:
-        if isinstance(message, Exception):
-            logger.error("Error: %s", message)
-            continue
+async def message_handler(
+    message: RequestResponder[types.ServerRequest, types.ClientResult]
+    | types.ServerNotification
+    | Exception,
+) -> None:
+    if isinstance(message, Exception):
+        logger.error("Error: %s", message)
+        return
 
-        logger.info("Received message from server: %s", message)
+    logger.info("Received message from server: %s", message)
 
 
 async def run_session(
     read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
     write_stream: MemoryObjectSendStream[JSONRPCMessage],
 ):
-    async with (
-        ClientSession(read_stream, write_stream) as session,
-        anyio.create_task_group() as tg,
-    ):
-        tg.start_soon(receive_loop, session)
-
+    async with ClientSession(
+        read_stream, write_stream, message_handler=message_handler
+    ) as session:
         logger.info("Initializing session")
         await session.initialize()
         logger.info("Initialized")