test_logging_callback.py

from typing import Any, Literal

import pytest

from mcp import Client, types
from mcp.server.mcpserver import Context, MCPServer
from mcp.shared.session import RequestResponder
from mcp.types import (
    LoggingMessageNotificationParams,
    TextContent,
)


class LoggingCollector:
    def __init__(self):
        self.log_messages: list[LoggingMessageNotificationParams] = []

    async def __call__(self, params: LoggingMessageNotificationParams) -> None:
        self.log_messages.append(params)


@pytest.mark.anyio
async def test_logging_callback():
    server = MCPServer("test")
    logging_collector = LoggingCollector()

    # Create a simple test tool
    @server.tool("test_tool")
    async def test_tool() -> bool:
        # The actual tool is very simple and just returns True
        return True

    # Create a function that can send a log notification
    @server.tool("test_tool_with_log")
    async def test_tool_with_log(
        message: str, level: Literal["debug", "info", "warning", "error"], logger: str, ctx: Context
    ) -> bool:
        """Send a log notification to the client."""
        await ctx.log(level=level, message=message, logger_name=logger)
        return True

    @server.tool("test_tool_with_log_extra")
    async def test_tool_with_log_extra(
        message: str,
        level: Literal["debug", "info", "warning", "error"],
        logger: str,
        extra_string: str,
        extra_dict: dict[str, Any],
        ctx: Context,
    ) -> bool:
        """Send a log notification to the client with extra fields."""
        await ctx.log(
            level=level,
            message=message,
            logger_name=logger,
            extra={
                "extra_string": extra_string,
                "extra_dict": extra_dict,
            },
        )
        return True

    # Create a message handler to catch exceptions
    async def message_handler(
        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
    ) -> None:
        if isinstance(message, Exception):  # pragma: no cover
            raise message

    async with Client(
        server,
        logging_callback=logging_collector,
        message_handler=message_handler,
    ) as client:
        # First verify our test tool works
        result = await client.call_tool("test_tool", {})
        assert result.is_error is False
        assert isinstance(result.content[0], TextContent)
        assert result.content[0].text == "true"

        # Now send a log message via our tool
        log_result = await client.call_tool(
            "test_tool_with_log",
            {
                "message": "Test log message",
                "level": "info",
                "logger": "test_logger",
            },
        )
        log_result_with_extra = await client.call_tool(
            "test_tool_with_log_extra",
            {
                "message": "Test log message",
                "level": "info",
                "logger": "test_logger",
                "extra_string": "example",
                "extra_dict": {"a": 1, "b": 2, "c": 3},
            },
        )
        assert log_result.is_error is False
        assert log_result_with_extra.is_error is False
        assert len(logging_collector.log_messages) == 2
        # Create meta object with related_request_id added dynamically
        log = logging_collector.log_messages[0]
        assert log.level == "info"
        assert log.logger == "test_logger"
        assert log.data == "Test log message"

        log_with_extra = logging_collector.log_messages[1]
        assert log_with_extra.level == "info"
        assert log_with_extra.logger == "test_logger"
        assert log_with_extra.data == {
            "message": "Test log message",
            "extra_string": "example",
            "extra_dict": {"a": 1, "b": 2, "c": 3},
        }