test_sse.py

"""Tests for the SSE client and server transports, driven entirely in process."""

import json
from collections.abc import AsyncGenerator
from typing import Any
from unittest.mock import AsyncMock, MagicMock, Mock, patch
from urllib.parse import urlparse

import anyio
import httpx
import pytest
from httpx_sse import ServerSentEvent
from inline_snapshot import snapshot
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import Response, StreamingResponse
from starlette.routing import Mount, Route

import mcp.client.sse
from mcp import types
from mcp.client.session import ClientSession
from mcp.client.sse import _extract_session_id_from_endpoint, sse_client
from mcp.server import Server, ServerRequestContext
from mcp.server.sse import SseServerTransport
from mcp.server.transport_security import TransportSecuritySettings
from mcp.shared._httpx_utils import McpHttpClientFactory
from mcp.shared.exceptions import MCPError
from mcp.shared.message import SessionMessage
from mcp.types import (
    CallToolRequestParams,
    CallToolResult,
    EmptyResult,
    Implementation,
    InitializeResult,
    JSONRPCResponse,
    ListToolsResult,
    PaginatedRequestParams,
    ReadResourceRequestParams,
    ReadResourceResult,
    ServerCapabilities,
    TextContent,
    TextResourceContents,
    Tool,
)
from tests.interaction.transports import StreamingASGITransport

SERVER_NAME = "test_server_for_SSE"

# The in-process app is mounted at this origin purely so URLs are well-formed; nothing listens here.
BASE_URL = "http://127.0.0.1:8000"


def in_process_client_factory(app: Starlette) -> McpHttpClientFactory:
    """An httpx_client_factory for sse_client whose clients are served in process by `app`."""

    def factory(
        headers: dict[str, str] | None = None,
        timeout: httpx.Timeout | None = None,
        auth: httpx.Auth | None = None,
    ) -> httpx.AsyncClient:
        # The SSE GET runs until it observes a disconnect, so the bridge must let the
        # application drain on close rather than cancelling it. follow_redirects matches
        # create_mcp_http_client, the factory this one stands in for.
        return httpx.AsyncClient(
            transport=StreamingASGITransport(app, cancel_on_close=False),
            base_url=BASE_URL,
            headers=headers,
            timeout=timeout,
            auth=auth,
            follow_redirects=True,
        )

    return factory


async def _handle_read_resource(ctx: ServerRequestContext, params: ReadResourceRequestParams) -> ReadResourceResult:
    uri = str(params.uri)
    parsed = urlparse(uri)
    if parsed.scheme == "foobar":
        return ReadResourceResult(
            contents=[TextResourceContents(uri=uri, text=f"Read {parsed.netloc}", mime_type="text/plain")]
        )
    raise MCPError(code=404, message="OOPS! no resource with that URI was found")


def make_app(server: Server) -> Starlette:
    """Mount `server` on a Starlette app exposing the SSE transport at /sse and /messages/."""
    # DNS-rebinding protection validates Host/Origin headers against a network attack that cannot
    # exist for an in-process app; the transport security behaviour itself is pinned by
    # tests/server/test_sse_security.py.
    sse = SseServerTransport(
        "/messages/", security_settings=TransportSecuritySettings(enable_dns_rebinding_protection=False)
    )

    async def handle_sse(request: Request) -> Response:
        async with sse.connect_sse(request.scope, request.receive, request._send) as (read_stream, write_stream):
            await server.run(read_stream, write_stream, server.create_initialization_options())
        return Response()

    return Starlette(
        routes=[
            Route("/sse", endpoint=handle_sse),
            Mount("/messages/", app=sse.handle_post_message),
        ]
    )


def make_server_app() -> Starlette:
    return make_app(Server(SERVER_NAME, on_read_resource=_handle_read_resource))


def make_failing_post_app() -> Starlette:
    """An SSE app that completes the handshake but fails every message POST with 503.

    The `/sse` stream announces a valid endpoint (so the client reaches the POST path) and
    stays open until the client disconnects; `/messages/` always returns 503. Used to drive
    the client's POST-error propagation path (#2110).
    """

    async def handle_sse(request: Request) -> Response:
        async def event_stream() -> AsyncGenerator[bytes, None]:
            yield b"event: endpoint\r\ndata: /messages/\r\n\r\n"
            # Hold the stream open with a single, branch-free wait (so coverage is
            # deterministic) that comfortably outlasts the in-process POST round trip.
            # The client tears the connection down as soon as it receives the error.
            await anyio.sleep(1.0)

        return StreamingResponse(event_stream(), media_type="text/event-stream")

    async def handle_message(request: Request) -> Response:
        return Response("upstream exploded", status_code=503)

    return Starlette(
        routes=[
            Route("/sse", endpoint=handle_sse),
            Route("/messages/", endpoint=handle_message, methods=["POST"]),
        ]
    )


@pytest.mark.anyio
async def test_sse_client_post_error_propagates_to_caller() -> None:
    """A non-2xx on the message POST surfaces to the caller via the read stream.

    Regression test for #2110: the error was previously swallowed by the post_writer task
    group and `read_stream.receive()` blocked forever.
    """
    factory = in_process_client_factory(make_failing_post_app())
    async with sse_client(f"{BASE_URL}/sse", httpx_client_factory=factory) as (read_stream, write_stream):
        await write_stream.send(SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")))
        with anyio.fail_after(10):
            item = await read_stream.receive()

    assert isinstance(item, httpx.HTTPStatusError)
    assert item.response.status_code == 503


@pytest.mark.anyio
async def test_raw_sse_connection() -> None:
    """The SSE GET responds 200 with an event-stream content type, announcing the session
    endpoint as its first event."""
    http_client = httpx.AsyncClient(
        transport=StreamingASGITransport(make_server_app(), cancel_on_close=False), base_url=BASE_URL
    )

    with anyio.fail_after(5):
        async with http_client, http_client.stream("GET", "/sse") as response:
            assert response.status_code == 200
            assert response.headers["content-type"] == "text/event-stream; charset=utf-8"

            lines = response.aiter_lines()
            assert await anext(lines) == "event: endpoint"
            assert (await anext(lines)).startswith("data: /messages/?session_id=")


@pytest.mark.anyio
async def test_sse_client_basic_connection() -> None:
    """A client initializes against, and pings, a server over the SSE transport."""
    factory = in_process_client_factory(make_server_app())
    async with sse_client(f"{BASE_URL}/sse", httpx_client_factory=factory) as streams:
        async with ClientSession(*streams) as session:
            result = await session.initialize()
            assert isinstance(result, InitializeResult)
            assert result.server_info.name == SERVER_NAME

            ping_result = await session.send_ping()
            assert isinstance(ping_result, EmptyResult)


@pytest.mark.anyio
async def test_sse_client_on_session_created() -> None:
    """The session-created callback receives the new session ID before sse_client yields."""
    factory = in_process_client_factory(make_server_app())
    captured: list[str] = []

    async with sse_client(
        f"{BASE_URL}/sse", httpx_client_factory=factory, on_session_created=captured.append
    ) as streams:
        async with ClientSession(*streams) as session:
            result = await session.initialize()
            assert isinstance(result, InitializeResult)
            # Callback fires when the endpoint event arrives, before sse_client yields.
            assert len(captured) == 1
            assert len(captured[0]) > 0


@pytest.mark.parametrize(
    "endpoint_url,expected",
    [
        ("/messages?sessionId=abc123", "abc123"),
        ("/messages?session_id=def456", "def456"),
        ("/messages?sessionId=abc&session_id=def", "abc"),
        ("/messages?other=value", None),
        ("/messages", None),
        ("", None),
    ],
)
def test_extract_session_id_from_endpoint(endpoint_url: str, expected: str | None) -> None:
    """The session ID is read from the endpoint URL's sessionId/session_id query parameters."""
    assert _extract_session_id_from_endpoint(endpoint_url) == expected


@pytest.mark.anyio
async def test_sse_client_on_session_created_not_called_when_no_session_id(monkeypatch: pytest.MonkeyPatch) -> None:
    """No session-created callback fires when the endpoint URL carries no session ID."""
    factory = in_process_client_factory(make_server_app())
    callback_mock = Mock()

    def mock_extract(url: str) -> None:
        return None

    monkeypatch.setattr(mcp.client.sse, "_extract_session_id_from_endpoint", mock_extract)

    async with sse_client(f"{BASE_URL}/sse", httpx_client_factory=factory, on_session_created=callback_mock) as streams:
        async with ClientSession(*streams) as session:
            result = await session.initialize()
            assert isinstance(result, InitializeResult)
            # Callback would have fired by now (endpoint event arrives before
            # sse_client yields); if it hasn't, it won't.
            callback_mock.assert_not_called()


@pytest.fixture
async def initialized_sse_client_session() -> AsyncGenerator[ClientSession, None]:
    factory = in_process_client_factory(make_server_app())
    async with sse_client(f"{BASE_URL}/sse", httpx_client_factory=factory) as streams:
        async with ClientSession(*streams) as session:
            await session.initialize()
            yield session


@pytest.mark.anyio
async def test_sse_client_happy_request_and_response(
    initialized_sse_client_session: ClientSession,
) -> None:
    """A resource read round-trips its arguments and the handler's content over SSE."""
    session = initialized_sse_client_session
    response = await session.read_resource(uri="foobar://should-work")
    assert len(response.contents) == 1
    assert isinstance(response.contents[0], TextResourceContents)
    assert response.contents[0].text == "Read should-work"


@pytest.mark.anyio
async def test_sse_client_exception_handling(
    initialized_sse_client_session: ClientSession,
) -> None:
    """A server-side MCPError reaches the client with its message intact."""
    session = initialized_sse_client_session
    with pytest.raises(MCPError, match="OOPS! no resource with that URI was found"):
        await session.read_resource(uri="xxx://will-not-work")


@pytest.mark.anyio
async def test_sse_client_basic_connection_mounted_app() -> None:
    """The SSE transport works unchanged when its app is mounted under a sub-path."""
    main_app = Starlette(routes=[Mount("/mounted_app", app=make_server_app())])
    factory = in_process_client_factory(main_app)

    async with sse_client(f"{BASE_URL}/mounted_app/sse", httpx_client_factory=factory) as streams:
        async with ClientSession(*streams) as session:
            result = await session.initialize()
            assert isinstance(result, InitializeResult)
            assert result.server_info.name == SERVER_NAME

            ping_result = await session.send_ping()
            assert isinstance(ping_result, EmptyResult)


async def _handle_context_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
    assert params.name in ("echo_headers", "echo_context")
    assert ctx.request is not None
    headers_info = dict(ctx.request.headers)

    if params.name == "echo_headers":
        return CallToolResult(content=[TextContent(type="text", text=json.dumps(headers_info))])

    assert params.arguments is not None
    context_data = {
        "request_id": params.arguments.get("request_id"),
        "headers": headers_info,
    }
    return CallToolResult(content=[TextContent(type="text", text=json.dumps(context_data))])


async def _handle_context_list_tools(
    ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListToolsResult:
    return ListToolsResult(
        tools=[
            Tool(
                name="echo_headers",
                description="Echoes request headers",
                input_schema={"type": "object", "properties": {}},
            ),
            Tool(
                name="echo_context",
                description="Echoes request context",
                input_schema={
                    "type": "object",
                    "properties": {"request_id": {"type": "string"}},
                    "required": ["request_id"],
                },
            ),
        ]
    )


def make_context_server_app() -> Starlette:
    return make_app(
        Server(
            "request_context_server",
            on_call_tool=_handle_context_call_tool,
            on_list_tools=_handle_context_list_tools,
        )
    )


@pytest.mark.anyio
async def test_request_context_propagation() -> None:
    """Custom HTTP headers on the SSE connection are visible to server handlers via ctx.request."""
    factory = in_process_client_factory(make_context_server_app())

    custom_headers = {
        "Authorization": "Bearer test-token",
        "X-Custom-Header": "test-value",
        "X-Trace-Id": "trace-123",
    }

    async with sse_client(f"{BASE_URL}/sse", httpx_client_factory=factory, headers=custom_headers) as streams:
        async with ClientSession(*streams) as session:
            result = await session.initialize()
            assert isinstance(result, InitializeResult)

            tool_result = await session.call_tool("echo_headers", {})

            assert len(tool_result.content) == 1
            content = tool_result.content[0]
            assert isinstance(content, TextContent)
            headers_data = json.loads(content.text)

            assert headers_data.get("authorization") == "Bearer test-token"
            assert headers_data.get("x-custom-header") == "test-value"
            assert headers_data.get("x-trace-id") == "trace-123"


@pytest.mark.anyio
async def test_request_context_isolation() -> None:
    """Each SSE connection's handlers see only that connection's request headers."""
    factory = in_process_client_factory(make_context_server_app())
    contexts: list[dict[str, Any]] = []

    # Connect three clients in turn, each with its own headers.
    for i in range(3):
        headers = {"X-Request-Id": f"request-{i}", "X-Custom-Value": f"value-{i}"}

        async with sse_client(f"{BASE_URL}/sse", httpx_client_factory=factory, headers=headers) as streams:
            async with ClientSession(*streams) as session:
                await session.initialize()

                tool_result = await session.call_tool("echo_context", {"request_id": f"request-{i}"})

                assert len(tool_result.content) == 1
                content = tool_result.content[0]
                assert isinstance(content, TextContent)
                contexts.append(json.loads(content.text))

    assert len(contexts) == 3
    for i, ctx in enumerate(contexts):
        assert ctx["request_id"] == f"request-{i}"
        assert ctx["headers"].get("x-request-id") == f"request-{i}"
        assert ctx["headers"].get("x-custom-value") == f"value-{i}"


def test_sse_message_id_coercion() -> None:
    """Previously, the `RequestId` would coerce a string that looked like an integer into an integer.

    See <https://github.com/modelcontextprotocol/python-sdk/pull/851> for more details.

    As per the JSON-RPC 2.0 specification, the id in the response object needs to be the same type as the id in the
    request object. In other words, we can't perform the coercion.

    See <https://www.jsonrpc.org/specification#response_object> for more details.
    """
    json_message = '{"jsonrpc": "2.0", "id": "123", "method": "ping", "params": null}'
    msg = types.JSONRPCRequest.model_validate_json(json_message)
    assert msg == snapshot(types.JSONRPCRequest(method="ping", jsonrpc="2.0", id="123"))

    json_message = '{"jsonrpc": "2.0", "id": 123, "method": "ping", "params": null}'
    msg = types.JSONRPCRequest.model_validate_json(json_message)
    assert msg == snapshot(types.JSONRPCRequest(method="ping", jsonrpc="2.0", id=123))


@pytest.mark.parametrize(
    "endpoint, expected_result",
    [
        # Valid endpoints - should normalize and work
        ("/messages/", "/messages/"),
        ("messages/", "/messages/"),
        ("/", "/"),
        # Invalid endpoints - should raise ValueError
        ("http://example.com/messages/", ValueError),
        ("//example.com/messages/", ValueError),
        ("ftp://example.com/messages/", ValueError),
        ("/messages/?param=value", ValueError),
        ("/messages/#fragment", ValueError),
    ],
)
def test_sse_server_transport_endpoint_validation(endpoint: str, expected_result: str | type[Exception]) -> None:
    """Test that SseServerTransport properly validates and normalizes endpoints."""
    if isinstance(expected_result, type):
        # Test invalid endpoints that should raise an exception
        with pytest.raises(expected_result, match="is not a relative path.*expecting a relative path"):
            SseServerTransport(endpoint)
    else:
        # Test valid endpoints that should normalize correctly
        sse = SseServerTransport(endpoint)
        assert sse._endpoint == expected_result
        assert sse._endpoint.startswith("/")


@pytest.mark.anyio
async def test_sse_client_handles_empty_keepalive_pings() -> None:
    """Test that SSE client properly handles empty data lines (keep-alive pings).

    Per the MCP spec (Streamable HTTP transport): "The server SHOULD immediately
    send an SSE event consisting of an event ID and an empty data field in order
    to prime the client to reconnect."

    This test mocks the SSE event stream to include empty "message" events and
    verifies the client skips them without crashing.
    """
    # Build a proper JSON-RPC response using types (not hardcoded strings)
    init_result = InitializeResult(
        protocol_version="2024-11-05",
        capabilities=ServerCapabilities(),
        server_info=Implementation(name="test", version="1.0"),
    )
    response = JSONRPCResponse(
        jsonrpc="2.0",
        id=1,
        result=init_result.model_dump(by_alias=True, exclude_none=True),
    )
    response_json = response.model_dump_json(by_alias=True, exclude_none=True)

    # Create mock SSE events using httpx_sse's ServerSentEvent
    async def mock_aiter_sse() -> AsyncGenerator[ServerSentEvent, None]:
        # First: endpoint event
        yield ServerSentEvent(event="endpoint", data="/messages/?session_id=abc123")
        # Empty data keep-alive ping - this is what we're testing
        yield ServerSentEvent(event="message", data="")
        # Real JSON-RPC response
        yield ServerSentEvent(event="message", data=response_json)

    mock_event_source = MagicMock()
    mock_event_source.aiter_sse.return_value = mock_aiter_sse()
    mock_event_source.response = MagicMock()
    mock_event_source.response.raise_for_status = MagicMock()

    mock_aconnect_sse = MagicMock()
    mock_aconnect_sse.__aenter__ = AsyncMock(return_value=mock_event_source)
    mock_aconnect_sse.__aexit__ = AsyncMock(return_value=None)

    mock_client = MagicMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)
    mock_client.post = AsyncMock(return_value=MagicMock(status_code=200, raise_for_status=MagicMock()))

    with (
        patch("mcp.client.sse.create_mcp_http_client", return_value=mock_client),
        patch("mcp.client.sse.aconnect_sse", return_value=mock_aconnect_sse),
    ):
        async with sse_client("http://test/sse") as (read_stream, _):
            # Read the message - should skip the empty one and get the real response
            msg = await read_stream.receive()
            # If we get here without error, the empty message was skipped successfully
            assert not isinstance(msg, Exception)
            assert isinstance(msg.message, types.JSONRPCResponse)
            assert msg.message.id == 1


@pytest.mark.anyio
async def test_sse_session_cleanup_on_disconnect() -> None:
    """Regression test for https://github.com/modelcontextprotocol/python-sdk/issues/1227

    When a client disconnects, the server should remove the session from
    _read_stream_writers. Without this cleanup, stale sessions accumulate and
    POST requests to disconnected sessions return 202 Accepted followed by a
    ClosedResourceError when the server tries to write to the dead stream.
    """
    factory = in_process_client_factory(make_server_app())
    captured: list[str] = []

    # Connect a client session, then disconnect
    async with sse_client(
        f"{BASE_URL}/sse", httpx_client_factory=factory, on_session_created=captured.append
    ) as streams:
        async with ClientSession(*streams) as session:
            await session.initialize()

    # After disconnect, POST to the stale session should return 404
    # (not 202 as it did before the fix)
    async with factory() as client:
        response = await client.post(
            f"/messages/?session_id={captured[0]}",
            json={"jsonrpc": "2.0", "method": "ping", "id": 99},
            headers={"Content-Type": "application/json"},
        )
        assert response.status_code == 404