test_progress_notifications.py

from typing import Any, cast
from unittest.mock import patch

import anyio
import pytest

import mcp.types as types
from mcp import Client
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.shared.context import RequestContext
from mcp.shared.message import SessionMessage
from mcp.shared.progress import progress
from mcp.shared.session import BaseSession, RequestResponder


@pytest.mark.anyio
async def test_bidirectional_progress_notifications():
    """Test that both client and server can send progress notifications."""
    # Create memory streams for client/server
    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](5)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](5)

    # Run a server session so we can send progress updates in tool
    async def run_server():
        # Create a server session
        async with ServerSession(
            client_to_server_receive,
            server_to_client_send,
            InitializationOptions(
                server_name="ProgressTestServer",
                server_version="0.1.0",
                capabilities=server.get_capabilities(NotificationOptions(), {}),
            ),
        ) as server_session:
            global serv_sesh

            serv_sesh = server_session
            async for message in server_session.incoming_messages:
                try:
                    await server._handle_message(message, server_session, {})
                except Exception as e:  # pragma: no cover
                    raise e

    # Track progress updates
    server_progress_updates: list[dict[str, Any]] = []
    client_progress_updates: list[dict[str, Any]] = []

    # Progress tokens
    server_progress_token = "server_token_123"
    client_progress_token = "client_token_456"

    # Create a server with progress capability
    server = Server(name="ProgressTestServer")

    # Register progress handler
    @server.progress_notification()
    async def handle_progress(
        progress_token: str | int,
        progress: float,
        total: float | None,
        message: str | None,
    ):
        server_progress_updates.append(
            {
                "token": progress_token,
                "progress": progress,
                "total": total,
                "message": message,
            }
        )

    # Register list tool handler
    @server.list_tools()
    async def handle_list_tools() -> list[types.Tool]:
        return [
            types.Tool(
                name="test_tool",
                description="A tool that sends progress notifications <o/",
                input_schema={},
            )
        ]

    # Register tool handler
    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[types.TextContent]:
        # Make sure we received a progress token
        if name == "test_tool":
            if arguments and "_meta" in arguments:
                progressToken = arguments["_meta"]["progressToken"]

                if not progressToken:  # pragma: no cover
                    raise ValueError("Empty progress token received")

                if progressToken != client_progress_token:  # pragma: no cover
                    raise ValueError("Server sending back incorrect progressToken")

                # Send progress notifications
                await serv_sesh.send_progress_notification(
                    progress_token=progressToken,
                    progress=0.25,
                    total=1.0,
                    message="Server progress 25%",
                )

                await serv_sesh.send_progress_notification(
                    progress_token=progressToken,
                    progress=0.5,
                    total=1.0,
                    message="Server progress 50%",
                )

                await serv_sesh.send_progress_notification(
                    progress_token=progressToken,
                    progress=1.0,
                    total=1.0,
                    message="Server progress 100%",
                )

            else:  # pragma: no cover
                raise ValueError("Progress token not sent.")

            return [types.TextContent(type="text", text="Tool executed successfully")]

        raise ValueError(f"Unknown tool: {name}")  # pragma: no cover

    # Client message handler to store progress notifications
    async def handle_client_message(
        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
    ) -> None:
        if isinstance(message, Exception):  # pragma: no cover
            raise message

        if isinstance(message, types.ServerNotification):  # pragma: no branch
            if isinstance(message.root, types.ProgressNotification):  # pragma: no branch
                params = message.root.params
                client_progress_updates.append(
                    {
                        "token": params.progress_token,
                        "progress": params.progress,
                        "total": params.total,
                        "message": params.message,
                    }
                )

    # Test using client
    async with (
        ClientSession(
            server_to_client_receive,
            client_to_server_send,
            message_handler=handle_client_message,
        ) as client_session,
        anyio.create_task_group() as tg,
    ):
        # Start the server in a background task
        tg.start_soon(run_server)

        # Initialize the client connection
        await client_session.initialize()

        # Call list_tools with progress token
        await client_session.list_tools()

        # Call test_tool with progress token
        await client_session.call_tool("test_tool", {"_meta": {"progressToken": client_progress_token}})

        # Send progress notifications from client to server
        await client_session.send_progress_notification(
            progress_token=server_progress_token,
            progress=0.33,
            total=1.0,
            message="Client progress 33%",
        )

        await client_session.send_progress_notification(
            progress_token=server_progress_token,
            progress=0.66,
            total=1.0,
            message="Client progress 66%",
        )

        await client_session.send_progress_notification(
            progress_token=server_progress_token,
            progress=1.0,
            total=1.0,
            message="Client progress 100%",
        )

        # Wait and exit
        await anyio.sleep(0.5)
        tg.cancel_scope.cancel()

    # Verify client received progress updates from server
    assert len(client_progress_updates) == 3
    assert client_progress_updates[0]["token"] == client_progress_token
    assert client_progress_updates[0]["progress"] == 0.25
    assert client_progress_updates[0]["message"] == "Server progress 25%"
    assert client_progress_updates[2]["progress"] == 1.0

    # Verify server received progress updates from client
    assert len(server_progress_updates) == 3
    assert server_progress_updates[0]["token"] == server_progress_token
    assert server_progress_updates[0]["progress"] == 0.33
    assert server_progress_updates[0]["message"] == "Client progress 33%"
    assert server_progress_updates[2]["progress"] == 1.0


@pytest.mark.anyio
async def test_progress_context_manager():
    """Test client using progress context manager for sending progress notifications."""
    # Create memory streams for client/server
    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](5)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](5)

    # Track progress updates
    server_progress_updates: list[dict[str, Any]] = []

    server = Server(name="ProgressContextTestServer")

    progress_token = None

    # Register progress handler
    @server.progress_notification()
    async def handle_progress(
        progress_token: str | int,
        progress: float,
        total: float | None,
        message: str | None,
    ):
        server_progress_updates.append(
            {"token": progress_token, "progress": progress, "total": total, "message": message}
        )

    # Run server session to receive progress updates
    async def run_server():
        # Create a server session
        async with ServerSession(
            client_to_server_receive,
            server_to_client_send,
            InitializationOptions(
                server_name="ProgressContextTestServer",
                server_version="0.1.0",
                capabilities=server.get_capabilities(NotificationOptions(), {}),
            ),
        ) as server_session:
            async for message in server_session.incoming_messages:
                try:
                    await server._handle_message(message, server_session, {})
                except Exception as e:  # pragma: no cover
                    raise e

    # Client message handler
    async def handle_client_message(
        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
    ) -> None:
        if isinstance(message, Exception):  # pragma: no cover
            raise message

    # run client session
    async with (
        ClientSession(
            server_to_client_receive,
            client_to_server_send,
            message_handler=handle_client_message,
        ) as client_session,
        anyio.create_task_group() as tg,
    ):
        tg.start_soon(run_server)

        await client_session.initialize()

        progress_token = "client_token_456"

        # Create request context
        meta = types.RequestParams.Meta(progress_token=progress_token)
        request_context = RequestContext(
            request_id="test-request",
            session=client_session,
            meta=meta,
            lifespan_context=None,
        )

        # cast for type checker
        typed_context = cast(RequestContext[BaseSession[Any, Any, Any, Any, Any], Any], request_context)

        # Utilize progress context manager
        with progress(typed_context, total=100) as p:
            await p.progress(10, message="Loading configuration...")
            await p.progress(30, message="Connecting to database...")
            await p.progress(40, message="Fetching data...")
            await p.progress(20, message="Processing results...")

        # Wait for all messages to be processed
        await anyio.sleep(0.5)
        tg.cancel_scope.cancel()

    # Verify progress updates were received by server
    assert len(server_progress_updates) == 4

    # first update
    assert server_progress_updates[0]["token"] == progress_token
    assert server_progress_updates[0]["progress"] == 10
    assert server_progress_updates[0]["total"] == 100
    assert server_progress_updates[0]["message"] == "Loading configuration..."

    # second update
    assert server_progress_updates[1]["token"] == progress_token
    assert server_progress_updates[1]["progress"] == 40
    assert server_progress_updates[1]["total"] == 100
    assert server_progress_updates[1]["message"] == "Connecting to database..."

    # third update
    assert server_progress_updates[2]["token"] == progress_token
    assert server_progress_updates[2]["progress"] == 80
    assert server_progress_updates[2]["total"] == 100
    assert server_progress_updates[2]["message"] == "Fetching data..."

    # final update
    assert server_progress_updates[3]["token"] == progress_token
    assert server_progress_updates[3]["progress"] == 100
    assert server_progress_updates[3]["total"] == 100
    assert server_progress_updates[3]["message"] == "Processing results..."


@pytest.mark.anyio
async def test_progress_callback_exception_logging():
    """Test that exceptions in progress callbacks are logged and \
        don't crash the session."""
    # Track logged warnings
    logged_errors: list[str] = []

    def mock_log_error(msg: str, *args: Any) -> None:
        logged_errors.append(msg % args if args else msg)

    # Create a progress callback that raises an exception
    async def failing_progress_callback(
        progress: float, total: float | None, message: str | None
    ) -> None:  # pragma: no cover
        raise ValueError("Progress callback failed!")

    # Create a server with a tool that sends progress notifications
    server = Server(name="TestProgressServer")

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: Any) -> list[types.TextContent]:
        if name == "progress_tool":
            # Send a progress notification
            await server.request_context.session.send_progress_notification(
                progress_token=server.request_context.request_id,
                progress=50.0,
                total=100.0,
                message="Halfway done",
            )
            return [types.TextContent(type="text", text="progress_result")]
        raise ValueError(f"Unknown tool: {name}")  # pragma: no cover

    @server.list_tools()
    async def handle_list_tools() -> list[types.Tool]:
        return [
            types.Tool(
                name="progress_tool",
                description="A tool that sends progress notifications",
                input_schema={},
            )
        ]

    # Test with mocked logging
    with patch("mcp.shared.session.logging.error", side_effect=mock_log_error):
        async with Client(server) as client:
            # Call tool with a failing progress callback
            result = await client.call_tool(
                "progress_tool",
                arguments={},
                progress_callback=failing_progress_callback,
            )

            # Verify the request completed successfully despite the callback failure
            assert len(result.content) == 1
            content = result.content[0]
            assert isinstance(content, types.TextContent)
            assert content.text == "progress_result"

            # Check that a warning was logged for the progress callback exception
            assert len(logged_errors) > 0
            assert any("Progress callback raised an exception" in warning for warning in logged_errors)