Streamable HTTP - improve usability, fast mcp and auth (#641)

tests/server/fastmcp/test_integration.py

@@ -15,6 +15,7 @@ import uvicorn
 
 from mcp.client.session import ClientSession
 from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamablehttp_client
 from mcp.server.fastmcp import FastMCP
 from mcp.types import InitializeResult, TextContent
 
@@ -33,6 +34,34 @@ def server_url(server_port: int) -> str:
     return f"http://127.0.0.1:{server_port}"
 
 
+@pytest.fixture
+def http_server_port() -> int:
+    """Get a free port for testing the StreamableHTTP server."""
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+@pytest.fixture
+def http_server_url(http_server_port: int) -> str:
+    """Get the StreamableHTTP server URL for testing."""
+    return f"http://127.0.0.1:{http_server_port}"
+
+
+@pytest.fixture
+def stateless_http_server_port() -> int:
+    """Get a free port for testing the stateless StreamableHTTP server."""
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+@pytest.fixture
+def stateless_http_server_url(stateless_http_server_port: int) -> str:
+    """Get the stateless StreamableHTTP server URL for testing."""
+    return f"http://127.0.0.1:{stateless_http_server_port}"
+
+
 # Create a function to make the FastMCP server app
 def make_fastmcp_app():
     """Create a FastMCP server without auth settings."""
@@ -51,6 +80,40 @@ def make_fastmcp_app():
     return mcp, app
 
 
+def make_fastmcp_streamable_http_app():
+    """Create a FastMCP server with StreamableHTTP transport."""
+    from starlette.applications import Starlette
+
+    mcp = FastMCP(name="NoAuthServer")
+
+    # Add a simple tool
+    @mcp.tool(description="A simple echo tool")
+    def echo(message: str) -> str:
+        return f"Echo: {message}"
+
+    # Create the StreamableHTTP app
+    app: Starlette = mcp.streamable_http_app()
+
+    return mcp, app
+
+
+def make_fastmcp_stateless_http_app():
+    """Create a FastMCP server with stateless StreamableHTTP transport."""
+    from starlette.applications import Starlette
+
+    mcp = FastMCP(name="StatelessServer", stateless_http=True)
+
+    # Add a simple tool
+    @mcp.tool(description="A simple echo tool")
+    def echo(message: str) -> str:
+        return f"Echo: {message}"
+
+    # Create the StreamableHTTP app
+    app: Starlette = mcp.streamable_http_app()
+
+    return mcp, app
+
+
 def run_server(server_port: int) -> None:
     """Run the server."""
     _, app = make_fastmcp_app()
@@ -63,6 +126,30 @@ def run_server(server_port: int) -> None:
     server.run()
 
 
+def run_streamable_http_server(server_port: int) -> None:
+    """Run the StreamableHTTP server."""
+    _, app = make_fastmcp_streamable_http_app()
+    server = uvicorn.Server(
+        config=uvicorn.Config(
+            app=app, host="127.0.0.1", port=server_port, log_level="error"
+        )
+    )
+    print(f"Starting StreamableHTTP server on port {server_port}")
+    server.run()
+
+
+def run_stateless_http_server(server_port: int) -> None:
+    """Run the stateless StreamableHTTP server."""
+    _, app = make_fastmcp_stateless_http_app()
+    server = uvicorn.Server(
+        config=uvicorn.Config(
+            app=app, host="127.0.0.1", port=server_port, log_level="error"
+        )
+    )
+    print(f"Starting stateless StreamableHTTP server on port {server_port}")
+    server.run()
+
+
 @pytest.fixture()
 def server(server_port: int) -> Generator[None, None, None]:
     """Start the server in a separate process and clean up after the test."""
@@ -94,6 +181,80 @@ def server(server_port: int) -> Generator[None, None, None]:
         print("Server process failed to terminate")
 
 
+@pytest.fixture()
+def streamable_http_server(http_server_port: int) -> Generator[None, None, None]:
+    """Start the StreamableHTTP server in a separate process."""
+    proc = multiprocessing.Process(
+        target=run_streamable_http_server, args=(http_server_port,), daemon=True
+    )
+    print("Starting StreamableHTTP server process")
+    proc.start()
+
+    # Wait for server to be running
+    max_attempts = 20
+    attempt = 0
+    print("Waiting for StreamableHTTP server to start")
+    while attempt < max_attempts:
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.connect(("127.0.0.1", http_server_port))
+                break
+        except ConnectionRefusedError:
+            time.sleep(0.1)
+            attempt += 1
+    else:
+        raise RuntimeError(
+            f"StreamableHTTP server failed to start after {max_attempts} attempts"
+        )
+
+    yield
+
+    print("Killing StreamableHTTP server")
+    proc.kill()
+    proc.join(timeout=2)
+    if proc.is_alive():
+        print("StreamableHTTP server process failed to terminate")
+
+
+@pytest.fixture()
+def stateless_http_server(
+    stateless_http_server_port: int,
+) -> Generator[None, None, None]:
+    """Start the stateless StreamableHTTP server in a separate process."""
+    proc = multiprocessing.Process(
+        target=run_stateless_http_server,
+        args=(stateless_http_server_port,),
+        daemon=True,
+    )
+    print("Starting stateless StreamableHTTP server process")
+    proc.start()
+
+    # Wait for server to be running
+    max_attempts = 20
+    attempt = 0
+    print("Waiting for stateless StreamableHTTP server to start")
+    while attempt < max_attempts:
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.connect(("127.0.0.1", stateless_http_server_port))
+                break
+        except ConnectionRefusedError:
+            time.sleep(0.1)
+            attempt += 1
+    else:
+        raise RuntimeError(
+            f"Stateless server failed to start after {max_attempts} attempts"
+        )
+
+    yield
+
+    print("Killing stateless StreamableHTTP server")
+    proc.kill()
+    proc.join(timeout=2)
+    if proc.is_alive():
+        print("Stateless StreamableHTTP server process failed to terminate")
+
+
 @pytest.mark.anyio
 async def test_fastmcp_without_auth(server: None, server_url: str) -> None:
     """Test that FastMCP works when auth settings are not provided."""
@@ -110,3 +271,55 @@ async def test_fastmcp_without_auth(server: None, server_url: str) -> None:
             assert len(tool_result.content) == 1
             assert isinstance(tool_result.content[0], TextContent)
             assert tool_result.content[0].text == "Echo: hello"
+
+
+@pytest.mark.anyio
+async def test_fastmcp_streamable_http(
+    streamable_http_server: None, http_server_url: str
+) -> None:
+    """Test that FastMCP works with StreamableHTTP transport."""
+    # Connect to the server using StreamableHTTP
+    async with streamablehttp_client(http_server_url + "/mcp") as (
+        read_stream,
+        write_stream,
+        _,
+    ):
+        # Create a session using the client streams
+        async with ClientSession(read_stream, write_stream) as session:
+            # Test initialization
+            result = await session.initialize()
+            assert isinstance(result, InitializeResult)
+            assert result.serverInfo.name == "NoAuthServer"
+
+            # Test that we can call tools without authentication
+            tool_result = await session.call_tool("echo", {"message": "hello"})
+            assert len(tool_result.content) == 1
+            assert isinstance(tool_result.content[0], TextContent)
+            assert tool_result.content[0].text == "Echo: hello"
+
+
+@pytest.mark.anyio
+async def test_fastmcp_stateless_streamable_http(
+    stateless_http_server: None, stateless_http_server_url: str
+) -> None:
+    """Test that FastMCP works with stateless StreamableHTTP transport."""
+    # Connect to the server using StreamableHTTP
+    async with streamablehttp_client(stateless_http_server_url + "/mcp") as (
+        read_stream,
+        write_stream,
+        _,
+    ):
+        async with ClientSession(read_stream, write_stream) as session:
+            result = await session.initialize()
+            assert isinstance(result, InitializeResult)
+            assert result.serverInfo.name == "StatelessServer"
+            tool_result = await session.call_tool("echo", {"message": "hello"})
+            assert len(tool_result.content) == 1
+            assert isinstance(tool_result.content[0], TextContent)
+            assert tool_result.content[0].text == "Echo: hello"
+
+            for i in range(3):
+                tool_result = await session.call_tool("echo", {"message": f"test_{i}"})
+                assert len(tool_result.content) == 1
+                assert isinstance(tool_result.content[0], TextContent)
+                assert tool_result.content[0].text == f"Echo: test_{i}"

tests/server/test_streamable_http_manager.py

@@ -0,0 +1,81 @@
+"""Tests for StreamableHTTPSessionManager."""
+
+import anyio
+import pytest
+
+from mcp.server.lowlevel import Server
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+
+
+@pytest.mark.anyio
+async def test_run_can_only_be_called_once():
+    """Test that run() can only be called once per instance."""
+    app = Server("test-server")
+    manager = StreamableHTTPSessionManager(app=app)
+
+    # First call should succeed
+    async with manager.run():
+        pass
+
+    # Second call should raise RuntimeError
+    with pytest.raises(RuntimeError) as excinfo:
+        async with manager.run():
+            pass
+
+    assert (
+        "StreamableHTTPSessionManager .run() can only be called once per instance"
+        in str(excinfo.value)
+    )
+
+
+@pytest.mark.anyio
+async def test_run_prevents_concurrent_calls():
+    """Test that concurrent calls to run() are prevented."""
+    app = Server("test-server")
+    manager = StreamableHTTPSessionManager(app=app)
+
+    errors = []
+
+    async def try_run():
+        try:
+            async with manager.run():
+                # Simulate some work
+                await anyio.sleep(0.1)
+        except RuntimeError as e:
+            errors.append(e)
+
+    # Try to run concurrently
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(try_run)
+        tg.start_soon(try_run)
+
+    # One should succeed, one should fail
+    assert len(errors) == 1
+    assert (
+        "StreamableHTTPSessionManager .run() can only be called once per instance"
+        in str(errors[0])
+    )
+
+
+@pytest.mark.anyio
+async def test_handle_request_without_run_raises_error():
+    """Test that handle_request raises error if run() hasn't been called."""
+    app = Server("test-server")
+    manager = StreamableHTTPSessionManager(app=app)
+
+    # Mock ASGI parameters
+    scope = {"type": "http", "method": "POST", "path": "/test"}
+
+    async def receive():
+        return {"type": "http.request", "body": b""}
+
+    async def send(message):
+        pass
+
+    # Should raise error because run() hasn't been called
+    with pytest.raises(RuntimeError) as excinfo:
+        await manager.handle_request(scope, receive, send)
+
+    assert "Task group is not initialized. Make sure to use run()." in str(
+        excinfo.value
+    )