Close the coverage gaps in the rewritten transport

Removes now-dead code (the attach take-over branch, redundant containment
around the notification handler, an unreachable None-connection arm, the
channel-closing loop in run() teardown) and adds tests for the behaviours
that were untested: terminating a session with a request in flight,
client-posted progress for a server-initiated request, containment of a
raising event store per request, concurrent POSTs sharing a request id, the
channel attach/detach identity guard, and the serve_loop driver.

Author: maxisbey

src/mcp/server/runner.py

@@ -477,9 +477,9 @@ async def serve_loop(
     """Drive ``server`` in handshake-only loop mode over a stream pair until the channel closes.
 
     Builds the loop-mode `JSONRPCDispatcher` + `Connection` and hands them to
-    `serve_connection`. The streamable-HTTP manager (which owns its lifespan
-    and serves the modern era on the single-exchange entry instead) calls
-    this; `Server.run` drives `serve_dual_era_loop`, which extends the same
+    `serve_connection`. For a transport that supplies a duplex message stream
+    pair but owns its own lifespan (so `Server.run`'s lifespan entry is not
+    wanted); `Server.run` drives `serve_dual_era_loop`, which extends the same
     dispatcher recipe (notably the `inline_methods={"initialize"}` rule) with
     era routing.
     """

src/mcp/server/streamable_http.py

@@ -224,9 +224,8 @@ def attach(self) -> MemoryObjectReceiveStream[EventMessage]:
         (the standalone GET stream); re-attaching after a detach is how a
         `Last-Event-ID` reconnect resumes a live stream.
         """
+        assert self._writer is None, "every attach site checks `attached` first"
         writer, reader = anyio.create_memory_object_stream[EventMessage](REQUEST_STREAM_BUFFER_SIZE)
-        if self._writer is not None:
-            self._writer.close()
         self._writer = writer
         return reader
 
@@ -262,8 +261,15 @@ def close(self) -> None:
         self.finished.set()
 
     def finish(self) -> None:
-        """Mark the request as finished even if no terminal frame was recorded."""
+        """The request is over: mark it finished and end any response still attached.
+
+        Normally the terminal frame already ended the response; this covers a
+        request whose terminal write never landed (e.g. the event store raised
+        mid-stream), so the client sees the stream close rather than hanging.
+        """
         self.finished.set()
+        if self._writer is not None:
+            self._detach(self._writer)
 
 
 @dataclass
@@ -513,6 +519,9 @@ async def run(self, *, task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STAT
         cancelling it (server shutdown or the idle timeout) cancels the
         handlers and tears the connection down.
         """
+        self._require_app()
+        connection = self._connection
+        assert connection is not None, "a session-bound transport always has a connection"
         try:
             async with anyio.create_task_group() as tg:
                 self._task_group = tg
@@ -521,15 +530,12 @@ async def run(self, *, task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STAT
                 tg.cancel_scope.cancel()
         finally:
             self._task_group = None
-            # End every still-open response stream and wake anything awaiting a
-            # client answer; runs on termination and on manager shutdown alike.
-            for channel in list(self._channels.values()):
-                channel.close()
-            self._channels.clear()
+            # By now every request handler has finished and released its own
+            # channel; end the standalone stream and wake anything awaiting a
+            # client answer (runs on termination and manager shutdown alike).
             self._standalone.close()
             self._corr.close()
-            if self._connection is not None:
-                await aclose_shielded(self._connection)
+            await aclose_shielded(connection)
 
     def _build_message_metadata(
         self, request: Request, request_id: RequestId, protocol_version: str
@@ -650,7 +656,14 @@ def _sse_headers(self) -> dict[str, str]:
         }
 
     async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
-        """Application entry point that handles all HTTP requests."""
+        """Application entry point that handles all HTTP requests.
+
+        Raises:
+            RuntimeError: The transport was constructed without a server to
+                dispatch to (`app`); it is created and driven by
+                `StreamableHTTPSessionManager`.
+        """
+        self._require_app()
         request = Request(scope, receive)
 
         # Validate request headers for DNS rebinding protection
@@ -807,8 +820,8 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re
             return
 
     def _session_runner(self) -> ServerRunner[Any]:
-        """The stateful session's handler kernel."""
-        assert self._runner is not None, "stateful session was built without a server"
+        """The stateful session's handler kernel; built with the transport's server binding."""
+        assert self._runner is not None
         return self._runner
 
     def _stateless_runner(self, request: Request) -> ServerRunner[Any]:
@@ -853,10 +866,10 @@ async def _deliver_client_message(
         )
 
         async def _run_notification() -> None:
+            # `on_notify` contains handler exceptions itself, so a crashing
+            # notification handler cannot take the session down.
             try:
                 await runner.on_notify(dctx, message.method, message.params)
-            except Exception:  # a crashing notification handler must not take the session down
-                logger.exception("notification handler for %r raised", message.method)
             finally:
                 if connection is not None:
                     await aclose_shielded(connection)

src/mcp/shared/jsonrpc_dispatcher.py

@@ -19,7 +19,6 @@
 import anyio
 import anyio.abc
 from mcp_types import (
-    CONNECTION_CLOSED,
     INTERNAL_ERROR,
     ErrorData,
     JSONRPCError,
@@ -51,7 +50,7 @@
     as_request_id,
     run_notify_intercept,
 )
-from mcp.shared.exceptions import MCPError, NoBackChannelError
+from mcp.shared.exceptions import NoBackChannelError
 from mcp.shared.message import (
     ClientMessageMetadata,
     MessageMetadata,
@@ -271,10 +270,9 @@ async def send_raw_request(
                 transport closed or the dispatcher shut down.
             RuntimeError: Called before `run()`.
         """
-        # Post-close sends get the same CONNECTION_CLOSED contract as in-flight waiters.
-        if self._corr.closed:
-            raise MCPError(code=CONNECTION_CLOSED, message="Connection closed")
-        if not self._running:
+        # Post-close sends get the same CONNECTION_CLOSED contract as in-flight
+        # waiters (raised by the correlator); only a never-run dispatcher is a usage error.
+        if not self._running and not self._corr.closed:
             raise RuntimeError("JSONRPCDispatcher.send_raw_request called before run()")
         plan = _plan_outbound(_related_request_id, opts)
         return await self._corr.call(

tests/server/test_runner.py

@@ -53,6 +53,7 @@
 )
 
 import mcp.server.runner
+from mcp.client.session import ClientSession
 from mcp.server.caching import CacheHint
 from mcp.server.connection import Connection, NotifyOnlyOutbound
 from mcp.server.context import ServerRequestContext
@@ -67,6 +68,7 @@
     aclose_shielded,
     serve_connection,
     serve_dual_era_loop,
+    serve_loop,
     serve_one,
 )
 from mcp.server.session import ServerSession
@@ -75,6 +77,7 @@
 from mcp.shared.dispatcher import CallOptions
 from mcp.shared.exceptions import MCPError, NoBackChannelError
 from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
+from mcp.shared.memory import create_client_server_memory_streams
 from mcp.shared.message import MessageMetadata, SessionMessage
 from mcp.shared.peer import dump_params
 from mcp.shared.transport_context import TransportContext
@@ -2056,3 +2059,22 @@ async def test_dual_era_client_propagates_body_exception_unwrapped(server: SrvT)
     with pytest.raises(RuntimeError, match="boom"):
         async with dual_era_client(server):
             raise RuntimeError("boom")
+
+
+@pytest.mark.anyio
+async def test_serve_loop_serves_a_handshake_connection_over_a_stream_pair(server: SrvT) -> None:
+    """`serve_loop`, the loop-mode driver for transports that own their own lifespan, round-trips
+    a handshake and a request over a duplex stream pair and returns when the channel closes."""
+    async with create_client_server_memory_streams() as ((client_read, client_write), (server_read, server_write)):
+        with anyio.fail_after(5):
+            async with anyio.create_task_group() as tg:  # pragma: no branch
+                tg.start_soon(
+                    partial(serve_loop, server, server_read, server_write, lifespan_state={}, session_id="loop-1")
+                )
+                async with ClientSession(client_read, client_write) as session:
+                    initialized = await session.initialize()
+                    tools = await session.list_tools()
+                    assert initialized.server_info.name == "test-server"
+                    assert [tool.name for tool in tools.tools] == ["t"]
+                # Closing the client's write side EOFs the loop and lets it return.
+                await client_write.aclose()