ruff

akash329d · akash329d · commit b2fce7dd1a24 · 2025-04-08T16:35:27.000-07:00
diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py
@@ -75,7 +75,7 @@ class Settings(BaseSettings, Generic[LifespanResultT]):
     port: int = 8000
     sse_path: str = "/sse"
     message_path: str = "/messages/"
-    
+
     # SSE message queue settings
     message_queue: Literal["memory", "redis"] = "memory"
     redis_url: str = "redis://localhost:6379/0"
@@ -488,20 +488,25 @@ def sse_app(self) -> Starlette:
         if self.settings.message_queue == "redis":
             try:
                 from mcp.server.message_queue import RedisMessageQueue
+
                 message_queue = RedisMessageQueue(
-                    redis_url=self.settings.redis_url, 
-                    prefix=self.settings.redis_prefix
+                    redis_url=self.settings.redis_url, prefix=self.settings.redis_prefix
                 )
                 logger.info(f"Using Redis message queue at {self.settings.redis_url}")
             except ImportError:
-                logger.error("Redis message queue requested but 'redis' package not installed. ")
+                logger.error(
+                    "Redis message queue requested but 'redis' package not installed. "
+                )
                 raise
         else:
             from mcp.server.message_queue import InMemoryMessageQueue
+
             message_queue = InMemoryMessageQueue()
             logger.info("Using in-memory message queue")
-            
-        sse = SseServerTransport(self.settings.message_path, message_queue=message_queue)
+
+        sse = SseServerTransport(
+            self.settings.message_path, message_queue=message_queue
+        )
 
         async def handle_sse(request: Request) -> None:
             async with sse.connect_sse(
diff --git a/src/mcp/server/message_queue/__init__.py b/src/mcp/server/message_queue/__init__.py
@@ -1,7 +1,8 @@
 """
 Message Queue Module for MCP Server
 
-This module implements queue interfaces for handling messages between clients and servers.
+This module implements queue interfaces for handling
+messages between clients and servers.
 """
 
 from mcp.server.message_queue.base import InMemoryMessageQueue, MessageQueue
@@ -12,4 +13,4 @@
 except ImportError:
     RedisMessageQueue = None
 
-__all__ = ["MessageQueue", "InMemoryMessageQueue", "RedisMessageQueue"]
+__all__ = ["MessageQueue", "InMemoryMessageQueue", "RedisMessageQueue"]
diff --git a/src/mcp/server/message_queue/base.py b/src/mcp/server/message_queue/base.py
@@ -1,9 +1,3 @@
-"""
-Base Message Queue Protocol and In-Memory Implementation
-
-This module defines the message queue protocol and provides a default in-memory implementation.
-"""
-
 import logging
 from typing import Protocol, runtime_checkable
 from uuid import UUID
@@ -16,57 +10,61 @@
 @runtime_checkable
 class MessageQueue(Protocol):
     """Abstract interface for an SSE message queue.
-    
-    This interface allows messages to be queued and processed by any SSE server instance,
+
+    This interface allows messages to be queued and processed by any SSE server instance
     enabling multiple servers to handle requests for the same session.
     """
-    
-    async def add_message(self, session_id: UUID, message: types.JSONRPCMessage | Exception) -> bool:
+
+    async def add_message(
+        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+    ) -> bool:
         """Add a message to the queue for the specified session.
-        
+
         Args:
             session_id: The UUID of the session this message is for
             message: The message to queue
-            
+
         Returns:
             bool: True if message was accepted, False if session not found
         """
         ...
-    
-    async def get_message(self, session_id: UUID, timeout: float = 0.1) -> types.JSONRPCMessage | Exception | None:
+
+    async def get_message(
+        self, session_id: UUID, timeout: float = 0.1
+    ) -> types.JSONRPCMessage | Exception | None:
         """Get the next message for the specified session.
-        
+
         Args:
             session_id: The UUID of the session to get messages for
             timeout: Maximum time to wait for a message, in seconds
-            
+
         Returns:
             The next message or None if no message is available
         """
         ...
-    
+
     async def register_session(self, session_id: UUID) -> None:
         """Register a new session with the queue.
-        
+
         Args:
             session_id: The UUID of the new session to register
         """
         ...
-    
+
     async def unregister_session(self, session_id: UUID) -> None:
         """Unregister a session when it's closed.
-        
+
         Args:
             session_id: The UUID of the session to unregister
         """
         ...
-        
+
     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
         """
@@ -75,54 +73,59 @@ async def session_exists(self, session_id: UUID) -> bool:
 
 class InMemoryMessageQueue:
     """Default in-memory implementation of the MessageQueue interface.
-    
-    This implementation keeps messages in memory for each session until they're retrieved.
+
+    This implementation keeps messages in memory for
+    each session until they're retrieved.
     """
-    
+
     def __init__(self) -> None:
         self._message_queues: dict[UUID, list[types.JSONRPCMessage | Exception]] = {}
         self._active_sessions: set[UUID] = set()
-    
-    async def add_message(self, session_id: UUID, message: types.JSONRPCMessage | Exception) -> bool:
+
+    async def add_message(
+        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+    ) -> bool:
         """Add a message to the queue for the specified session."""
         if session_id not in self._active_sessions:
             logger.warning(f"Message received for unknown session {session_id}")
             return False
-            
+
         if session_id not in self._message_queues:
             self._message_queues[session_id] = []
-        
+
         self._message_queues[session_id].append(message)
         logger.debug(f"Added message to queue for session {session_id}")
         return True
-    
-    async def get_message(self, session_id: UUID, timeout: float = 0.1) -> types.JSONRPCMessage | Exception | None:
+
+    async def get_message(
+        self, session_id: UUID, timeout: float = 0.1
+    ) -> types.JSONRPCMessage | Exception | None:
         """Get the next message for the specified session."""
         if session_id not in self._active_sessions:
             return None
-            
+
         queue = self._message_queues.get(session_id, [])
         if not queue:
             return None
-            
+
         message = queue.pop(0)
         if not queue:  # Clean up empty queue
             del self._message_queues[session_id]
-            
+
         return message
-    
+
     async def register_session(self, session_id: UUID) -> None:
         """Register a new session with the queue."""
         self._active_sessions.add(session_id)
         logger.debug(f"Registered session {session_id}")
-    
+
     async def unregister_session(self, session_id: UUID) -> None:
         """Unregister a session when it's closed."""
         self._active_sessions.discard(session_id)
         if session_id in self._message_queues:
             del self._message_queues[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
+        return session_id in self._active_sessions
diff --git a/src/mcp/server/message_queue/redis.py b/src/mcp/server/message_queue/redis.py
@@ -1,9 +1,3 @@
-"""
-Redis Message Queue Module for MCP Server
-
-This module implements a Redis-backed message queue for handling messages between clients and servers.
-"""
-
 import json
 import logging
 from uuid import UUID
@@ -23,14 +17,16 @@
 
 class RedisMessageQueue:
     """Redis implementation of the MessageQueue interface.
-    
+
     This implementation uses Redis lists to store messages for each session.
     Redis provides persistence and allows multiple servers to share the same queue.
     """
-    
-    def __init__(self, redis_url: str = "redis://localhost:6379/0", prefix: str = "mcp:queue:") -> None:
+
+    def __init__(
+        self, redis_url: str = "redis://localhost:6379/0", prefix: str = "mcp:queue:"
+    ) -> None:
         """Initialize Redis message queue.
-        
+
         Args:
             redis_url: Redis connection string
             prefix: Key prefix for Redis keys to avoid collisions
@@ -39,79 +35,87 @@ def __init__(self, redis_url: str = "redis://localhost:6379/0", prefix: str = "m
         self._prefix = prefix
         self._active_sessions_key = f"{prefix}active_sessions"
         logger.debug(f"Initialized Redis message queue with URL: {redis_url}")
-    
+
     def _session_queue_key(self, session_id: UUID) -> str:
         """Get the Redis key for a session's message queue."""
         return f"{self._prefix}session:{session_id.hex}"
-    
-    async def add_message(self, session_id: UUID, message: types.JSONRPCMessage | Exception) -> bool:
+
+    async def add_message(
+        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+    ) -> bool:
         """Add a message to the queue for the specified session."""
         # Check if session exists
         if not await self.session_exists(session_id):
             logger.warning(f"Message received for unknown session {session_id}")
             return False
-        
+
         # Serialize the message
         if isinstance(message, Exception):
             # For exceptions, store them as special format
-            data = json.dumps({
-                "_exception": True,
-                "type": type(message).__name__,
-                "message": str(message)
-            })
+            data = json.dumps(
+                {
+                    "_exception": True,
+                    "type": type(message).__name__,
+                    "message": str(message),
+                }
+            )
         else:
             data = message.model_dump_json(by_alias=True, exclude_none=True)
-        
+
         # Push to the right side of the list (queue)
         await self._redis.rpush(self._session_queue_key(session_id), data)
         logger.debug(f"Added message to Redis queue for session {session_id}")
         return True
-    
-    async def get_message(self, session_id: UUID, timeout: float = 0.1) -> types.JSONRPCMessage | Exception | None:
+
+    async def get_message(
+        self, session_id: UUID, timeout: float = 0.1
+    ) -> types.JSONRPCMessage | Exception | None:
         """Get the next message for the specified session."""
         # Check if session exists
         if not await self.session_exists(session_id):
             return None
-        
+
         # Pop from the left side of the list (queue)
         # Use BLPOP with timeout to avoid busy waiting
         result = await self._redis.blpop([self._session_queue_key(session_id)], timeout)
-        
+
         if not result:
             return None
-        
+
         # result is a tuple of (key, value)
         _, data = result
-        
+
         # Deserialize the message
         json_data = json.loads(data)
-        
+
         # Check if it's an exception
         if isinstance(json_data, dict) and json_data.get("_exception"):
             # Reconstitute a generic exception
             return Exception(f"{json_data['type']}: {json_data['message']}")
-        
+
         # Regular message
         try:
             return types.JSONRPCMessage.model_validate_json(data)
         except Exception as e:
             logger.error(f"Failed to deserialize message: {e}")
             return None
-    
+
     async def register_session(self, session_id: UUID) -> None:
         """Register a new session with the queue."""
         # Add session ID to the set of active sessions
         await self._redis.sadd(self._active_sessions_key, session_id.hex)
         logger.debug(f"Registered session {session_id} in Redis")
-    
+
     async def unregister_session(self, session_id: UUID) -> None:
         """Unregister a session when it's closed."""
         # Remove session ID from active sessions
         await self._redis.srem(self._active_sessions_key, session_id.hex)
         # Delete the session's message queue
         await self._redis.delete(self._session_queue_key(session_id))
         logger.debug(f"Unregistered session {session_id} from Redis")
-    
+
     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))
+        return bool(
+            await self._redis.sismember(self._active_sessions_key, session_id.hex)
+        )
diff --git a/src/mcp/server/sse.py b/src/mcp/server/sse.py
@@ -66,14 +66,16 @@ class SseServerTransport:
     _endpoint: str
     _message_queue: MessageQueue
 
-    def __init__(self, endpoint: str, message_queue: MessageQueue | None = None) -> None:
+    def __init__(
+        self, endpoint: str, message_queue: MessageQueue | None = None
+    ) -> None:
         """
         Creates a new SSE server transport, which will direct the client to POST
         messages to the relative or absolute URL given.
-        
+
         Args:
             endpoint: The endpoint URL for SSE connections
-            message_queue: Optional message queue to use. If None, creates an InMemoryMessageQueue.
+            message_queue: Optional message queue to use
         """
 
         super().__init__()
@@ -107,6 +109,7 @@ async def connect_sse(self, scope: Scope, receive: Receive, send: Send):
         ](0)
 
         message_polling_active = True
+
         async def poll_queue():
             """Background task to poll for messages in the queue"""
             logger.debug(f"Starting queue polling for session {session_id}")
