feat: expand InitializationState with explicit lifecycle state machine

Expand the InitializationState enum with new states (Stateless, Closing, Closed, Error) and add centralized transition validation via a _VALID_TRANSITIONS table and _transition_state() method.

Key changes:

- Add Stateless state for sessions that skip the initialization handshake

- Add Closing/Closed states for orderly shutdown tracking

- Add Error state for unrecoverable failures with recovery paths

- Add _transition_state() method that validates transitions against a table

- Add initialization_state property (read-only) and is_initialized property

- Override __aexit__ to transition through Closing -> Closed on exit

- Update _received_request and _received_notification to use new APIs

- Add comprehensive test suite (20 tests) covering all state transitions

Github-Issue: #1691

Author: dgenio

src/mcp/server/session.py

@@ -37,7 +37,9 @@ async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]:
 be instantiated directly by users of the MCP framework.
 """
 
+import logging
 from enum import Enum
+from types import TracebackType
 from typing import Any, TypeVar
 
 import anyio
@@ -58,11 +60,58 @@ async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]:
 )
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
+logger = logging.getLogger(__name__)
+
 
 class InitializationState(Enum):
+    """Represents the lifecycle states of a server session.
+
+    State transitions:
+        NotInitialized -> Initializing -> Initialized -> Closing -> Closed
+        Stateless -> Closing -> Closed
+        Any state -> Error (on unrecoverable failure)
+    """
+
     NotInitialized = 1
     Initializing = 2
     Initialized = 3
+    Stateless = 4
+    Closing = 5
+    Closed = 6
+    Error = 7
+
+
+# Valid state transitions: maps each state to the set of states it can transition to.
+_VALID_TRANSITIONS: dict[InitializationState, set[InitializationState]] = {
+    InitializationState.NotInitialized: {
+        InitializationState.Initializing,
+        InitializationState.Closing,
+        InitializationState.Error,
+    },
+    InitializationState.Initializing: {
+        InitializationState.Initialized,
+        InitializationState.Closing,
+        InitializationState.Error,
+    },
+    InitializationState.Initialized: {
+        InitializationState.Initializing,  # re-initialization
+        InitializationState.Closing,
+        InitializationState.Error,
+    },
+    InitializationState.Stateless: {
+        InitializationState.Closing,
+        InitializationState.Error,
+    },
+    InitializationState.Closing: {
+        InitializationState.Closed,
+        InitializationState.Error,
+    },
+    InitializationState.Closed: set(),
+    InitializationState.Error: {
+        InitializationState.Closing,
+        InitializationState.Closed,
+    },
+}
 
 
 ServerSessionT = TypeVar("ServerSessionT", bound="ServerSession")
@@ -81,7 +130,7 @@ class ServerSession(
         types.ClientNotification,
     ]
 ):
-    _initialized: InitializationState = InitializationState.NotInitialized
+    _initialization_state: InitializationState = InitializationState.NotInitialized
     _client_params: types.InitializeRequestParams | None = None
     _experimental_features: ExperimentalServerSessionFeatures | None = None
 
@@ -93,16 +142,47 @@ def __init__(
         stateless: bool = False,
     ) -> None:
         super().__init__(read_stream, write_stream, types.ClientRequest, types.ClientNotification)
-        self._initialization_state = (
-            InitializationState.Initialized if stateless else InitializationState.NotInitialized
-        )
+        self._initialization_state = InitializationState.Stateless if stateless else InitializationState.NotInitialized
 
         self._init_options = init_options
         self._incoming_message_stream_writer, self._incoming_message_stream_reader = anyio.create_memory_object_stream[
             ServerRequestResponder
         ](0)
         self._exit_stack.push_async_callback(lambda: self._incoming_message_stream_reader.aclose())
 
+    @property
+    def initialization_state(self) -> InitializationState:
+        """Return the current initialization state of the session."""
+        return self._initialization_state
+
+    @property
+    def is_initialized(self) -> bool:
+        """Check whether the session is ready to process requests.
+
+        Returns True when the session has completed initialization handshake
+        (Initialized) or is operating in stateless mode (Stateless).
+        """
+        return self._initialization_state in (
+            InitializationState.Initialized,
+            InitializationState.Stateless,
+        )
+
+    def _transition_state(self, new_state: InitializationState) -> None:
+        """Transition the session to a new state, validating the transition.
+
+        Args:
+            new_state: The target state to transition to.
+
+        Raises:
+            RuntimeError: If the transition is not valid from the current state.
+        """
+        current = self._initialization_state
+        valid_targets = _VALID_TRANSITIONS.get(current, set())
+        if new_state not in valid_targets:
+            raise RuntimeError(f"Invalid session state transition: {current.name} -> {new_state.name}")
+        logger.debug("Session state transition: %s -> %s", current.name, new_state.name)
+        self._initialization_state = new_state
+
     @property
     def client_params(self) -> types.InitializeRequestParams | None:
         return self._client_params  # pragma: no cover
@@ -160,11 +240,34 @@ async def _receive_loop(self) -> None:
         async with self._incoming_message_stream_writer:
             await super()._receive_loop()
 
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool | None:
+        """Clean up the session with proper state transitions."""
+        try:
+            if self._initialization_state not in (
+                InitializationState.Closed,
+                InitializationState.Closing,
+            ):
+                self._transition_state(InitializationState.Closing)
+        except RuntimeError:
+            logger.debug("Could not transition to Closing from %s", self._initialization_state.name)
+        try:
+            return await super().__aexit__(exc_type, exc_val, exc_tb)
+        finally:
+            try:
+                self._transition_state(InitializationState.Closed)
+            except RuntimeError:
+                logger.debug("Could not transition to Closed from %s", self._initialization_state.name)
+
     async def _received_request(self, responder: RequestResponder[types.ClientRequest, types.ServerResult]):
         match responder.request.root:
             case types.InitializeRequest(params=params):
                 requested_version = params.protocolVersion
-                self._initialization_state = InitializationState.Initializing
+                self._transition_state(InitializationState.Initializing)
                 self._client_params = params
                 with responder:
                     await responder.respond(
@@ -184,22 +287,24 @@ async def _received_request(self, responder: RequestResponder[types.ClientReques
                             )
                         )
                     )
-                self._initialization_state = InitializationState.Initialized
+                self._transition_state(InitializationState.Initialized)
             case types.PingRequest():
                 # Ping requests are allowed at any time
                 pass
             case _:
-                if self._initialization_state != InitializationState.Initialized:
+                if not self.is_initialized:
                     raise RuntimeError("Received request before initialization was complete")
 
     async def _received_notification(self, notification: types.ClientNotification) -> None:
         # Need this to avoid ASYNC910
         await anyio.lowlevel.checkpoint()
         match notification.root:
             case types.InitializedNotification():
-                self._initialization_state = InitializationState.Initialized
+                # Only transition if not already initialized (e.g. stateless mode)
+                if self._initialization_state == InitializationState.Initializing:
+                    self._transition_state(InitializationState.Initialized)
             case _:
-                if self._initialization_state != InitializationState.Initialized:  # pragma: no cover
+                if not self.is_initialized:  # pragma: no cover
                     raise RuntimeError("Received notification before initialization was complete")
 
     async def send_log_message(