Add message queue for SSE messages POST endpoint

README.md

@@ -385,6 +385,30 @@ app.router.routes.append(Host('mcp.acme.corp', app=mcp.sse_app()))
 
 For more information on mounting applications in Starlette, see the [Starlette documentation](https://www.starlette.io/routing/#submounting-routes).
 
+#### Message Queue Options
+
+By default, the SSE server uses an in-memory message queue for incoming POST messages. For production deployments or distributed scenarios, you can use Redis:
+
+```python
+# Using the built-in Redis message queue
+from mcp.server.fastmcp import FastMCP
+from mcp.server.message_queue import RedisMessageQueue
+
+# Create a Redis message queue
+redis_queue = RedisMessageQueue(
+    redis_url="redis://localhost:6379/0", prefix="mcp:pubsub:"
+)
+
+# Pass the message queue instance to the server
+mcp = FastMCP("My App", message_queue=redis_queue)
+```
+
+To use Redis, add the Redis dependency:
+
+```bash
+uv add "mcp[redis]"
+```
+
 ## Examples
 
 ### Echo Server

examples/fastmcp/unicode_example.py

@@ -4,8 +4,14 @@
 """
 
 from mcp.server.fastmcp import FastMCP
+from mcp.server.message_queue import RedisMessageQueue
 
-mcp = FastMCP()
+# Create a Redis message queue
+redis_queue = RedisMessageQueue(
+    redis_url="redis://localhost:6379/0", prefix="mcp:pubsub:"
+)
+
+mcp = FastMCP(message_queue=redis_queue)
 
 
 @mcp.tool(
@@ -61,4 +67,4 @@ def multilingual_hello() -> str:
 
 
 if __name__ == "__main__":
-    mcp.run()
+    mcp.run(transport="sse")

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
 rich = ["rich>=13.9.4"]
 cli = ["typer>=0.12.4", "python-dotenv>=1.0.0"]
 ws = ["websockets>=15.0.1"]
+redis = ["redis>=5.2.1", "types-redis>=4.6.0.20241004"]
 
 [project.scripts]
 mcp = "mcp.cli:app [cli]"

src/mcp/server/fastmcp/server.py

@@ -33,6 +33,7 @@
 from mcp.server.lowlevel.server import LifespanResultT
 from mcp.server.lowlevel.server import Server as MCPServer
 from mcp.server.lowlevel.server import lifespan as default_lifespan
+from mcp.server.message_queue import MessageQueue
 from mcp.server.session import ServerSession, ServerSessionT
 from mcp.server.sse import SseServerTransport
 from mcp.server.stdio import stdio_server
@@ -76,6 +77,11 @@ class Settings(BaseSettings, Generic[LifespanResultT]):
     sse_path: str = "/sse"
     message_path: str = "/messages/"
 
+    # SSE message queue settings
+    message_queue: MessageQueue | None = Field(
+        None, description="Custom message queue instance"
+    )
+
     # resource settings
     warn_on_duplicate_resources: bool = True
 
@@ -479,7 +485,16 @@ async def run_sse_async(self) -> None:
 
     def sse_app(self) -> Starlette:
         """Return an instance of the SSE server app."""
-        sse = SseServerTransport(self.settings.message_path)
+        message_queue = self.settings.message_queue
+        if message_queue is None:
+            from mcp.server.message_queue import InMemoryMessageQueue
+
+            message_queue = InMemoryMessageQueue()
+            logger.info("Using default in-memory message queue")
+
+        sse = SseServerTransport(
+            self.settings.message_path, message_queue=message_queue
+        )
 
         async def handle_sse(request: Request) -> None:
             async with sse.connect_sse(

src/mcp/server/message_queue/__init__.py

@@ -0,0 +1,16 @@
+"""
+Message Queue Module for MCP Server
+
+This module implements queue interfaces for handling
+messages between clients and servers.
+"""
+
+from mcp.server.message_queue.base import InMemoryMessageQueue, MessageQueue
+
+# Try to import Redis implementation if available
+try:
+    from mcp.server.message_queue.redis import RedisMessageQueue
+except ImportError:
+    RedisMessageQueue = None
+
+__all__ = ["MessageQueue", "InMemoryMessageQueue", "RedisMessageQueue"]

src/mcp/server/message_queue/base.py

@@ -0,0 +1,103 @@
+import logging
+from collections.abc import Awaitable, Callable
+from contextlib import asynccontextmanager
+from typing import Protocol, runtime_checkable
+from uuid import UUID
+
+import mcp.types as types
+
+logger = logging.getLogger(__name__)
+
+MessageCallback = Callable[[types.JSONRPCMessage | Exception], Awaitable[None]]
+
+
+@runtime_checkable
+class MessageQueue(Protocol):
+    """Abstract interface for SSE messaging.
+
+    This interface allows messages to be published to sessions and callbacks to be
+    registered for message handling, enabling multiple servers to handle requests.
+    """
+
+    async def publish_message(
+        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+    ) -> bool:
+        """Publish a message for the specified session.
+
+        Args:
+            session_id: The UUID of the session this message is for
+            message: The message to publish
+
+        Returns:
+            bool: True if message was published, False if session not found
+        """
+        ...
+
+    @asynccontextmanager
+    async def active_for_request(self, session_id: UUID, callback: MessageCallback):
+        """Request-scoped context manager that ensures the listener is active.
+
+        Args:
+            session_id: The UUID of the session to activate
+            callback: Async callback function to handle messages for this session
+        """
+        yield
+
+    async def session_exists(self, session_id: UUID) -> bool:
+        """Check if a session exists.
+
+        Args:
+            session_id: The UUID of the session to check
+
+        Returns:
+            bool: True if the session is active, False otherwise
+        """
+        ...
+
+
+class InMemoryMessageQueue:
+    """Default in-memory implementation of the MessageQueue interface.
+
+    This implementation immediately calls registered callbacks when messages
+    are received.
+    """
+
+    def __init__(self) -> None:
+        self._callbacks: dict[UUID, MessageCallback] = {}
+        self._active_sessions: set[UUID] = set()
+
+    async def publish_message(
+        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+    ) -> bool:
+        """Publish a message for the specified session."""
+        if not await self.session_exists(session_id):
+            logger.warning(f"Message received for unknown session {session_id}")
+            return False
+
+        # Call the callback directly if registered
+        if session_id in self._callbacks:
+            await self._callbacks[session_id](message)
+            logger.debug(f"Called callback for session {session_id}")
+        else:
+            logger.warning(f"No callback registered for session {session_id}")
+
+        return True
+
+    @asynccontextmanager
+    async def active_for_request(self, session_id: UUID, callback: MessageCallback):
+        """Request-scoped context manager that ensures the listener is active."""
+        self._active_sessions.add(session_id)
+        self._callbacks[session_id] = callback
+        logger.debug(f"Registered session {session_id} with callback")
+
+        try:
+            yield
+        finally:
+            self._active_sessions.discard(session_id)
+            if session_id in self._callbacks:
+                del self._callbacks[session_id]
+            logger.debug(f"Unregistered session {session_id}")
+
+    async def session_exists(self, session_id: UUID) -> bool:
+        """Check if a session exists."""
+        return session_id in self._active_sessions

src/mcp/server/message_queue/redis.py

@@ -0,0 +1,143 @@
+import json
+import logging
+from contextlib import asynccontextmanager
+from typing import Any, cast
+from uuid import UUID
+
+import anyio
+from anyio import CapacityLimiter, lowlevel
+
+import mcp.types as types
+from mcp.server.message_queue.base import MessageCallback
+
+try:
+    import redis.asyncio as redis
+except ImportError:
+    raise ImportError(
+        "Redis support requires the 'redis' package. "
+        "Install it with: 'uv add redis' or 'uv add \"mcp[redis]\"'"
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class RedisMessageQueue:
+    """Redis implementation of the MessageQueue interface using pubsub.
+
+    This implementation uses Redis pubsub for real-time message distribution across
+    multiple servers handling the same sessions.
+    """
+
+    def __init__(
+        self, redis_url: str = "redis://localhost:6379/0", prefix: str = "mcp:pubsub:"
+    ) -> None:
+        """Initialize Redis message queue.
+
+        Args:
+            redis_url: Redis connection string
+            prefix: Key prefix for Redis channels to avoid collisions
+        """
+        self._redis = redis.from_url(redis_url, decode_responses=True)  # type: ignore
+        self._pubsub = self._redis.pubsub(ignore_subscribe_messages=True)  # type: ignore
+        self._prefix = prefix
+        self._active_sessions_key = f"{prefix}active_sessions"
+        self._callbacks: dict[UUID, MessageCallback] = {}
+        # Ensures only one polling task runs at a time for message handling
+        self._limiter = CapacityLimiter(1)
+        logger.debug(f"Initialized Redis message queue with URL: {redis_url}")
+
+    def _session_channel(self, session_id: UUID) -> str:
+        """Get the Redis channel for a session."""
+        return f"{self._prefix}session:{session_id.hex}"
+
+    @asynccontextmanager
+    async def active_for_request(self, session_id: UUID, callback: MessageCallback):
+        """Request-scoped context manager that ensures the listener task is running."""
+        await self._redis.sadd(self._active_sessions_key, session_id.hex)
+        self._callbacks[session_id] = callback
+        channel = self._session_channel(session_id)
+        await self._pubsub.subscribe(channel)  # type: ignore
+
+        logger.debug(f"Registered session {session_id} in Redis with callback")
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(self._listen_for_messages)
+            try:
+                yield
+            finally:
+                tg.cancel_scope.cancel()
+                await self._pubsub.unsubscribe(channel)  # type: ignore
+                await self._redis.srem(self._active_sessions_key, session_id.hex)
+                del self._callbacks[session_id]
+                logger.debug(f"Unregistered session {session_id} from Redis")
+
+    async def _listen_for_messages(self) -> None:
+        """Background task that listens for messages on subscribed channels."""
+        async with self._limiter:
+            while True:
+                await lowlevel.checkpoint()
+                message: None | dict[str, Any] = await self._pubsub.get_message(  # type: ignore
+                    ignore_subscribe_messages=True,
+                    timeout=None,  # type: ignore
+                )
+                if message is None:
+                    continue
+
+                # Extract session ID from channel name
+                channel: str = cast(str, message["channel"])
+                if not channel.startswith(self._prefix):
+                    continue
+
+                session_hex = channel.split(":")[-1]
+                try:
+                    session_id = UUID(hex=session_hex)
+                except ValueError:
+                    logger.error(f"Invalid session channel: {channel}")
+                    continue
+
+                data: str = cast(str, message["data"])
+                msg: None | types.JSONRPCMessage | Exception = None
+                try:
+                    json_data = json.loads(data)
+                    if isinstance(json_data, dict):
+                        json_dict: dict[str, Any] = json_data
+                        if json_dict.get("_exception", False):
+                            msg = Exception(
+                                f"{json_dict['type']}: {json_dict['message']}"
+                            )
+                        else:
+                            msg = types.JSONRPCMessage.model_validate_json(data)
+
+                    if msg and session_id in self._callbacks:
+                        await self._callbacks[session_id](msg)
+                except Exception as e:
+                    logger.error(f"Failed to process message: {e}")
+
+    async def publish_message(
+        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+    ) -> bool:
+        """Publish a message for the specified session."""
+        if not await self.session_exists(session_id):
+            logger.warning(f"Message received for unknown session {session_id}")
+            return False
+
+        if isinstance(message, Exception):
+            data = json.dumps(
+                {
+                    "_exception": True,
+                    "type": type(message).__name__,
+                    "message": str(message),
+                }
+            )
+        else:
+            data = message.model_dump_json()
+
+        channel = self._session_channel(session_id)
+        await self._redis.publish(channel, data)  # type: ignore[attr-defined]
+        logger.debug(f"Published message to Redis channel for session {session_id}")
+        return True
+
+    async def session_exists(self, session_id: UUID) -> bool:
+        """Check if a session exists."""
+        return bool(
+            await self._redis.sismember(self._active_sessions_key, session_id.hex)  # type: ignore[attr-defined]
+        )