testing

Author: akash329d

examples/servers/simple-prompt/mcp_simple_prompt/server.py

@@ -92,7 +92,11 @@ async def get_prompt(
         from starlette.applications import Starlette
         from starlette.routing import Mount, Route
 
-        sse = SseServerTransport("/messages/")
+        from mcp.server.message_queue.redis import RedisMessageDispatch
+
+        message_dispatch = RedisMessageDispatch("redis://localhost:6379/0")
+
+        sse = SseServerTransport("/messages/", message_dispatch=message_dispatch)
 
         async def handle_sse(request):
             async with sse.connect_sse(

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "sse-starlette>=1.6.1",
     "pydantic-settings>=2.5.2",
     "uvicorn>=0.23.1; sys_platform != 'emscripten'",
+    "fakeredis==2.28.1",
 ]
 
 [project.optional-dependencies]

src/mcp/server/fastmcp/server.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations as _annotations
 
+from email import message
 import inspect
 import json
 import re
@@ -507,14 +508,22 @@ async def handle_sse(request: Request) -> None:
                     streams[0],
                     streams[1],
                     self._mcp_server.create_initialization_options(),
-                )
+                )            
+
+        @asynccontextmanager
+        async def lifespan(app: Starlette):
+            try:
+                yield
+            finally:
+                await message_dispatch.close()
 
         return Starlette(
             debug=self.settings.debug,
             routes=[
                 Route(self.settings.sse_path, endpoint=handle_sse),
                 Mount(self.settings.message_path, app=sse.handle_post_message),
             ],
+            lifespan=lifespan,
         )
 
     async def list_prompts(self) -> list[MCPPrompt]:

src/mcp/server/message_queue/base.py

@@ -56,6 +56,10 @@ async def session_exists(self, session_id: UUID) -> bool:
         """
         ...
 
+    async def close(self) -> None:
+        """Close the message dispatch."""
+        ...
+
 
 class InMemoryMessageDispatch:
     """Default in-memory implementation of the MessageDispatch interface.
@@ -106,3 +110,7 @@ async def subscribe(self, session_id: UUID, callback: MessageCallback):
     async def session_exists(self, session_id: UUID) -> bool:
         """Check if a session exists."""
         return session_id in self._callbacks
+
+    async def close(self) -> None:
+        """Close the message dispatch."""
+        pass

src/mcp/server/message_queue/redis.py

@@ -48,6 +48,10 @@ def __init__(
         self._limiter = CapacityLimiter(1)
         logger.debug(f"Redis message dispatch initialized: {redis_url}")
 
+    async def close(self):
+        await self._pubsub.aclose() # type: ignore
+        await self._redis.aclose() # type: ignore
+
     def _session_channel(self, session_id: UUID) -> str:
         """Get the Redis channel for a session."""
         return f"{self._prefix}session:{session_id.hex}"

src/mcp/server/sse.py

@@ -49,6 +49,7 @@ async def handle_sse(request):
 from mcp.server.message_queue import InMemoryMessageDispatch, MessageDispatch
 
 logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG)
 
 
 class SseServerTransport:
@@ -184,4 +185,4 @@ async def handle_post_message(
         logger.debug(f"Publishing message for session {session_id}: {message}")
         response = Response("Accepted", status_code=202)
         await response(scope, receive, send)
-        await self._message_dispatch.publish_message(session_id, message)
+        await self._message_dispatch.publish_message(session_id, message)

tests/server/message_queue/__init__.py

@@ -0,0 +1 @@
+# Message queue tests module