logging improvements

Author: akash329d

src/mcp/server/message_queue/__init__.py

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

src/mcp/server/message_queue/base.py

@@ -3,6 +3,7 @@
 from contextlib import asynccontextmanager
 from typing import Protocol, runtime_checkable
 from uuid import UUID
+from pydantic import ValidationError
 
 import mcp.types as types
 
@@ -20,13 +21,13 @@ class MessageDispatch(Protocol):
     """
 
     async def publish_message(
-        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+        self, session_id: UUID, message: types.JSONRPCMessage | str
     ) -> 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
+            message: The message to publish (JSONRPCMessage or str for invalid JSON)
 
         Returns:
             bool: True if message was published, False if session not found
@@ -67,31 +68,40 @@ def __init__(self) -> None:
         # We don't need a separate _active_sessions set since _callbacks already tracks this
 
     async def publish_message(
-        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+        self, session_id: UUID, message: types.JSONRPCMessage | str
     ) -> bool:
         """Publish a message for the specified session."""
         if session_id not in self._callbacks:
-            logger.warning(f"Message received for unknown session {session_id}")
+            logger.warning(f"Message dropped: unknown session {session_id}")
             return False
-
-        # Call the callback directly
-        await self._callbacks[session_id](message)
-        logger.debug(f"Called callback for session {session_id}")
-
+        
+        # For string messages, attempt parsing and recreate original ValidationError if invalid
+        if isinstance(message, str):
+            try:
+                callback_argument = types.JSONRPCMessage.model_validate_json(message)
+            except ValidationError as exc:
+                callback_argument = exc
+        else:
+            callback_argument = message
+        
+        # Call the callback with either valid message or recreated ValidationError
+        await self._callbacks[session_id](callback_argument)
+        
+        logger.debug(f"Message dispatched to session {session_id}")
         return True
 
     @asynccontextmanager
     async def subscribe(self, session_id: UUID, callback: MessageCallback):
         """Request-scoped context manager that subscribes to messages for a session."""
         self._callbacks[session_id] = callback
-        logger.debug(f"Registered session {session_id} with callback")
+        logger.debug(f"Subscribing to messages for session {session_id}")
 
         try:
             yield
         finally:
             if session_id in self._callbacks:
                 del self._callbacks[session_id]
-            logger.debug(f"Unregistered session {session_id}")
+            logger.debug(f"Unsubscribed from session {session_id}")
 
     async def session_exists(self, session_id: UUID) -> bool:
         """Check if a session exists."""

src/mcp/server/message_queue/redis.py

@@ -1,8 +1,8 @@
-import json
 import logging
 from contextlib import asynccontextmanager
 from typing import Any, cast
 from uuid import UUID
+from pydantic import ValidationError
 
 import anyio
 from anyio import CapacityLimiter, lowlevel
@@ -44,7 +44,7 @@ def __init__(
         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 dispatch with URL: {redis_url}")
+        logger.debug(f"Redis message dispatch initialized: {redis_url}")
 
     def _session_channel(self, session_id: UUID) -> str:
         """Get the Redis channel for a session."""
@@ -58,7 +58,7 @@ async def subscribe(self, session_id: UUID, callback: MessageCallback):
         channel = self._session_channel(session_id)
         await self._pubsub.subscribe(channel)  # type: ignore
 
-        logger.debug(f"Registered session {session_id} in Redis with callback")
+        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:
@@ -68,7 +68,7 @@ async def subscribe(self, session_id: UUID, callback: MessageCallback):
                 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")
+                logger.debug(f"Unsubscribed from Redis channel for session {session_id}")
 
     async def _listen_for_messages(self) -> None:
         """Background task that listens for messages on subscribed channels."""
@@ -85,61 +85,49 @@ async def _listen_for_messages(self) -> None:
                 # Extract session ID from channel name
                 channel: str = cast(str, message["channel"])
                 if not channel.startswith(self._prefix):
+                    logger.debug(f"Ignoring message from non-MCP channel: {channel}")
                     continue
 
                 session_hex = channel.split(":")[-1]
                 try:
                     session_id = UUID(hex=session_hex)
                 except ValueError:
-                    logger.error(f"Invalid session channel: {channel}")
+                    logger.error(f"Received message for invalid session channel: {channel}")
                     continue
 
                 data: str = cast(str, message["data"])
-                msg: None | types.JSONRPCMessage | Exception = None
                 try:
-                    json_data = json.loads(data)
-                    if not isinstance(json_data, dict):
-                        logger.error(f"Received non-dict JSON data: {type(json_data)}")
+                    if session_id not in self._callbacks:
+                        logger.warning(f"Message dropped: no callback for session {session_id}")
                         continue
 
-                    json_dict: dict[str, Any] = json_data
-                    if json_dict.get("_exception", False):
-                        msg = Exception(
-                            f"{json_dict['type']}: {json_dict['message']}"
-                        )
-                    else:
+                    # Try to parse as valid message or recreate original ValidationError
+                    try:
                         msg = types.JSONRPCMessage.model_validate_json(data)
-
-                    if msg:
-                        if session_id in self._callbacks:
-                            await self._callbacks[session_id](msg)
-                        else:
-                            logger.warning(f"No callback registered for session {session_id}")
+                        await self._callbacks[session_id](msg)
+                    except ValidationError as exc:
+                        # Pass the identical validation error that would have occurred originally
+                        await self._callbacks[session_id](exc)
                 except Exception as e:
-                    logger.error(f"Failed to process message: {e}")
+                    logger.error(f"Error processing message for session {session_id}: {e}")
 
     async def publish_message(
-        self, session_id: UUID, message: types.JSONRPCMessage | Exception
+        self, session_id: UUID, message: types.JSONRPCMessage | str
     ) -> 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}")
+            logger.warning(f"Message dropped: unknown session {session_id}")
             return False
 
-        if isinstance(message, Exception):
-            data = json.dumps(
-                {
-                    "_exception": True,
-                    "type": type(message).__name__,
-                    "message": str(message),
-                }
-            )
+        # Pass raw JSON strings directly, preserving validation errors
+        if isinstance(message, str):
+            data = 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}")
+        logger.debug(f"Message published to Redis channel for session {session_id}")
         return True
 
     async def session_exists(self, session_id: UUID) -> bool:

src/mcp/server/sse.py

@@ -46,7 +46,7 @@ async def handle_sse(request):
 from starlette.types import Receive, Scope, Send
 
 import mcp.types as types
-from mcp.server.message_queue import InMemoryMessageQueue, MessageQueue
+from mcp.server.message_queue import InMemoryMessageDispatch, MessageDispatch
 
 logger = logging.getLogger(__name__)
 
@@ -64,23 +64,23 @@ class SseServerTransport:
     """
 
     _endpoint: str
-    _message_queue: MessageQueue
+    _message_dispatch: MessageDispatch
 
     def __init__(
-        self, endpoint: str, message_queue: MessageQueue | None = None
+        self, endpoint: str, message_dispatch: MessageDispatch | 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
+            message_dispatch: Optional message dispatch to use
         """
 
         super().__init__()
         self._endpoint = endpoint
-        self._message_dispatch = message_queue or InMemoryMessageDispatch()
+        self._message_dispatch = message_dispatch or InMemoryMessageDispatch()
         logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}")
 
     @asynccontextmanager
@@ -137,7 +137,7 @@ async def sse_writer():
             logger.debug("Starting SSE response task")
             tg.start_soon(response, scope, receive, send)
 
-            async with self._message_queue.active_for_request(
+            async with self._message_dispatch.subscribe(
                 session_id, message_callback
             ):
                 logger.debug("Yielding read and write streams")
@@ -163,7 +163,7 @@ async def handle_post_message(
             response = Response("Invalid session ID", status_code=400)
             return await response(scope, receive, send)
 
-        if not await self._message_queue.session_exists(session_id):
+        if not await self._message_dispatch.session_exists(session_id):
             logger.warning(f"Could not find session for ID: {session_id}")
             response = Response("Could not find session", status_code=404)
             return await response(scope, receive, send)
@@ -178,10 +178,12 @@ async def handle_post_message(
             logger.error(f"Failed to parse message: {err}")
             response = Response("Could not parse message", status_code=400)
             await response(scope, receive, send)
-            await self._message_queue.publish_message(session_id, err)
+            # Pass raw JSON string through dispatch; original ValidationError will be recreated when 
+            # the receiver tries to validate the same invalid JSON
+            await self._message_dispatch.publish_message(session_id, body.decode())
             return
 
         logger.debug(f"Publishing message for session {session_id}: {message}")
         response = Response("Accepted", status_code=202)
         await response(scope, receive, send)
-        await self._message_queue.publish_message(session_id, message)
+        await self._message_dispatch.publish_message(session_id, message)