Adding run_ws_async to FastMCP

[aksheyd]

Instead of having to create a class on top of FastMCP to add a run_ws_async() func to connect the websocket_server to uvicorn over a Starlette app (and add test cases), I was wondering if this could be added to the sdk innately?

The solution would entail adding "ws" as a transport option to FastMCP, adding a run_ws_async() func that launches a uvicorn server. Then, a ws_app() func to create a starlette app which binds the FastMCP read/write streams with the websocket_server read/write streams. (Very similar to how run_sse_async() is done).

An alternate would be to just let the user create their own methods? But, I feel like it would be nice to have direct support from the sdk.

I have working code and some tests for this, and I can start a PR (if this would be a viable feature to add). If you want to take a look a the code, I can add it in the comments.

Cheers!

[SinOverCos]

Related: #302

I'm not sure if Anthropic plans to add the support for websockets themselves but I am facing a similar issue where communication over websockets would be much more preferred.

[aksheyd]

Seems like they aren't going in that direction. If it's of any help, here's how I implemented WebSocket support for FastMCP:

from typing import Literal

import anyio
import uvicorn
import uvicorn.config
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP as FastMCPBase
from mcp.server.websocket import websocket_server
from starlette.applications import Starlette
from starlette.responses import PlainTextResponse
from starlette.routing import Route, WebSocketRoute

load_dotenv()


class FastMCP(FastMCPBase):
    """Extended FastMCP server with websocket integration."""

    def run(self, transport: Literal["stdio", "sse", "ws"] = "stdio") -> None:
        """Run the FastMCP server. Note this is a synchronous function.

        Args:
            transport: Transport protocol to use ("stdio", "sse", or "ws")
        """
        TRANSPORTS = Literal["stdio", "sse", "ws"]
        if transport not in TRANSPORTS.__args__:  # type: ignore
            raise ValueError(f"Unknown transport: {transport}")

        if transport == "stdio":
            anyio.run(self.run_stdio_async)
        elif transport == "ws":
            anyio.run(self.run_ws_async)
        else:
            anyio.run(self.run_sse_async)

    async def run_ws_async(self) -> None:
        """Run the server using WebSocket transport."""
        starlette_app = self.ws_app()

        config = uvicorn.Config(
            app=starlette_app,
            host=self.settings.host,
            port=self.settings.port,
            log_level="debug",
        )

        server = uvicorn.Server(config)

        await server.serve()

    def ws_app(self):
        """Return an instance of the Websocket server app."""

        async def handle_ws(websocket):
            async with websocket_server(
                websocket.scope, websocket.receive, websocket.send
            ) as (ws_read_stream, ws_write_stream):
                await self._mcp_server.run(
                    ws_read_stream,
                    ws_write_stream,
                    self._mcp_server.create_initialization_options(),
                    raise_exceptions=self.settings.debug,
                )

        async def homepage(request):
            return PlainTextResponse("Hello MCP WebSocket!")

        app = Starlette(
            debug=self.settings.debug,
            routes=[
                Route("/", endpoint=homepage),
                WebSocketRoute("/ws", endpoint=handle_ws),
            ],
        )

        return app


if __name__ == "__main__":
    client = FastMCP("websocket")
    client.run(transport="ws")

[Kludex]

I personally think WebSocket should be the only transport, but since it's not supported in the specification yet, we will not support on the Python SDK yet.