Skip to content

Commit a3d5a83

Browse files
ShivamShivam
authored andcommitted
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
1 parent 528abfa commit a3d5a83

File tree

2 files changed

+60
-0
lines changed

2 files changed

+60
-0
lines changed

src/mcp/client/session.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ def __init__(
132132
self._message_handler = message_handler or _default_message_handler
133133
self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
134134
self._server_capabilities: types.ServerCapabilities | None = None
135+
self._server_info: types.Implementation | None = None
135136
self._experimental_features: ExperimentalClientFeatures | None = None
136137

137138
# Experimental: Task handlers (use defaults if not provided)
@@ -186,6 +187,7 @@ async def initialize(self) -> types.InitializeResult:
186187
raise RuntimeError(f"Unsupported protocol version from the server: {result.protocol_version}")
187188

188189
self._server_capabilities = result.capabilities
190+
self._server_info = result.server_info
189191

190192
await self.send_notification(types.InitializedNotification())
191193

@@ -198,6 +200,13 @@ def get_server_capabilities(self) -> types.ServerCapabilities | None:
198200
"""
199201
return self._server_capabilities
200202

203+
def get_server_info(self) -> types.Implementation | None:
204+
"""Return the server info received during initialization.
205+
206+
Returns None if the session has not been initialized yet.
207+
"""
208+
return self._server_info
209+
201210
@property
202211
def experimental(self) -> ExperimentalClientFeatures:
203212
"""Experimental APIs for tasks and other features.

tests/client/test_session.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -607,6 +607,57 @@ async def mock_server():
607607
assert capabilities.tools.list_changed is False
608608

609609

610+
@pytest.mark.anyio
611+
async def test_get_server_info():
612+
"""Test that get_server_info returns None before init and server info after"""
613+
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
614+
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
615+
616+
expected_server_info = Implementation(name="test-server", version="2.0.0")
617+
618+
async def mock_server():
619+
session_message = await client_to_server_receive.receive()
620+
jsonrpc_request = session_message.message
621+
assert isinstance(jsonrpc_request, JSONRPCRequest)
622+
623+
result = InitializeResult(
624+
protocol_version=LATEST_PROTOCOL_VERSION,
625+
capabilities=ServerCapabilities(),
626+
server_info=expected_server_info,
627+
)
628+
629+
async with server_to_client_send:
630+
await server_to_client_send.send(
631+
SessionMessage(
632+
JSONRPCResponse(
633+
jsonrpc="2.0",
634+
id=jsonrpc_request.id,
635+
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
636+
)
637+
)
638+
)
639+
await client_to_server_receive.receive()
640+
641+
async with (
642+
ClientSession(server_to_client_receive, client_to_server_send) as session,
643+
anyio.create_task_group() as tg,
644+
client_to_server_send,
645+
client_to_server_receive,
646+
server_to_client_send,
647+
server_to_client_receive,
648+
):
649+
assert session.get_server_info() is None
650+
651+
tg.start_soon(mock_server)
652+
await session.initialize()
653+
654+
server_info = session.get_server_info()
655+
assert server_info is not None
656+
assert server_info == expected_server_info
657+
assert server_info.name == "test-server"
658+
assert server_info.version == "2.0.0"
659+
660+
610661
@pytest.mark.anyio
611662
@pytest.mark.parametrize(argnames="meta", argvalues=[None, {"toolMeta": "value"}])
612663
async def test_client_tool_call_with_meta(meta: RequestParamsMeta | None):

0 commit comments

Comments
 (0)