prefecthq--fastmcp/tests/client/test_streamable_http.py

import asyncio
import json
import sys
from contextlib import suppress
from unittest.mock import AsyncMock, call

import pytest
from mcp import MCPError
from mcp_types import TextResourceContents

from fastmcp import Context
from fastmcp.client import Client
from fastmcp.client.transports import StreamableHttpTransport
from fastmcp.server.dependencies import get_http_request
from fastmcp.server.server import FastMCP
from fastmcp.utilities.tests import run_server_async


def create_test_server() -> FastMCP:
    """Create a FastMCP server with tools, resources, and prompts."""
    server = FastMCP("TestServer")

    @server.tool
    def greet(name: str) -> str:
        """Greet someone by name."""
        return f"Hello, {name}!"

    @server.tool
    async def elicit(ctx: Context) -> str:
        """Elicit a response from the user."""
        result = await ctx.elicit("What is your name?", response_type=str)

        if result.action == "accept":
            return f"You said your name was: {result.data}!"
        else:
            return "No name provided"

    @server.tool
    def add(a: int, b: int) -> int:
        """Add two numbers together."""
        return a + b

    @server.tool
    async def sleep(seconds: float) -> str:
        """Sleep for a given number of seconds."""
        await asyncio.sleep(seconds)
        return f"Slept for {seconds} seconds"

    @server.tool
    async def greet_with_progress(name: str, ctx: Context) -> str:
        """Report progress for a greeting."""
        await ctx.report_progress(0.5, 1.0, "Greeting in progress")
        await ctx.report_progress(0.75, 1.0, "Almost there!")
        return f"Hello, {name}!"

    @server.resource(uri="data://users")
    async def get_users() -> str:
        import json

        return json.dumps(["Alice", "Bob", "Charlie"])

    @server.resource(uri="data://user/{user_id}")
    async def get_user(user_id: str) -> str:
        import json

        return json.dumps({"id": user_id, "name": f"User {user_id}", "active": True})

    @server.resource(uri="request://headers")
    async def get_headers() -> str:
        import json

        request = get_http_request()
        return json.dumps(dict(request.headers))

    @server.prompt
    def welcome(name: str) -> str:
        """Example greeting prompt."""
        return f"Welcome to FastMCP, {name}!"

    return server


@pytest.fixture
async def streamable_http_server(request):
    """Start a test server and return its URL."""
    import fastmcp

    stateless_http = getattr(request, "param", False)
    if stateless_http:
        fastmcp.settings.stateless_http = True

    server = create_test_server()
    async with run_server_async(server) as url:
        yield url

    if stateless_http:
        fastmcp.settings.stateless_http = False


@pytest.fixture
async def streamable_http_server_with_streamable_http_alias():
    """Test that the "streamable-http" transport alias works."""
    server = create_test_server()
    async with run_server_async(server, transport="streamable-http") as url:
        yield url


@pytest.fixture
async def nested_server():
    """Test nested server mounts with Starlette."""
    import uvicorn
    from starlette.applications import Starlette
    from starlette.routing import Mount

    from fastmcp.utilities.http import find_available_port

    mcp_server = create_test_server()
    mcp_app = mcp_server.http_app(path="/final/mcp")

    # Nest the app under multiple mounts to test URL resolution
    inner = Starlette(routes=[Mount("/nest-inner", app=mcp_app)])
    outer = Starlette(
        routes=[Mount("/nest-outer", app=inner)], lifespan=mcp_app.lifespan
    )

    # Run uvicorn with the nested ASGI app
    port = find_available_port()

    config = uvicorn.Config(
        app=outer,
        host="127.0.0.1",
        port=port,
        log_level="critical",
        ws="websockets-sansio",
        timeout_graceful_shutdown=0,
    )

    uvicorn_server = uvicorn.Server(config)
    server_task = asyncio.create_task(uvicorn_server.serve())
    await asyncio.sleep(0.1)

    yield f"http://127.0.0.1:{port}/nest-outer/nest-inner/final/mcp"

    # Graceful shutdown - required for uvicorn 0.39+ due to context isolation
    uvicorn_server.should_exit = True
    with suppress(asyncio.CancelledError, asyncio.TimeoutError):
        await asyncio.wait_for(server_task, timeout=2.0)


async def test_ping(streamable_http_server: str):
    """Test pinging the server."""
    async with Client(
        transport=StreamableHttpTransport(streamable_http_server)
    ) as client:
        result = await client.ping()
        assert result is True


async def test_ping_with_streamable_http_alias(
    streamable_http_server_with_streamable_http_alias: str,
):
    """Test pinging the server."""
    async with Client(
        transport=StreamableHttpTransport(
            streamable_http_server_with_streamable_http_alias
        )
    ) as client:
        result = await client.ping()
        assert result is True


async def test_http_headers(streamable_http_server: str):
    """Test getting HTTP headers from the server."""
    async with Client(
        transport=StreamableHttpTransport(
            streamable_http_server, headers={"X-DEMO-HEADER": "ABC"}
        )
    ) as client:
        raw_result = await client.read_resource("request://headers")
        assert isinstance(raw_result[0], TextResourceContents)
        json_result = json.loads(raw_result[0].text)
        assert "x-demo-header" in json_result
        assert json_result["x-demo-header"] == "ABC"


async def test_session_id_callback(streamable_http_server: str):
    """Test getting mcp-session-id from the transport."""
    transport = StreamableHttpTransport(streamable_http_server)
    assert transport.get_session_id() is None
    async with Client(transport=transport):
        session_id = transport.get_session_id()
        assert session_id is not None


@pytest.mark.parametrize("streamable_http_server", [True, False], indirect=True)
async def test_greet_with_progress_tool(streamable_http_server: str):
    """Test calling the greet tool."""
    progress_handler = AsyncMock(return_value=None)

    async with Client(
        transport=StreamableHttpTransport(streamable_http_server),
        progress_handler=progress_handler,
    ) as client:
        result = await client.call_tool("greet_with_progress", {"name": "Alice"})
        assert result.data == "Hello, Alice!"

        progress_handler.assert_has_calls(
            [
                call(0.5, 1.0, "Greeting in progress"),
                call(0.75, 1.0, "Almost there!"),
            ]
        )


@pytest.mark.parametrize("streamable_http_server", [True, False], indirect=True)
async def test_elicitation_tool(streamable_http_server: str, request):
    """Test calling the elicitation tool in both stateless and stateful modes."""

    async def elicitation_handler(message, response_type, params, ctx):
        return {"value": "Alice"}

    stateless_http = request.node.callspec.params.get("streamable_http_server", False)
    if stateless_http:
        pytest.xfail("Elicitation is not supported in stateless HTTP mode")

    async with Client(
        transport=StreamableHttpTransport(streamable_http_server),
        elicitation_handler=elicitation_handler,
    ) as client:
        result = await client.call_tool("elicit")
        assert result.data == "You said your name was: Alice!"


@pytest.mark.parametrize("streamable_http_server", [True], indirect=True)
async def test_stateless_http_rejects_get_sse(streamable_http_server: str):
    """Stateless servers should reject GET SSE requests with 405."""
    import httpx

    async with httpx.AsyncClient() as http_client:
        response = await http_client.get(streamable_http_server)
        assert response.status_code == 405


@pytest.mark.parametrize("streamable_http_server", [True], indirect=True)
async def test_stateless_http_still_accepts_post(streamable_http_server: str):
    """Stateless servers should still handle POST requests normally."""
    async with Client(
        transport=StreamableHttpTransport(streamable_http_server)
    ) as client:
        result = await client.call_tool("greet", {"name": "World"})
        assert result.data == "Hello, World!"


async def test_nested_streamable_http_server_resolves_correctly(nested_server: str):
    """Test patch for https://github.com/modelcontextprotocol/python-sdk/pull/659"""
    async with Client(transport=StreamableHttpTransport(nested_server)) as client:
        result = await client.ping()
        assert result is True


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="Timeout tests are flaky on Windows. Timeouts *are* supported but the tests are unreliable.",
)
class TestTimeout:
    async def test_timeout(self, streamable_http_server: str):
        # note this transport behaves differently than others and raises
        # MCPError from the *client* context
        with pytest.raises(MCPError, match="timed out"):
            async with Client(
                transport=StreamableHttpTransport(streamable_http_server),
                timeout=0.02,
            ) as client:
                await client.call_tool("sleep", {"seconds": 0.05})

    async def test_timeout_tool_call(self, streamable_http_server: str):
        async with Client(
            transport=StreamableHttpTransport(streamable_http_server),
        ) as client:
            with pytest.raises(MCPError):
                await client.call_tool("sleep", {"seconds": 0.2}, timeout=0.1)

    async def test_timeout_tool_call_overrides_client_timeout(
        self, streamable_http_server: str
    ):
        async with Client(
            transport=StreamableHttpTransport(streamable_http_server),
            timeout=2,
        ) as client:
            with pytest.raises(MCPError):
                await client.call_tool("sleep", {"seconds": 0.2}, timeout=0.1)