Drop RootModel from JSONRPCMessage (#1908)
This commit is contained in:
parent
dcc9b4f8e0
commit
f4672c5084
2 changed files with 314 additions and 476 deletions
|
|
@ -66,7 +66,7 @@ async def test_request_id_match() -> None:
|
|||
jsonrpc="2.0",
|
||||
)
|
||||
|
||||
await client_writer.send(SessionMessage(JSONRPCMessage(root=init_req)))
|
||||
await client_writer.send(SessionMessage(init_req))
|
||||
response = await server_reader.receive() # Get init response but don't need to check it
|
||||
|
||||
# Send initialized notification
|
||||
|
|
@ -75,12 +75,12 @@ async def test_request_id_match() -> None:
|
|||
params=NotificationParams().model_dump(by_alias=True, exclude_none=True),
|
||||
jsonrpc="2.0",
|
||||
)
|
||||
await client_writer.send(SessionMessage(JSONRPCMessage(root=initialized_notification)))
|
||||
await client_writer.send(SessionMessage(initialized_notification))
|
||||
|
||||
# Send ping request with custom ID
|
||||
ping_request = JSONRPCRequest(id=custom_request_id, method="ping", params={}, jsonrpc="2.0")
|
||||
|
||||
await client_writer.send(SessionMessage(JSONRPCMessage(root=ping_request)))
|
||||
await client_writer.send(SessionMessage(ping_request))
|
||||
|
||||
# Read response
|
||||
response = await server_reader.receive()
|
||||
|
|
@ -88,8 +88,8 @@ async def test_request_id_match() -> None:
|
|||
# Verify response ID matches request ID
|
||||
assert isinstance(response, SessionMessage)
|
||||
assert isinstance(response.message, JSONRPCMessage)
|
||||
assert isinstance(response.message.root, JSONRPCResponse)
|
||||
assert response.message.root.id == custom_request_id, "Response ID should match request ID"
|
||||
assert isinstance(response.message, JSONRPCResponse)
|
||||
assert response.message.id == custom_request_id, "Response ID should match request ID"
|
||||
|
||||
# Cancel server task
|
||||
tg.cancel_scope.cancel()
|
||||
|
|
|
|||
|
|
@ -1,21 +1,13 @@
|
|||
# Claude Debug
|
||||
"""Test for HackerOne vulnerability report #3156202 - malformed input DOS."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import (
|
||||
INVALID_PARAMS,
|
||||
JSONRPCError,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
ServerCapabilities,
|
||||
)
|
||||
from mcp.types import INVALID_PARAMS, JSONRPCError, JSONRPCMessage, JSONRPCRequest, ServerCapabilities
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
@ -37,7 +29,7 @@ async def test_malformed_initialize_request_does_not_crash_server():
|
|||
)
|
||||
|
||||
# Wrap in session message
|
||||
request_message = SessionMessage(message=JSONRPCMessage(malformed_request))
|
||||
request_message = SessionMessage(message=malformed_request)
|
||||
|
||||
# Start a server session
|
||||
async with ServerSession(
|
||||
|
|
@ -58,7 +50,7 @@ async def test_malformed_initialize_request_does_not_crash_server():
|
|||
# Check that we received an error response instead of a crash
|
||||
try:
|
||||
response_message = write_receive_stream.receive_nowait()
|
||||
response = response_message.message.root
|
||||
response = response_message.message
|
||||
|
||||
# Verify it's a proper JSON-RPC error response
|
||||
assert isinstance(response, JSONRPCError)
|
||||
|
|
@ -75,14 +67,14 @@ async def test_malformed_initialize_request_does_not_crash_server():
|
|||
method="tools/call",
|
||||
# params=None # Missing required params
|
||||
)
|
||||
another_request_message = SessionMessage(message=JSONRPCMessage(another_malformed_request))
|
||||
another_request_message = SessionMessage(message=another_malformed_request)
|
||||
|
||||
await read_send_stream.send(another_request_message)
|
||||
await anyio.sleep(0.1)
|
||||
|
||||
# Should get another error response, not a crash
|
||||
second_response_message = write_receive_stream.receive_nowait()
|
||||
second_response = second_response_message.message.root
|
||||
second_response = second_response_message.message
|
||||
|
||||
assert isinstance(second_response, JSONRPCError)
|
||||
assert second_response.id == "test_id_2"
|
||||
|
|
@ -125,7 +117,7 @@ async def test_multiple_concurrent_malformed_requests():
|
|||
method="initialize",
|
||||
# params=None # Missing required params
|
||||
)
|
||||
request_message = SessionMessage(message=JSONRPCMessage(malformed_request))
|
||||
request_message = SessionMessage(message=malformed_request)
|
||||
malformed_requests.append(request_message)
|
||||
|
||||
# Send all requests
|
||||
|
|
@ -136,11 +128,11 @@ async def test_multiple_concurrent_malformed_requests():
|
|||
await anyio.sleep(0.2)
|
||||
|
||||
# Verify we get error responses for all requests
|
||||
error_responses: list[Any] = []
|
||||
error_responses: list[JSONRPCMessage] = []
|
||||
try:
|
||||
while True:
|
||||
response_message = write_receive_stream.receive_nowait()
|
||||
error_responses.append(response_message.message.root)
|
||||
error_responses.append(response_message.message)
|
||||
except anyio.WouldBlock:
|
||||
pass # No more messages
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue