mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 14:54:24 +01:00
Support custom client info throughout client APIs (#474)
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -38,9 +38,13 @@ async def message_handler(
|
|||||||
async def run_session(
|
async def run_session(
|
||||||
read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
|
read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
|
||||||
write_stream: MemoryObjectSendStream[JSONRPCMessage],
|
write_stream: MemoryObjectSendStream[JSONRPCMessage],
|
||||||
|
client_info: types.Implementation | None = None,
|
||||||
):
|
):
|
||||||
async with ClientSession(
|
async with ClientSession(
|
||||||
read_stream, write_stream, message_handler=message_handler
|
read_stream,
|
||||||
|
write_stream,
|
||||||
|
message_handler=message_handler,
|
||||||
|
client_info=client_info,
|
||||||
) as session:
|
) as session:
|
||||||
logger.info("Initializing session")
|
logger.info("Initializing session")
|
||||||
await session.initialize()
|
await session.initialize()
|
||||||
|
|||||||
@@ -10,6 +10,8 @@ from mcp.shared.context import RequestContext
|
|||||||
from mcp.shared.session import BaseSession, RequestResponder
|
from mcp.shared.session import BaseSession, RequestResponder
|
||||||
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
||||||
|
|
||||||
|
DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
|
||||||
|
|
||||||
|
|
||||||
class SamplingFnT(Protocol):
|
class SamplingFnT(Protocol):
|
||||||
async def __call__(
|
async def __call__(
|
||||||
@@ -97,6 +99,7 @@ class ClientSession(
|
|||||||
list_roots_callback: ListRootsFnT | None = None,
|
list_roots_callback: ListRootsFnT | None = None,
|
||||||
logging_callback: LoggingFnT | None = None,
|
logging_callback: LoggingFnT | None = None,
|
||||||
message_handler: MessageHandlerFnT | None = None,
|
message_handler: MessageHandlerFnT | None = None,
|
||||||
|
client_info: types.Implementation | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
super().__init__(
|
super().__init__(
|
||||||
read_stream,
|
read_stream,
|
||||||
@@ -105,6 +108,7 @@ class ClientSession(
|
|||||||
types.ServerNotification,
|
types.ServerNotification,
|
||||||
read_timeout_seconds=read_timeout_seconds,
|
read_timeout_seconds=read_timeout_seconds,
|
||||||
)
|
)
|
||||||
|
self._client_info = client_info or DEFAULT_CLIENT_INFO
|
||||||
self._sampling_callback = sampling_callback or _default_sampling_callback
|
self._sampling_callback = sampling_callback or _default_sampling_callback
|
||||||
self._list_roots_callback = list_roots_callback or _default_list_roots_callback
|
self._list_roots_callback = list_roots_callback or _default_list_roots_callback
|
||||||
self._logging_callback = logging_callback or _default_logging_callback
|
self._logging_callback = logging_callback or _default_logging_callback
|
||||||
@@ -130,7 +134,7 @@ class ClientSession(
|
|||||||
experimental=None,
|
experimental=None,
|
||||||
roots=roots,
|
roots=roots,
|
||||||
),
|
),
|
||||||
clientInfo=types.Implementation(name="mcp", version="0.1.0"),
|
clientInfo=self._client_info,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
),
|
),
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ from typing import Any
|
|||||||
import anyio
|
import anyio
|
||||||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||||
|
|
||||||
|
import mcp.types as types
|
||||||
from mcp.client.session import (
|
from mcp.client.session import (
|
||||||
ClientSession,
|
ClientSession,
|
||||||
ListRootsFnT,
|
ListRootsFnT,
|
||||||
@@ -65,6 +66,7 @@ async def create_connected_server_and_client_session(
|
|||||||
list_roots_callback: ListRootsFnT | None = None,
|
list_roots_callback: ListRootsFnT | None = None,
|
||||||
logging_callback: LoggingFnT | None = None,
|
logging_callback: LoggingFnT | None = None,
|
||||||
message_handler: MessageHandlerFnT | None = None,
|
message_handler: MessageHandlerFnT | None = None,
|
||||||
|
client_info: types.Implementation | None = None,
|
||||||
raise_exceptions: bool = False,
|
raise_exceptions: bool = False,
|
||||||
) -> AsyncGenerator[ClientSession, None]:
|
) -> AsyncGenerator[ClientSession, None]:
|
||||||
"""Creates a ClientSession that is connected to a running MCP server."""
|
"""Creates a ClientSession that is connected to a running MCP server."""
|
||||||
@@ -95,6 +97,7 @@ async def create_connected_server_and_client_session(
|
|||||||
list_roots_callback=list_roots_callback,
|
list_roots_callback=list_roots_callback,
|
||||||
logging_callback=logging_callback,
|
logging_callback=logging_callback,
|
||||||
message_handler=message_handler,
|
message_handler=message_handler,
|
||||||
|
client_info=client_info,
|
||||||
) as client_session:
|
) as client_session:
|
||||||
await client_session.initialize()
|
await client_session.initialize()
|
||||||
yield client_session
|
yield client_session
|
||||||
|
|||||||
@@ -2,7 +2,7 @@ import anyio
|
|||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
import mcp.types as types
|
import mcp.types as types
|
||||||
from mcp.client.session import ClientSession
|
from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
|
||||||
from mcp.shared.session import RequestResponder
|
from mcp.shared.session import RequestResponder
|
||||||
from mcp.types import (
|
from mcp.types import (
|
||||||
LATEST_PROTOCOL_VERSION,
|
LATEST_PROTOCOL_VERSION,
|
||||||
@@ -111,3 +111,131 @@ async def test_client_session_initialize():
|
|||||||
# Check that the client sent the initialized notification
|
# Check that the client sent the initialized notification
|
||||||
assert initialized_notification
|
assert initialized_notification
|
||||||
assert isinstance(initialized_notification.root, InitializedNotification)
|
assert isinstance(initialized_notification.root, InitializedNotification)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_client_session_custom_client_info():
|
||||||
|
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
|
||||||
|
JSONRPCMessage
|
||||||
|
](1)
|
||||||
|
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
|
||||||
|
JSONRPCMessage
|
||||||
|
](1)
|
||||||
|
|
||||||
|
custom_client_info = Implementation(name="test-client", version="1.2.3")
|
||||||
|
received_client_info = None
|
||||||
|
|
||||||
|
async def mock_server():
|
||||||
|
nonlocal received_client_info
|
||||||
|
|
||||||
|
jsonrpc_request = await client_to_server_receive.receive()
|
||||||
|
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||||
|
request = ClientRequest.model_validate(
|
||||||
|
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||||
|
)
|
||||||
|
assert isinstance(request.root, InitializeRequest)
|
||||||
|
received_client_info = request.root.params.clientInfo
|
||||||
|
|
||||||
|
result = ServerResult(
|
||||||
|
InitializeResult(
|
||||||
|
protocolVersion=LATEST_PROTOCOL_VERSION,
|
||||||
|
capabilities=ServerCapabilities(),
|
||||||
|
serverInfo=Implementation(name="mock-server", version="0.1.0"),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
async with server_to_client_send:
|
||||||
|
await server_to_client_send.send(
|
||||||
|
JSONRPCMessage(
|
||||||
|
JSONRPCResponse(
|
||||||
|
jsonrpc="2.0",
|
||||||
|
id=jsonrpc_request.root.id,
|
||||||
|
result=result.model_dump(
|
||||||
|
by_alias=True, mode="json", exclude_none=True
|
||||||
|
),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
# Receive initialized notification
|
||||||
|
await client_to_server_receive.receive()
|
||||||
|
|
||||||
|
async with (
|
||||||
|
ClientSession(
|
||||||
|
server_to_client_receive,
|
||||||
|
client_to_server_send,
|
||||||
|
client_info=custom_client_info,
|
||||||
|
) 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)
|
||||||
|
await session.initialize()
|
||||||
|
|
||||||
|
# Assert that the custom client info was sent
|
||||||
|
assert received_client_info == custom_client_info
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_client_session_default_client_info():
|
||||||
|
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
|
||||||
|
JSONRPCMessage
|
||||||
|
](1)
|
||||||
|
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
|
||||||
|
JSONRPCMessage
|
||||||
|
](1)
|
||||||
|
|
||||||
|
received_client_info = None
|
||||||
|
|
||||||
|
async def mock_server():
|
||||||
|
nonlocal received_client_info
|
||||||
|
|
||||||
|
jsonrpc_request = await client_to_server_receive.receive()
|
||||||
|
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||||
|
request = ClientRequest.model_validate(
|
||||||
|
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||||
|
)
|
||||||
|
assert isinstance(request.root, InitializeRequest)
|
||||||
|
received_client_info = request.root.params.clientInfo
|
||||||
|
|
||||||
|
result = ServerResult(
|
||||||
|
InitializeResult(
|
||||||
|
protocolVersion=LATEST_PROTOCOL_VERSION,
|
||||||
|
capabilities=ServerCapabilities(),
|
||||||
|
serverInfo=Implementation(name="mock-server", version="0.1.0"),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
async with server_to_client_send:
|
||||||
|
await server_to_client_send.send(
|
||||||
|
JSONRPCMessage(
|
||||||
|
JSONRPCResponse(
|
||||||
|
jsonrpc="2.0",
|
||||||
|
id=jsonrpc_request.root.id,
|
||||||
|
result=result.model_dump(
|
||||||
|
by_alias=True, mode="json", exclude_none=True
|
||||||
|
),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
# Receive initialized notification
|
||||||
|
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)
|
||||||
|
await session.initialize()
|
||||||
|
|
||||||
|
# Assert that the default client info was sent
|
||||||
|
assert received_client_info == DEFAULT_CLIENT_INFO
|
||||||
|
|||||||
Reference in New Issue
Block a user