Move incoming message stream from BaseSession to ServerSession

Fixes GitHub issue #201 by moving the incoming message stream and related methods from BaseSession to ServerSession where they are actually needed. This change follows the principle of having functionality only where it's required.

GitHub-Issue:#201

🤖 Generated with [Claude Code](https://claude.ai/code)

Author: dsp-ant

src/mcp/server/session.py

@@ -61,6 +61,12 @@ class InitializationState(Enum):
 
 ServerSessionT = TypeVar("ServerSessionT", bound="ServerSession")
 
+ServerRequestResponder = (
+    RequestResponder[types.ClientRequest, types.ServerResult]
+    | types.ClientNotification
+    | Exception
+)
+
 
 class ServerSession(
     BaseSession[
@@ -85,6 +91,15 @@ def __init__(
         )
         self._initialization_state = InitializationState.NotInitialized
         self._init_options = init_options
+        self._incoming_message_stream_writer, self._incoming_message_stream_reader = (
+            anyio.create_memory_object_stream[ServerRequestResponder]()
+        )
+        self._exit_stack.push_async_callback(
+            lambda: self._incoming_message_stream_reader.aclose()
+        )
+        self._exit_stack.push_async_callback(
+            lambda: self._incoming_message_stream_writer.aclose()
+        )
 
     @property
     def client_params(self) -> types.InitializeRequestParams | None:
@@ -291,3 +306,12 @@ async def send_prompt_list_changed(self) -> None:
                 )
             )
         )
+
+    async def _handle_incoming(self, req: ServerRequestResponder) -> None:
+        return await self._incoming_message_stream_writer.send(req)
+
+    @property
+    def incoming_messages(
+        self,
+    ) -> MemoryObjectReceiveStream[ServerRequestResponder]:
+        return self._incoming_message_stream_reader

src/mcp/shared/session.py

@@ -182,19 +182,6 @@ def __init__(
         self._in_flight = {}
 
         self._exit_stack = AsyncExitStack()
-        self._incoming_message_stream_writer, self._incoming_message_stream_reader = (
-            anyio.create_memory_object_stream[
-                RequestResponder[ReceiveRequestT, SendResultT]
-                | ReceiveNotificationT
-                | Exception
-            ]()
-        )
-        self._exit_stack.push_async_callback(
-            lambda: self._incoming_message_stream_reader.aclose()
-        )
-        self._exit_stack.push_async_callback(
-            lambda: self._incoming_message_stream_writer.aclose()
-        )
 
     async def __aenter__(self) -> Self:
         self._task_group = anyio.create_task_group()
@@ -300,11 +287,10 @@ async def _receive_loop(self) -> None:
         async with (
             self._read_stream,
             self._write_stream,
-            self._incoming_message_stream_writer,
         ):
             async for message in self._read_stream:
                 if isinstance(message, Exception):
-                    await self._incoming_message_stream_writer.send(message)
+                    await self._handle_incoming(message)
                 elif isinstance(message.root, JSONRPCRequest):
                     validated_request = self._receive_request_type.model_validate(
                         message.root.model_dump(
@@ -325,7 +311,7 @@ async def _receive_loop(self) -> None:
                     self._in_flight[responder.request_id] = responder
                     await self._received_request(responder)
                     if not responder._completed:
-                        await self._incoming_message_stream_writer.send(responder)
+                        await self._handle_incoming(responder)
 
                 elif isinstance(message.root, JSONRPCNotification):
                     try:
@@ -341,9 +327,7 @@ async def _receive_loop(self) -> None:
                                 await self._in_flight[cancelled_id].cancel()
                         else:
                             await self._received_notification(notification)
-                            await self._incoming_message_stream_writer.send(
-                                notification
-                            )
+                            await self._handle_incoming(notification)
                     except Exception as e:
                         # For other validation errors, log and continue
                         logging.warning(
@@ -355,7 +339,7 @@ async def _receive_loop(self) -> None:
                     if stream:
                         await stream.send(message.root)
                     else:
-                        await self._incoming_message_stream_writer.send(
+                        await self._handle_incoming(
                             RuntimeError(
                                 "Received response with an unknown "
                                 f"request ID: {message}"
@@ -387,12 +371,11 @@ async def send_progress_notification(
         processed.
         """
 
-    @property
-    def incoming_messages(
+    async def _handle_incoming(
         self,
-    ) -> MemoryObjectReceiveStream[
-        RequestResponder[ReceiveRequestT, SendResultT]
+        req: RequestResponder[ReceiveRequestT, SendResultT]
         | ReceiveNotificationT
-        | Exception
-    ]:
-        return self._incoming_message_stream_reader
+        | Exception,
+    ) -> None:
+        """A generic handler for incoming messages. Overwritten by subclasses."""
+        await anyio.lowlevel.checkpoint()