Rework to fix POST blocking issue

akash329d · akash329d · commit 0bfd800dfce2 · 2025-04-21T13:27:59.000-07:00
diff --git a/src/mcp/server/message_queue/base.py b/src/mcp/server/message_queue/base.py
@@ -4,7 +4,7 @@
 from typing import Protocol, runtime_checkable
 from uuid import UUID
 
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 
 import mcp.types as types
 
@@ -13,6 +13,18 @@
 MessageCallback = Callable[[types.JSONRPCMessage | Exception], Awaitable[None]]
 
 
+class MessageWrapper(BaseModel):
+    message_id: str
+    payload: str
+    
+    def get_json_rpc_message(self) -> types.JSONRPCMessage | ValidationError:
+        """Parse the payload into a JSONRPCMessage or return ValidationError."""
+        try:
+            return types.JSONRPCMessage.model_validate_json(self.payload)
+        except ValidationError as exc:
+            return exc
+
+
 @runtime_checkable
 class MessageDispatch(Protocol):
     """Abstract interface for SSE message dispatching.
@@ -35,6 +47,25 @@ async def publish_message(
         """
         ...
 
+    async def publish_message_sync(
+        self, session_id: UUID, message: types.JSONRPCMessage | str, timeout: float = 30.0
+    ) -> bool:
+        """Publish a message for the specified session and wait for consumption confirmation.
+
+        This method blocks until the message has been fully consumed by the subscriber,
+        or until the timeout is reached.
+
+        Args:
+            session_id: The UUID of the session this message is for
+            message: The message to publish (JSONRPCMessage or str for invalid JSON)
+            timeout: Maximum time to wait for consumption in seconds
+
+        Returns:
+            bool: True if message was published and consumed, False otherwise
+        """
+        # Default implementation falls back to standard publish
+        return await self.publish_message(session_id, message)
+
     @asynccontextmanager
     async def subscribe(self, session_id: UUID, callback: MessageCallback):
         """Request-scoped context manager that subscribes to messages for a session.
@@ -90,6 +121,18 @@ async def publish_message(
 
         logger.debug(f"Message dispatched to session {session_id}")
         return True
+        
+    async def publish_message_sync(
+        self, session_id: UUID, message: types.JSONRPCMessage | str, timeout: float = 30.0
+    ) -> bool:
+        """Publish a message for the specified session and wait for consumption.
+        
+        For InMemoryMessageDispatch, this is the same as publish_message since
+        the callback is executed synchronously.
+        """
+        # For in-memory dispatch, the message is processed immediately
+        # so we can just call the regular publish method
+        return await self.publish_message(session_id, message)
 
     @asynccontextmanager
     async def subscribe(self, session_id: UUID, callback: MessageCallback):
diff --git a/src/mcp/server/message_queue/redis.py b/src/mcp/server/message_queue/redis.py
@@ -1,14 +1,14 @@
 import logging
 from contextlib import asynccontextmanager
 from typing import Any, cast
-from uuid import UUID
+from uuid import UUID, uuid4
 
 import anyio
-from anyio import CapacityLimiter, lowlevel
-from pydantic import ValidationError
+from anyio import CancelScope, CapacityLimiter, Event, lowlevel
+from anyio.abc import TaskGroup
 
 import mcp.types as types
-from mcp.server.message_queue.base import MessageCallback
+from mcp.server.message_queue.base import MessageCallback, MessageWrapper
 
 try:
     import redis.asyncio as redis
@@ -42,98 +42,212 @@ def __init__(
         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)
+        self._ack_events: dict[str, Event] = {}
+
         logger.debug(f"Redis message dispatch initialized: {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}"
 
+    def _ack_channel(self, session_id: UUID) -> str:
+        """Get the acknowledgment channel for a session."""
+        return f"{self._prefix}ack:{session_id.hex}"
+
     @asynccontextmanager
     async def subscribe(self, session_id: UUID, callback: MessageCallback):
         """Request-scoped context manager that subscribes to messages for a session."""
         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"Subscribing to Redis channel for session {session_id}")
-        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"Unsubscribed from Redis channel: {session_id}")
-
-    async def _listen_for_messages(self) -> None:
+
+        session_channel = self._session_channel(session_id)
+        ack_channel = self._ack_channel(session_id)
+
+        await self._pubsub.subscribe(session_channel)  # type: ignore
+        await self._pubsub.subscribe(ack_channel)  # type: ignore
+
+        logger.debug(f"Subscribing to Redis channels for session {session_id}")
+
+        # Two nested task groups ensure proper cleanup: the inner one cancels the
+        # listener, while the outer one allows any handlers to complete before exiting.
+        async with anyio.create_task_group() as tg_handler:
+            async with anyio.create_task_group() as tg:
+                tg.start_soon(self._listen_for_messages, tg_handler)
+                try:
+                    yield
+                finally:
+                    tg.cancel_scope.cancel()
+                    await self._pubsub.unsubscribe(session_channel)  # type: ignore
+                    await self._pubsub.unsubscribe(ack_channel)  # type: ignore
+                    await self._redis.srem(self._active_sessions_key, session_id.hex)
+                    del self._callbacks[session_id]
+                    logger.debug(
+                        f"Unsubscribed from Redis channels for session {session_id}"
+                    )
+
+    async def _listen_for_messages(self, tg_handler: TaskGroup) -> 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
-
-                channel: str = cast(str, message["channel"])
-                expected_prefix = f"{self._prefix}session:"
+                # Shield message retrieval from cancellation to ensure no messages are
+                # lost when a session disconnects during processing.
+                with CancelScope(shield=True):
+                    redis_message: (  # type: ignore
+                        None | dict[str, Any]
+                    ) = await self._pubsub.get_message(  # type: ignore
+                        ignore_subscribe_messages=True,
+                        timeout=0.1,  # type: ignore
+                    )
+                    if redis_message is None:
+                        continue
 
-                if not channel.startswith(expected_prefix):
-                    logger.debug(f"Ignoring message from non-MCP channel: {channel}")
-                    continue
+                    channel: str = cast(str, redis_message["channel"])
+                    data: str = cast(str, redis_message["data"])
 
-                session_hex = channel[len(expected_prefix) :]
-                try:
-                    session_id = UUID(hex=session_hex)
-                    expected_channel = self._session_channel(session_id)
-                    if channel != expected_channel:
-                        logger.error(f"Channel format mismatch: {channel}")
+                    # Handle acknowledgment messages
+                    if channel.startswith(f"{self._prefix}ack:"):
+                        tg_handler.start_soon(self._handle_ack_message, channel, data)
                         continue
-                except ValueError:
-                    logger.error(f"Invalid UUID in channel: {channel}")
-                    continue
 
-                data: str = cast(str, message["data"])
-                try:
-                    if session_id not in self._callbacks:
-                        logger.warning(f"Message dropped: no callback for {session_id}")
+                    # Handle session messages
+                    elif channel.startswith(f"{self._prefix}session:"):
+                        tg_handler.start_soon(
+                            self._handle_session_message, channel, data
+                        )
                         continue
 
-                    # Try to parse as valid message or recreate original ValidationError
-                    try:
-                        msg = types.JSONRPCMessage.model_validate_json(data)
-                        await self._callbacks[session_id](msg)
-                    except ValidationError as exc:
-                        # Pass the identical validation error that would have occurred
-                        await self._callbacks[session_id](exc)
-                except Exception as e:
-                    logger.error(f"Error processing message for {session_id}: {e}")
+                    # Ignore other channels
+                    else:
+                        logger.debug(
+                            f"Ignoring message from non-MCP channel: {channel}"
+                        )
+
+    async def _handle_ack_message(self, channel: str, data: str) -> None:
+        """Handle acknowledgment messages received on ack channels."""
+        ack_prefix = f"{self._prefix}ack:"
+        if not channel.startswith(ack_prefix):
+            return
+
+        # Validate channel format exactly matches our expected format
+        session_hex = channel[len(ack_prefix) :]
+        try:
+            # Validate this is a valid UUID hex and channel has correct format
+            session_id = UUID(hex=session_hex)
+            expected_channel = self._ack_channel(session_id)
+            if channel != expected_channel:
+                logger.error(
+                    f"Channel mismatch: got {channel}, expected {expected_channel}"
+                )
+                return
+        except ValueError:
+            logger.error(f"Invalid UUID hex in ack channel: {channel}")
+            return
+
+        # Extract message ID from data
+        message_id = data.strip()
+        if message_id in self._ack_events:
+            logger.debug(f"Received acknowledgment for message: {message_id}")
+            self._ack_events[message_id].set()
+
+    async def _handle_session_message(self, channel: str, data: str) -> None:
+        """Handle regular messages received on session channels."""
+        session_prefix = f"{self._prefix}session:"
+        if not channel.startswith(session_prefix):
+            return
+
+        session_hex = channel[len(session_prefix) :]
+        try:
+            session_id = UUID(hex=session_hex)
+            expected_channel = self._session_channel(session_id)
+            if channel != expected_channel:
+                logger.error(
+                    f"Channel mismatch: got {channel}, expected {expected_channel}"
+                )
+                return
+        except ValueError:
+            logger.error(f"Invalid UUID hex in session channel: {channel}")
+            return
+
+        if session_id not in self._callbacks:
+            logger.warning(f"Message dropped: no callback for {session_id}")
+            return
+
+        try:
+            wrapper = MessageWrapper.model_validate_json(data)
+            result = wrapper.get_json_rpc_message()
+            await self._callbacks[session_id](result)
+            await self._send_acknowledgment(session_id, wrapper.message_id)
+
+        except Exception as e:
+            logger.error(f"Error processing message for {session_id}: {e}")
+
+    async def _send_acknowledgment(self, session_id: UUID, message_id: str) -> None:
+        """Send an acknowledgment for a message that was successfully processed."""
+        ack_channel = self._ack_channel(session_id)
+        await self._redis.publish(ack_channel, message_id)  # type: ignore
+        logger.debug(
+            f"Sent acknowledgment for message {message_id} to session {session_id}"
+        )
 
     async def publish_message(
-        self, session_id: UUID, message: types.JSONRPCMessage | str
-    ) -> bool:
+        self,
+        session_id: UUID,
+        message: types.JSONRPCMessage | str,
+        message_id: str | None = None,
+    ) -> str | None:
         """Publish a message for the specified session."""
         if not await self.session_exists(session_id):
             logger.warning(f"Message dropped: unknown session {session_id}")
-            return False
+            return None
 
         # Pass raw JSON strings directly, preserving validation errors
+        message_id = message_id or str(uuid4())
         if isinstance(message, str):
-            data = message
+            wrapper = MessageWrapper(message_id=message_id, payload=message)
         else:
-            data = message.model_dump_json()
+            wrapper = MessageWrapper(
+                message_id=message_id, payload=message.model_dump_json()
+            )
 
         channel = self._session_channel(session_id)
-        await self._redis.publish(channel, data)  # type: ignore[attr-defined]
-        logger.debug(f"Message published to Redis channel for session {session_id}")
-        return True
+        await self._redis.publish(channel, wrapper.model_dump_json())  # type: ignore
+        logger.debug(
+            f"Message {message_id} published to Redis channel for session {session_id}"
+        )
+        return message_id
+
+    async def publish_message_sync(
+        self,
+        session_id: UUID,
+        message: types.JSONRPCMessage | str,
+        timeout: float = 120.0,
+    ) -> bool:
+        """Publish a message and wait for acknowledgment of processing."""
+        message_id = str(uuid4())
+        ack_event = Event()
+        self._ack_events[message_id] = ack_event
+
+        try:
+            published_id = await self.publish_message(session_id, message, message_id)
+            if published_id is None:
+                return False
+
+            with anyio.fail_after(timeout):
+                await ack_event.wait()
+                logger.debug(f"Received acknowledgment for message {message_id}")
+                return True
+
+        except TimeoutError:
+            logger.warning(
+                f"Timed out waiting for acknowledgment of message {message_id}"
+            )
+            return False
+
+        finally:
+            if message_id in self._ack_events:
+                del self._ack_events[message_id]
 
     async def session_exists(self, session_id: UUID) -> bool:
         """Check if a session exists."""
diff --git a/src/mcp/server/sse.py b/src/mcp/server/sse.py
@@ -174,14 +174,24 @@ async def handle_post_message(
             logger.debug(f"Validated client message: {message}")
         except ValidationError as err:
             logger.error(f"Failed to parse message: {err}")
+            # Still publish the invalid message, but using synchronized version
             response = Response("Could not parse message", status_code=400)
             await response(scope, receive, send)
             # Pass raw JSON string; receiver will recreate identical ValidationError
             # when parsing the same invalid JSON
-            await self._message_dispatch.publish_message(session_id, body.decode())
+            await self._message_dispatch.publish_message_sync(session_id, body.decode())
             return
 
         logger.debug(f"Publishing message for session {session_id}: {message}")
-        response = Response("Accepted", status_code=202)
+        
+        # Use sync publish to block until the message is processed
+        result = await self._message_dispatch.publish_message_sync(session_id, message)
+        
+        if result:
+            # Message was successfully processed
+            response = Response("OK", status_code=200)
+        else:
+            # Message timed out or failed to be processed
+            response = Response("Processing Timeout", status_code=504)
+            
         await response(scope, receive, send)
-        await self._message_dispatch.publish_message(session_id, message)
