subscriptions.py

"""Client-side `subscriptions/listen` driver (2026-07-28, SEP-2575).

`listen()` opens the stream as an async context manager: entering waits for
the server's acknowledgment, iteration yields typed change events, a graceful
server close ends the loop, and an abrupt drop raises `SubscriptionLost`.
There is no replay and no automatic re-listen: a client that re-opens a
subscription refetches what it depends on.
"""

from __future__ import annotations

from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
from contextlib import asynccontextmanager
from itertools import count
from typing import TYPE_CHECKING, Literal

import anyio
import mcp_types as types
from mcp_types.version import MODERN_PROTOCOL_VERSIONS

from mcp.shared.direct_dispatcher import DirectDispatcher
from mcp.shared.dispatcher import CallOptions
from mcp.shared.exceptions import MCPError
from mcp.shared.subscriptions import (
    PromptsListChanged,
    ResourcesListChanged,
    ResourceUpdated,
    ServerEvent,
    ToolsListChanged,
    event_matches,
)

if TYPE_CHECKING:
    from mcp.client.session import ClientSession

__all__ = [
    "ListenNotSupportedError",
    "OnEvent",
    "PromptsListChanged",
    "ResourceUpdated",
    "ResourcesListChanged",
    "ServerEvent",
    "Subscription",
    "SubscriptionLost",
    "ToolsListChanged",
    "listen",
]

_listen_ids = count(1)
"""Process-wide `listen-N` sequence: string ids can never collide with a dispatcher's minted ints."""

_MAX_PENDING_EVENTS = 1024
"""Backlog backstop: the spec allows sub-resource URIs, so distinct pending
`ResourceUpdated` events are unbounded; overflowing this cap settles the
subscription lost rather than growing client memory."""

_SubscriptionEnd = Literal["graceful", "lost", "local"]


class ListenNotSupportedError(RuntimeError):
    """`subscriptions/listen` requires a 2026-07-28 connection."""

    def __init__(self, negotiated_version: str | None) -> None:
        self.negotiated_version = negotiated_version
        super().__init__(
            f"subscriptions/listen is not available at protocol version {negotiated_version!r}; it requires "
            "2026-07-28. On earlier versions use subscribe_resource() and the change notifications delivered "
            "through message_handler."
        )


class SubscriptionLost(RuntimeError):
    """The stream ended without the server's graceful close; re-listen and refetch."""


class ListenRoute:
    """Package-internal demux state for one listen stream, fed synchronously in receive order by the session."""

    def __init__(self) -> None:
        self.honored: types.SubscriptionFilter | None = None
        self.acked = anyio.Event()
        self.error: MCPError | None = None
        self.end: _SubscriptionEnd | None = None
        self._honored_uris: frozenset[str] = frozenset()
        self._pending: dict[ServerEvent, None] = {}
        self._wake = anyio.Event()

    def set_acked(self, honored: types.SubscriptionFilter) -> None:
        """Record the acknowledged filter; the first ack wins."""
        if not self.acked.is_set():
            self.honored = honored
            self._honored_uris = frozenset(honored.resource_subscriptions or ())
            self.acked.set()

    def deliver(self, event: ServerEvent) -> None:
        """Queue an event within the honored filter, deduplicated against the backlog.

        Any `ResourceUpdated` is admitted once URI subscriptions were honored at
        all: the spec allows the stamped URI to be a sub-resource of a subscribed one.
        """
        if self.end is not None or self.honored is None:
            return
        if isinstance(event, ResourceUpdated):
            admitted = bool(self._honored_uris)
        else:
            admitted = event_matches(self.honored, self._honored_uris, event)
        if not admitted or event in self._pending:
            return
        if len(self._pending) >= _MAX_PENDING_EVENTS:
            self.settle(
                "lost",
                error=MCPError(
                    types.INTERNAL_ERROR,
                    f"subscription backlog exceeded {_MAX_PENDING_EVENTS} unconsumed events; re-listen and refetch",
                ),
            )
            return
        self._pending[event] = None
        self._wake.set()

    def settle(self, end: _SubscriptionEnd, error: MCPError | None = None) -> None:
        """Record the stream's end; the first reason wins and wakes both waiters."""
        if self.end is None:
            self.end = end
            self.error = error
            self.acked.set()
            self._wake.set()

    async def next_event(self) -> ServerEvent | _SubscriptionEnd:
        """Peek the next pending event, or the stream's end once the backlog drains.

        A "local" end short-circuits the backlog; the other endings drain it first,
        so a graceful close never swallows events that preceded it.
        """
        while True:
            # Snapshot the wake event before checking state so a deliver landing after the checks cannot be missed.
            wake = self._wake
            if self.end == "local":
                return self.end
            if self._pending:
                return next(iter(self._pending))
            if self.end is not None:
                return self.end
            await wake.wait()
            self._wake = anyio.Event()

    def consume(self, event: ServerEvent) -> None:
        """Remove a peeked event from the backlog."""
        self._pending.pop(event, None)


OnEvent = Callable[[ServerEvent], Awaitable[None]]
"""Per-event barrier awaited before a `Subscription` returns each event to its consumer."""


class Subscription:
    """One open `subscriptions/listen` stream: an async iterator of typed events.

    Produced by `listen()` / `Client.listen()`, not constructed directly.
    """

    def __init__(
        self,
        route: ListenRoute,
        subscription_id: types.RequestId,
        honored: types.SubscriptionFilter,
        on_event: OnEvent | None = None,
    ):
        self._route = route
        self._on_event = on_event
        self.subscription_id = subscription_id
        """The listen request's JSON-RPC id, stamped into every frame's `_meta`."""
        self.honored = honored
        """The subset of the requested filter the server agreed to deliver."""

    def __aiter__(self) -> Subscription:
        return self

    async def __anext__(self) -> ServerEvent:
        """Yield the next change event; the loop ends when the stream does.

        Raises:
            SubscriptionLost: the stream dropped without the server's graceful close.
        """
        outcome = await self._route.next_event()
        if isinstance(outcome, str):
            if outcome == "lost":
                raise SubscriptionLost(
                    f"subscription {self.subscription_id!r} ended without the server's graceful close;"
                    " re-listen and refetch"
                ) from self._route.error
            raise StopAsyncIteration
        if self._on_event is not None:
            # The event stays pending while the barrier runs: a cancellation or a
            # raising barrier leaves it for the next anext instead of dropping it.
            await self._on_event(outcome)
        self._route.consume(outcome)
        return outcome


@asynccontextmanager
async def listen(
    session: ClientSession,
    *,
    tools_list_changed: bool = False,
    prompts_list_changed: bool = False,
    resources_list_changed: bool = False,
    resource_subscriptions: Sequence[str] = (),
    on_event: OnEvent | None = None,
) -> AsyncIterator[Subscription]:
    """Open one `subscriptions/listen` stream on `session` (2026-07-28 only).

    Entering sends the request and returns once the server's acknowledgment
    arrives; exiting ends the subscription. `on_event` is awaited before each
    event is returned - the seam `Client.listen` uses to finish cache eviction
    before the consumer can refetch.

    Raises:
        ListenNotSupportedError: negotiated version predates 2026-07-28.
        MCPError: the server rejected the request, or the connection failed pre-ack.
        SubscriptionLost: the stream ended before it was acknowledged.
        TimeoutError: the session's read timeout elapsed before the acknowledgment.
    """
    if session.protocol_version not in MODERN_PROTOCOL_VERSIONS:
        raise ListenNotSupportedError(session.protocol_version)
    if isinstance(resource_subscriptions, str):
        raise TypeError("resource_subscriptions takes a sequence of URIs, not a bare string")
    request = types.SubscriptionsListenRequest(
        params=types.SubscriptionsListenRequestParams(
            notifications=types.SubscriptionFilter(
                tools_list_changed=tools_list_changed or None,
                prompts_list_changed=prompts_list_changed or None,
                resources_list_changed=resources_list_changed or None,
                resource_subscriptions=list(resource_subscriptions) or None,
            )
        )
    )
    task_group = session._task_group  # pyright: ignore[reportPrivateUsage]
    if task_group is None:
        raise RuntimeError("listen() requires an entered session")
    request_id: types.RequestId = f"listen-{next(_listen_ids)}"
    data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
    opts: CallOptions = {"request_id": request_id}
    session._stamp(data, opts)  # pyright: ignore[reportPrivateUsage]
    dispatcher = session._dispatcher  # pyright: ignore[reportPrivateUsage]
    driver_scope = anyio.CancelScope()
    driver_done = anyio.Event()

    async def drive() -> None:
        # Deliberately no result timeout: the response arrives when the stream ends.
        try:
            with driver_scope:
                try:
                    await dispatcher.send_raw_request(data["method"], data.get("params"), opts)
                except MCPError as error:
                    route.settle("lost", error=error)
                    return
                except ValueError as error:
                    # A raw request id collided with our minted listen id: fail this subscription
                    # and release the route in this same slice, so it cannot consume the raw caller's ack.
                    session._unregister_listen_route(request_id)  # pyright: ignore[reportPrivateUsage]
                    route.settle("lost", error=MCPError(types.INTERNAL_ERROR, str(error)))
                    return
                # A result, whatever its body, is the spec's graceful close; with no prior ack
                # it opens the subscription already closed.
                route.set_acked(types.SubscriptionFilter())
                route.settle("graceful")
        finally:
            driver_done.set()

    # Register the demux route before the request is written so the ack cannot race it.
    route = session._register_listen_route(request_id)  # pyright: ignore[reportPrivateUsage]
    try:
        task_group.start_soon(drive)
        try:
            with anyio.fail_after(session._session_read_timeout_seconds):  # pyright: ignore[reportPrivateUsage]
                await route.acked.wait()
            if route.honored is None:
                # Only reachable on failure paths: a graceful no-ack result acked an empty filter in drive().
                if route.error is not None:
                    raise route.error
                raise SubscriptionLost(f"subscription {request_id!r} ended before it was acknowledged")
            yield Subscription(route, request_id, route.honored, on_event)
        finally:
            route.settle("local")
            driver_scope.cancel()
            # Only direct drivers own handler cleanup; remote courtesy writes remain session-owned.
            if isinstance(dispatcher, DirectDispatcher):
                with anyio.move_on_after(5, shield=True):
                    await driver_done.wait()
    finally:
        session._unregister_listen_route(request_id)  # pyright: ignore[reportPrivateUsage]