Drop RootModel from JSONRPCMessage (#1908)
This commit is contained in:
parent
dcc9b4f8e0
commit
f4672c5084
4 changed files with 314 additions and 476 deletions
|
|
@ -82,13 +82,11 @@ async def test_lowlevel_server_lifespan():
|
|||
)
|
||||
await send_stream1.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -96,27 +94,16 @@ async def test_lowlevel_server_lifespan():
|
|||
response = response.message
|
||||
|
||||
# Send initialized notification
|
||||
await send_stream1.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await send_stream1.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized")))
|
||||
|
||||
# Call the tool to verify lifespan context
|
||||
await send_stream1.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -125,8 +112,8 @@ async def test_lowlevel_server_lifespan():
|
|||
response = await receive_stream2.receive()
|
||||
response = response.message
|
||||
assert isinstance(response, JSONRPCMessage)
|
||||
assert isinstance(response.root, JSONRPCResponse)
|
||||
assert response.root.result["content"][0]["text"] == "true"
|
||||
assert isinstance(response, JSONRPCResponse)
|
||||
assert response.result["content"][0]["text"] == "true"
|
||||
|
||||
# Cancel server task
|
||||
tg.cancel_scope.cancel()
|
||||
|
|
@ -162,13 +149,7 @@ async def test_fastmcp_server_lifespan():
|
|||
return True
|
||||
|
||||
# Run server in background task
|
||||
async with (
|
||||
anyio.create_task_group() as tg,
|
||||
send_stream1,
|
||||
receive_stream1,
|
||||
send_stream2,
|
||||
receive_stream2,
|
||||
):
|
||||
async with anyio.create_task_group() as tg, send_stream1, receive_stream1, send_stream2, receive_stream2:
|
||||
|
||||
async def run_server():
|
||||
await server._mcp_server.run(
|
||||
|
|
@ -188,13 +169,11 @@ async def test_fastmcp_server_lifespan():
|
|||
)
|
||||
await send_stream1.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -202,27 +181,16 @@ async def test_fastmcp_server_lifespan():
|
|||
response = response.message
|
||||
|
||||
# Send initialized notification
|
||||
await send_stream1.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await send_stream1.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized")))
|
||||
|
||||
# Call the tool to verify lifespan context
|
||||
await send_stream1.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -231,8 +199,8 @@ async def test_fastmcp_server_lifespan():
|
|||
response = await receive_stream2.receive()
|
||||
response = response.message
|
||||
assert isinstance(response, JSONRPCMessage)
|
||||
assert isinstance(response.root, JSONRPCResponse)
|
||||
assert response.root.result["content"][0]["text"] == "true"
|
||||
assert isinstance(response, JSONRPCResponse)
|
||||
assert response.result["content"][0]["text"] == "true"
|
||||
|
||||
# Cancel server task
|
||||
tg.cancel_scope.cancel()
|
||||
|
|
|
|||
|
|
@ -169,25 +169,23 @@ async def test_server_session_initialize_with_older_protocol_version():
|
|||
# Send initialization request with older protocol version (2024-11-05)
|
||||
await client_to_server_send.send(
|
||||
SessionMessage(
|
||||
types.JSONRPCMessage(
|
||||
types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=types.InitializeRequestParams(
|
||||
protocol_version="2024-11-05",
|
||||
capabilities=types.ClientCapabilities(),
|
||||
client_info=types.Implementation(name="test-client", version="1.0.0"),
|
||||
).model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=types.InitializeRequestParams(
|
||||
protocol_version="2024-11-05",
|
||||
capabilities=types.ClientCapabilities(),
|
||||
client_info=types.Implementation(name="test-client", version="1.0.0"),
|
||||
).model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for the initialize response
|
||||
init_response_message = await server_to_client_receive.receive()
|
||||
assert isinstance(init_response_message.message.root, types.JSONRPCResponse)
|
||||
result_data = init_response_message.message.root.result
|
||||
assert isinstance(init_response_message.message, types.JSONRPCResponse)
|
||||
result_data = init_response_message.message.result
|
||||
init_result = types.InitializeResult.model_validate(result_data)
|
||||
|
||||
# Check that the server responded with the requested protocol version
|
||||
|
|
@ -196,14 +194,7 @@ async def test_server_session_initialize_with_older_protocol_version():
|
|||
|
||||
# Send initialized notification
|
||||
await client_to_server_send.send(
|
||||
SessionMessage(
|
||||
types.JSONRPCMessage(
|
||||
types.JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
)
|
||||
)
|
||||
SessionMessage(types.JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized"))
|
||||
)
|
||||
|
||||
async with (
|
||||
|
|
@ -256,24 +247,14 @@ async def test_ping_request_before_initialization():
|
|||
nonlocal ping_response_received, ping_response_id
|
||||
|
||||
# Send ping request before any initialization
|
||||
await client_to_server_send.send(
|
||||
SessionMessage(
|
||||
types.JSONRPCMessage(
|
||||
types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=42,
|
||||
method="ping",
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await client_to_server_send.send(SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=42, method="ping")))
|
||||
|
||||
# Wait for the ping response
|
||||
ping_response_message = await server_to_client_receive.receive()
|
||||
assert isinstance(ping_response_message.message.root, types.JSONRPCResponse)
|
||||
assert isinstance(ping_response_message.message, types.JSONRPCResponse)
|
||||
|
||||
ping_response_received = True
|
||||
ping_response_id = ping_response_message.message.root.id
|
||||
ping_response_id = ping_response_message.message.id
|
||||
|
||||
async with (
|
||||
client_to_server_send,
|
||||
|
|
@ -493,22 +474,14 @@ async def test_other_requests_blocked_before_initialization():
|
|||
|
||||
# Try to send a non-ping request before initialization
|
||||
await client_to_server_send.send(
|
||||
SessionMessage(
|
||||
types.JSONRPCMessage(
|
||||
types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="prompts/list",
|
||||
)
|
||||
)
|
||||
)
|
||||
SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=1, method="prompts/list"))
|
||||
)
|
||||
|
||||
# Wait for the error response
|
||||
error_message = await server_to_client_receive.receive()
|
||||
if isinstance(error_message.message.root, types.JSONRPCError): # pragma: no branch
|
||||
if isinstance(error_message.message, types.JSONRPCError): # pragma: no branch
|
||||
error_response_received = True
|
||||
error_code = error_message.message.root.error.code
|
||||
error_code = error_message.message.error.code
|
||||
|
||||
async with (
|
||||
client_to_server_send,
|
||||
|
|
|
|||
|
|
@ -87,54 +87,35 @@ async def test_request_immediately_after_initialize_response():
|
|||
# Step 1: Send InitializeRequest
|
||||
await client_to_server_send.send(
|
||||
SessionMessage(
|
||||
types.JSONRPCMessage(
|
||||
types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=types.InitializeRequestParams(
|
||||
protocol_version=types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=types.ClientCapabilities(),
|
||||
client_info=types.Implementation(name="test-client", version="1.0.0"),
|
||||
).model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=types.InitializeRequestParams(
|
||||
protocol_version=types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=types.ClientCapabilities(),
|
||||
client_info=types.Implementation(name="test-client", version="1.0.0"),
|
||||
).model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Step 2: Wait for InitializeResult
|
||||
init_msg = await server_to_client_receive.receive()
|
||||
assert isinstance(init_msg.message.root, types.JSONRPCResponse)
|
||||
assert isinstance(init_msg.message, types.JSONRPCResponse)
|
||||
|
||||
# Step 3: Immediately send tools/list BEFORE InitializedNotification
|
||||
# This is the race condition scenario
|
||||
await client_to_server_send.send(
|
||||
SessionMessage(
|
||||
types.JSONRPCMessage(
|
||||
types.JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/list",
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await client_to_server_send.send(SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=2, method="tools/list")))
|
||||
|
||||
# Step 4: Check the response
|
||||
tools_msg = await server_to_client_receive.receive()
|
||||
if isinstance(tools_msg.message.root, types.JSONRPCError): # pragma: no cover
|
||||
error_received = tools_msg.message.root.error.message
|
||||
if isinstance(tools_msg.message, types.JSONRPCError): # pragma: no cover
|
||||
error_received = tools_msg.message.error.message
|
||||
|
||||
# Step 5: Send InitializedNotification
|
||||
await client_to_server_send.send(
|
||||
SessionMessage(
|
||||
types.JSONRPCMessage(
|
||||
types.JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
)
|
||||
)
|
||||
SessionMessage(types.JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized"))
|
||||
)
|
||||
|
||||
async with (
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import pytest
|
|||
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
|
||||
from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse, jsonrpc_message_adapter
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
@ -14,8 +14,8 @@ async def test_stdio_server():
|
|||
stdout = io.StringIO()
|
||||
|
||||
messages = [
|
||||
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
|
||||
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
|
||||
JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"),
|
||||
JSONRPCResponse(jsonrpc="2.0", id=2, result={}),
|
||||
]
|
||||
|
||||
for message in messages:
|
||||
|
|
@ -37,13 +37,13 @@ async def test_stdio_server():
|
|||
|
||||
# Verify received messages
|
||||
assert len(received_messages) == 2
|
||||
assert received_messages[0] == JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
|
||||
assert received_messages[1] == JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={}))
|
||||
assert received_messages[0] == JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
|
||||
assert received_messages[1] == JSONRPCResponse(jsonrpc="2.0", id=2, result={})
|
||||
|
||||
# Test sending responses from the server
|
||||
responses = [
|
||||
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")),
|
||||
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})),
|
||||
JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"),
|
||||
JSONRPCResponse(jsonrpc="2.0", id=4, result={}),
|
||||
]
|
||||
|
||||
async with write_stream:
|
||||
|
|
@ -55,7 +55,7 @@ async def test_stdio_server():
|
|||
output_lines = stdout.readlines()
|
||||
assert len(output_lines) == 2
|
||||
|
||||
received_responses = [JSONRPCMessage.model_validate_json(line.strip()) for line in output_lines]
|
||||
received_responses = [jsonrpc_message_adapter.validate_json(line.strip()) for line in output_lines]
|
||||
assert len(received_responses) == 2
|
||||
assert received_responses[0] == JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"))
|
||||
assert received_responses[1] == JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={}))
|
||||
assert received_responses[0] == JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
|
||||
assert received_responses[1] == JSONRPCResponse(jsonrpc="2.0", id=4, result={})
|
||||
|
|
|
|||
Loading…
Reference in a new issue