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

Wrap JSONRPC messages with SessionMessage for metadata support (#590)

Browse source
This commit is contained in:
ihrpr 2025-05-02 14:29:00 +01:00 • committed by GitHub
commit da0cf22355
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 286 additions and 173 deletions
Download patch file Download diff file Expand all files Collapse all files

73
tests/client/test_session.py
View file

@ -3,6 +3,7 @@ import pytest
import mcp.types as types
from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from mcp.types import (
LATEST_PROTOCOL_VERSION,
@ -24,10 +25,10 @@ from mcp.types import (
@pytest.mark.anyio
async def test_client_session_initialize():
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
JSONRPCMessage
SessionMessage
](1)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
JSONRPCMessage
SessionMessage
](1)
initialized_notification = None
@ -35,7 +36,8 @@ async def test_client_session_initialize():
async def mock_server():
nonlocal initialized_notification
jsonrpc_request = await client_to_server_receive.receive()
session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
@ -59,17 +61,20 @@ async def test_client_session_initialize():
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
),
SessionMessage(
JSONRPCMessage(
JSONRPCResponse(
jsonrpc="2.0",
id=jsonrpc_request.root.id,
result=result.model_dump(
by_alias=True, mode="json", exclude_none=True
),
)
)
)
)
jsonrpc_notification = await client_to_server_receive.receive()
session_notification = await client_to_server_receive.receive()
jsonrpc_notification = session_notification.message
assert isinstance(jsonrpc_notification.root, JSONRPCNotification)
initialized_notification = ClientNotification.model_validate(
jsonrpc_notification.model_dump(
@ -116,10 +121,10 @@ async def test_client_session_initialize():
@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
SessionMessage
](1)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
JSONRPCMessage
SessionMessage
](1)
custom_client_info = Implementation(name="test-client", version="1.2.3")
@ -128,7 +133,8 @@ async def test_client_session_custom_client_info():
async def mock_server():
nonlocal received_client_info
jsonrpc_request = await client_to_server_receive.receive()
session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
@ -146,13 +152,15 @@ async def test_client_session_custom_client_info():
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
),
SessionMessage(
JSONRPCMessage(
JSONRPCResponse(
jsonrpc="2.0",
id=jsonrpc_request.root.id,
result=result.model_dump(
by_alias=True, mode="json", exclude_none=True
),
)
)
)
)
@ -181,10 +189,10 @@ async def test_client_session_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
SessionMessage
](1)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
JSONRPCMessage
SessionMessage
](1)
received_client_info = None
@ -192,7 +200,8 @@ async def test_client_session_default_client_info():
async def mock_server():
nonlocal received_client_info
jsonrpc_request = await client_to_server_receive.receive()
session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
@ -210,13 +219,15 @@ async def test_client_session_default_client_info():
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
),
SessionMessage(
JSONRPCMessage(
JSONRPCResponse(
jsonrpc="2.0",
id=jsonrpc_request.root.id,
result=result.model_dump(
by_alias=True, mode="json", exclude_none=True
),
)
)
)
)

6
tests/client/test_stdio.py
View file

@ -3,6 +3,7 @@ import shutil
import pytest
from mcp.client.stdio import StdioServerParameters, stdio_client
from mcp.shared.message import SessionMessage
from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
tee: str = shutil.which("tee") # type: ignore
@ -22,7 +23,8 @@ async def test_stdio_client():
async with write_stream:
for message in messages:
await write_stream.send(message)
session_message = SessionMessage(message)
await write_stream.send(session_message)
read_messages = []
async with read_stream:
@ -30,7 +32,7 @@ async def test_stdio_client():
if isinstance(message, Exception):
raise message
read_messages.append(message)
read_messages.append(message.message)
if len(read_messages) == 2:
break