feat: add get_session_state() method to ClientSession

This method extracts serializable session state that can be stored
in external storage for distributed deployments.

Related: #2111

Author: peter-luminova

src/mcp/client/session.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import uuid
 from typing import Any, Protocol
 
 import anyio.lowlevel
@@ -13,6 +14,7 @@
 from mcp.shared._context import RequestContext
 from mcp.shared.message import SessionMessage
 from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder
+from mcp.shared.session_state import SessionState
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 from mcp.types._types import RequestParamsMeta
 
@@ -132,6 +134,9 @@ def __init__(
         self._message_handler = message_handler or _default_message_handler
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
         self._server_capabilities: types.ServerCapabilities | None = None
+        self._server_info: types.Implementation | None = None
+        self._initialized_sent: bool = False
+        self._session_id: str = str(uuid.uuid4())
         self._experimental_features: ExperimentalClientFeatures | None = None
 
         # Experimental: Task handlers (use defaults if not provided)
@@ -186,8 +191,10 @@ async def initialize(self) -> types.InitializeResult:
             raise RuntimeError(f"Unsupported protocol version from the server: {result.protocol_version}")
 
         self._server_capabilities = result.capabilities
+        self._server_info = result.server_info
 
         await self.send_notification(types.InitializedNotification())
+        self._initialized_sent = True
 
         return result
 
@@ -198,6 +205,34 @@ def get_server_capabilities(self) -> types.ServerCapabilities | None:
         """
         return self._server_capabilities
 
+    def get_session_state(self) -> SessionState:
+        """Extract a serializable snapshot of the current session state.
+
+        This allows the session state to be stored in external storage
+        (Redis, database, etc.) for distributed deployments.
+
+        Returns:
+            A SessionState object containing the serializable state
+        """
+        from mcp.shared.version import LATEST_PROTOCOL_VERSION
+
+        return SessionState(
+            session_id=self._session_id,
+            protocol_version=LATEST_PROTOCOL_VERSION,
+            next_request_id=self._request_id,
+            server_capabilities=(
+                self._server_capabilities.model_dump(by_alias=True, mode="json", exclude_none=True)
+                if self._server_capabilities
+                else None
+            ),
+            server_info=(
+                self._server_info.model_dump(by_alias=True, mode="json", exclude_none=True)
+                if self._server_info
+                else None
+            ),
+            initialized_sent=self._initialized_sent,
+        )
+
     @property
     def experimental(self) -> ExperimentalClientFeatures:
         """Experimental APIs for tasks and other features.

tests/client/test_session.py

@@ -706,3 +706,74 @@ async def mock_server():
         await session.initialize()
 
         await session.call_tool(name=mocked_tool.name, arguments={"foo": "bar"}, meta=meta)
+
+
+@pytest.mark.anyio
+async def test_client_session_get_state():
+    """Test that get_session_state() returns a valid SessionState."""
+    from mcp.shared.session_state import SessionState
+
+    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)
+
+    async def mock_server():
+        session_message = await client_to_server_receive.receive()
+        jsonrpc_request = session_message.message
+        assert isinstance(jsonrpc_request, JSONRPCRequest)
+
+        result = InitializeResult(
+            protocol_version=LATEST_PROTOCOL_VERSION,
+            capabilities=ServerCapabilities(
+                logging=None,
+                resources=None,
+                tools=None,
+                experimental=None,
+                prompts=None,
+            ),
+            server_info=Implementation(name="mock-server", version="0.1.0"),
+        )
+
+        async with server_to_client_send:
+            await server_to_client_send.send(
+                SessionMessage(
+                    JSONRPCResponse(
+                        jsonrpc="2.0",
+                        id=jsonrpc_request.id,
+                        result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
+                    )
+                )
+            )
+            await client_to_server_receive.receive()
+
+    async with (
+        ClientSession(
+            server_to_client_receive,
+            client_to_server_send,
+        ) as session,
+        anyio.create_task_group() as tg,
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        tg.start_soon(mock_server)
+
+        # Initialize the session
+        await session.initialize()
+
+        # Get session state
+        state = session.get_session_state()
+
+        # Verify the state
+        assert isinstance(state, SessionState)
+        assert state.session_id is not None
+        assert state.protocol_version == LATEST_PROTOCOL_VERSION
+        assert state.next_request_id == 1  # After initialize request
+        assert state.server_capabilities is not None
+        assert state.server_info is not None
+        assert state.initialized_sent is True
+
+        # Verify it's serializable
+        json_str = state.model_dump_json()
+        assert json_str is not None
+