prefecthq--fastmcp/tests/server/middleware/test_response_limiting.py

"""Tests for ResponseLimitingMiddleware."""

import pytest
from mcp_types import ImageContent, TextContent
from pydantic import BaseModel

from fastmcp import Client, FastMCP
from fastmcp.server.middleware.response_limiting import ResponseLimitingMiddleware
from fastmcp.tools.base import ToolResult


# Regression test model for #3717
class Answer(BaseModel):
    text: str


class TestResponseLimitingMiddleware:
    """Tests for ResponseLimitingMiddleware."""

    @pytest.fixture
    def mcp_server(self) -> FastMCP:
        """Create a basic MCP server for testing."""
        return FastMCP("test-server")

    async def test_response_under_limit_passes_unchanged(self, mcp_server: FastMCP):
        """Test that responses under the limit pass through unchanged."""
        mcp_server.add_middleware(ResponseLimitingMiddleware(max_size=1_000_000))

        @mcp_server.tool()
        def small_tool() -> ToolResult:
            return ToolResult(content=[TextContent(type="text", text="hello world")])

        async with Client(mcp_server) as client:
            result = await client.call_tool("small_tool", {})
            assert len(result.content) == 1
            assert result.content[0].text == "hello world"

    async def test_response_over_limit_is_truncated(self, mcp_server: FastMCP):
        """Test that responses over the limit are truncated."""
        mcp_server.add_middleware(ResponseLimitingMiddleware(max_size=500))

        @mcp_server.tool()
        def large_tool() -> ToolResult:
            return ToolResult(content=[TextContent(type="text", text="x" * 10_000)])

        async with Client(mcp_server) as client:
            result = await client.call_tool("large_tool", {})
            assert len(result.content) == 1
            assert "[Response truncated due to size limit]" in result.content[0].text
            # Verify truncated result fits within limit
            assert len(result.content[0].text.encode("utf-8")) < 500

    async def test_tool_filtering(self, mcp_server: FastMCP):
        """Test that tool filtering only applies to specified tools."""
        mcp_server.add_middleware(
            ResponseLimitingMiddleware(max_size=100, tools=["limited_tool"])
        )

        @mcp_server.tool()
        def limited_tool() -> ToolResult:
            return ToolResult(content=[TextContent(type="text", text="x" * 10_000)])

        @mcp_server.tool()
        def unlimited_tool() -> ToolResult:
            return ToolResult(content=[TextContent(type="text", text="y" * 10_000)])

        async with Client(mcp_server) as client:
            # Limited tool should be truncated
            result = await client.call_tool("limited_tool", {})
            assert "[Response truncated" in result.content[0].text

            # Unlimited tool should pass through
            result = await client.call_tool("unlimited_tool", {})
            assert "y" * 100 in result.content[0].text

    async def test_empty_tools_list_limits_nothing(self, mcp_server: FastMCP):
        """Test that empty tools list means no tools are limited."""
        mcp_server.add_middleware(ResponseLimitingMiddleware(max_size=100, tools=[]))

        @mcp_server.tool()
        def any_tool() -> ToolResult:
            return ToolResult(content=[TextContent(type="text", text="x" * 10_000)])

        async with Client(mcp_server) as client:
            result = await client.call_tool("any_tool", {})
            # Should NOT be truncated
            assert "[Response truncated" not in result.content[0].text

    async def test_custom_truncation_suffix(self, mcp_server: FastMCP):
        """Test that custom truncation suffix is applied."""
        mcp_server.add_middleware(
            ResponseLimitingMiddleware(max_size=200, truncation_suffix="\n[CUT]")
        )

        @mcp_server.tool()
        def large_tool() -> ToolResult:
            return ToolResult(content=[TextContent(type="text", text="x" * 10_000)])

        async with Client(mcp_server) as client:
            result = await client.call_tool("large_tool", {})
            assert "[CUT]" in result.content[0].text

    async def test_multiple_text_blocks_combined(self, mcp_server: FastMCP):
        """Test that multiple text blocks are combined when truncating."""
        mcp_server.add_middleware(ResponseLimitingMiddleware(max_size=300))

        @mcp_server.tool()
        def multi_block() -> ToolResult:
            return ToolResult(
                content=[
                    TextContent(type="text", text="First: " + "a" * 500),
                    TextContent(type="text", text="Second: " + "b" * 500),
                ]
            )

        async with Client(mcp_server) as client:
            result = await client.call_tool("multi_block", {})
            # Both blocks should be joined and truncated
            assert len(result.content) == 1
            assert "[Response truncated" in result.content[0].text

    async def test_binary_only_content_serialized(self, mcp_server: FastMCP):
        """Test that binary-only responses fall back to serialized content."""
        mcp_server.add_middleware(ResponseLimitingMiddleware(max_size=200))

        @mcp_server.tool()
        def binary_tool() -> ToolResult:
            return ToolResult(
                content=[
                    ImageContent(type="image", data="x" * 10_000, mime_type="image/png")
                ]
            )

        async with Client(mcp_server) as client:
            result = await client.call_tool("binary_tool", {})
            # Should be truncated (using serialized fallback)
            assert len(result.content) == 1
            assert "[Response truncated" in result.content[0].text

    async def test_default_max_size_is_1mb(self):
        """Test that the default max size is 1MB."""
        middleware = ResponseLimitingMiddleware()
        assert middleware.max_size == 1_000_000

    def test_invalid_max_size_raises(self):
        """Test that zero or negative max_size raises ValueError."""
        with pytest.raises(ValueError, match="max_size must be positive"):
            ResponseLimitingMiddleware(max_size=0)
        with pytest.raises(ValueError, match="max_size must be positive"):
            ResponseLimitingMiddleware(max_size=-100)

    async def test_truncation_does_not_break_output_schema_tools(
        self, mcp_server: FastMCP
    ):
        """Truncating a tool with outputSchema must not cause validation errors.

        Regression test for #3717: the MCP SDK rejects truncated results
        from tools with outputSchema because structured_content is dropped.
        We verify the server returns a successful (non-error) truncated result.
        """
        mcp_server.add_middleware(ResponseLimitingMiddleware(max_size=1_000))

        @mcp_server.tool()
        def big_answer() -> Answer:
            return Answer(text="x" * 2_000)

        result = await mcp_server.call_tool("big_answer", {})
        first = result.content[0]
        assert isinstance(first, TextContent)
        assert "[Response truncated" in first.text

    def test_utf8_truncation_preserves_characters(self):
        """Test that UTF-8 truncation doesn't break multi-byte characters."""
        middleware = ResponseLimitingMiddleware(max_size=100)
        # Text with multi-byte characters (emoji)
        text = "Hello 🌍 World 🎉 Test " * 100
        result = middleware._truncate_to_result(text)
        # Should not raise and should be valid UTF-8
        content = result.content[0]
        assert isinstance(content, TextContent)
        content.text.encode("utf-8")