test_subscriptions.py

"""`docs/{handlers,client}/subscriptions.md`: every claim the two pages make, proved against the real SDK."""

from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any

import anyio
import mcp_types as types
import pytest
from trio.testing import MockClock

from docs_src.subscriptions import (
    tutorial001,
    tutorial002,
    tutorial003,
    tutorial004_anyio,
    tutorial004_asyncio,
    tutorial004_trio,
    tutorial005,
    tutorial006,
)
from mcp import Client
from mcp.client.subscriptions import Subscription
from mcp.server.auth.middleware.auth_context import auth_context_var
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
from mcp.server.auth.provider import AccessToken
from mcp.server.lowlevel import Server
from mcp.server.subscriptions import SUBSCRIPTION_ID_META_KEY, ListenHandler, ToolsListChanged
from mcp.shared.exceptions import MCPError

# See test_index.py for why this is a per-module mark and not a conftest hook.
pytestmark = [pytest.mark.anyio, pytest.mark.filterwarnings("error::mcp.MCPDeprecationWarning")]


@pytest.fixture(autouse=True)
def _module_runner_lease() -> None:
    """Opt out of the shared per-module event loop: this module parametrizes `anyio_backend`."""


class _Stream:
    """Collects listen-stream notifications and lets tests await arrival counts."""

    def __init__(self) -> None:
        self.received: list[types.ServerNotification] = []
        self._arrival = anyio.Event()

    async def handler(
        self,
        message: object,
    ) -> None:
        # The only messages these connections produce are the stream's frames.
        assert isinstance(
            message,
            types.SubscriptionsAcknowledgedNotification
            | types.ResourceUpdatedNotification
            | types.ToolListChangedNotification,
        ), message
        self.received.append(message)
        self._arrival.set()
        self._arrival = anyio.Event()

    async def wait_for(self, count: int) -> None:
        with anyio.fail_after(5):
            while len(self.received) < count:
                await self._arrival.wait()


class _Output:
    """Counts completed prints so tests wait for client output, not server-side reads."""

    def __init__(self) -> None:
        self.count = 0
        self._bump = anyio.Event()

    def __call__(self, *values: str | list[str]) -> None:
        print(*values)
        self.count += 1
        self._bump.set()
        self._bump = anyio.Event()

    async def wait_for(self, count: int) -> None:
        with anyio.fail_after(5):
            while self.count < count:
                await self._bump.wait()


def _listen_request(**fields: Any) -> types.SubscriptionsListenRequest:
    return types.SubscriptionsListenRequest(
        params=types.SubscriptionsListenRequestParams(notifications=types.SubscriptionFilter(**fields))
    )


@pytest.fixture(autouse=True)
def _fresh_server_state() -> Any:
    """Each test starts from an all-unfinished board and the base tool set.

    The tutorials mutate module state deliberately (that is what publishes events), so the
    board contents and the `enable_reports` registration have to be undone between tests.
    """
    boards = {name: dict(tasks) for name, tasks in tutorial001.BOARDS.items()}
    lowlevel_board = dict(tutorial002.BOARD)
    tools = dict(tutorial001.mcp._tool_manager._tools)  # pyright: ignore[reportPrivateUsage]
    yield
    tutorial001.BOARDS.clear()
    tutorial001.BOARDS.update(boards)
    tutorial002.BOARD.clear()
    tutorial002.BOARD.update(lowlevel_board)
    tutorial001.mcp._tool_manager._tools.clear()  # pyright: ignore[reportPrivateUsage]
    tutorial001.mcp._tool_manager._tools.update(tools)  # pyright: ignore[reportPrivateUsage]


async def test_publishes_reach_the_stream_filtered_and_tagged() -> None:
    """tutorial001: the full arc - ack first, exact-URI filtering, list_changed
    leading to a refreshed tool list, and client-side close."""
    stream = _Stream()
    async with Client(tutorial001.mcp, mode="2026-07-28", message_handler=stream.handler) as client:
        async with anyio.create_task_group() as tg:

            async def listen() -> None:
                await client.session.send_request(
                    _listen_request(tools_list_changed=True, resource_subscriptions=["board://sprint"]),
                    types.SubscriptionsListenResult,
                )

            tg.start_soon(listen)
            await stream.wait_for(1)

            ack = stream.received[0]
            assert isinstance(ack, types.SubscriptionsAcknowledgedNotification)
            assert ack.params.notifications == types.SubscriptionFilter(
                tools_list_changed=True, resource_subscriptions=["board://sprint"]
            )
            assert ack.params.meta is not None and SUBSCRIPTION_ID_META_KEY in ack.params.meta

            # An edit to a URI the stream did not subscribe to stays silent...
            await client.call_tool("complete_task", {"board": "backlog", "task": "tidy docs"})
            # ...and the subscribed URI delivers, tagged with the same subscription id.
            await client.call_tool("complete_task", {"board": "sprint", "task": "design"})
            await stream.wait_for(2)
            updated = stream.received[1]
            assert isinstance(updated, types.ResourceUpdatedNotification)
            assert updated.params.uri == "board://sprint"
            assert updated.params.meta == ack.params.meta

            await client.call_tool("enable_reports", {})
            await stream.wait_for(3)
            assert isinstance(stream.received[2], types.ToolListChangedNotification)

            # The client ends the stream by closing it - cancel the parked request.
            tg.cancel_scope.cancel()

        # The list_changed told us to re-fetch: the new tool is there, and the
        # session outlives the closed stream.
        tools = await client.list_tools()
        assert "sprint_report" in {tool.name for tool in tools.tools}
        contents = (await client.read_resource("board://sprint")).contents[0]
        assert isinstance(contents, types.TextResourceContents)
        assert contents.text == "[x] design\n[ ] build\n[ ] ship"


async def test_publish_with_no_subscribers_is_a_no_op() -> None:
    """tutorial001: publishing to an idle server does nothing and breaks nothing."""
    async with Client(tutorial001.mcp, mode="2026-07-28") as client:
        result = await client.call_tool("complete_task", {"board": "sprint", "task": "design"})
        assert result.is_error is not True


async def test_lowlevel_composition_serves_the_same_stream() -> None:
    """tutorial002: bus + ListenHandler on the lowlevel Server is the same machinery."""
    stream = _Stream()
    async with Client(tutorial002.server, mode="2026-07-28", message_handler=stream.handler) as client:
        tools = await client.list_tools()
        assert [tool.name for tool in tools.tools] == ["complete_task"]

        async with anyio.create_task_group() as tg:

            async def listen() -> None:
                await client.session.send_request(
                    _listen_request(resource_subscriptions=["board://sprint"]),
                    types.SubscriptionsListenResult,
                )

            tg.start_soon(listen)
            await stream.wait_for(1)

            await client.call_tool("complete_task", {"task": "design"})
            await stream.wait_for(2)
            updated = stream.received[1]
            assert isinstance(updated, types.ResourceUpdatedNotification)
            assert updated.params.uri == "board://sprint"

            # The bus you constructed is also the publish surface outside a
            # request; an unrequested kind never reaches this stream.
            await tutorial002.bus.publish(ToolsListChanged())
            await client.call_tool("complete_task", {"task": "build"})
            await stream.wait_for(3)
            assert isinstance(stream.received[2], types.ResourceUpdatedNotification)

            tg.cancel_scope.cancel()


async def test_follow_board_prints_the_refetched_board_and_the_new_tool_list(
    capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
) -> None:
    """tutorial003: each event drives a refetch - the board reprints, and a tools change reprints the tool names."""
    output = _Output()
    monkeypatch.setattr(tutorial003, "print", output, raising=False)
    listening = anyio.Event()
    async with Client(tutorial001.mcp) as client:
        listen = client.listen

        @asynccontextmanager
        async def listen_when_ready(
            *, tools_list_changed: bool, resource_subscriptions: list[str]
        ) -> AsyncIterator[Subscription]:
            async with listen(
                tools_list_changed=tools_list_changed, resource_subscriptions=resource_subscriptions
            ) as sub:
                listening.set()
                yield sub

        monkeypatch.setattr(client, "listen", listen_when_ready)
        async with anyio.create_task_group() as tg:
            tg.start_soon(tutorial003.follow_board, client)
            with anyio.fail_after(5):
                await listening.wait()
            await client.call_tool("complete_task", {"board": "sprint", "task": "design"})
            await output.wait_for(1)
            await client.call_tool("enable_reports", {})
            await output.wait_for(2)
            tg.cancel_scope.cancel()

    printed = capsys.readouterr().out
    assert "[x] design\n[ ] build\n[ ] ship" in printed
    assert "sprint_report" in printed


EMPTY_BOARD = "[ ] design\n[ ] build\n[ ] ship"
FINISHED_BOARD = "[x] design\n[x] build\n[x] ship"


def _assert_snapshot_then_current_board(printed: str) -> None:
    """The snapshot taken inside the open subscription came first, and the watcher ended up current.

    How many times the watcher printed is deliberately not asserted: identical events that pile up
    unconsumed coalesce, so a fast main flow can turn three completions into one refetch. What the
    stream guarantees is that no change after the acknowledgment is missed.
    """
    assert printed.startswith(EMPTY_BOARD), printed
    assert printed.strip().endswith(FINISHED_BOARD), printed


@pytest.mark.parametrize("anyio_backend", [pytest.param("asyncio", id="asyncio")])
async def test_the_asyncio_watcher_runs_beside_the_main_flow(capsys: pytest.CaptureFixture[str]) -> None:
    """tutorial004 (asyncio tab): run_sprint opens the subscription, snapshots the board, then a watcher
    task reprints it while the main flow keeps calling tools.

    The example connects over HTTP; the in-memory client here is the maintainer-side stand-in."""
    async with Client(tutorial001.mcp) as client:
        with anyio.fail_after(5):
            await tutorial004_asyncio.run_sprint(client)
    _assert_snapshot_then_current_board(capsys.readouterr().out)


@pytest.mark.parametrize("anyio_backend", [pytest.param("trio", id="trio")])
async def test_the_trio_watcher_runs_beside_the_main_flow(capsys: pytest.CaptureFixture[str]) -> None:
    """tutorial004 (trio tab): the same shape as the asyncio tab, with a nursery owning the watcher."""
    async with Client(tutorial001.mcp) as client:
        with anyio.fail_after(5):
            await tutorial004_trio.run_sprint(client)
    _assert_snapshot_then_current_board(capsys.readouterr().out)


async def test_the_anyio_watcher_runs_beside_the_main_flow(capsys: pytest.CaptureFixture[str]) -> None:
    """tutorial004 (anyio tab): the same shape again, with a task group owning the watcher."""
    async with Client(tutorial001.mcp) as client:
        with anyio.fail_after(5):
            await tutorial004_anyio.run_sprint(client)
    _assert_snapshot_then_current_board(capsys.readouterr().out)


@pytest.mark.parametrize(
    "anyio_backend",
    [pytest.param(("trio", {"clock": MockClock(autojump_threshold=0)}), id="trio-mockclock")],
)
async def test_the_follower_re_listens_after_the_stream_ends(
    capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
) -> None:
    """tutorial005: a graceful server close ends one stream; the loop backs off, re-listens, and refetches.

    Runs on trio's autojumping MockClock so the loop's backoff sleep takes no wall-clock time.
    """
    output = _Output()
    monkeypatch.setattr(tutorial005, "print", output, raising=False)
    handler = ListenHandler(tutorial002.bus)
    server = Server(
        "sprint-board",
        on_read_resource=tutorial002.read_resource,
        on_list_tools=tutorial002.list_tools,
        on_call_tool=tutorial002.call_tool,
        on_subscriptions_listen=handler,
    )

    async with Client(server) as client:
        async with anyio.create_task_group() as tg:
            tg.start_soon(tutorial005.keep_following, client)
            # First stream: print the entry snapshot, then the board after the event.
            await output.wait_for(1)
            await client.call_tool("complete_task", {"task": "design"})
            await output.wait_for(2)

            # The mock clock jumps the backoff; the third print is the new stream's snapshot.
            handler.close()
            await output.wait_for(3)
            await client.call_tool("complete_task", {"task": "build"})
            await output.wait_for(4)
            tg.cancel_scope.cancel()

    printed = capsys.readouterr().out
    assert "[x] design\n[ ] build" in printed  # first stream, after design
    assert "[x] design\n[x] build" in printed  # second stream, after build


def _signed_in_as(subject: str) -> Any:
    """Stand in for the auth middleware: put this user's token in the auth context."""
    token = AccessToken(token="demo", client_id="docs-client", scopes=[], subject=subject)
    return auth_context_var.set(AuthenticatedUser(token))


async def test_the_middleware_refuses_a_listen_the_caller_could_not_read() -> None:
    """tutorial006: one `can_access` gates both `resources/read` and `subscriptions/listen`.

    Alice may read (and so watch) the report, and is refused the payroll file on both
    paths - the listen refusal is in-band, before any acknowledgment.
    """
    reset = _signed_in_as("alice")
    try:
        async with Client(tutorial006.mcp, mode="2026-07-28") as client:
            async with client.listen(resource_subscriptions=["files://report.pdf"]) as sub:
                assert sub.honored.resource_subscriptions == ["files://report.pdf"]
            with pytest.raises(MCPError) as listen_error:
                async with client.listen(resource_subscriptions=["files://report.pdf", "files://payroll.csv"]):
                    pass  # pragma: no cover - the refusal precedes the stream
            assert listen_error.value.error.message == "not permitted to watch the requested resources"
            with pytest.raises(MCPError):
                await client.read_resource("files://payroll.csv")
    finally:
        auth_context_var.reset(reset)