connection.py

"""`Connection` - per-client connection state and the standalone outbound channel.

Always present on `Context` (never `None`), even in stateless deployments.
Holds peer info, 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).

Construct via the factories: `Connection.from_envelope` for the 2026-era
single-exchange path (born ready, no back-channel) and `Connection.for_loop`
for the handshake-driven loop path. Both populate `protocol_version` so the
kernel reads it as a fact.

`notify` is best-effort: it never raises. If there's no standalone channel
or the stream has been dropped, the notification is debug-logged and silently
discarded - server-initiated notifications are inherently advisory.
`send_raw_request` raises `NoBackChannelError` when there's no channel; `ping`
is the only spec-sanctioned standalone request.
"""

from __future__ import annotations

import logging
from collections.abc import Mapping
from contextlib import AsyncExitStack
from typing import Any, Final, TypeVar, get_args, overload

import anyio
from mcp_types import (
    LOG_LEVEL_META_KEY,
    ClientCapabilities,
    CreateMessageRequest,
    CreateMessageResult,
    ElicitRequest,
    ElicitResult,
    EmptyResult,
    Implementation,
    InitializeRequestParams,
    ListRootsRequest,
    ListRootsResult,
    LoggingLevel,
    PingRequest,
    Request,
)
from mcp_types import methods as _methods
from mcp_types.version import LATEST_HANDSHAKE_VERSION, MODERN_PROTOCOL_VERSIONS
from pydantic import BaseModel, ValidationError
from typing_extensions import deprecated

from mcp.shared.dispatcher import CallOptions, Outbound
from mcp.shared.exceptions import MCPDeprecationWarning, NoBackChannelError
from mcp.shared.peer import Meta, dump_params

__all__ = ["Connection"]

logger = logging.getLogger(__name__)
# `Connection.log`'s `logger` parameter (public API, the spec's logger-name
# field) shadows the module logger inside that method; this alias keeps the
# module logger reachable there.
_logger = logger

_LOG_LEVELS: Final[tuple[LoggingLevel, ...]] = get_args(LoggingLevel)
"""Severity-ascending, from the `LoggingLevel` literal's declaration order (the
RFC 5424 scale) - the literal is the single source of the ordering."""

_ALL_LOG_LEVELS: Final[frozenset[LoggingLevel]] = frozenset(_LOG_LEVELS)


def allowed_log_levels(protocol_version: str, meta: Mapping[str, Any] | None) -> frozenset[LoggingLevel]:
    """The `notifications/message` levels deliverable for one inbound request.

    2026-07-28+ makes log delivery a per-request opt-in (server/utilities/
    logging): the client sets the reserved `io.modelcontextprotocol/logLevel`
    `_meta` key, absent means no levels - the server MUST NOT send - and
    present means that level and above. An unrecognized value reads as absent;
    spec methods already reject a malformed value at surface validation
    before any handler runs, so that arm only serves custom methods, where
    dropping is the safe direction. Connection-scoped emitters pass
    `meta=None`: `logging/setLevel` is gone at 2026 and log delivery is
    request-scoped only, so they deliver nothing. Handshake versions keep
    their `logging/setLevel`-era semantics: every level may be sent, filtering
    is the application's `logging/setLevel` handler's job as before.
    """
    if protocol_version not in MODERN_PROTOCOL_VERSIONS:
        return _ALL_LOG_LEVELS
    requested = (meta or {}).get(LOG_LEVEL_META_KEY)
    if requested not in _LOG_LEVELS:
        return frozenset()
    return frozenset(_LOG_LEVELS[_LOG_LEVELS.index(requested) :])


ResultT = TypeVar("ResultT", bound=BaseModel)

# Result types for the spec's server-to-client request set, used by
# `Connection.send_request` to infer the result type. If the spec's request
# set grows substantially, consider declaring the result mapping on the
# request types themselves (a `__mcp_result__` ClassVar read via a structural
# protocol) so this table and the overload ladder don't need maintaining.
_RESULT_FOR: dict[type[Request[Any, Any]], type[BaseModel]] = {
    CreateMessageRequest: CreateMessageResult,
    ElicitRequest: ElicitResult,
    ListRootsRequest: ListRootsResult,
    PingRequest: EmptyResult,
}


_ModelT = TypeVar("_ModelT", bound=BaseModel)


def _typed(model: type[_ModelT], raw: Any) -> _ModelT | None:
    """Validate a raw envelope value into a typed model.

    A missing, null or mis-shaped value falls through to `ValidationError`
    and is treated as not supplied so the request still routes. Spec methods
    are separately re-validated by the kernel's per-version params surface,
    which types the reserved `_meta` keys strictly.
    """
    try:
        return model.model_validate(raw, by_name=False)
    except ValidationError:
        return None


def _notification_params(payload: dict[str, Any] | None, meta: Meta | None) -> dict[str, Any] | None:
    if not meta:
        return payload
    out = dict(payload or {})
    out["_meta"] = meta
    return out


class _NoChannelOutbound:
    """Connection-scoped `Outbound` for the no-back-channel case.

    The structural answer to "this connection cannot push to its peer":
    `send_raw_request` raises `NoBackChannelError`; `notify` drops with a
    debug log. `Connection.from_envelope` installs this so the modern
    single-exchange path never needs a mode flag - the channel itself says no.
    """

    async def send_raw_request(
        self,
        method: str,
        params: Mapping[str, Any] | None,
        opts: CallOptions | None = None,
    ) -> dict[str, Any]:
        raise NoBackChannelError(method)

    async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None:
        logger.debug("dropped %s: no standalone channel", method)


_NO_CHANNEL = _NoChannelOutbound()


class NotifyOnlyOutbound(_NoChannelOutbound):
    """Connection-scoped `Outbound` that forwards notifications and refuses requests.

    Installed by `serve_dual_era_loop` for modern (2026-07-28+) connections
    over duplex stream transports: the pipe is real, so server notifications
    ride it, but the modern protocol forbids server-initiated JSON-RPC
    requests, so `send_raw_request` (inherited) refuses by construction.
    """

    def __init__(self, outbound: Outbound) -> None:
        self._outbound = outbound

    async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None:
        await self._outbound.notify(method, params, opts)


class Connection:
    """Per-client connection state and standalone-stream `Outbound`.

    Construct via `from_envelope` (modern single-exchange: born ready, no
    back-channel) or `for_loop` (handshake-driven: ready once the client's
    `notifications/initialized` arrives). Either way `protocol_version` is
    populated at construction.
    """

    outbound: Outbound
    """The connection-scoped channel for server-initiated messages."""

    session_id: str | None

    client_capabilities: ClientCapabilities | None
    """The capabilities the peer declared: the handshake's on the loop path,
    the request envelope's on the modern path. `None` when none were declared.
    Kept in lockstep with `client_params` by its setter, and settable on its
    own for the modern envelope, where capabilities are required but client
    info is optional (spec PR #3002) - capability checks must not depend on the
    peer having identified itself."""

    protocol_version: str
    """The protocol version this connection speaks. Populated at construction
    by the factory and overwritten by `_handle_initialize` once the handshake
    commits on the loop path."""

    initialized: anyio.Event
    """Set when `notifications/initialized` arrives (matches TS `oninitialized`);
    the point from which the spec permits server-initiated requests beyond
    ping/logging. Pre-set on connections built via `from_envelope`."""

    state: dict[str, Any]
    """Per-connection scratch state; persists across requests on this connection."""

    exit_stack: AsyncExitStack
    """Per-connection teardown, unwound LIFO (shielded) when the connection
    closes. Push cleanup from handlers or middleware; exceptions are logged
    and swallowed."""

    def __init__(
        self,
        outbound: Outbound,
        *,
        protocol_version: str,
        session_id: str | None = None,
        client_params: InitializeRequestParams | None = None,
    ) -> None:
        self.outbound = outbound
        self.protocol_version = protocol_version
        self.session_id = session_id
        self.client_capabilities = None
        self.client_params = client_params
        self.initialized = anyio.Event()
        self.state = {}
        self.exit_stack = AsyncExitStack()

    @property
    def client_params(self) -> InitializeRequestParams | None:
        """The full `initialize` request params, or the equivalent built from the
        2026-era envelope. `None` when no client info was supplied."""
        return self._client_params

    @client_params.setter
    def client_params(self, value: InitializeRequestParams | None) -> None:
        # Assignment is the sync point: recording full client params (the
        # handshake commit, or a modern envelope carrying client info) also
        # records the capabilities fact, so the two can never drift. Clearing
        # to `None` leaves `client_capabilities` alone - the modern envelope
        # declares capabilities without client info.
        self._client_params = value
        if value is not None:
            self.client_capabilities = value.capabilities

    @classmethod
    def from_envelope(
        cls,
        protocol_version: str,
        client_info: Any,
        client_capabilities: Any,
        *,
        outbound: Outbound = _NO_CHANNEL,
    ) -> Connection:
        """A born-ready connection populated from a request's `_meta` envelope.

        `protocol_version` must be an already-validated version string - the
        inbound classification ladder owns rejecting non-string or unsupported
        values. `client_info` and `client_capabilities` are the raw envelope
        values: this constructor owns turning them into connection identity,
        identically on every modern entry, so a mis-shaped value degrades to
        not-supplied rather than failing the request. `initialized` is set,
        well-formed capabilities are recorded as `client_capabilities` (client
        info is optional per spec PR #3002, so capability checks never depend on
        it), and the full `client_params` is additionally synthesized when
        client info was supplied too. `outbound` defaults to the no-channel
        sentinel for the single-exchange HTTP path; duplex modern transports
        (e.g. stdio) pass a notify-only wrapper around the dispatcher so
        server notifications ride the pipe while server-initiated requests
        stay refused.
        """
        info = _typed(Implementation, client_info)
        capabilities = _typed(ClientCapabilities, client_capabilities)
        client_params = None
        if info is not None and capabilities is not None:
            client_params = InitializeRequestParams(
                protocol_version=protocol_version,
                capabilities=capabilities,
                client_info=info,
            )
        connection = cls(outbound, protocol_version=protocol_version, client_params=client_params)
        connection.client_capabilities = capabilities
        connection.initialized.set()
        return connection

    @classmethod
    def for_loop(
        cls,
        outbound: Outbound,
        *,
        session_id: str | None = None,
        protocol_version_hint: str | None = None,
    ) -> Connection:
        """A connection for the handshake-driven loop path.

        Not born-ready: `initialized` is set later by the kernel when
        `notifications/initialized` arrives. `protocol_version` is seeded from
        the transport hint (or `LATEST_HANDSHAKE_VERSION`) so it's never `None`;
        the handshake overwrites it once negotiated.
        """
        return cls(
            outbound,
            protocol_version=protocol_version_hint if protocol_version_hint is not None else LATEST_HANDSHAKE_VERSION,
            session_id=session_id,
        )

    @property
    def has_standalone_channel(self) -> bool:
        """Whether this connection has a real back-channel for server-initiated
        messages. Derived from `outbound` - the no-channel sentinel is the only
        case that doesn't.

        Channel presence, not request permission: a modern (2026-07-28+)
        duplex connection has a channel that carries notifications while
        `send_raw_request` still refuses, because the protocol forbids
        server-initiated requests."""
        return self.outbound is not _NO_CHANNEL

    @property
    def initialize_accepted(self) -> bool:
        """True once the inbound request gate is open: `initialize` recorded the
        peer info, or the handshake completed outright (born-ready, or a bare
        `notifications/initialized`). Derived, never stored."""
        return self.client_params is not None or self.initialized.is_set()

    async def send_raw_request(
        self,
        method: str,
        params: Mapping[str, Any] | None,
        opts: CallOptions | None = None,
    ) -> dict[str, Any]:
        """Send a raw request on the standalone stream.

        Low-level `Outbound` channel. Prefer the typed `send_request` or the
        convenience methods below; use this directly only for off-spec
        messages. `opts` carries per-call `timeout` / `on_progress` /
        resumption hints; see `CallOptions`.

        Raises:
            MCPError: The peer responded with an error.
            NoBackChannelError: no back-channel for server-initiated requests -
                `has_standalone_channel` is `False`, or a modern (2026-07-28+)
                connection, where the protocol forbids them.
        """
        return await self.outbound.send_raw_request(method, params, opts)

    @overload
    async def send_request(
        self, req: CreateMessageRequest, *, opts: CallOptions | None = None
    ) -> CreateMessageResult: ...
    @overload
    async def send_request(self, req: ElicitRequest, *, opts: CallOptions | None = None) -> ElicitResult: ...
    @overload
    async def send_request(self, req: ListRootsRequest, *, opts: CallOptions | None = None) -> ListRootsResult: ...
    @overload
    async def send_request(self, req: PingRequest, *, opts: CallOptions | None = None) -> EmptyResult: ...
    @overload
    async def send_request(
        self, req: Request[Any, Any], *, result_type: type[ResultT], opts: CallOptions | None = None
    ) -> ResultT: ...
    async def send_request(
        self,
        req: Request[Any, Any],
        *,
        result_type: type[BaseModel] | None = None,
        opts: CallOptions | None = None,
    ) -> BaseModel:
        """Send a typed server-to-client request and return its typed result.

        For spec request types the result type is inferred. For custom requests
        pass `result_type=` explicitly.

        Raises:
            MCPError: The peer responded with an error.
            NoBackChannelError: No back-channel for server-initiated requests.
            pydantic.ValidationError: The peer's result does not match the expected result type.
            KeyError: `result_type` omitted for a non-spec request type.
        """
        raw = await self.send_raw_request(req.method, dump_params(req.params), opts)
        if req.method in _methods.MONOLITH_REQUESTS:
            try:
                _methods.validate_client_result(req.method, self.protocol_version, raw)
            except KeyError:
                pass
        cls = result_type if result_type is not None else _RESULT_FOR[type(req)]
        return cls.model_validate(raw, by_name=False)

    async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None:
        """Send a best-effort notification on the standalone stream.

        Never raises. If there's no standalone channel or the stream is broken,
        the notification is dropped and debug-logged.
        """
        try:
            await self.outbound.notify(method, params, opts)
        except (anyio.BrokenResourceError, anyio.ClosedResourceError):
            logger.debug("dropped %s: standalone stream closed", method)

    async def ping(self, *, meta: Meta | None = None, opts: CallOptions | None = None) -> None:
        """Send a `ping` request on the standalone stream.

        Raises:
            MCPError: The peer responded with an error.
            NoBackChannelError: no back-channel for server-initiated requests -
                `has_standalone_channel` is `False`, or a modern (2026-07-28+)
                connection, where the protocol forbids them.
        """
        await self.send_raw_request("ping", dump_params(None, meta), opts)

    @deprecated("The logging capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning)
    async def log(self, level: LoggingLevel, data: Any, logger: str | None = None, *, meta: Meta | None = None) -> None:
        """Send a `notifications/message` log entry on the standalone stream. Best-effort.

        On 2026-07-28+ connections this never sends: log delivery is a
        per-request opt-in that rides the requesting stream (`ctx.log`,
        `ctx.session.send_log_message`), and the standalone stream is
        forbidden from carrying `notifications/message`, so the entry is
        debug-logged and dropped.
        """
        if level not in allowed_log_levels(self.protocol_version, None):
            _logger.debug("dropped notifications/message: no connection-wide log delivery at %s", self.protocol_version)
            return
        params: dict[str, Any] = {"level": level, "data": data}
        if logger is not None:
            params["logger"] = logger
        await self.notify("notifications/message", _notification_params(params, meta))

    async def send_tool_list_changed(self, *, meta: Meta | None = None) -> None:
        await self.notify("notifications/tools/list_changed", _notification_params(None, meta))

    async def send_prompt_list_changed(self, *, meta: Meta | None = None) -> None:
        await self.notify("notifications/prompts/list_changed", _notification_params(None, meta))

    async def send_resource_list_changed(self, *, meta: Meta | None = None) -> None:
        await self.notify("notifications/resources/list_changed", _notification_params(None, meta))

    async def send_resource_updated(self, uri: str, *, meta: Meta | None = None) -> None:
        await self.notify("notifications/resources/updated", _notification_params({"uri": uri}, meta))

    def check_capability(self, capability: ClientCapabilities) -> bool:
        """Return whether the connected client declared the given capability.

        Returns `False` when no capabilities have been recorded.
        """
        # TODO(L53): redesign - mirrors v1 ServerSession.check_client_capability
        # verbatim for parity.
        if self.client_capabilities is None:
            return False
        have = self.client_capabilities
        if capability.roots is not None:
            if have.roots is None:
                return False
            if capability.roots.list_changed and not have.roots.list_changed:
                return False
        if capability.sampling is not None:
            if have.sampling is None:
                return False
            if capability.sampling.context is not None and have.sampling.context is None:
                return False
            if capability.sampling.tools is not None and have.sampling.tools is None:
                return False
        if capability.elicitation is not None and have.elicitation is None:
            return False
        if capability.experimental is not None:
            if have.experimental is None:
                return False
            for k, v in capability.experimental.items():
                if k not in have.experimental or have.experimental[k] != v:
                    return False
        if capability.extensions is not None:
            # SEP-2133: an extension is supported when the client declares its
            # identifier. Settings are negotiated per-extension (the client may
            # advertise more than the server asks for), so presence - not value
            # equality - is the meaningful check.
            if have.extensions is None:
                return False
            for identifier in capability.extensions:
                if identifier not in have.extensions:
                    return False
        return True