feat: Connection.state + exit_stack; ctx.session_id/headers; TransportContext.headers

Per-connection state without a connection_lifespan CM or a second Server
generic. Stateless is the default deployment, where a per-connection
lifespan would wrap a single request; the enter-late mechanics it would
need (race init vs dispatcher-done, ready-gate) were more machinery than
the use case warrants.

- Connection.session_id: str | None — set by the mount via
  ServerRunner(session_id=...); per-connection, not per-message
- Connection.state: dict[str, Any] — scratch that persists across
  requests; handlers/middleware read and write freely
- Connection.exit_stack: AsyncExitStack — handlers/middleware push CMs
  or callbacks for per-connection teardown; ServerRunner.run() unwinds
  it (shielded) in a finally after dispatcher.run() returns
- TransportContext.headers: Mapping[str, str] | None on the base —
  populated by HTTP transports, None on stdio
- Context.session_id / Context.headers convenience properties
- create_direct_dispatcher_pair(headers=...) and
  connected_runner(session_id=..., headers=...) for tests

Author: maxisbey

src/mcp/server/connection.py

@@ -1,9 +1,10 @@
 """`Connection` — per-client connection state and the standalone outbound channel.
 
 Always present on `Context` (never ``None``), even in stateless deployments.
-Holds peer info populated at ``initialize`` time, the per-connection lifespan
-output, and an `Outbound` for the standalone stream (the SSE GET stream in
-streamable HTTP, or the single duplex stream in stdio).
+Holds peer info populated at ``initialize`` time, per-connection scratch
+``state`` and an ``exit_stack`` for teardown, and an `Outbound` for the
+standalone stream (the SSE GET stream in streamable HTTP, or the single duplex
+stream in stdio).
 
 `notify` is best-effort: it never raises. If there's no standalone channel
 (stateless HTTP) or the stream has been dropped, the notification is
@@ -14,6 +15,7 @@
 
 import logging
 from collections.abc import Mapping
+from contextlib import AsyncExitStack
 from typing import Any
 
 import anyio
@@ -44,17 +46,27 @@ class Connection(TypedServerRequestMixin):
     ``None`` until ``initialize`` completes; ``initialized`` is set then.
     """
 
-    def __init__(self, outbound: Outbound, *, has_standalone_channel: bool) -> None:
+    def __init__(self, outbound: Outbound, *, has_standalone_channel: bool, session_id: str | None = None) -> None:
         self._outbound = outbound
         self.has_standalone_channel = has_standalone_channel
+        self.session_id: str | None = session_id
 
         self.client_info: Implementation | None = None
         self.client_capabilities: ClientCapabilities | None = None
         self.protocol_version: str | None = None
         self.initialized: anyio.Event = anyio.Event()
-        # TODO: make this generic (Connection[StateT]) once connection_lifespan
-        # wiring lands in ServerRunner.
-        self.state: Any = None
+
+        self.state: dict[str, Any] = {}
+        """Per-connection scratch state. Handlers and middleware may read and
+        write freely; persists across requests on this connection."""
+
+        self.exit_stack: AsyncExitStack = AsyncExitStack()
+        """Cleanup stack unwound by `ServerRunner` when the connection closes.
+
+        Push context managers (``await exit_stack.enter_async_context(...)``)
+        or callbacks (``exit_stack.push_async_callback(...)``) from handlers or
+        middleware to register per-connection teardown. Unwound LIFO after
+        `dispatcher.run()` returns, shielded from cancellation."""
 
     async def send_raw_request(
         self,

src/mcp/server/context.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass
 from typing import Any, Generic, Protocol
 
@@ -69,6 +69,23 @@ def connection(self) -> Connection:
         """The per-client `Connection` for this request's connection."""
         return self._connection
 
+    @property
+    def session_id(self) -> str | None:
+        """The transport's session id for this connection, when one exists.
+
+        Convenience for ``ctx.connection.session_id``. ``None`` on stdio and
+        stateless HTTP.
+        """
+        return self._connection.session_id
+
+    @property
+    def headers(self) -> Mapping[str, str] | None:
+        """Request headers carried by this message, when the transport has them.
+
+        Convenience for ``ctx.transport.headers``. ``None`` on stdio.
+        """
+        return self.transport.headers
+
     async def log(self, level: LoggingLevel, data: Any, logger: str | None = None, *, meta: Meta | None = None) -> None:
         """Send a request-scoped ``notifications/message`` log entry.
 

src/mcp/server/runner.py

@@ -116,6 +116,7 @@ class ServerRunner(Generic[LifespanT]):
     dispatcher: Dispatcher[TransportContext]
     lifespan_state: LifespanT
     has_standalone_channel: bool
+    session_id: str | None = None
     stateless: bool = False
     dispatch_middleware: list[DispatchMiddleware] = field(default_factory=list[DispatchMiddleware])
 
@@ -124,17 +125,25 @@ class ServerRunner(Generic[LifespanT]):
 
     def __post_init__(self) -> None:
         self._initialized = self.stateless
-        self.connection = Connection(self.dispatcher, has_standalone_channel=self.has_standalone_channel)
+        self.connection = Connection(
+            self.dispatcher, has_standalone_channel=self.has_standalone_channel, session_id=self.session_id
+        )
 
     async def run(self, *, task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED) -> None:
         """Drive the dispatcher until the underlying channel closes.
 
         Composes `dispatch_middleware` over `_on_request` and hands the result
         to `dispatcher.run()`. ``task_status.started()`` is forwarded so callers
         can ``await tg.start(runner.run)`` and resume once the dispatcher is
-        ready to accept requests.
+        ready to accept requests. Once the dispatcher exits,
+        `connection.exit_stack` is unwound (shielded) so any per-connection
+        cleanup registered by handlers or middleware runs to completion.
         """
-        await self.dispatcher.run(self._compose_on_request(), self._on_notify, task_status=task_status)
+        try:
+            await self.dispatcher.run(self._compose_on_request(), self._on_notify, task_status=task_status)
+        finally:
+            with anyio.CancelScope(shield=True):
+                await self.connection.exit_stack.aclose()
 
     def _compose_on_request(self) -> OnRequest:
         """Wrap `_on_request` in `dispatch_middleware`, outermost-first.

src/mcp/shared/direct_dispatcher.py

@@ -162,18 +162,20 @@ async def _dispatch_notify(self, method: str, params: Mapping[str, Any] | None)
 def create_direct_dispatcher_pair(
     *,
     can_send_request: bool = True,
+    headers: Mapping[str, str] | None = None,
 ) -> tuple[DirectDispatcher, DirectDispatcher]:
     """Create two `DirectDispatcher` instances wired to each other.
 
     Args:
         can_send_request: Sets `TransportContext.can_send_request` on both
             sides. Pass ``False`` to simulate a transport with no back-channel.
+        headers: Sets `TransportContext.headers` on both sides.
 
     Returns:
         A ``(left, right)`` pair. Conventionally ``left`` is the client side
         and ``right`` is the server side, but the wiring is symmetric.
     """
-    ctx = TransportContext(kind=DIRECT_TRANSPORT_KIND, can_send_request=can_send_request)
+    ctx = TransportContext(kind=DIRECT_TRANSPORT_KIND, can_send_request=can_send_request, headers=headers)
     left = DirectDispatcher(ctx)
     right = DirectDispatcher(ctx)
     left.connect_to(right)

src/mcp/shared/transport_context.py

@@ -6,6 +6,7 @@
 dispatcher (`ServerRunner`, `Context`, user handlers) read its concrete fields.
 """
 
+from collections.abc import Mapping
 from dataclasses import dataclass
 
 __all__ = ["TransportContext"]
@@ -28,3 +29,10 @@ class TransportContext:
     stdio, SSE, and stateful streamable HTTP. When ``False``,
     `DispatchContext.send_raw_request` raises `NoBackChannelError`.
     """
+
+    headers: Mapping[str, str] | None = None
+    """Request headers carried by this message, when the transport has them.
+
+    Populated by HTTP-based transports; ``None`` on stdio. Handlers should
+    None-check before use.
+    """

tests/server/test_runner.py

@@ -6,8 +6,8 @@
 under test.
 """
 
-from collections.abc import AsyncIterator
-from contextlib import asynccontextmanager
+from collections.abc import AsyncIterator, Mapping
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from typing import Any, cast
 
 import anyio
@@ -76,6 +76,8 @@ async def connected_runner(
     initialized: bool = True,
     stateless: bool = False,
     has_standalone_channel: bool = True,
+    session_id: str | None = None,
+    headers: Mapping[str, str] | None = None,
     dispatch_middleware: list[DispatchMiddleware] | None = None,
 ) -> AsyncIterator[tuple[DirectDispatcher, ServerRunner[dict[str, Any]]]]:
     """Yield ``(client, runner)`` running over an in-memory dispatcher pair.
@@ -85,12 +87,13 @@ async def connected_runner(
     ``initialized`` is true the helper performs the real ``initialize`` request
     before yielding, so tests start past the init-gate via the public path.
     """
-    client, server_d = create_direct_dispatcher_pair()
+    client, server_d = create_direct_dispatcher_pair(headers=headers)
     runner = ServerRunner(
         server=server,
         dispatcher=server_d,
         lifespan_state={},
         has_standalone_channel=has_standalone_channel,
+        session_id=session_id,
         stateless=stateless,
         dispatch_middleware=dispatch_middleware or [],
     )
@@ -380,3 +383,100 @@ async def failing(ctx: Any, params: Any) -> Any:
     [event] = [e for e in span.events if e.name == "exception"]
     assert event.attributes is not None
     assert event.attributes["exception.type"] == "ValueError"
+
+
+@pytest.mark.anyio
+async def test_connection_state_persists_across_requests_on_same_connection(server: SrvT) -> None:
+    async def count(ctx: Any, params: PaginatedRequestParams | None) -> ListToolsResult:
+        ctx.connection.state["n"] = ctx.connection.state.get("n", 0) + 1
+        return ListToolsResult(tools=[])
+
+    server.add_request_handler("tools/list", PaginatedRequestParams, count)
+    async with connected_runner(server) as (client, runner):
+        await client.send_raw_request("tools/list", None)
+        await client.send_raw_request("tools/list", None)
+        assert runner.connection.state == {"n": 2}
+
+
+@pytest.mark.anyio
+async def test_connection_exit_stack_runs_pushed_callback_after_close(server: SrvT) -> None:
+    cleaned: list[str] = []
+
+    async def push(ctx: Any, params: PaginatedRequestParams | None) -> ListToolsResult:
+        async def _cleanup() -> None:
+            cleaned.append("done")
+
+        ctx.connection.exit_stack.push_async_callback(_cleanup)
+        return ListToolsResult(tools=[])
+
+    server.add_request_handler("tools/list", PaginatedRequestParams, push)
+    async with connected_runner(server) as (client, _runner):
+        await client.send_raw_request("tools/list", None)
+        assert cleaned == []
+    assert cleaned == ["done"]
+
+
+@pytest.mark.anyio
+async def test_connection_exit_stack_unwinds_entered_context_manager_after_close(server: SrvT) -> None:
+    events: list[str] = []
+
+    class _Tracker(AbstractAsyncContextManager[str]):
+        async def __aenter__(self) -> str:
+            events.append("enter")
+            return "resource"
+
+        async def __aexit__(self, *exc: object) -> None:
+            events.append("exit")
+
+    async def acquire(ctx: Any, params: PaginatedRequestParams | None) -> ListToolsResult:
+        res = await ctx.connection.exit_stack.enter_async_context(_Tracker())
+        ctx.connection.state["res"] = res
+        return ListToolsResult(tools=[])
+
+    server.add_request_handler("tools/list", PaginatedRequestParams, acquire)
+    async with connected_runner(server) as (client, runner):
+        await client.send_raw_request("tools/list", None)
+        assert events == ["enter"]
+        assert runner.connection.state["res"] == "resource"
+    assert events == ["enter", "exit"]
+
+
+@pytest.mark.anyio
+async def test_connection_exit_stack_runs_callbacks_lifo_after_handler_error(server: SrvT) -> None:
+    cleaned: list[int] = []
+
+    async def push_then_fail(ctx: Any, params: PaginatedRequestParams | None) -> ListToolsResult:
+        for i in (1, 2, 3):
+            ctx.connection.exit_stack.push_async_callback(_append, i)
+        raise RuntimeError("boom")
+
+    async def _append(i: int) -> None:
+        cleaned.append(i)
+
+    server.add_request_handler("tools/list", PaginatedRequestParams, push_then_fail)
+    async with connected_runner(server) as (client, _runner):
+        with pytest.raises(MCPError) as ei:
+            await client.send_raw_request("tools/list", None)
+        assert ei.value.error.code == INTERNAL_ERROR
+        assert cleaned == []
+    assert cleaned == [3, 2, 1]
+
+
+@pytest.mark.anyio
+async def test_context_session_id_and_headers_expose_connection_and_transport(server: SrvT) -> None:
+    async with connected_runner(server, session_id="sess-abc", headers={"authorization": "Bearer t"}) as (client, _r):
+        await client.send_raw_request("tools/list", None)
+    [ctx] = _seen_ctx
+    assert ctx.session_id == "sess-abc"
+    assert ctx.session_id == ctx.connection.session_id
+    assert ctx.headers == {"authorization": "Bearer t"}
+    assert ctx.headers is ctx.transport.headers
+
+
+@pytest.mark.anyio
+async def test_context_session_id_and_headers_default_none(server: SrvT) -> None:
+    async with connected_runner(server) as (client, _r):
+        await client.send_raw_request("tools/list", None)
+    [ctx] = _seen_ctx
+    assert ctx.session_id is None
+    assert ctx.headers is None