fix: address review feedback and add receive middleware test

- Move 'import inspect' to top of file
- Pre-compute whether middleware is async using inspect.iscoroutinefunction()
  instead of checking on every message
- Add test for receive_middleware to fix coverage

Author: jerome3o-anthropic

src/mcp/shared/session.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 from collections.abc import Awaitable, Callable
 from contextlib import AsyncExitStack
@@ -209,22 +210,24 @@ def __init__(
         self._progress_callbacks = {}
         self._response_routers = []
         self._exit_stack = AsyncExitStack()
-        self._send_middleware = send_middleware or []
-        self._receive_middleware = receive_middleware or []
+        # Pre-compute whether each middleware is async to avoid checking on every message
+        self._send_middleware: list[tuple[MessageMiddleware, bool]] = [
+            (m, inspect.iscoroutinefunction(m)) for m in (send_middleware or [])
+        ]
+        self._receive_middleware: list[tuple[MessageMiddleware, bool]] = [
+            (m, inspect.iscoroutinefunction(m)) for m in (receive_middleware or [])
+        ]
 
     async def _apply_middleware(
-        self, message: JSONRPCMessage, middleware_list: list[MessageMiddleware]
+        self, message: JSONRPCMessage, middleware_list: list[tuple[MessageMiddleware, bool]]
     ) -> JSONRPCMessage:
         """Apply a list of middleware functions to a message."""
-        import inspect
-
-        for middleware in middleware_list:
+        for middleware, is_async in middleware_list:
             result = middleware(message)
-            if inspect.isawaitable(result):
-                message = await result
-            else:
-                message = result  # type: ignore[assignment]
-        return message
+            if is_async:
+                result = await result  # type: ignore[misc]
+            message = result  # type: ignore[assignment]
+        return message  # type: ignore[return-value]
 
     def add_response_router(self, router: ResponseRouter) -> None:
         """

tests/client/test_session.py

@@ -909,3 +909,67 @@ async def mock_server():
         await session.initialize()
 
     assert middleware_called
+
+
+@pytest.mark.anyio
+async def test_client_session_receive_middleware():
+    """Test that receive middleware can transform incoming messages."""
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
+
+    middleware_called = False
+    received_response = None
+
+    def receive_transform(message: JSONRPCMessage) -> JSONRPCMessage:
+        """Middleware that observes incoming messages."""
+        nonlocal middleware_called, received_response
+        middleware_called = True
+        if isinstance(message.root, JSONRPCResponse):
+            received_response = message.root
+        return message
+
+    async def mock_server():
+        session_message = await client_to_server_receive.receive()
+        jsonrpc_request = session_message.message
+        assert isinstance(jsonrpc_request.root, JSONRPCRequest)
+
+        result = ServerResult(
+            InitializeResult(
+                protocolVersion=LATEST_PROTOCOL_VERSION,
+                capabilities=ServerCapabilities(),
+                serverInfo=Implementation(name="mock-server", version="0.1.0"),
+            )
+        )
+
+        async with server_to_client_send:
+            await server_to_client_send.send(
+                SessionMessage(
+                    JSONRPCMessage(
+                        JSONRPCResponse(
+                            jsonrpc="2.0",
+                            id=jsonrpc_request.root.id,
+                            result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
+                        )
+                    )
+                )
+            )
+            await client_to_server_receive.receive()
+
+    async with (
+        ClientSession(
+            server_to_client_receive,
+            client_to_server_send,
+            receive_middleware=[receive_transform],
+        ) as session,
+        anyio.create_task_group() as tg,
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        tg.start_soon(mock_server)
+        await session.initialize()
+
+    # Verify receive middleware was called and saw the response
+    assert middleware_called
+    assert received_response is not None