fix: InMemoryTransport EOFs the server instead of cancelling; coverage cleanup

src/mcp/client/_memory.py: replace tg.cancel_scope.cancel() with stream
aclose(). Cancelling here throws CancelledError into the host (test's)
task; on CPython 3.11 (gh-106749) coro.throw() drops 'call' trace
events for the outer await chain, underflowing coverage's CTracer past
the test frame. Post-swap the dispatcher's empty-at-EOF inner task group
takes a one-tick fast path, so the join no longer suspends a second time
to heal via .send(). EOF teardown is equivalent (the dispatcher's run()
cancels its own in-flight handlers on read-stream EOF).

Also bundled (coverage cleanup):
- shared/session.py: delete server-only BaseSession code now unreachable
  after the swap (RequestResponder.cancel/_cancel_scope/_on_complete,
  CancelledNotification handling, _in_flight dict, deferred-respond
  arm). ClientSession is the only remaining subclass.
- tests/client/test_session.py: add tests for the client-reachable
  BaseSession paths (malformed inbound request -> INVALID_PARAMS,
  sampling callback raises -> INVALID_PARAMS, progress callback
  exception swallowed, malformed notification dropped, transport
  exception forwarded to message_handler).
- tests/shared/test_session.py: drop test_in_flight_requests_cleared
  (asserts the deleted _in_flight dict).
- tests/shared/test_dispatcher.py: add contract test for ValidationError
  -> INVALID_PARAMS (covers DirectDispatcher's arm).
- tests/server/test_validation.py: cover the previous-message-has-no-
  tool-use branch.
- examples/everything-server: _add_request_handler (deleted) ->
  add_request_handler with explicit params type.

Author: maxisbey

examples/servers/everything-server/mcp_everything_server/server.py

@@ -417,9 +417,15 @@ async def handle_unsubscribe(ctx: ServerRequestContext, params: UnsubscribeReque
     return EmptyResult()
 
 
-mcp._lowlevel_server._add_request_handler("logging/setLevel", handle_set_logging_level)  # pyright: ignore[reportPrivateUsage]
-mcp._lowlevel_server._add_request_handler("resources/subscribe", handle_subscribe)  # pyright: ignore[reportPrivateUsage]
-mcp._lowlevel_server._add_request_handler("resources/unsubscribe", handle_unsubscribe)  # pyright: ignore[reportPrivateUsage]
+mcp._lowlevel_server.add_request_handler(  # pyright: ignore[reportPrivateUsage]
+    "logging/setLevel", SetLevelRequestParams, handle_set_logging_level
+)
+mcp._lowlevel_server.add_request_handler(  # pyright: ignore[reportPrivateUsage]
+    "resources/subscribe", SubscribeRequestParams, handle_subscribe
+)
+mcp._lowlevel_server.add_request_handler(  # pyright: ignore[reportPrivateUsage]
+    "resources/unsubscribe", UnsubscribeRequestParams, handle_unsubscribe
+)
 
 
 @mcp.completion()

src/mcp/client/_memory.py

@@ -62,7 +62,15 @@ async def _connect(self) -> AsyncIterator[TransportStreams]:
                 try:
                     yield client_read, client_write
                 finally:
-                    tg.cancel_scope.cancel()
+                    # EOF the server (and our own read side) instead of
+                    # cancelling. The dispatcher's run() cancels its own
+                    # in-flight handlers on read-stream EOF, so cleanup is
+                    # equivalent. Cancelling here would `coro.throw()` into the
+                    # host task, which on CPython 3.11 (gh-106749) drops
+                    # `'call'` trace events for the outer await chain and
+                    # desyncs coverage's CTracer past the test frame.
+                    await client_write.aclose()
+                    await server_write.aclose()
 
     async def __aenter__(self) -> TransportStreams:
         """Connect to the server and return streams for communication."""

src/mcp/shared/session.py

@@ -2,7 +2,6 @@
 
 import contextvars
 import logging
-from collections.abc import Callable
 from contextlib import AsyncExitStack
 from types import TracebackType
 from typing import Any, Generic, Protocol, TypeVar
@@ -21,7 +20,6 @@
     CONNECTION_CLOSED,
     INVALID_PARAMS,
     REQUEST_TIMEOUT,
-    CancelledNotification,
     ClientNotification,
     ClientRequest,
     ClientResult,
@@ -80,7 +78,6 @@ def __init__(
         request_meta: RequestParamsMeta | None,
         request: ReceiveRequestT,
         session: BaseSession[SendRequestT, SendNotificationT, SendResultT, ReceiveRequestT, ReceiveNotificationT],
-        on_complete: Callable[[RequestResponder[ReceiveRequestT, SendResultT]], Any],
         message_metadata: MessageMetadata = None,
         context: contextvars.Context | None = None,
     ) -> None:
@@ -91,15 +88,10 @@ def __init__(
         self.context = context
         self._session = session
         self._completed = False
-        self._cancel_scope = anyio.CancelScope()
-        self._on_complete = on_complete
         self._entered = False  # Track if we're in a context manager
 
     def __enter__(self) -> RequestResponder[ReceiveRequestT, SendResultT]:
-        """Enter the context manager, enabling request cancellation tracking."""
         self._entered = True
-        self._cancel_scope = anyio.CancelScope()
-        self._cancel_scope.__enter__()
         return self
 
     def __exit__(
@@ -108,15 +100,7 @@ def __exit__(
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        """Exit the context manager, performing cleanup and notifying completion."""
-        try:
-            if self._completed:
-                self._on_complete(self)
-        finally:
-            self._entered = False
-            if not self._cancel_scope:  # pragma: no cover
-                raise RuntimeError("No active cancel scope")
-            self._cancel_scope.__exit__(exc_type, exc_val, exc_tb)
+        self._entered = False
 
     async def respond(self, response: SendResultT | ErrorData) -> None:
         """Send a response for this request.
@@ -130,37 +114,11 @@ async def respond(self, response: SendResultT | ErrorData) -> None:
         if not self._entered:  # pragma: no cover
             raise RuntimeError("RequestResponder must be used as a context manager")
         assert not self._completed, "Request already responded to"
-
-        if not self.cancelled:  # pragma: no branch
-            self._completed = True
-
-            await self._session._send_response(  # type: ignore[reportPrivateUsage]
-                request_id=self.request_id, response=response
-            )
-
-    async def cancel(self) -> None:
-        """Cancel this request and mark it as completed."""
-        if not self._entered:  # pragma: no cover
-            raise RuntimeError("RequestResponder must be used as a context manager")
-        if not self._cancel_scope:  # pragma: no cover
-            raise RuntimeError("No active cancel scope")
-
-        self._cancel_scope.cancel()
-        self._completed = True  # Mark as completed so it's removed from in_flight
-        # Send an error response to indicate cancellation
+        self._completed = True
         await self._session._send_response(  # type: ignore[reportPrivateUsage]
-            request_id=self.request_id,
-            response=ErrorData(code=0, message="Request cancelled"),
+            request_id=self.request_id, response=response
         )
 
-    @property
-    def in_flight(self) -> bool:  # pragma: no cover
-        return not self._completed and not self.cancelled
-
-    @property
-    def cancelled(self) -> bool:
-        return self._cancel_scope.cancel_called
-
 
 class BaseSession(
     Generic[
@@ -180,7 +138,6 @@ class BaseSession(
 
     _response_streams: dict[RequestId, MemoryObjectSendStream[JSONRPCResponse | JSONRPCError]]
     _request_id: int
-    _in_flight: dict[RequestId, RequestResponder[ReceiveRequestT, SendResultT]]
     _progress_callbacks: dict[RequestId, ProgressFnT]
 
     def __init__(
@@ -195,7 +152,6 @@ def __init__(
         self._response_streams = {}
         self._request_id = 0
         self._session_read_timeout_seconds = read_timeout_seconds
-        self._in_flight = {}
         self._progress_callbacks = {}
         self._exit_stack = AsyncExitStack()
 
@@ -347,15 +303,10 @@ async def _handle_session_message(message: SessionMessage) -> None:
                                 request_meta=validated_request.params.meta if validated_request.params else None,
                                 request=validated_request,
                                 session=self,
-                                on_complete=lambda r: self._in_flight.pop(r.request_id, None),
                                 message_metadata=message.metadata,
                                 context=sender_context,
                             )
-                            self._in_flight[responder.request_id] = responder
                             await self._received_request(responder)
-
-                            if not responder._completed:  # type: ignore[reportPrivateUsage]
-                                await self._handle_incoming(responder)
                         except Exception:
                             # For request validation errors, send a proper JSON-RPC error
                             # response instead of crashing the server
@@ -375,33 +326,28 @@ async def _handle_session_message(message: SessionMessage) -> None:
                                 message.message.model_dump(by_alias=True, mode="json", exclude_none=True),
                                 by_name=False,
                             )
-                            # Handle cancellation notifications
-                            if isinstance(notification, CancelledNotification):
-                                cancelled_id = notification.params.request_id
-                                if cancelled_id in self._in_flight:  # pragma: no branch
-                                    await self._in_flight[cancelled_id].cancel()
-                            else:
-                                # Handle progress notifications callback
-                                if isinstance(notification, ProgressNotification):
-                                    progress_token = notification.params.progress_token
-                                    # If there is a progress callback for this token,
-                                    # call it with the progress information
-                                    if progress_token in self._progress_callbacks:
-                                        callback = self._progress_callbacks[progress_token]
-                                        try:
-                                            await callback(
-                                                notification.params.progress,
-                                                notification.params.total,
-                                                notification.params.message,
-                                            )
-                                        except Exception:
-                                            logging.exception("Progress callback raised an exception")
-                                await self._received_notification(notification)
-                                await self._handle_incoming(notification)
+                            # Handle progress notifications callback
+                            if isinstance(notification, ProgressNotification):
+                                progress_token = notification.params.progress_token
+                                # If there is a progress callback for this token,
+                                # call it with the progress information
+                                if progress_token in self._progress_callbacks:
+                                    callback = self._progress_callbacks[progress_token]
+                                    try:
+                                        await callback(
+                                            notification.params.progress,
+                                            notification.params.total,
+                                            notification.params.message,
+                                        )
+                                    except Exception:
+                                        logging.exception("Progress callback raised an exception")
+                            await self._received_notification(notification)
+                            await self._handle_incoming(notification)
                         except Exception:
                             # For other validation errors, log and continue
-                            logging.warning(  # pragma: no cover
-                                f"Failed to validate notification:. Message was: {message.message}",
+                            logging.warning(
+                                "Failed to validate notification: %s",
+                                message.message,
                                 exc_info=True,
                             )
                     else:  # Response or error

tests/client/test_session.py

@@ -1,6 +1,11 @@
 from __future__ import annotations
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
 import anyio
+import anyio.streams.memory
 import pytest
 
 from mcp import types
@@ -10,12 +15,14 @@
 from mcp.shared.session import RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 from mcp.types import (
+    INVALID_PARAMS,
     LATEST_PROTOCOL_VERSION,
     CallToolResult,
     Implementation,
     InitializedNotification,
     InitializeRequest,
     InitializeResult,
+    JSONRPCError,
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
@@ -26,6 +33,29 @@
     client_request_adapter,
 )
 
+_SendToClient = anyio.streams.memory.MemoryObjectSendStream[SessionMessage | Exception]
+_RecvFromClient = anyio.streams.memory.MemoryObjectReceiveStream[SessionMessage]
+
+
+@asynccontextmanager
+async def raw_client_session(
+    **kwargs: Any,
+) -> AsyncIterator[tuple[ClientSession, _SendToClient, _RecvFromClient]]:
+    """Yield `(session, send_to_client, recv_from_client)` with the receive loop running.
+
+    `send_to_client` accepts `SessionMessage | Exception` so tests can inject
+    transport-level exceptions. No initialize handshake is performed.
+    """
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage](32)
+    async with ClientSession(s2c_recv, c2s_send, **kwargs) as session:
+        try:
+            with anyio.fail_after(5):
+                yield session, s2c_send, c2s_recv
+        finally:
+            s2c_send.close()
+            c2s_recv.close()
+
 
 @pytest.mark.anyio
 async def test_client_session_initialize():
@@ -705,3 +735,96 @@ async def mock_server():
         await session.initialize()
 
         await session.call_tool(name=mocked_tool.name, arguments={"foo": "bar"}, meta=meta)
+
+
+@pytest.mark.anyio
+async def test_receive_loop_answers_malformed_inbound_request_with_invalid_params():
+    """A request that fails ServerRequest validation gets an INVALID_PARAMS error response."""
+    async with raw_client_session() as (_session, to_client, from_client):
+        await to_client.send(
+            SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=7, method="sampling/createMessage", params={"broken": 1}))
+        )
+        out = await from_client.receive()
+    assert isinstance(out.message, JSONRPCError)
+    assert out.message.id == 7
+    assert out.message.error.code == INVALID_PARAMS
+
+
+@pytest.mark.anyio
+async def test_receive_loop_answers_invalid_params_when_sampling_callback_raises():
+    """Same boundary catches exceptions from the request handler itself."""
+
+    async def boom(ctx: object, params: object) -> types.CreateMessageResult:
+        raise RuntimeError("sampling boom")
+
+    params = types.CreateMessageRequestParams(
+        messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hi"))],
+        max_tokens=10,
+    ).model_dump(by_alias=True, mode="json", exclude_none=True)
+    async with raw_client_session(sampling_callback=boom) as (_session, to_client, from_client):
+        await to_client.send(
+            SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=8, method="sampling/createMessage", params=params))
+        )
+        out = await from_client.receive()
+    assert isinstance(out.message, JSONRPCError)
+    assert out.message.error.code == INVALID_PARAMS
+
+
+@pytest.mark.anyio
+async def test_receive_loop_logs_and_drops_malformed_notification(caplog: pytest.LogCaptureFixture):
+    """A notification that fails ServerNotification validation is logged and dropped."""
+    seen: list[object] = []
+    delivered = anyio.Event()
+
+    async def handler(msg: object) -> None:
+        seen.append(msg)
+        delivered.set()
+
+    async with raw_client_session(message_handler=handler) as (_session, to_client, _):
+        await to_client.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="not/a/spec/notification")))
+        # Follow with a valid notification so we know the loop is still alive.
+        await to_client.send(
+            SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/tools/list_changed"))
+        )
+        await delivered.wait()
+    assert isinstance(seen[0], types.ToolListChangedNotification)
+    assert "Failed to validate notification" in caplog.text
+
+
+@pytest.mark.anyio
+async def test_receive_loop_forwards_transport_exception_to_message_handler():
+    seen: list[object] = []
+    delivered = anyio.Event()
+
+    async def handler(msg: object) -> None:
+        seen.append(msg)
+        delivered.set()
+
+    async with raw_client_session(message_handler=handler) as (_session, to_client, _):
+        exc = ValueError("bad bytes")
+        await to_client.send(exc)
+        await delivered.wait()
+    assert seen == [exc]
+
+
+@pytest.mark.anyio
+async def test_receive_loop_swallows_progress_callback_exception(caplog: pytest.LogCaptureFixture):
+    delivered = anyio.Event()
+
+    async def boom(progress: float, total: float | None, message: str | None) -> None:
+        raise RuntimeError("progress boom")
+
+    async def handler(msg: object) -> None:
+        delivered.set()
+
+    async with raw_client_session(message_handler=handler) as (session, to_client, _):
+        # Register the callback under a known token without sending a request.
+        session._progress_callbacks[42] = boom  # pyright: ignore[reportPrivateUsage]
+        params = {"progressToken": 42, "progress": 0.5}
+        await to_client.send(
+            SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/progress", params=params))
+        )
+        # The progress notification also reaches the message handler after the
+        # callback runs, so this fires once the callback's exception is handled.
+        await delivered.wait()
+    assert "Progress callback raised an exception" in caplog.text

tests/server/test_validation.py

@@ -120,6 +120,16 @@ def test_validate_tool_use_result_messages_raises_when_tool_result_without_previ
         validate_tool_use_result_messages(messages)
 
 
+def test_validate_tool_use_result_messages_raises_when_previous_message_has_no_tool_use() -> None:
+    """Raises when tool_result follows a message that has content but no tool_use."""
+    messages = [
+        SamplingMessage(role="assistant", content=TextContent(type="text", text="just text")),
+        SamplingMessage(role="user", content=ToolResultContent(type="tool_result", tool_use_id="tool-1")),
+    ]
+    with pytest.raises(ValueError, match="do not match any tool_use in the previous message"):
+        validate_tool_use_result_messages(messages)
+
+
 def test_validate_tool_use_result_messages_raises_when_tool_result_ids_dont_match_tool_use() -> None:
     """Raises when tool_result IDs don't match tool_use IDs."""
     messages = [