modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: Pass through and expose additional parameters in ClientSessionGroup.call_tool and .connect_to_server (#1576)

Browse source
This commit is contained in:
inaku 2025-11-16 23:57:43 +08:00 • committed by GitHub
commit a357380cfa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 101 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

86
src/mcp/client/session_group.py
View file

@ -11,20 +11,23 @@ hook.
import contextlib
import logging
from collections.abc import Callable
from dataclasses import dataclass
from datetime import timedelta
from types import TracebackType
from typing import Any, TypeAlias
from typing import Any, TypeAlias, overload
import anyio
from pydantic import BaseModel
from typing_extensions import Self
from typing_extensions import Self, deprecated
import mcp
from mcp import types
from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
from mcp.client.sse import sse_client
from mcp.client.stdio import StdioServerParameters
from mcp.client.streamable_http import streamablehttp_client
from mcp.shared.exceptions import McpError
from mcp.shared.session import ProgressFnT
class SseServerParameters(BaseModel):
@ -65,6 +68,21 @@ class StreamableHttpParameters(BaseModel):
ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters
# Use dataclass instead of pydantic BaseModel
# because pydantic BaseModel cannot handle Protocol fields.
@dataclass
class ClientSessionParameters:
"""Parameters for establishing a client session to an MCP server."""
read_timeout_seconds: timedelta | 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
client_info: types.Implementation | None = None
class ClientSessionGroup:
"""Client for managing connections to multiple MCP servers.
@ -172,11 +190,49 @@ class ClientSessionGroup:
"""Returns the tools as a dictionary of names to tools."""
return self._tools
async def call_tool(self, name: str, args: dict[str, Any]) -> types.CallToolResult:
@overload
async def call_tool(
self,
name: str,
arguments: dict[str, Any],
read_timeout_seconds: timedelta | None = None,
progress_callback: ProgressFnT | None = None,
*,
meta: dict[str, Any] | None = None,
) -> types.CallToolResult: ...
@overload
@deprecated("The 'args' parameter is deprecated. Use 'arguments' instead.")
async def call_tool(
self,
name: str,
*,
args: dict[str, Any],
read_timeout_seconds: timedelta | None = None,
progress_callback: ProgressFnT | None = None,
meta: dict[str, Any] | None = None,
) -> types.CallToolResult: ...
async def call_tool(
self,
name: str,
arguments: dict[str, Any] | None = None,
read_timeout_seconds: timedelta | None = None,
progress_callback: ProgressFnT | None = None,
*,
meta: dict[str, Any] | None = None,
args: dict[str, Any] | None = None,
) -> types.CallToolResult:
"""Executes a tool given its name and arguments."""
session = self._tool_to_session[name]
session_tool_name = self.tools[name].name
return await session.call_tool(session_tool_name, args)
return await session.call_tool(
session_tool_name,
arguments if args is None else args,
read_timeout_seconds=read_timeout_seconds,
progress_callback=progress_callback,
meta=meta,
)
async def disconnect_from_server(self, session: mcp.ClientSession) -> None:
"""Disconnects from a single MCP server."""
@ -225,13 +281,16 @@ class ClientSessionGroup:
async def connect_to_server(
self,
server_params: ServerParameters,
session_params: ClientSessionParameters | None = None,
) -> mcp.ClientSession:
"""Connects to a single MCP server."""
server_info, session = await self._establish_session(server_params)
server_info, session = await self._establish_session(server_params, session_params or ClientSessionParameters())
return await self.connect_with_session(server_info, session)
async def _establish_session(
self, server_params: ServerParameters
self,
server_params: ServerParameters,
session_params: ClientSessionParameters,
) -> tuple[types.Implementation, mcp.ClientSession]:
"""Establish a client session to an MCP server."""
@ -259,7 +318,20 @@ class ClientSessionGroup:
)
read, write, _ = await session_stack.enter_async_context(client)
session = await session_stack.enter_async_context(mcp.ClientSession(read, write))
session = await session_stack.enter_async_context(
mcp.ClientSession(
read,
write,
read_timeout_seconds=session_params.read_timeout_seconds,
sampling_callback=session_params.sampling_callback,
elicitation_callback=session_params.elicitation_callback,
list_roots_callback=session_params.list_roots_callback,
logging_callback=session_params.logging_callback,
message_handler=session_params.message_handler,
client_info=session_params.client_info,
)
)
result = await session.initialize()
# Session successfully initialized.

26
tests/client/test_session_group.py
View file

@ -5,7 +5,12 @@ import pytest
import mcp
from mcp import types
from mcp.client.session_group import ClientSessionGroup, SseServerParameters, StreamableHttpParameters
from mcp.client.session_group import (
ClientSessionGroup,
ClientSessionParameters,
SseServerParameters,
StreamableHttpParameters,
)
from mcp.client.stdio import StdioServerParameters
from mcp.shared.exceptions import McpError
@ -62,7 +67,7 @@ class TestClientSessionGroup:
# --- Test Execution ---
result = await mcp_session_group.call_tool(
name="server1-my_tool",
args={
arguments={
"name": "value1",
"args": {},
},
@ -73,6 +78,9 @@ class TestClientSessionGroup:
mock_session.call_tool.assert_called_once_with(
"my_tool",
{"name": "value1", "args": {}},
read_timeout_seconds=None,
progress_callback=None,
meta=None,
)
async def test_connect_to_server(self, mock_exit_stack: contextlib.AsyncExitStack):
@ -329,7 +337,7 @@ class TestClientSessionGroup:
(
returned_server_info,
returned_session,
) = await group._establish_session(server_params_instance)
) = await group._establish_session(server_params_instance, ClientSessionParameters())
# --- Assertions ---
# 1. Assert the correct specific client function was called
@ -357,7 +365,17 @@ class TestClientSessionGroup:
mock_client_cm_instance.__aenter__.assert_awaited_once()
# 2. Assert ClientSession was called correctly
mock_ClientSession_class.assert_called_once_with(mock_read_stream, mock_write_stream)
mock_ClientSession_class.assert_called_once_with(
mock_read_stream,
mock_write_stream,
read_timeout_seconds=None,
sampling_callback=None,
elicitation_callback=None,
list_roots_callback=None,
logging_callback=None,
message_handler=None,
client_info=None,
)
mock_raw_session_cm.__aenter__.assert_awaited_once()
mock_entered_session.initialize.assert_awaited_once()