feat: add from_session_state() classmethod to ClientSession

This classmethod creates a new ClientSession from a previously saved
SessionState, enabling distributed deployments.

Related: #2111

Author: peter-luminova

src/mcp/client/session.py

@@ -233,6 +233,67 @@ def get_session_state(self) -> SessionState:
             initialized_sent=self._initialized_sent,
         )
 
+    @classmethod
+    def from_session_state(
+        cls,
+        state: SessionState,
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
+        write_stream: MemoryObjectSendStream[SessionMessage],
+        client_info: types.Implementation | None = None,
+        read_timeout_seconds: float | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
+        list_roots_callback: ListRootsFnT | None = None,
+        logging_callback: LoggingFnT | None = None,
+        message_handler: MessageHandlerFnT | None = None,
+        *,
+        sampling_capabilities: types.SamplingCapability | None = None,
+        experimental_task_handlers: ExperimentalTaskHandlers | None = None,
+    ) -> ClientSession:
+        """Create a new ClientSession from a previously saved SessionState.
+
+        This restores session context from external storage, allowing
+        distributed instances to continue a session.
+
+        Args:
+            state: The SessionState to restore from
+            read_stream: The read stream for receiving messages
+            write_stream: The write stream for sending messages
+            client_info: Optional client info (defaults to DEFAULT_CLIENT_INFO)
+            read_timeout_seconds: Optional read timeout for this session
+
+        Returns:
+            A new ClientSession instance with the restored state
+        """
+        # Create session with default initialization
+        session = cls(
+            read_stream=read_stream,
+            write_stream=write_stream,
+            client_info=client_info,
+            read_timeout_seconds=read_timeout_seconds,
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+            list_roots_callback=list_roots_callback,
+            logging_callback=logging_callback,
+            message_handler=message_handler,
+            sampling_capabilities=sampling_capabilities,
+            experimental_task_handlers=experimental_task_handlers,
+        )
+
+        # Restore the state
+        session._session_id = state.session_id
+        session._request_id = state.next_request_id
+
+        if state.server_capabilities:
+            session._server_capabilities = types.ServerCapabilities.model_validate(state.server_capabilities)
+
+        if state.server_info:
+            session._server_info = types.Implementation.model_validate(state.server_info)
+
+        session._initialized_sent = state.initialized_sent
+
+        return session
+
     @property
     def experimental(self) -> ExperimentalClientFeatures:
         """Experimental APIs for tasks and other features.

tests/client/test_session.py

@@ -777,3 +777,55 @@ async def mock_server():
         json_str = state.model_dump_json()
         assert json_str is not None
 
+
+@pytest.mark.anyio
+async def test_client_session_from_state():
+    """Test that from_session_state() creates a valid session."""
+    from mcp.shared.session_state import SessionState
+
+    # Create a session state
+    state = SessionState(
+        session_id="test-session-from-state",
+        protocol_version=LATEST_PROTOCOL_VERSION,
+        next_request_id=5,
+        server_capabilities={
+            "tools": {},
+            "resources": {},
+            "prompts": None,
+            "logging": None,
+            "experimental": None,
+        },
+        server_info={"name": "test-server", "version": "1.0.0"},
+        initialized_sent=True,
+    )
+
+    # Create streams
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
+
+    # Create session from state
+    session = ClientSession.from_session_state(
+        state,
+        server_to_client_receive,
+        client_to_server_send,
+    )
+
+    # Verify the session was created with the correct state
+    assert session._session_id == "test-session-from-state"
+    assert session._request_id == 5  # Continues from saved state
+    assert session._server_capabilities is not None
+    assert session._initialized_sent is True
+    assert session._server_info is not None
+    assert session._server_info.name == "test-server"
+    assert session._server_info.version == "1.0.0"
+
+    # Clean up streams
+    async with (
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        pass
+
+