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

This commit is contained in:
ihrpr
2025-05-02 14:29:00 +01:00
committed by GitHub
parent 3978c6e1b9
commit da0cf22355
22 changed files with 286 additions and 173 deletions

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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()