mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-23 00:34:27 +01:00
Wrap JSONRPC messages with SessionMessage for metadata support (#590)
This commit is contained in:
@@ -10,6 +10,7 @@ from pydantic import TypeAdapter
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
from mcp.server.lowlevel.server import NotificationOptions, Server
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import (
|
||||
ClientCapabilities,
|
||||
Implementation,
|
||||
@@ -82,41 +83,49 @@ async def test_lowlevel_server_lifespan():
|
||||
clientInfo=Implementation(name="test-client", version="0.1.0"),
|
||||
)
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
response = await receive_stream2.receive()
|
||||
response = response.message
|
||||
|
||||
# Send initialized notification
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Call the tool to verify lifespan context
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Get response and verify
|
||||
response = await receive_stream2.receive()
|
||||
response = response.message
|
||||
assert response.root.result["content"][0]["text"] == "true"
|
||||
|
||||
# Cancel server task
|
||||
@@ -178,41 +187,49 @@ async def test_fastmcp_server_lifespan():
|
||||
clientInfo=Implementation(name="test-client", version="0.1.0"),
|
||||
)
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
response = await receive_stream2.receive()
|
||||
response = response.message
|
||||
|
||||
# Send initialized notification
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Call the tool to verify lifespan context
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Get response and verify
|
||||
response = await receive_stream2.receive()
|
||||
response = response.message
|
||||
assert response.root.result["content"][0]["text"] == "true"
|
||||
|
||||
# Cancel server task
|
||||
|
||||
@@ -8,10 +8,10 @@ from mcp.server import Server
|
||||
from mcp.server.lowlevel import NotificationOptions
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.types import (
|
||||
ClientResult,
|
||||
JSONRPCMessage,
|
||||
ServerNotification,
|
||||
ServerRequest,
|
||||
Tool,
|
||||
@@ -46,10 +46,10 @@ async def test_lowlevel_server_tool_annotations():
|
||||
]
|
||||
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
|
||||
JSONRPCMessage
|
||||
SessionMessage
|
||||
](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
|
||||
JSONRPCMessage
|
||||
SessionMessage
|
||||
](10)
|
||||
|
||||
# Message handler for client
|
||||
|
||||
@@ -7,11 +7,11 @@ from mcp.server import Server
|
||||
from mcp.server.lowlevel import NotificationOptions
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.types import (
|
||||
ClientNotification,
|
||||
InitializedNotification,
|
||||
JSONRPCMessage,
|
||||
PromptsCapability,
|
||||
ResourcesCapability,
|
||||
ServerCapabilities,
|
||||
@@ -21,10 +21,10 @@ from mcp.types import (
|
||||
@pytest.mark.anyio
|
||||
async def test_server_session_initialize():
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
|
||||
JSONRPCMessage
|
||||
SessionMessage
|
||||
](1)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
|
||||
JSONRPCMessage
|
||||
SessionMessage
|
||||
](1)
|
||||
|
||||
# Create a message handler to catch exceptions
|
||||
|
||||
@@ -4,6 +4,7 @@ import anyio
|
||||
import pytest
|
||||
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
|
||||
|
||||
|
||||
@@ -29,7 +30,7 @@ async def test_stdio_server():
|
||||
async for message in read_stream:
|
||||
if isinstance(message, Exception):
|
||||
raise message
|
||||
received_messages.append(message)
|
||||
received_messages.append(message.message)
|
||||
if len(received_messages) == 2:
|
||||
break
|
||||
|
||||
@@ -50,7 +51,8 @@ async def test_stdio_server():
|
||||
|
||||
async with write_stream:
|
||||
for response in responses:
|
||||
await write_stream.send(response)
|
||||
session_message = SessionMessage(response)
|
||||
await write_stream.send(session_message)
|
||||
|
||||
stdout.seek(0)
|
||||
output_lines = stdout.readlines()
|
||||
|
||||
Reference in New Issue
Block a user