runner

async def aclose_shielded(connection: Connection) -> None:
    """Unwind ``connection.exit_stack`` under a shielded, bounded scope.

    Called from a driver's ``finally``: the shield lets per-connection cleanup
    callbacks run even when the driver itself is being cancelled, the
    `_EXIT_STACK_CLOSE_TIMEOUT` bound stops a hung callback wedging shutdown,
    and a raising callback is logged-and-swallowed so it never masks the
    driver's own exception.
    """
    with anyio.move_on_after(_EXIT_STACK_CLOSE_TIMEOUT, shield=True) as scope:
        try:
            await connection.exit_stack.aclose()
        except Exception:
            logger.exception("connection exit_stack cleanup raised")
    if scope.cancelled_caught:
        logger.warning(
            "connection exit_stack cleanup exceeded %s seconds; abandoning remaining callbacks",
            _EXIT_STACK_CLOSE_TIMEOUT,
        )

def modern_error_data(exc: Exception) -> ErrorData:
    """Map a modern request's handler exception to its wire `ErrorData`.

    The exception-to-wire fact shared by the modern entries (the
    single-exchange HTTP path and the dual-era stream loop), so an identical
    modern request fails identically on every transport: `MCPError` and
    `ValidationError` map via the shared `handler_exception_to_error_data`
    ladder; anything else is logged server-side and surfaced as a generic
    INTERNAL_ERROR so handler internals never reach the wire.
    """
    error = handler_exception_to_error_data(exc)
    if error is not None:
        return error
    logger.exception("modern request handler raised")
    return ErrorData(code=INTERNAL_ERROR, message="Internal server error")

def modern_on_request(server: Server[LifespanT], lifespan_state: LifespanT) -> OnRequest:
    """Return an `OnRequest` callback that serves each call via `serve_one` with a fresh per-request `Connection`.

    Wire this into the server side of a `DirectDispatcher` peer-pair to drive an
    in-process server on the modern per-request-envelope path (each request
    carries protocol version, client info, and capabilities in `params._meta`;
    no `initialize` handshake). The dispatch context is wrapped in the
    server-requests denial, so the modern prohibition on server-initiated
    JSON-RPC requests holds on this entry like on the others. Like `serve_one`,
    this raises whatever the handler chain raises - the dispatcher owns the
    exception-to-error mapping.
    """

    async def handle(
        dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
    ) -> dict[str, Any]:
        meta = (params or {}).get("_meta", {})
        connection = Connection.from_envelope(
            meta.get(PROTOCOL_VERSION_META_KEY, LATEST_MODERN_VERSION),
            meta.get(CLIENT_INFO_META_KEY),
            meta.get(CLIENT_CAPABILITIES_META_KEY),
        )
        return await serve_one(
            server,
            NoServerRequestsDispatchContext(dctx),
            method,
            params,
            connection=connection,
            lifespan_state=lifespan_state,
        )

    return handle

async def serve_connection(
    server: Server[LifespanT],
    dispatcher: Dispatcher[Any],
    *,
    connection: Connection,
    lifespan_state: LifespanT,
    init_options: InitializationOptions | None = None,
    task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
) -> None:
    """Drive ``dispatcher`` until the underlying channel closes.

    The loop-mode driver: builds the kernel, hands `on_request`/`on_notify`
    to `dispatcher.run()`, and tears down `connection.exit_stack` (shielded)
    on the way out. The entry constructs the `Connection`; this only consumes
    it.
    """
    runner = ServerRunner(server, connection, lifespan_state, init_options=init_options)
    try:
        await dispatcher.run(runner.on_request, runner.on_notify, task_status=task_status)
    finally:
        await aclose_shielded(connection)

async def serve_dual_era_loop(
    server: Server[LifespanT],
    read_stream: ReadStream[SessionMessage | Exception],
    write_stream: WriteStream[SessionMessage],
    *,
    lifespan_state: LifespanT,
    session_id: str | None = None,
    init_options: InitializationOptions | None = None,
    raise_exceptions: bool = False,
) -> None:
    """Drive `server` over a duplex stream pair, serving both protocol eras.

    The stream-pair counterpart of the modern HTTP entry's era router. Era is
    a property of the connection, decided by how the client opens it, and
    mid-stream switching is undefined - so the first era-distinctive message
    to SUCCEED locks the connection (matching the typescript-sdk):

    - A successful `initialize` locks legacy: the connection behaves exactly
      like `serve_loop` for its lifetime, and modern envelope traffic is then
      rejected with INVALID_REQUEST. `initialize` never routes modern - the
      method is legacy-distinctive by definition - even when a confused
      client stamps the envelope triple on it.
    - A request carrying the modern `_meta` envelope triple - or
      `server/discover`, a modern-only method - is classified
      (`classify_inbound_request`) and served single-exchange via `serve_one`
      with a born-ready per-request `Connection`, the same dispatch model as
      the modern HTTP entry. The first such request to succeed locks the
      connection modern; a later `initialize` is then rejected with
      UNSUPPORTED_PROTOCOL_VERSION naming the modern versions.

    Modern connections push notifications over the duplex pipe but refuse
    server-initiated requests on both channels (the modern protocol forbids
    them). A request that fails - rejected classification, malformed envelope
    content, unknown method - never locks either era, so a failed probe
    leaves the legacy handshake available: released auto-negotiating clients
    fall back on any error code except -32022, and that code is only emitted
    for genuine version negotiation or for `initialize` on an
    already-modern connection.

    The modern lock commits on the FIRST client-visible success frame the
    classified request produces - a request-scoped notification (progress and
    the `subscriptions/listen` ack ride those) or its result - immediately
    before that frame is written, with no checkpoint in between. One
    definition for every method: a plain request locks at its response
    exactly as before, and a streaming request (`subscriptions/listen`, whose
    response arrives only at close) locks at its ack, so a pipelined
    `initialize` behind the ack can never hijack a live stream back to
    legacy. For the inline methods (`initialize`, `server/discover`) the
    dispatch completes before the next frame is read, so the canonical
    probe-then-go flow is race-free; a pinned-modern client that pipelines
    frames ahead of its first response should expect envelope-less
    notifications sent in that window to be dropped. The lock settles exactly
    once: a request from the other era that was already in flight when the
    lock committed may still complete and its response stands, but the era
    does not move; and a peer cancel observed before the first frame's write
    begins prevents the lock (the commit re-checks `cancel_requested` at fire
    time).

    One residual corner is accepted, deliberately: the lock commits
    immediately before the frame's transport write, and that write's own
    checkpoint is the first point a pending peer cancel can land - so a
    cancel arriving DURING the write cancels the frame away after the lock
    already committed, leaving the connection locked modern with ZERO frames
    delivered. Only mid-handler frames (the listen ack, progress) have this
    window; a plain request's response cannot be cancelled away (its
    in-flight entry is removed before the response write starts, with no
    checkpoint in between, so a peer cancel no longer reaches it). The
    converse exists too: a handler that shields its own send past a pending
    cancel delivers a frame WITHOUT the lock - `era_settles` declined at fire
    time - but that frame is output the handler forced after the cancel said
    stop, not a state this loop produces. The
    orphaned lock is self-consistent and client-recoverable: only a
    validly-classified modern envelope reaches this path, so the peer already
    committed to modern; a follow-up `initialize` is answered with
    UNSUPPORTED_PROTOCOL_VERSION (-32022) naming the modern versions, which
    released auto-negotiating clients treat as modern evidence and re-probe;
    and the connection keeps serving modern traffic. Closing the corner would
    require committing the lock only after the write is known delivered,
    which reopens the worse race this design exists to prevent (an
    `initialize` pipelined behind an ack the client already read, flipping a
    live stream legacy).

    Classified-modern requests are also governed by the 2026 cancellation
    rule for their whole lifetime: after an inbound `notifications/cancelled`
    the server sends nothing further for that request - no result, no error
    (the transport spec's MUST NOT). Once the connection is LOCKED modern the
    rule widens from classification-scoped to connection-scoped: every
    request - including one rejected before classification succeeds
    (envelope-less, malformed envelope) - carries the silence commitment, so
    a rejection error computed after the peer's cancel never reaches the wire
    (the rejection itself, absent a cancel, is of course still answered).
    Legacy-era requests keep the byte-identical "Request cancelled" answer
    released 2025 clients expect - including a legacy-classified request
    still in flight when the modern lock commits (the straggler carve-out
    covers its cancel answer, not just its response) - as does a request that
    fails classification on a still-unlocked connection: an unclassifiable
    request has not proven modern semantics, so the legacy answer stands
    there.
    """
    dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
        read_stream,
        write_stream,
        raise_handler_exceptions=raise_exceptions,
        # `initialize` inline for the same pipelining reason as `serve_loop`;
        # `server/discover` inline so the modern era lock commits before the
        # next pipelined message is read.
        inline_methods=frozenset({"initialize", "server/discover"}),
    )
    loop_connection = Connection.for_loop(dispatcher, session_id=session_id)
    loop_runner = ServerRunner(server, loop_connection, lifespan_state, init_options=init_options)
    standalone_outbound = NotifyOnlyOutbound(dispatcher)
    era: Literal["unlocked", "legacy", "modern"] = "unlocked"
    modern_version = LATEST_MODERN_VERSION

    def era_settles(dctx: DispatchContext[TransportContext]) -> bool:
        # The one definition of "this request may lock the era": it produced
        # a client-visible success on a still-unlocked connection. The lock is
        # monotone - the first success wins, so a straggling request from the
        # other era can never overwrite a committed lock. A pending peer
        # cancel means the client already abandoned this request and will
        # never read the frame as a success: no lock.
        return era == "unlocked" and not dctx.cancel_requested.is_set()

    def commit_modern(dctx: DispatchContext[TransportContext], version: str) -> None:
        # The success-frame observer for classified-modern requests: invoked
        # immediately before each non-error frame's write, with no checkpoint
        # in between - by the time the client can read any modern output, the
        # era is locked and a pipelined `initialize` is rejected instead of
        # hijacking the connection. `era_settles` re-checks `cancel_requested`
        # at fire time; the cancel-during-the-write residue is the accepted
        # corner documented on `serve_dual_era_loop`.
        nonlocal era, modern_version
        if era_settles(dctx):
            era, modern_version = "modern", version

    async def serve_modern(
        dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
    ) -> dict[str, Any]:
        route = classify_inbound_request({"method": method, "params": params})
        if isinstance(route, InboundLadderRejection):
            raise MCPError(code=route.code, message=route.message, data=route.data)
        # Classification succeeded: 2026 wire semantics govern this request's
        # lifecycle from here. On a still-unlocked connection this is the
        # silence rule's exact scope - a request that fails classification
        # above has not proven modern semantics and keeps the legacy cancel
        # answer - and the era lock rides the request's first success frame.
        # Both facts attach synchronously - no checkpoint since the
        # dispatcher entered the handler scope - so a peer cancel can never
        # interleave and observe the legacy defaults. Once the connection is
        # locked, neither attaches here: `on_request` owns the
        # connection-scoped silence fact, and the lock can never move again.
        if era == "unlocked":
            suppress_cancel_answer(dctx)
            observe_success_frames(dctx, partial(commit_modern, dctx, route.protocol_version))
        connection = Connection.from_envelope(
            route.protocol_version,
            route.client_info,
            route.client_capabilities,
            outbound=standalone_outbound,
        )
        try:
            return await serve_one(
                server,
                NoServerRequestsDispatchContext(dctx),
                method,
                params,
                connection=connection,
                lifespan_state=lifespan_state,
            )
        except (MCPError, ValidationError):
            # The dispatcher's shared ladder maps these to the same wire error
            # the modern HTTP entry produces.
            raise
        except Exception as exc:
            if raise_exceptions:
                raise
            error = modern_error_data(exc)
            raise MCPError(code=error.code, message=error.message, data=error.data) from exc

    async def on_request(
        dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
    ) -> dict[str, Any]:
        nonlocal era
        if era == "legacy":
            if _has_modern_envelope(params):
                raise MCPError(
                    code=INVALID_REQUEST,
                    message="connection is locked to the legacy handshake era; "
                    "modern envelope requests are not accepted",
                )
            # Bare modern-only methods (e.g. `server/discover`) fall through to
            # the loop runner's per-version surface validation - the same
            # METHOD_NOT_FOUND a handshake-only server produced, byte for byte.
            return await loop_runner.on_request(dctx, method, params)
        if era == "modern":
            # Connection-scoped silence: once locked modern, EVERY request is
            # governed by the 2026 cancellation rule, rejections included - a
            # rejection error is a legitimate answer, but nothing may follow
            # the peer's `notifications/cancelled`. Pre-classification rejects
            # run on a spawned task, so an adjacent cancel can precede them
            # under trio's randomized scheduling (asyncio's FIFO always starts
            # the reject's answer write first - the legitimate
            # already-answering case; inline `initialize` has no window at
            # all). See the cancellation paragraph on `serve_dual_era_loop`.
            suppress_cancel_answer(dctx)
            if method == "initialize":
                raise MCPError(
                    code=UNSUPPORTED_PROTOCOL_VERSION,
                    message="connection already negotiated a modern protocol version",
                    data=_initialize_after_modern_data(params),
                )
            return await serve_modern(dctx, method, params)
        # Unlocked. `initialize` is legacy-distinctive by definition (the
        # method does not exist at modern versions), so it takes the handshake
        # path even when the envelope triple is stamped on it.
        if method != "initialize" and (method == "server/discover" or _has_modern_envelope(params)):
            return await serve_modern(dctx, method, params)
        result = await loop_runner.on_request(dctx, method, params)
        if method == "initialize" and era_settles(dctx):
            # Lock only on success: a failed handshake leaves both eras open.
            era = "legacy"
        return result

    async def on_notify(dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None:
        if era != "modern":
            return await loop_runner.on_notify(dctx, method, params)
        # The envelope is request-only, so notifications inherit the
        # connection's locked version.
        connection = Connection.from_envelope(modern_version, None, None, outbound=standalone_outbound)
        notify_runner = ServerRunner(server, connection, lifespan_state)
        try:
            await notify_runner.on_notify(NoServerRequestsDispatchContext(dctx), method, params)
        finally:
            await aclose_shielded(connection)

    try:
        await dispatcher.run(on_request, on_notify)
    finally:
        await aclose_shielded(loop_connection)

async def serve_loop(
    server: Server[LifespanT],
    read_stream: ReadStream[SessionMessage | Exception],
    write_stream: WriteStream[SessionMessage],
    *,
    lifespan_state: LifespanT,
    session_id: str | None = None,
    init_options: InitializationOptions | None = None,
    raise_exceptions: bool = False,
) -> None:
    """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
    dispatcher recipe (notably the `inline_methods={"initialize"}` rule) with
    era routing.
    """
    dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
        read_stream,
        write_stream,
        raise_handler_exceptions=raise_exceptions,
        # Handle `initialize` inline so a client that pipelines it with the
        # next request (spec: SHOULD NOT, not MUST NOT) sees the initialized
        # state instead of failing the init-gate.
        inline_methods=frozenset({"initialize"}),
    )
    connection = Connection.for_loop(dispatcher, session_id=session_id)
    await serve_connection(
        server, dispatcher, connection=connection, lifespan_state=lifespan_state, init_options=init_options
    )

async def serve_one(
    server: Server[LifespanT],
    dctx: DispatchContext[TransportContext],
    method: str,
    params: Mapping[str, Any] | None,
    *,
    connection: Connection,
    lifespan_state: LifespanT,
) -> dict[str, Any]:
    """Handle a single request ``(method, params)`` and return its result dict.

    The single-exchange driver: builds the kernel, runs `on_request` once under
    `dctx`, and tears down `connection.exit_stack` (shielded) on the way out.
    The entry constructs the (born-ready) `Connection` and the `dctx`; this
    only consumes them.

    Raises whatever the handler chain raises (`MCPError` / `ValidationError` /
    unmapped); callers own the exception-to-wire mapping.
    """
    runner = ServerRunner(server, connection, lifespan_state)
    try:
        return await runner.on_request(dctx, method, params)
    finally:
        await aclose_shielded(connection)