merge fixes

Author: akash329d

src/mcp/server/fastmcp/server.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations as _annotations
 
-from email import message
 import inspect
 import re
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
@@ -689,7 +688,10 @@ async def sse_endpoint(request: Request) -> None:
 
         # Create Starlette app with routes and middleware
         return Starlette(
-            debug=self.settings.debug, routes=routes, middleware=middleware, lifespan=lifespan
+            debug=self.settings.debug,
+            routes=routes,
+            middleware=middleware,
+            lifespan=lifespan,
         )
 
     async def list_prompts(self) -> list[MCPPrompt]:

src/mcp/server/message_queue/base.py

@@ -6,11 +6,11 @@
 
 from pydantic import ValidationError
 
-import mcp.types as types
+from mcp.shared.message import SessionMessage
 
 logger = logging.getLogger(__name__)
 
-MessageCallback = Callable[[types.JSONRPCMessage | Exception], Awaitable[None]]
+MessageCallback = Callable[[SessionMessage | Exception], Awaitable[None]]
 
 
 @runtime_checkable
@@ -22,7 +22,7 @@ class MessageDispatch(Protocol):
     """
 
     async def publish_message(
-        self, session_id: UUID, message: types.JSONRPCMessage | str
+        self, session_id: UUID, message: SessionMessage | str
     ) -> bool:
         """Publish a message for the specified session.
 
@@ -72,7 +72,7 @@ def __init__(self) -> None:
         self._callbacks: dict[UUID, MessageCallback] = {}
 
     async def publish_message(
-        self, session_id: UUID, message: types.JSONRPCMessage | str
+        self, session_id: UUID, message: SessionMessage | str
     ) -> bool:
         """Publish a message for the specified session."""
         if session_id not in self._callbacks:
@@ -82,7 +82,7 @@ async def publish_message(
         # Parse string messages or recreate original ValidationError
         if isinstance(message, str):
             try:
-                callback_argument = types.JSONRPCMessage.model_validate_json(message)
+                callback_argument = SessionMessage.model_validate_json(message)
             except ValidationError as exc:
                 callback_argument = exc
         else:

src/mcp/server/sse.py

@@ -201,9 +201,10 @@ async def handle_post_message(
             await self._message_dispatch.publish_message(session_id, body.decode())
             return
 
-        session_message = SessionMessage(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)
-        logger.debug(f"Sending session message to writer: {session_message}")
+        await self._message_dispatch.publish_message(
+            session_id, SessionMessage(message=message)
+        )
+        logger.debug(f"Sending session message to writer: {message}")

src/mcp/shared/message.py

@@ -6,7 +6,8 @@
 """
 
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass
+
+from pydantic import BaseModel
 
 from mcp.types import JSONRPCMessage, RequestId
 
@@ -15,8 +16,7 @@
 ResumptionTokenUpdateCallback = Callable[[ResumptionToken], Awaitable[None]]
 
 
-@dataclass
-class ClientMessageMetadata:
+class ClientMessageMetadata(BaseModel):
     """Metadata specific to client messages."""
 
     resumption_token: ResumptionToken | None = None
@@ -25,8 +25,7 @@ class ClientMessageMetadata:
     )
 
 
-@dataclass
-class ServerMessageMetadata:
+class ServerMessageMetadata(BaseModel):
     """Metadata specific to server messages."""
 
     related_request_id: RequestId | None = None
@@ -35,8 +34,7 @@ class ServerMessageMetadata:
 MessageMetadata = ClientMessageMetadata | ServerMessageMetadata | None
 
 
-@dataclass
-class SessionMessage:
+class SessionMessage(BaseModel):
     """A message with specific metadata for transport-specific features."""
 
     message: JSONRPCMessage