Skip to content

Commit 10c5af8

Browse files
akash329dakash329d
committed
initial
1 parent c2ca8e0 commit 10c5af8

File tree

6 files changed

+321
-15
lines changed

6 files changed

+321
-15
lines changed

‎pyproject.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ dependencies = [
3636
rich = ["rich>=13.9.4"]
3737
cli = ["typer>=0.12.4", "python-dotenv>=1.0.0"]
3838
ws = ["websockets>=15.0.1"]
39+
redis = ["redis>=4.5.0"]
3940

4041
[project.scripts]
4142
mcp = "mcp.cli:app [cli]"
@@ -57,7 +58,7 @@ dev = [
5758
docs = [
5859
"mkdocs>=1.6.1",
5960
"mkdocs-glightbox>=0.4.0",
60-
"mkdocs-material[imaging]>=9.5.45",
61+
"mkdocs-material[imaging]>=9.5.18",
6162
"mkdocstrings-python>=1.12.2",
6263
]
6364

‎src/mcp/server/fastmcp/server.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,11 @@ class Settings(BaseSettings, Generic[LifespanResultT]):
7575
port: int = 8000
7676
sse_path: str = "/sse"
7777
message_path: str = "/messages/"
78+
79+
# SSE message queue settings
80+
message_queue: Literal["memory", "redis"] = "memory"
81+
redis_url: str = "redis://localhost:6379/0"
82+
redis_prefix: str = "mcp:queue:"
7883

7984
# resource settings
8085
warn_on_duplicate_resources: bool = True
@@ -479,7 +484,24 @@ async def run_sse_async(self) -> None:
479484

480485
def sse_app(self) -> Starlette:
481486
"""Return an instance of the SSE server app."""
482-
sse = SseServerTransport(self.settings.message_path)
487+
message_queue = None
488+
if self.settings.message_queue == "redis":
489+
try:
490+
from mcp.server.message_queue import RedisMessageQueue
491+
message_queue = RedisMessageQueue(
492+
redis_url=self.settings.redis_url,
493+
prefix=self.settings.redis_prefix
494+
)
495+
logger.info(f"Using Redis message queue at {self.settings.redis_url}")
496+
except ImportError:
497+
logger.error("Redis message queue requested but 'redis' package not installed. ")
498+
raise
499+
else:
500+
from mcp.server.message_queue import InMemoryMessageQueue
501+
message_queue = InMemoryMessageQueue()
502+
logger.info("Using in-memory message queue")
503+
504+
sse = SseServerTransport(self.settings.message_path, message_queue=message_queue)
483505

484506
async def handle_sse(request: Request) -> None:
485507
async with sse.connect_sse(

‎src/mcp/server/message_queue/__init__.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
"""
2+
Message Queue Module for MCP Server
3+
4+
This module implements queue interfaces for handling messages between clients and servers.
5+
"""
6+
7+
from mcp.server.message_queue.base import InMemoryMessageQueue, MessageQueue
8+
9+
# Try to import Redis implementation if available
10+
try:
11+
from mcp.server.message_queue.redis import RedisMessageQueue
12+
except ImportError:
13+
RedisMessageQueue = None
14+
15+
__all__ = ["MessageQueue", "InMemoryMessageQueue", "RedisMessageQueue"]

‎src/mcp/server/message_queue/base.py‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
"""
2+
Base Message Queue Protocol and In-Memory Implementation
3+
4+
This module defines the message queue protocol and provides a default in-memory implementation.
5+
"""
6+
7+
import logging
8+
from typing import Protocol, runtime_checkable
9+
from uuid import UUID
10+
11+
import mcp.types as types
12+
13+
logger = logging.getLogger(__name__)
14+
15+
16+
@runtime_checkable
17+
class MessageQueue(Protocol):
18+
"""Abstract interface for an SSE message queue.
19+
20+
This interface allows messages to be queued and processed by any SSE server instance,
21+
enabling multiple servers to handle requests for the same session.
22+
"""
23+
24+
async def add_message(self, session_id: UUID, message: types.JSONRPCMessage | Exception) -> bool:
25+
"""Add a message to the queue for the specified session.
26+
27+
Args:
28+
session_id: The UUID of the session this message is for
29+
message: The message to queue
30+
31+
Returns:
32+
bool: True if message was accepted, False if session not found
33+
"""
34+
...
35+
36+
async def get_message(self, session_id: UUID, timeout: float = 0.1) -> types.JSONRPCMessage | Exception | None:
37+
"""Get the next message for the specified session.
38+
39+
Args:
40+
session_id: The UUID of the session to get messages for
41+
timeout: Maximum time to wait for a message, in seconds
42+
43+
Returns:
44+
The next message or None if no message is available
45+
"""
46+
...
47+
48+
async def register_session(self, session_id: UUID) -> None:
49+
"""Register a new session with the queue.
50+
51+
Args:
52+
session_id: The UUID of the new session to register
53+
"""
54+
...
55+
56+
async def unregister_session(self, session_id: UUID) -> None:
57+
"""Unregister a session when it's closed.
58+
59+
Args:
60+
session_id: The UUID of the session to unregister
61+
"""
62+
...
63+
64+
async def session_exists(self, session_id: UUID) -> bool:
65+
"""Check if a session exists.
66+
67+
Args:
68+
session_id: The UUID of the session to check
69+
70+
Returns:
71+
bool: True if the session is active, False otherwise
72+
"""
73+
...
74+
75+
76+
class InMemoryMessageQueue:
77+
"""Default in-memory implementation of the MessageQueue interface.
78+
79+
This implementation keeps messages in memory for each session until they're retrieved.
80+
"""
81+
82+
def __init__(self) -> None:
83+
self._message_queues: dict[UUID, list[types.JSONRPCMessage | Exception]] = {}
84+
self._active_sessions: set[UUID] = set()
85+
86+
async def add_message(self, session_id: UUID, message: types.JSONRPCMessage | Exception) -> bool:
87+
"""Add a message to the queue for the specified session."""
88+
if session_id not in self._active_sessions:
89+
logger.warning(f"Message received for unknown session {session_id}")
90+
return False
91+
92+
if session_id not in self._message_queues:
93+
self._message_queues[session_id] = []
94+
95+
self._message_queues[session_id].append(message)
96+
logger.debug(f"Added message to queue for session {session_id}")
97+
return True
98+
99+
async def get_message(self, session_id: UUID, timeout: float = 0.1) -> types.JSONRPCMessage | Exception | None:
100+
"""Get the next message for the specified session."""
101+
if session_id not in self._active_sessions:
102+
return None
103+
104+
queue = self._message_queues.get(session_id, [])
105+
if not queue:
106+
return None
107+
108+
message = queue.pop(0)
109+
if not queue: # Clean up empty queue
110+
del self._message_queues[session_id]
111+
112+
return message
113+
114+
async def register_session(self, session_id: UUID) -> None:
115+
"""Register a new session with the queue."""
116+
self._active_sessions.add(session_id)
117+
logger.debug(f"Registered session {session_id}")
118+
119+
async def unregister_session(self, session_id: UUID) -> None:
120+
"""Unregister a session when it's closed."""
121+
self._active_sessions.discard(session_id)
122+
if session_id in self._message_queues:
123+
del self._message_queues[session_id]
124+
logger.debug(f"Unregistered session {session_id}")
125+
126+
async def session_exists(self, session_id: UUID) -> bool:
127+
"""Check if a session exists."""
128+
return session_id in self._active_sessions

‎src/mcp/server/message_queue/redis.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
"""
2+
Redis Message Queue Module for MCP Server
3+
4+
This module implements a Redis-backed message queue for handling messages between clients and servers.
5+
"""
6+
7+
import json
8+
import logging
9+
from uuid import UUID
10+
11+
import mcp.types as types
12+
13+
try:
14+
import redis.asyncio as redis
15+
except ImportError:
16+
raise ImportError(
17+
"Redis support requires the 'redis' package. "
18+
"Install it with: 'uv add redis' or 'uv add \"mcp[redis]\"'"
19+
)
20+
21+
logger = logging.getLogger(__name__)
22+
23+
24+
class RedisMessageQueue:
25+
"""Redis implementation of the MessageQueue interface.
26+
27+
This implementation uses Redis lists to store messages for each session.
28+
Redis provides persistence and allows multiple servers to share the same queue.
29+
"""
30+
31+
def __init__(self, redis_url: str = "redis://localhost:6379/0", prefix: str = "mcp:queue:") -> None:
32+
"""Initialize Redis message queue.
33+
34+
Args:
35+
redis_url: Redis connection string
36+
prefix: Key prefix for Redis keys to avoid collisions
37+
"""
38+
self._redis = redis.Redis.from_url(redis_url, decode_responses=True)
39+
self._prefix = prefix
40+
self._active_sessions_key = f"{prefix}active_sessions"
41+
logger.debug(f"Initialized Redis message queue with URL: {redis_url}")
42+
43+
def _session_queue_key(self, session_id: UUID) -> str:
44+
"""Get the Redis key for a session's message queue."""
45+
return f"{self._prefix}session:{session_id.hex}"
46+
47+
async def add_message(self, session_id: UUID, message: types.JSONRPCMessage | Exception) -> bool:
48+
"""Add a message to the queue for the specified session."""
49+
# Check if session exists
50+
if not await self.session_exists(session_id):
51+
logger.warning(f"Message received for unknown session {session_id}")
52+
return False
53+
54+
# Serialize the message
55+
if isinstance(message, Exception):
56+
# For exceptions, store them as special format
57+
data = json.dumps({
58+
"_exception": True,
59+
"type": type(message).__name__,
60+
"message": str(message)
61+
})
62+
else:
63+
data = message.model_dump_json(by_alias=True, exclude_none=True)
64+
65+
# Push to the right side of the list (queue)
66+
await self._redis.rpush(self._session_queue_key(session_id), data)
67+
logger.debug(f"Added message to Redis queue for session {session_id}")
68+
return True
69+
70+
async def get_message(self, session_id: UUID, timeout: float = 0.1) -> types.JSONRPCMessage | Exception | None:
71+
"""Get the next message for the specified session."""
72+
# Check if session exists
73+
if not await self.session_exists(session_id):
74+
return None
75+
76+
# Pop from the left side of the list (queue)
77+
# Use BLPOP with timeout to avoid busy waiting
78+
result = await self._redis.blpop([self._session_queue_key(session_id)], timeout)
79+
80+
if not result:
81+
return None
82+
83+
# result is a tuple of (key, value)
84+
_, data = result
85+
86+
# Deserialize the message
87+
json_data = json.loads(data)
88+
89+
# Check if it's an exception
90+
if isinstance(json_data, dict) and json_data.get("_exception"):
91+
# Reconstitute a generic exception
92+
return Exception(f"{json_data['type']}: {json_data['message']}")
93+
94+
# Regular message
95+
try:
96+
return types.JSONRPCMessage.model_validate_json(data)
97+
except Exception as e:
98+
logger.error(f"Failed to deserialize message: {e}")
99+
return None
100+
101+
async def register_session(self, session_id: UUID) -> None:
102+
"""Register a new session with the queue."""
103+
# Add session ID to the set of active sessions
104+
await self._redis.sadd(self._active_sessions_key, session_id.hex)
105+
logger.debug(f"Registered session {session_id} in Redis")
106+
107+
async def unregister_session(self, session_id: UUID) -> None:
108+
"""Unregister a session when it's closed."""
109+
# Remove session ID from active sessions
110+
await self._redis.srem(self._active_sessions_key, session_id.hex)
111+
# Delete the session's message queue
112+
await self._redis.delete(self._session_queue_key(session_id))
113+
logger.debug(f"Unregistered session {session_id} from Redis")
114+
115+
async def session_exists(self, session_id: UUID) -> bool:
116+
"""Check if a session exists."""
117+
return bool(await self._redis.sismember(self._active_sessions_key, session_id.hex))

0 commit comments

Comments
 (0)