Merge branch 'main' into fix/session-manager-resilience

Author: soby

examples/servers/simple-streamablehttp-stateless/mcp_simple_streamablehttp_stateless/server.py

@@ -41,7 +41,7 @@ def main(
     app = Server("mcp-streamable-http-stateless-demo")
 
     @app.call_tool()
-    async def call_tool(name: str, arguments: dict) -> list[types.Content]:
+    async def call_tool(name: str, arguments: dict) -> list[types.ContentBlock]:
         ctx = app.request_context
         interval = arguments.get("interval", 1.0)
         count = arguments.get("count", 5)

examples/servers/simple-streamablehttp/mcp_simple_streamablehttp/server.py

@@ -45,7 +45,7 @@ def main(
     app = Server("mcp-streamable-http-demo")
 
     @app.call_tool()
-    async def call_tool(name: str, arguments: dict) -> list[types.Content]:
+    async def call_tool(name: str, arguments: dict) -> list[types.ContentBlock]:
         ctx = app.request_context
         interval = arguments.get("interval", 1.0)
         count = arguments.get("count", 5)

examples/servers/simple-tool/mcp_simple_tool/server.py

@@ -7,7 +7,7 @@
 
 async def fetch_website(
     url: str,
-) -> list[types.Content]:
+) -> list[types.ContentBlock]:
     headers = {
         "User-Agent": "MCP Test Server (github.com/modelcontextprotocol/python-sdk)"
     }
@@ -29,7 +29,7 @@ def main(port: int, transport: str) -> int:
     app = Server("mcp-website-fetcher")
 
     @app.call_tool()
-    async def fetch_tool(name: str, arguments: dict) -> list[types.Content]:
+    async def fetch_tool(name: str, arguments: dict) -> list[types.ContentBlock]:
         if name != "fetch":
             raise ValueError(f"Unknown tool: {name}")
         if "url" not in arguments:

src/mcp/server/fastmcp/prompts/base.py

@@ -7,16 +7,16 @@
 import pydantic_core
 from pydantic import BaseModel, Field, TypeAdapter, validate_call
 
-from mcp.types import Content, TextContent
+from mcp.types import ContentBlock, TextContent
 
 
 class Message(BaseModel):
     """Base class for all prompt messages."""
 
     role: Literal["user", "assistant"]
-    content: Content
+    content: ContentBlock
 
-    def __init__(self, content: str | Content, **kwargs: Any):
+    def __init__(self, content: str | ContentBlock, **kwargs: Any):
         if isinstance(content, str):
             content = TextContent(type="text", text=content)
         super().__init__(content=content, **kwargs)
@@ -27,7 +27,7 @@ class UserMessage(Message):
 
     role: Literal["user", "assistant"] = "user"
 
-    def __init__(self, content: str | Content, **kwargs: Any):
+    def __init__(self, content: str | ContentBlock, **kwargs: Any):
         super().__init__(content=content, **kwargs)
 
 
@@ -36,7 +36,7 @@ class AssistantMessage(Message):
 
     role: Literal["user", "assistant"] = "assistant"
 
-    def __init__(self, content: str | Content, **kwargs: Any):
+    def __init__(self, content: str | ContentBlock, **kwargs: Any):
         super().__init__(content=content, **kwargs)
 
 

src/mcp/server/fastmcp/server.py

@@ -50,10 +50,11 @@
 from mcp.server.stdio import stdio_server
 from mcp.server.streamable_http import EventStore
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared.context import LifespanContextT, RequestContext, RequestT
 from mcp.types import (
     AnyFunction,
-    Content,
+    ContentBlock,
     GetPromptResult,
     TextContent,
     ToolAnnotations,
@@ -118,6 +119,9 @@ class Settings(BaseSettings, Generic[LifespanResultT]):
 
     auth: AuthSettings | None = None
 
+    # Transport security settings (DNS rebinding protection)
+    transport_security: TransportSecuritySettings | None = None
+
 
 def lifespan_wrapper(
     app: FastMCP,
@@ -256,7 +260,7 @@ def get_context(self) -> Context[ServerSession, object, Request]:
             request_context = None
         return Context(request_context=request_context, fastmcp=self)
 
-    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[Content]:
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[ContentBlock]:
         """Call a tool by name with arguments."""
         context = self.get_context()
         result = await self._tool_manager.call_tool(name, arguments, context=context)
@@ -674,6 +678,7 @@ def sse_app(self, mount_path: str | None = None) -> Starlette:
 
         sse = SseServerTransport(
             normalized_message_endpoint,
+            security_settings=self.settings.transport_security,
         )
 
         async def handle_sse(scope: Scope, receive: Receive, send: Send):
@@ -779,6 +784,7 @@ def streamable_http_app(self) -> Starlette:
                 event_store=self._event_store,
                 json_response=self.settings.json_response,
                 stateless=self.settings.stateless_http,  # Use the stateless setting
+                security_settings=self.settings.transport_security,
             )
 
         # Create the ASGI handler
@@ -872,12 +878,12 @@ async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None) -
 
 def _convert_to_content(
     result: Any,
-) -> Sequence[Content]:
+) -> Sequence[ContentBlock]:
     """Convert a result to a sequence of content objects."""
     if result is None:
         return []
 
-    if isinstance(result, Content):
+    if isinstance(result, ContentBlock):
         return [result]
 
     if isinstance(result, Image):

src/mcp/server/lowlevel/server.py

@@ -384,7 +384,7 @@ def call_tool(self):
         def decorator(
             func: Callable[
                 ...,
-                Awaitable[Iterable[types.Content]],
+                Awaitable[Iterable[types.ContentBlock]],
             ],
         ):
             logger.debug("Registering handler for CallToolRequest")

src/mcp/server/sse.py

@@ -52,6 +52,10 @@ async def handle_sse(request):
 from starlette.types import Receive, Scope, Send
 
 import mcp.types as types
+from mcp.server.transport_security import (
+    TransportSecurityMiddleware,
+    TransportSecuritySettings,
+)
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 
 logger = logging.getLogger(__name__)
@@ -71,16 +75,22 @@ class SseServerTransport:
 
     _endpoint: str
     _read_stream_writers: dict[UUID, MemoryObjectSendStream[SessionMessage | Exception]]
+    _security: TransportSecurityMiddleware
 
-    def __init__(self, endpoint: str) -> None:
+    def __init__(self, endpoint: str, security_settings: TransportSecuritySettings | 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 relative or absolute URL for POST messages.
+            security_settings: Optional security settings for DNS rebinding protection.
         """
 
         super().__init__()
         self._endpoint = endpoint
         self._read_stream_writers = {}
+        self._security = TransportSecurityMiddleware(security_settings)
         logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}")
 
     @asynccontextmanager
@@ -89,6 +99,13 @@ async def connect_sse(self, scope: Scope, receive: Receive, send: Send):
             logger.error("connect_sse received non-HTTP request")
             raise ValueError("connect_sse can only handle HTTP requests")
 
+        # Validate request headers for DNS rebinding protection
+        request = Request(scope, receive)
+        error_response = await self._security.validate_request(request, is_post=False)
+        if error_response:
+            await error_response(scope, receive, send)
+            raise ValueError("Request validation failed")
+
         logger.debug("Setting up SSE connection")
         read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
         read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
@@ -160,6 +177,11 @@ async def handle_post_message(self, scope: Scope, receive: Receive, send: Send)
         logger.debug("Handling POST message")
         request = Request(scope, receive)
 
+        # Validate request headers for DNS rebinding protection
+        error_response = await self._security.validate_request(request, is_post=True)
+        if error_response:
+            return await error_response(scope, receive, send)
+
         session_id_param = request.query_params.get("session_id")
         if session_id_param is None:
             logger.warning("Received request without session_id")

src/mcp/server/streamable_http.py

@@ -24,6 +24,10 @@
 from starlette.responses import Response
 from starlette.types import Receive, Scope, Send
 
+from mcp.server.transport_security import (
+    TransportSecurityMiddleware,
+    TransportSecuritySettings,
+)
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 from mcp.types import (
@@ -130,12 +134,14 @@ class StreamableHTTPServerTransport:
     _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception] | None = None
     _write_stream: MemoryObjectSendStream[SessionMessage] | None = None
     _write_stream_reader: MemoryObjectReceiveStream[SessionMessage] | None = None
+    _security: TransportSecurityMiddleware
 
     def __init__(
         self,
         mcp_session_id: str | None,
         is_json_response_enabled: bool = False,
         event_store: EventStore | None = None,
+        security_settings: TransportSecuritySettings | None = None,
     ) -> None:
         """
         Initialize a new StreamableHTTP server transport.
@@ -148,6 +154,7 @@ def __init__(
             event_store: Event store for resumability support. If provided,
                         resumability will be enabled, allowing clients to
                         reconnect and resume messages.
+            security_settings: Optional security settings for DNS rebinding protection.
 
         Raises:
             ValueError: If the session ID contains invalid characters.
@@ -158,6 +165,7 @@ def __init__(
         self.mcp_session_id = mcp_session_id
         self.is_json_response_enabled = is_json_response_enabled
         self._event_store = event_store
+        self._security = TransportSecurityMiddleware(security_settings)
         self._request_streams: dict[
             RequestId,
             tuple[
@@ -251,6 +259,14 @@ async def _clean_up_memory_streams(self, request_id: RequestId) -> None:
     async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Application entry point that handles all HTTP requests"""
         request = Request(scope, receive)
+
+        # Validate request headers for DNS rebinding protection
+        is_post = request.method == "POST"
+        error_response = await self._security.validate_request(request, is_post=is_post)
+        if error_response:
+            await error_response(scope, receive, send)
+            return
+
         if self._terminated:
             # If the session has been terminated, return 404 Not Found
             response = self._create_error_response(

src/mcp/server/streamable_http_manager.py

@@ -22,6 +22,7 @@
     EventStore,
     StreamableHTTPServerTransport,
 )
+from mcp.server.transport_security import TransportSecuritySettings
 
 logger = logging.getLogger(__name__)
 
@@ -59,11 +60,13 @@ def __init__(
         event_store: EventStore | None = None,
         json_response: bool = False,
         stateless: bool = False,
+        security_settings: TransportSecuritySettings | None = None,
     ):
         self.app = app
         self.event_store = event_store
         self.json_response = json_response
         self.stateless = stateless
+        self.security_settings = security_settings
 
         # Session tracking (only used if not stateless)
         self._session_creation_lock = anyio.Lock()
@@ -161,6 +164,7 @@ async def _handle_stateless_request(
             mcp_session_id=None,  # No session tracking in stateless mode
             is_json_response_enabled=self.json_response,
             event_store=None,  # No event store in stateless mode
+            security_settings=self.security_settings,
         )
 
         # Start server in a new task
@@ -219,6 +223,7 @@ async def _handle_stateful_request(
                     mcp_session_id=new_session_id,
                     is_json_response_enabled=self.json_response,
                     event_store=self.event_store,  # May be None (no resumability)
+                    security_settings=self.security_settings,
                 )
 
                 assert http_transport.mcp_session_id is not None