draft

Author: swathimshetty19

src/mcp/__init__.py

@@ -13,7 +13,7 @@
     CompleteRequest,
     CreateMessageRequest,
     CreateMessageResult,
-    CreateMessageResultWithTools,
+    #CreateMessageResultWithTools,
     ErrorData,
     GetPromptRequest,
     GetPromptResult,
@@ -43,7 +43,7 @@
     ResourceUpdatedNotification,
     RootsCapability,
     SamplingCapability,
-    SamplingContent,
+    #SamplingContent,
     SamplingContextCapability,
     SamplingMessage,
     SamplingMessageContentBlock,

src/mcp/server/fastmcp/server.py

@@ -64,6 +64,7 @@
 from mcp.server.streamable_http import EventStore
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.server.transport_security import TransportSecuritySettings
+from mcp.server.checkpoint import CheckpointBackend
 from mcp.shared.context import LifespanContextT, RequestContext, RequestT
 from mcp.types import Annotations, AnyFunction, ContentBlock, GetPromptResult, Icon, ToolAnnotations
 from mcp.types import Prompt as MCPPrompt
@@ -173,6 +174,7 @@ def __init__(  # noqa: PLR0913
         lifespan: (Callable[[FastMCP[LifespanResultT]], AbstractAsyncContextManager[LifespanResultT]] | None) = None,
         auth: AuthSettings | None = None,
         transport_security: TransportSecuritySettings | None = None,
+        checkpoint_backend: CheckpointBackend | None = None,
     ):
         # Auto-enable DNS rebinding protection for localhost (IPv4 and IPv6)
         if transport_security is None and host in ("127.0.0.1", "localhost", "::1"):
@@ -230,6 +232,7 @@ def __init__(  # noqa: PLR0913
         if auth_server_provider and not token_verifier:  # pragma: no cover
             self._token_verifier = ProviderTokenVerifier(auth_server_provider)
         self._event_store = event_store
+        self._checkpoint_backend = checkpoint_backend
         self._retry_interval = retry_interval
         self._custom_starlette_routes: list[Route] = []
         self.dependencies = self.settings.dependencies
@@ -275,6 +278,11 @@ def session_manager(self) -> StreamableHTTPSessionManager:
                 "to avoid unnecessary initialization."
             )
         return self._session_manager  # pragma: no cover
+    
+    @property
+    def checkpoint_backend(self) -> CheckpointBackend | None:
+        """Return the checkpoint backend (if any) attached to this server."""
+        return self._checkpoint_backend
 
     def run(
         self,

src/mcp/server/lowlevel/server.py

@@ -94,6 +94,7 @@ async def main():
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.session import RequestResponder
 from mcp.shared.tool_name_validation import validate_and_warn_tool_name
+from mcp.server.checkpoint import CheckpointBackend
 
 logger = logging.getLogger(__name__)
 
@@ -146,6 +147,9 @@ def __init__(
             [Server[LifespanResultT, RequestT]],
             AbstractAsyncContextManager[LifespanResultT],
         ] = lifespan,
+        *,
+        stateless: bool = False,
+        checkpoint_backend: CheckpointBackend | None = None,
     ):
         self.name = name
         self.version = version
@@ -159,6 +163,8 @@ def __init__(
         self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
         self._tool_cache: dict[str, types.Tool] = {}
         self._experimental_handlers: ExperimentalHandlers | None = None
+        self._stateless = stateless
+        self._checkpoint_backend = checkpoint_backend
         logger.debug("Initializing server %r", name)
 
     def create_initialization_options(
@@ -650,6 +656,7 @@ async def run(
                     write_stream,
                     initialization_options,
                     stateless=stateless,
+                    checkpoint_backend=self._checkpoint_backend,
                 )
             )
 

src/mcp/server/session.py

@@ -90,6 +90,8 @@ def __init__(
         # NEW: optional checkpoint backend for servers that support
         # checkpoint/create, checkpoint/validate, checkpoint/resume,
         # checkpoint/delete at a higher level (e.g., FastMCP).
+        *,
+        stateless: bool = False,
         checkpoint_backend: "CheckpointBackend | None" = None,
     ) -> None:
         super().__init__(
@@ -110,6 +112,9 @@ def __init__(
             lambda: self._incoming_message_stream_writer.aclose()
         )
 
+        # Preserve original stateless behaviour (if anything uses it)
+        self._stateless = stateless
+
         # Store optional checkpoint backend. This does not itself handle any
         # checkpoint RPCs; it just makes the backend available to higher-level
         # server code via ctx.session.checkpoint_backend.

src/mcp/types.py

@@ -5,6 +5,7 @@
 from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel
 from pydantic.networks import AnyUrl, UrlConstraints
 from typing_extensions import deprecated
+from typing import Any
 
 """
 Model Context Protocol bindings for Python