runner

@dataclass
class ServerRunner(Generic[LifespanT]):
    """Per-connection handler kernel. One instance per client connection."""

    server: Server[LifespanT]
    connection: Connection
    lifespan_state: LifespanT
    _: KW_ONLY
    init_options: InitializationOptions | None = None
    """`InitializeResult` payload. Defaults to `server.create_initialization_options()`."""

    @cached_property
    def on_request(self) -> OnRequest:
        return self._on_request

    @cached_property
    def on_notify(self) -> OnNotify:
        return self._on_notify

    async def _on_request(
        self,
        dctx: DispatchContext[TransportContext],
        method: str,
        params: Mapping[str, Any] | None,
    ) -> dict[str, Any]:
        meta = _extract_meta(params)
        version = self.connection.protocol_version
        ctx = self._make_context(dctx, method, params, meta, version)

        async def _inner(ctx: ServerRequestContext[LifespanT, Any]) -> HandlerResult:
            # Read method/params off `ctx` so a middleware that rewrote them via
            # `call_next(replace(ctx, ...))` reaches lookup and the handler.
            method, params = ctx.method, ctx.params
            # Pinned compat: spec methods are surface-validated before lookup,
            # so malformed params are INVALID_PARAMS even with no handler
            # registered. Custom methods miss the monolith map and fall through
            # to `entry.params_type` exactly as before.
            if method in _methods.SPEC_CLIENT_METHODS:
                try:
                    _methods.validate_client_request(method, version, params)
                except KeyError:
                    raise MCPError(code=METHOD_NOT_FOUND, message="Method not found", data=method) from None
            # TODO(L29): the 2026-07-28 spec drops the handshake; this branch and
            # the gate become a per-version legacy path then. Initialize runs inline
            # (read loop parked), so awaiting the peer anywhere on this path deadlocks.
            if method == "initialize":
                return self._serialize(method, version, self._handle_initialize(params))
            # Methods without a handler are METHOD_NOT_FOUND regardless of
            # initialization state: JSON-RPC 2.0 reserves -32601 for "not
            # available on this server", and clients probing a server before
            # the handshake key off that code. The init gate below therefore
            # only ever applies to methods the server actually serves.
            entry = self.server.get_request_handler(method)
            if entry is None:
                raise MCPError(code=METHOD_NOT_FOUND, message="Method not found", data=method)
            if not self.connection.initialize_accepted and method not in _INIT_EXEMPT:
                # Pinned compat: the same error shape the union validation produced.
                raise MCPError(code=INVALID_PARAMS, message="Invalid request parameters", data="")
            # Absent params validate as {} (required fields still reject), so
            # the handler receives the model with its defaults, never None.
            typed_params = entry.params_type.model_validate({} if params is None else params, by_name=False)
            result = await entry.handler(ctx, typed_params)
            if isinstance(result, ErrorData):
                # Raise inside the chain so middleware observes the failure.
                raise MCPError.from_error_data(result)
            # Shape for the wire inside the chain so the OpenTelemetry span (the
            # outermost middleware) records a failing handler return shape too.
            return self._serialize(method, version, result)

        call = self._compose_server_middleware(_inner)
        # `_inner` already produced the wire dict; a middleware that short-circuited
        # without `call_next` is trusted to return its own well-formed result -
        # including its response envelope. The pipeline never patches it up after
        # the fact.
        result = _dump_result(await call(ctx))
        if method == "initialize":
            # Commit only on chain success, so a middleware veto leaves no state.
            # Race-free: the read loop is parked until this call returns.
            # TODO: this re-reads the wire `params`, so a middleware that rewrote
            # `ctx.params` (or `ctx.method`, or short-circuited without `call_next`)
            # can leave `connection.protocol_version` out of step with the
            # `InitializeResult` `_inner` produced. Resolve when `initialize` becomes
            # a built-in handler so commit and result derive from one negotiation.
            self.connection.client_params, self.connection.protocol_version = self._negotiate_initialize(params)
        return result

    async def _on_notify(
        self,
        dctx: DispatchContext[TransportContext],
        method: str,
        params: Mapping[str, Any] | None,
    ) -> None:
        meta = _extract_meta(params)
        version = self.connection.protocol_version
        ctx = self._make_context(dctx, method, params, meta, version)

        async def _inner(ctx: ServerRequestContext[LifespanT, Any]) -> None:
            method, params = ctx.method, ctx.params
            if method in _methods.SPEC_CLIENT_NOTIFICATION_METHODS:
                try:
                    _methods.validate_client_notification(method, version, params)
                except KeyError:
                    logger.debug("dropped %r: not defined at %s", method, version)
                    return
                except ValidationError:
                    logger.warning("dropped %r: malformed params", method)
                    return
            if method == "notifications/initialized":
                # Surface validation above already rejected a malformed body, so
                # commit; fall through so a registered handler observes an
                # initialized connection.
                self.connection.initialized.set()
            elif not self.connection.initialize_accepted:
                logger.debug("dropped %s: received before initialization", method)
                return
            entry = self.server.get_notification_handler(method)
            if entry is None:
                logger.debug("no handler for notification %s", method)
                return
            # Same absent-params contract as requests.
            try:
                typed_params = entry.params_type.model_validate({} if params is None else params, by_name=False)
            except ValidationError:
                logger.warning("dropped %r: malformed params", method)
                return
            await entry.handler(ctx, typed_params)

        call = self._compose_server_middleware(_inner)
        try:
            await call(ctx)
        except Exception:
            # A crashing handler must not cancel the dispatcher's task group;
            # middleware saw the raise out of call_next() first.
            logger.exception("notification handler for %r raised", method)

    def _compose_server_middleware(self, inner: CallNext) -> CallNext:
        """Wrap `inner` in `Server.middleware`, outermost-first.

        Shared by `_on_request` and `_on_notify` so the same middleware chain
        observes every inbound message. The composed callable takes the `ctx`
        at call time, so a middleware can rewrite it for the rest of the chain.
        """
        call = inner
        for middleware in reversed(self.server.middleware):
            call = partial(_apply_middleware, middleware, call)
        return call

    def _make_context(
        self,
        dctx: DispatchContext[TransportContext],
        method: str,
        params: Mapping[str, Any] | None,
        meta: RequestParamsMeta | None,
        protocol_version: str,
    ) -> ServerRequestContext[LifespanT, Any]:
        # TODO(L54): remove for Context rework. Reads the SHTTP per-request
        # data off the raw `dctx.message_metadata` carrier; replace with the
        # per-transport context once that lands.
        md = dctx.message_metadata
        if isinstance(md, ServerMessageMetadata):
            request = md.request_context
            close_sse_stream = md.close_sse_stream
            close_standalone_sse_stream = md.close_standalone_sse_stream
        else:
            request = close_sse_stream = close_standalone_sse_stream = None
        # Per-request session: `dctx` is the request-scoped channel (auto-threads
        # its own request_id on streamable HTTP); the standalone channel is read
        # off `connection.outbound`. `related_request_id` on the public API selects.
        session = ServerSession(dctx, self.connection)
        return ServerRequestContext(
            session=session,
            lifespan_context=self.lifespan_state,
            method=method,
            params=params,
            request_id=dctx.request_id,
            meta=meta,
            protocol_version=protocol_version,
            request=request,
            close_sse_stream=close_sse_stream,
            close_standalone_sse_stream=close_standalone_sse_stream,
        )

    def _serialize(self, method: str, version: str, result: HandlerResult) -> dict[str, Any]:
        """Shape a handler result into its wire form: the outbound counterpart
        of the inbound classification ladder.

        One pass owns the whole response envelope, in order: cache hints fill
        `ttlMs`/`cacheScope` the handler left unset, core-vocabulary spec-method
        results are validated and sieved by the per-version surface (a claimed
        extension `resultType` shape is the extension's to own), and 2026-era
        results get the `serverInfo` `_meta` stamp (spec #3002). Runs inside the
        middleware chain so the OpenTelemetry span observes a failing return
        shape (unsupported type, malformed spec result) as an error rather
        than closing on a request that the client sees fail - and so a
        middleware that short-circuits without `call_next` owns its result,
        envelope included.
        """
        # MRTR carve-out: `input_required` interim results, typed or mapping, never get hints.
        if (hint := self.server.cache_hints.get(method)) is not None:
            if isinstance(result, CacheableResult):
                result = apply_cache_hint(result, hint)
            elif isinstance(result, Mapping) and not _methods.is_input_required(result):
                # Hint keys first so wire keys the handler set win, matching `apply_cache_hint` precedence.
                result = {"ttlMs": hint.ttl_ms, "cacheScope": hint.scope, **result}
        dumped = _dump_result(result)
        # A modern-era extension `resultType` (outside the core vocabulary) marks
        # a claimed shape owned by the extension that defined it: the per-version
        # surface doesn't describe it, so the sieve applies to core results only.
        # Legacy connections sieve everything - claimed shapes are 2026-era
        # vocabulary and cannot be delivered on a legacy wire (mirrors the
        # client-side ResultClaim rule).
        # TODO(L56): reject extension resultType values unless the corresponding
        # extension is in this request's _meta clientCapabilities.extensions; the
        # explicit MUST-reject is client-side (basic/index.mdx ResultType), this enforces it proactively.
        result_type = dumped.get("resultType")
        core_shape = (
            version not in MODERN_PROTOCOL_VERSIONS
            or not isinstance(result_type, str)
            or result_type in CORE_RESULT_TYPES
        )
        if method in _methods.SPEC_CLIENT_METHODS and core_shape:
            try:
                dumped = _methods.serialize_server_result(method, version, dumped)
            except ValidationError:
                # Server bug, not client fault. Detail stays in the server log:
                # pydantic messages echo the result body.
                logger.exception("handler for %r returned an invalid result", method)
                raise MCPError(code=INTERNAL_ERROR, message="Handler returned an invalid result") from None
        if version in MODERN_PROTOCOL_VERSIONS and dumped.get("resultType") is None:
            # Spec 2026-07-28: `Result.resultType` is required - servers MUST
            # include it (the absent-means-complete bridge is for clients of
            # older servers only). The sieve guarantees it for core methods;
            # this covers everything else: custom methods, extension methods,
            # and empty results.
            dumped["resultType"] = "complete"
        return self._stamp_server_info(version, dumped)

    def _stamp_server_info(self, version: str, result: dict[str, Any]) -> dict[str, Any]:
        """Fill the `serverInfo` `_meta` stamp on a 2026-era result (spec #3002).

        A handler-authored value wins; an explicit `null` reads as absent and
        is stamped over, mirroring the request-side `clientInfo` posture (a
        `null` is not a valid `Implementation`, so presence means a value). A
        non-mapping `_meta` is the handler's to own, and handshake-era results
        are never stamped. `result` is
        pipeline-owned (`_dump_result` copies dicts; the spec-method sieve
        re-dumps), but `_meta` may still be the handler's object, so the stamp
        replaces it rather than writing into it. `server_info_stamp` is a
        fresh dict per access, so the response never aliases server state.
        """
        if version not in MODERN_PROTOCOL_VERSIONS:
            return result
        raw_meta = result.get("_meta")
        if raw_meta is None:
            result["_meta"] = {SERVER_INFO_META_KEY: self.server.server_info_stamp}
        elif isinstance(raw_meta, dict):
            meta = cast("dict[str, Any]", raw_meta)
            if meta.get(SERVER_INFO_META_KEY) is None:
                result["_meta"] = {**meta, SERVER_INFO_META_KEY: self.server.server_info_stamp}
        return result

    @staticmethod
    def _negotiate_initialize(params: Mapping[str, Any] | None) -> tuple[InitializeRequestParams, str]:
        """Validate `initialize` params and pick the protocol version."""
        init = InitializeRequestParams.model_validate(params or {}, by_name=False)
        requested = init.protocol_version
        negotiated = requested if requested in HANDSHAKE_PROTOCOL_VERSIONS else LATEST_HANDSHAKE_VERSION
        return init, negotiated

    def _handle_initialize(self, params: Mapping[str, Any] | None) -> InitializeResult:
        """Build the `initialize` result; state commits later in `_on_request`."""
        _, negotiated = self._negotiate_initialize(params)
        opts = self.init_options if self.init_options is not None else self.server.create_initialization_options()
        return InitializeResult(
            protocol_version=negotiated,
            capabilities=opts.capabilities,
            server_info=Implementation(
                name=opts.server_name,
                title=opts.title,
                description=opts.description,
                version=opts.server_version,
                website_url=opts.website_url,
                icons=opts.icons,
            ),
            instructions=opts.instructions,
        )