Drop RootModel from JSONRPCMessage (#1908)
This commit is contained in:
parent
dcc9b4f8e0
commit
f4672c5084
3 changed files with 314 additions and 476 deletions
|
|
@ -43,35 +43,33 @@ class StreamSpyCollection:
|
|||
def get_client_requests(self, method: str | None = None) -> list[JSONRPCRequest]: # pragma: no cover
|
||||
"""Get client-sent requests, optionally filtered by method."""
|
||||
return [
|
||||
req.message.root
|
||||
req.message
|
||||
for req in self.client.sent_messages
|
||||
if isinstance(req.message.root, JSONRPCRequest) and (method is None or req.message.root.method == method)
|
||||
if isinstance(req.message, JSONRPCRequest) and (method is None or req.message.method == method)
|
||||
]
|
||||
|
||||
def get_server_requests(self, method: str | None = None) -> list[JSONRPCRequest]: # pragma: no cover
|
||||
"""Get server-sent requests, optionally filtered by method."""
|
||||
return [ # pragma: no cover
|
||||
req.message.root
|
||||
req.message
|
||||
for req in self.server.sent_messages
|
||||
if isinstance(req.message.root, JSONRPCRequest) and (method is None or req.message.root.method == method)
|
||||
if isinstance(req.message, JSONRPCRequest) and (method is None or req.message.method == method)
|
||||
]
|
||||
|
||||
def get_client_notifications(self, method: str | None = None) -> list[JSONRPCNotification]: # pragma: no cover
|
||||
"""Get client-sent notifications, optionally filtered by method."""
|
||||
return [
|
||||
notif.message.root
|
||||
notif.message
|
||||
for notif in self.client.sent_messages
|
||||
if isinstance(notif.message.root, JSONRPCNotification)
|
||||
and (method is None or notif.message.root.method == method)
|
||||
if isinstance(notif.message, JSONRPCNotification) and (method is None or notif.message.method == method)
|
||||
]
|
||||
|
||||
def get_server_notifications(self, method: str | None = None) -> list[JSONRPCNotification]: # pragma: no cover
|
||||
"""Get server-sent notifications, optionally filtered by method."""
|
||||
return [
|
||||
notif.message.root
|
||||
notif.message
|
||||
for notif in self.server.sent_messages
|
||||
if isinstance(notif.message.root, JSONRPCNotification)
|
||||
and (method is None or notif.message.root.method == method)
|
||||
if isinstance(notif.message, JSONRPCNotification) and (method is None or notif.message.method == method)
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -18,7 +18,6 @@ from mcp.types import (
|
|||
InitializedNotification,
|
||||
InitializeRequest,
|
||||
InitializeResult,
|
||||
JSONRPCMessage,
|
||||
JSONRPCNotification,
|
||||
JSONRPCRequest,
|
||||
JSONRPCResponse,
|
||||
|
|
@ -41,7 +40,7 @@ async def test_client_session_initialize():
|
|||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -65,18 +64,16 @@ async def test_client_session_initialize():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
session_notification = await client_to_server_receive.receive()
|
||||
jsonrpc_notification = session_notification.message
|
||||
assert isinstance(jsonrpc_notification.root, JSONRPCNotification)
|
||||
assert isinstance(jsonrpc_notification, JSONRPCNotification)
|
||||
initialized_notification = ClientNotification.model_validate(
|
||||
jsonrpc_notification.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -128,7 +125,7 @@ async def test_client_session_custom_client_info():
|
|||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -146,12 +143,10 @@ async def test_client_session_custom_client_info():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -189,7 +184,7 @@ async def test_client_session_default_client_info():
|
|||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -207,12 +202,10 @@ async def test_client_session_default_client_info():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -220,10 +213,7 @@ async def test_client_session_default_client_info():
|
|||
await client_to_server_receive.receive()
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
) as session,
|
||||
ClientSession(server_to_client_receive, client_to_server_send) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
|
|
@ -247,7 +237,7 @@ async def test_client_session_version_negotiation_success():
|
|||
async def mock_server():
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -268,12 +258,10 @@ async def test_client_session_version_negotiation_success():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -281,10 +269,7 @@ async def test_client_session_version_negotiation_success():
|
|||
await client_to_server_receive.receive()
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
) as session,
|
||||
ClientSession(server_to_client_receive, client_to_server_send) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
|
|
@ -309,7 +294,7 @@ async def test_client_session_version_negotiation_failure():
|
|||
async def mock_server():
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -327,21 +312,16 @@ async def test_client_session_version_negotiation_failure():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
) as session,
|
||||
ClientSession(server_to_client_receive, client_to_server_send) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
|
|
@ -368,7 +348,7 @@ async def test_client_capabilities_default():
|
|||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -386,12 +366,10 @@ async def test_client_capabilities_default():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -399,10 +377,7 @@ async def test_client_capabilities_default():
|
|||
await client_to_server_receive.receive()
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
) as session,
|
||||
ClientSession(server_to_client_receive, client_to_server_send) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
|
|
@ -446,7 +421,7 @@ async def test_client_capabilities_with_custom_callbacks():
|
|||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -464,12 +439,10 @@ async def test_client_capabilities_with_custom_callbacks():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -529,7 +502,7 @@ async def test_client_capabilities_with_sampling_tools():
|
|||
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -547,12 +520,10 @@ async def test_client_capabilities_with_sampling_tools():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -600,7 +571,7 @@ async def test_get_server_capabilities():
|
|||
async def mock_server():
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -617,12 +588,10 @@ async def test_get_server_capabilities():
|
|||
async with server_to_client_send:
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -669,7 +638,7 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
|
|||
# Receive initialization request from client
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
|
@ -686,12 +655,10 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
|
|||
# Answer initialization request
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -702,14 +669,14 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
|
|||
# Wait for the client to send a 'tools/call' request
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
|
||||
assert jsonrpc_request.root.method == "tools/call"
|
||||
assert jsonrpc_request.method == "tools/call"
|
||||
|
||||
if meta is not None:
|
||||
assert jsonrpc_request.root.params
|
||||
assert "_meta" in jsonrpc_request.root.params
|
||||
assert jsonrpc_request.root.params["_meta"] == meta
|
||||
assert jsonrpc_request.params
|
||||
assert "_meta" in jsonrpc_request.params
|
||||
assert jsonrpc_request.params["_meta"] == meta
|
||||
|
||||
result = ServerResult(
|
||||
CallToolResult(content=[TextContent(type="text", text="Called successfully")], is_error=False)
|
||||
|
|
@ -718,12 +685,10 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
|
|||
# Send the tools/call result
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -732,20 +697,18 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
|
|||
# The client requires this step to validate the tool output schema
|
||||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
|
||||
assert jsonrpc_request.root.method == "tools/list"
|
||||
assert jsonrpc_request.method == "tools/list"
|
||||
|
||||
result = types.ListToolsResult(tools=[mocked_tool])
|
||||
|
||||
await server_to_client_send.send(
|
||||
SessionMessage(
|
||||
JSONRPCMessage(
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.root.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
JSONRPCResponse(
|
||||
jsonrpc="2.0",
|
||||
id=jsonrpc_request.id,
|
||||
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -753,10 +716,7 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
|
|||
server_to_client_send.close()
|
||||
|
||||
async with (
|
||||
ClientSession(
|
||||
server_to_client_receive,
|
||||
client_to_server_send,
|
||||
) as session,
|
||||
ClientSession(server_to_client_receive, client_to_server_send) as session,
|
||||
anyio.create_task_group() as tg,
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
|
|
|
|||
|
|
@ -47,8 +47,8 @@ async def test_stdio_client():
|
|||
async with stdio_client(server_parameters) as (read_stream, write_stream):
|
||||
# Test sending and receiving messages
|
||||
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={}),
|
||||
]
|
||||
|
||||
async with write_stream:
|
||||
|
|
@ -67,8 +67,8 @@ async def test_stdio_client():
|
|||
break
|
||||
|
||||
assert len(read_messages) == 2
|
||||
assert read_messages[0] == JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
|
||||
assert read_messages[1] == JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={}))
|
||||
assert read_messages[0] == JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
|
||||
assert read_messages[1] == JSONRPCResponse(jsonrpc="2.0", id=2, result={})
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
|
|||
Loading…
Reference in a new issue