From a3d5a83e9435f077c5a5fea7f3249245b098e4e7 Mon Sep 17 00:00:00 2001 From: Shivam Date: Thu, 5 Mar 2026 01:12:35 +0530 Subject: [PATCH] feat: store server_info on ClientSession during initialization Store the server_info from InitializeResult on the session (alongside the already-stored server_capabilities) and expose it via get_server_info(), so callers no longer need to capture the return value of initialize() to access server metadata. Closes #1018 --- src/mcp/client/session.py | 9 +++++++ tests/client/test_session.py | 51 ++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index a0ca751bd..53615aa13 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -132,6 +132,7 @@ 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._experimental_features: ExperimentalClientFeatures | None = None # Experimental: Task handlers (use defaults if not provided) @@ -186,6 +187,7 @@ 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()) @@ -198,6 +200,13 @@ def get_server_capabilities(self) -> types.ServerCapabilities | None: """ return self._server_capabilities + def get_server_info(self) -> types.Implementation | None: + """Return the server info received during initialization. + + Returns None if the session has not been initialized yet. + """ + return self._server_info + @property def experimental(self) -> ExperimentalClientFeatures: """Experimental APIs for tasks and other features. diff --git a/tests/client/test_session.py b/tests/client/test_session.py index d6d13e273..002889733 100644 --- a/tests/client/test_session.py +++ b/tests/client/test_session.py @@ -607,6 +607,57 @@ async def mock_server(): assert capabilities.tools.list_changed is False +@pytest.mark.anyio +async def test_get_server_info(): + """Test that get_server_info returns None before init and server info after""" + 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) + + expected_server_info = Implementation(name="test-server", version="2.0.0") + + 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(), + server_info=expected_server_info, + ) + + 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, + ): + assert session.get_server_info() is None + + tg.start_soon(mock_server) + await session.initialize() + + server_info = session.get_server_info() + assert server_info is not None + assert server_info == expected_server_info + assert server_info.name == "test-server" + assert server_info.version == "2.0.0" + + @pytest.mark.anyio @pytest.mark.parametrize(argnames="meta", argvalues=[None, {"toolMeta": "value"}]) async def test_client_tool_call_with_meta(meta: RequestParamsMeta | None):