Max Isbey
25fb05f416
Lowlevel plumbing for Multi Round-Trip Requests: Types: - IncompleteResult with result_type discriminator, input_requests, request_state - InputRequest/InputResponse unions (elicitation, sampling, roots) - input_responses + request_state fields on RequestParams Server (lowlevel): - on_call_tool return widened to include IncompleteResult Session: - send_request accepts TypeAdapter (overload) for union result parsing - call_tool_mrtr() returns CallToolResult | IncompleteResult - call_tool() stays narrow, raises on IncompleteResult with migration hint Client: - call_tool() drives MRTR retry loop internally — dispatches embedded input requests to elicitation/sampling/list_roots callbacks, retries with collected responses + echoed request_state - max_mrtr_rounds bound (default 8) The client-side delta from today's code is zero: elicitation_callback is the same function whether it fires from SSE push or MRTR retry.
531 lines
21 KiB
Python
531 lines
21 KiB
Python
"""Tests for the unified Client class."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import patch
|
|
|
|
import anyio
|
|
import pytest
|
|
from inline_snapshot import snapshot
|
|
from pydantic import FileUrl
|
|
|
|
from mcp import MCPError, types
|
|
from mcp.client._memory import InMemoryTransport
|
|
from mcp.client.client import Client
|
|
from mcp.client.context import ClientRequestContext
|
|
from mcp.server import Server, ServerRequestContext
|
|
from mcp.server.mcpserver import MCPServer
|
|
from mcp.types import (
|
|
CallToolResult,
|
|
EmptyResult,
|
|
GetPromptResult,
|
|
ListPromptsResult,
|
|
ListResourcesResult,
|
|
ListResourceTemplatesResult,
|
|
ListToolsResult,
|
|
Prompt,
|
|
PromptArgument,
|
|
PromptMessage,
|
|
PromptsCapability,
|
|
ReadResourceResult,
|
|
Resource,
|
|
ResourcesCapability,
|
|
ServerCapabilities,
|
|
TextContent,
|
|
TextResourceContents,
|
|
Tool,
|
|
ToolsCapability,
|
|
)
|
|
|
|
pytestmark = pytest.mark.anyio
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_server() -> Server:
|
|
"""Create a simple MCP server for testing."""
|
|
|
|
async def handle_list_resources(
|
|
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
|
|
) -> ListResourcesResult:
|
|
return ListResourcesResult(
|
|
resources=[Resource(uri="memory://test", name="Test Resource", description="A test resource")]
|
|
)
|
|
|
|
async def handle_subscribe_resource(ctx: ServerRequestContext, params: types.SubscribeRequestParams) -> EmptyResult:
|
|
return EmptyResult()
|
|
|
|
async def handle_unsubscribe_resource(
|
|
ctx: ServerRequestContext, params: types.UnsubscribeRequestParams
|
|
) -> EmptyResult:
|
|
return EmptyResult()
|
|
|
|
async def handle_set_logging_level(ctx: ServerRequestContext, params: types.SetLevelRequestParams) -> EmptyResult:
|
|
return EmptyResult()
|
|
|
|
async def handle_completion(ctx: ServerRequestContext, params: types.CompleteRequestParams) -> types.CompleteResult:
|
|
return types.CompleteResult(completion=types.Completion(values=[]))
|
|
|
|
return Server(
|
|
name="test_server",
|
|
on_list_resources=handle_list_resources,
|
|
on_subscribe_resource=handle_subscribe_resource,
|
|
on_unsubscribe_resource=handle_unsubscribe_resource,
|
|
on_set_logging_level=handle_set_logging_level,
|
|
on_completion=handle_completion,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def app() -> MCPServer:
|
|
"""Create an MCPServer server for testing."""
|
|
server = MCPServer("test")
|
|
|
|
@server.tool()
|
|
def greet(name: str) -> str:
|
|
"""Greet someone by name."""
|
|
return f"Hello, {name}!"
|
|
|
|
@server.resource("test://resource")
|
|
def test_resource() -> str:
|
|
"""A test resource."""
|
|
return "Test content"
|
|
|
|
@server.prompt()
|
|
def greeting_prompt(name: str) -> str:
|
|
"""A greeting prompt."""
|
|
return f"Please greet {name} warmly."
|
|
|
|
return server
|
|
|
|
|
|
async def test_client_is_initialized(app: MCPServer):
|
|
"""Test that the client is initialized after entering context."""
|
|
async with Client(app) as client:
|
|
assert client.initialize_result.capabilities == snapshot(
|
|
ServerCapabilities(
|
|
experimental={},
|
|
prompts=PromptsCapability(list_changed=False),
|
|
resources=ResourcesCapability(subscribe=False, list_changed=False),
|
|
tools=ToolsCapability(list_changed=False),
|
|
)
|
|
)
|
|
assert client.initialize_result.server_info.name == "test"
|
|
|
|
|
|
async def test_client_with_simple_server(simple_server: Server):
|
|
"""Test that from_server works with a basic Server instance."""
|
|
async with Client(simple_server) as client:
|
|
resources = await client.list_resources()
|
|
assert resources == snapshot(
|
|
ListResourcesResult(
|
|
resources=[Resource(name="Test Resource", uri="memory://test", description="A test resource")]
|
|
)
|
|
)
|
|
|
|
|
|
async def test_client_send_ping(app: MCPServer):
|
|
async with Client(app) as client:
|
|
result = await client.send_ping()
|
|
assert result == snapshot(EmptyResult())
|
|
|
|
|
|
async def test_client_list_tools(app: MCPServer):
|
|
async with Client(app) as client:
|
|
result = await client.list_tools()
|
|
assert result == snapshot(
|
|
ListToolsResult(
|
|
tools=[
|
|
Tool(
|
|
name="greet",
|
|
description="Greet someone by name.",
|
|
input_schema={
|
|
"properties": {"name": {"title": "Name", "type": "string"}},
|
|
"required": ["name"],
|
|
"title": "greetArguments",
|
|
"type": "object",
|
|
},
|
|
output_schema={
|
|
"properties": {"result": {"title": "Result", "type": "string"}},
|
|
"required": ["result"],
|
|
"title": "greetOutput",
|
|
"type": "object",
|
|
},
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
|
|
async def test_client_call_tool(app: MCPServer):
|
|
async with Client(app) as client:
|
|
result = await client.call_tool("greet", {"name": "World"})
|
|
assert result == snapshot(
|
|
CallToolResult(
|
|
content=[TextContent(text="Hello, World!")],
|
|
structured_content={"result": "Hello, World!"},
|
|
)
|
|
)
|
|
|
|
|
|
async def test_read_resource(app: MCPServer):
|
|
"""Test reading a resource."""
|
|
async with Client(app) as client:
|
|
result = await client.read_resource("test://resource")
|
|
assert result == snapshot(
|
|
ReadResourceResult(
|
|
contents=[TextResourceContents(uri="test://resource", mime_type="text/plain", text="Test content")]
|
|
)
|
|
)
|
|
|
|
|
|
async def test_read_resource_error_propagates():
|
|
"""MCPError raised by a server handler propagates to the client with its code intact."""
|
|
|
|
async def handle_read_resource(
|
|
ctx: ServerRequestContext, params: types.ReadResourceRequestParams
|
|
) -> ReadResourceResult:
|
|
raise MCPError(code=404, message="no resource with that URI was found")
|
|
|
|
server = Server("test", on_read_resource=handle_read_resource)
|
|
async with Client(server) as client:
|
|
with pytest.raises(MCPError) as exc_info:
|
|
await client.read_resource("unknown://example")
|
|
assert exc_info.value.error.code == 404
|
|
|
|
|
|
async def test_get_prompt(app: MCPServer):
|
|
"""Test getting a prompt."""
|
|
async with Client(app) as client:
|
|
result = await client.get_prompt("greeting_prompt", {"name": "Alice"})
|
|
assert result == snapshot(
|
|
GetPromptResult(
|
|
description="A greeting prompt.",
|
|
messages=[PromptMessage(role="user", content=TextContent(text="Please greet Alice warmly."))],
|
|
)
|
|
)
|
|
|
|
|
|
def test_client_session_property_before_enter(app: MCPServer):
|
|
"""Test that accessing session before context manager raises RuntimeError."""
|
|
client = Client(app)
|
|
with pytest.raises(RuntimeError, match="Client must be used within an async context manager"):
|
|
client.session
|
|
|
|
|
|
async def test_client_reentry_raises_runtime_error(app: MCPServer):
|
|
"""Test that reentering a client raises RuntimeError."""
|
|
async with Client(app) as client:
|
|
with pytest.raises(RuntimeError, match="Client is already entered"):
|
|
await client.__aenter__()
|
|
|
|
|
|
async def test_client_send_progress_notification():
|
|
"""Test sending progress notification."""
|
|
received_from_client = None
|
|
event = anyio.Event()
|
|
|
|
async def handle_progress(ctx: ServerRequestContext, params: types.ProgressNotificationParams) -> None:
|
|
nonlocal received_from_client
|
|
received_from_client = {"progress_token": params.progress_token, "progress": params.progress}
|
|
event.set()
|
|
|
|
server = Server(name="test_server", on_progress=handle_progress)
|
|
|
|
async with Client(server) as client:
|
|
await client.send_progress_notification(progress_token="token123", progress=50.0)
|
|
await event.wait()
|
|
assert received_from_client == snapshot({"progress_token": "token123", "progress": 50.0})
|
|
|
|
|
|
async def test_client_subscribe_resource(simple_server: Server):
|
|
async with Client(simple_server) as client:
|
|
result = await client.subscribe_resource("memory://test")
|
|
assert result == snapshot(EmptyResult())
|
|
|
|
|
|
async def test_client_unsubscribe_resource(simple_server: Server):
|
|
async with Client(simple_server) as client:
|
|
result = await client.unsubscribe_resource("memory://test")
|
|
assert result == snapshot(EmptyResult())
|
|
|
|
|
|
async def test_client_set_logging_level(simple_server: Server):
|
|
"""Test setting logging level."""
|
|
async with Client(simple_server) as client:
|
|
result = await client.set_logging_level("debug")
|
|
assert result == snapshot(EmptyResult())
|
|
|
|
|
|
async def test_client_list_resources_with_params(app: MCPServer):
|
|
"""Test listing resources with params parameter."""
|
|
async with Client(app) as client:
|
|
result = await client.list_resources()
|
|
assert result == snapshot(
|
|
ListResourcesResult(
|
|
resources=[
|
|
Resource(
|
|
name="test_resource",
|
|
uri="test://resource",
|
|
description="A test resource.",
|
|
mime_type="text/plain",
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
|
|
async def test_client_list_resource_templates(app: MCPServer):
|
|
"""Test listing resource templates with params parameter."""
|
|
async with Client(app) as client:
|
|
result = await client.list_resource_templates()
|
|
assert result == snapshot(ListResourceTemplatesResult(resource_templates=[]))
|
|
|
|
|
|
async def test_list_prompts(app: MCPServer):
|
|
"""Test listing prompts with params parameter."""
|
|
async with Client(app) as client:
|
|
result = await client.list_prompts()
|
|
assert result == snapshot(
|
|
ListPromptsResult(
|
|
prompts=[
|
|
Prompt(
|
|
name="greeting_prompt",
|
|
description="A greeting prompt.",
|
|
arguments=[PromptArgument(name="name", required=True)],
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
|
|
async def test_complete_with_prompt_reference(simple_server: Server):
|
|
"""Test getting completions for a prompt argument."""
|
|
async with Client(simple_server) as client:
|
|
ref = types.PromptReference(type="ref/prompt", name="test_prompt")
|
|
result = await client.complete(ref=ref, argument={"name": "arg", "value": "test"})
|
|
assert result == snapshot(types.CompleteResult(completion=types.Completion(values=[])))
|
|
|
|
|
|
def test_client_with_url_initializes_streamable_http_transport():
|
|
with patch("mcp.client.client.streamable_http_client") as mock:
|
|
_ = Client("http://localhost:8000/mcp")
|
|
mock.assert_called_once_with("http://localhost:8000/mcp")
|
|
|
|
|
|
async def test_client_uses_transport_directly(app: MCPServer):
|
|
transport = InMemoryTransport(app)
|
|
async with Client(transport) as client:
|
|
result = await client.call_tool("greet", {"name": "Transport"})
|
|
assert result == snapshot(
|
|
CallToolResult(
|
|
content=[TextContent(text="Hello, Transport!")],
|
|
structured_content={"result": "Hello, Transport!"},
|
|
)
|
|
)
|
|
|
|
|
|
# ─── MRTR (SEP-2322) ────────────────────────────────────────────────────────
|
|
|
|
|
|
async def _mrtr_list_tools(
|
|
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
|
|
) -> types.ListToolsResult:
|
|
return types.ListToolsResult(tools=[])
|
|
|
|
|
|
async def test_mrtr_single_round_elicitation():
|
|
"""Server returns IncompleteResult with one elicitation; Client drives retry transparently."""
|
|
|
|
async def on_call_tool(
|
|
ctx: ServerRequestContext, params: types.CallToolRequestParams
|
|
) -> types.CallToolResult | types.IncompleteResult:
|
|
units = params.input_responses.get("units") if params.input_responses else None
|
|
if units is None:
|
|
return types.IncompleteResult(
|
|
input_requests={
|
|
"units": types.ElicitRequest(
|
|
params=types.ElicitRequestFormParams(
|
|
message="Which units?",
|
|
requested_schema={"type": "object", "properties": {"u": {"type": "string"}}},
|
|
)
|
|
)
|
|
},
|
|
)
|
|
u = units["content"]["u"]
|
|
location = params.arguments["location"] if params.arguments else "?"
|
|
return types.CallToolResult(content=[TextContent(text=f"Weather in {location}: 22°{u}")])
|
|
|
|
server = Server("mrtr-test", on_call_tool=on_call_tool, on_list_tools=_mrtr_list_tools)
|
|
|
|
async def elicitation_cb(context: ClientRequestContext, params: types.ElicitRequestParams) -> types.ElicitResult:
|
|
return types.ElicitResult(action="accept", content={"u": "C"})
|
|
|
|
async with Client(server, elicitation_callback=elicitation_cb) as client:
|
|
result = await client.call_tool("weather", {"location": "Tokyo"})
|
|
assert result == snapshot(CallToolResult(content=[TextContent(text="Weather in Tokyo: 22°C")]))
|
|
|
|
|
|
async def test_mrtr_multi_round_with_request_state():
|
|
"""Two-round elicitation accumulating state in request_state (the ADO-rules SEP example)."""
|
|
|
|
async def on_call_tool(
|
|
ctx: ServerRequestContext, params: types.CallToolRequestParams
|
|
) -> types.CallToolResult | types.IncompleteResult:
|
|
responses = params.input_responses or {}
|
|
state = params.request_state
|
|
|
|
if "resolution" not in responses and state is None:
|
|
return types.IncompleteResult(
|
|
input_requests={
|
|
"resolution": types.ElicitRequest(
|
|
params=types.ElicitRequestFormParams(message="Resolution?", requested_schema={})
|
|
)
|
|
},
|
|
)
|
|
|
|
if state is None:
|
|
resolution = responses["resolution"]["content"]["r"]
|
|
return types.IncompleteResult(
|
|
input_requests={
|
|
"dup": types.ElicitRequest(
|
|
params=types.ElicitRequestFormParams(message="Duplicate of?", requested_schema={})
|
|
)
|
|
},
|
|
request_state=f"resolution={resolution}",
|
|
)
|
|
|
|
dup = responses["dup"]["content"]["id"]
|
|
return types.CallToolResult(content=[TextContent(text=f"{state} dup={dup}")])
|
|
|
|
server = Server("mrtr-test", on_call_tool=on_call_tool, on_list_tools=_mrtr_list_tools)
|
|
|
|
answers = {"Resolution?": {"r": "Duplicate"}, "Duplicate of?": {"id": "4301"}}
|
|
|
|
async def elicitation_cb(context: ClientRequestContext, params: types.ElicitRequestParams) -> types.ElicitResult:
|
|
assert isinstance(params, types.ElicitRequestFormParams)
|
|
return types.ElicitResult(action="accept", content=dict(answers[params.message]))
|
|
|
|
async with Client(server, elicitation_callback=elicitation_cb) as client:
|
|
result = await client.call_tool("update_item", {})
|
|
assert result == snapshot(CallToolResult(content=[TextContent(text="resolution=Duplicate dup=4301")]))
|
|
|
|
|
|
async def test_mrtr_round_limit_exceeded():
|
|
"""Server never converges → Client raises after max_mrtr_rounds."""
|
|
|
|
async def on_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.IncompleteResult:
|
|
return types.IncompleteResult(request_state="spin")
|
|
|
|
server = Server("mrtr-test", on_call_tool=on_call_tool, on_list_tools=_mrtr_list_tools)
|
|
|
|
async with Client(server, max_mrtr_rounds=3) as client:
|
|
with pytest.raises(RuntimeError, match="exceeded 3 rounds"):
|
|
await client.call_tool("stuck", {})
|
|
|
|
|
|
async def test_mrtr_elicitation_without_callback_raises():
|
|
"""IncompleteResult with elicitation but no callback → clear error."""
|
|
|
|
async def on_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.IncompleteResult:
|
|
return types.IncompleteResult(
|
|
input_requests={
|
|
"ask": types.ElicitRequest(params=types.ElicitRequestFormParams(message="?", requested_schema={}))
|
|
},
|
|
)
|
|
|
|
server = Server("mrtr-test", on_call_tool=on_call_tool, on_list_tools=_mrtr_list_tools)
|
|
|
|
async with Client(server) as client:
|
|
with pytest.raises(RuntimeError, match="no elicitation_callback"):
|
|
await client.call_tool("ask", {})
|
|
|
|
|
|
async def test_mrtr_sampling_input_request():
|
|
"""IncompleteResult with a sampling input request is dispatched to sampling_callback."""
|
|
|
|
async def on_call_tool(
|
|
ctx: ServerRequestContext, params: types.CallToolRequestParams
|
|
) -> types.CallToolResult | types.IncompleteResult:
|
|
if params.input_responses and "q" in params.input_responses:
|
|
answer = params.input_responses["q"]["content"]["text"]
|
|
return types.CallToolResult(content=[TextContent(text=answer)])
|
|
return types.IncompleteResult(
|
|
input_requests={
|
|
"q": types.CreateMessageRequest(
|
|
params=types.CreateMessageRequestParams(
|
|
messages=[
|
|
types.SamplingMessage(role="user", content=types.TextContent(text="Capital of France?"))
|
|
],
|
|
max_tokens=50,
|
|
)
|
|
)
|
|
},
|
|
)
|
|
|
|
server = Server("mrtr-test", on_call_tool=on_call_tool, on_list_tools=_mrtr_list_tools)
|
|
|
|
async def sampling_cb(
|
|
context: ClientRequestContext, params: types.CreateMessageRequestParams
|
|
) -> types.CreateMessageResult:
|
|
return types.CreateMessageResult(
|
|
role="assistant", content=types.TextContent(text="Paris"), model="test", stop_reason="endTurn"
|
|
)
|
|
|
|
async with Client(server, sampling_callback=sampling_cb) as client:
|
|
result = await client.call_tool("ask", {})
|
|
assert result == snapshot(CallToolResult(content=[TextContent(text="Paris")]))
|
|
|
|
|
|
async def test_mrtr_list_roots_input_request():
|
|
"""IncompleteResult with a roots/list input request is dispatched to list_roots_callback."""
|
|
|
|
async def on_call_tool(
|
|
ctx: ServerRequestContext, params: types.CallToolRequestParams
|
|
) -> types.CallToolResult | types.IncompleteResult:
|
|
if params.input_responses and "roots" in params.input_responses:
|
|
n = len(params.input_responses["roots"]["roots"])
|
|
return types.CallToolResult(content=[TextContent(text=f"saw {n} roots")])
|
|
return types.IncompleteResult(input_requests={"roots": types.ListRootsRequest()})
|
|
|
|
server = Server("mrtr-test", on_call_tool=on_call_tool, on_list_tools=_mrtr_list_tools)
|
|
|
|
async def list_roots_cb(context: ClientRequestContext) -> types.ListRootsResult:
|
|
return types.ListRootsResult(roots=[types.Root(uri=FileUrl("file:///a")), types.Root(uri=FileUrl("file:///b"))])
|
|
|
|
async with Client(server, list_roots_callback=list_roots_cb) as client:
|
|
result = await client.call_tool("scan", {})
|
|
assert result == snapshot(CallToolResult(content=[TextContent(text="saw 2 roots")]))
|
|
|
|
|
|
async def test_mrtr_callback_returns_error_data():
|
|
"""Callback returning ErrorData surfaces as RuntimeError."""
|
|
|
|
async def on_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.IncompleteResult:
|
|
return types.IncompleteResult(
|
|
input_requests={
|
|
"ask": types.ElicitRequest(params=types.ElicitRequestFormParams(message="?", requested_schema={}))
|
|
},
|
|
)
|
|
|
|
server = Server("mrtr-test", on_call_tool=on_call_tool, on_list_tools=_mrtr_list_tools)
|
|
|
|
async def elicitation_cb(context: ClientRequestContext, params: types.ElicitRequestParams) -> types.ErrorData:
|
|
return types.ErrorData(code=-1, message="user closed dialog")
|
|
|
|
async with Client(server, elicitation_callback=elicitation_cb) as client:
|
|
with pytest.raises(RuntimeError, match="user closed dialog"):
|
|
await client.call_tool("ask", {})
|
|
|
|
|
|
async def test_session_call_tool_raises_on_incomplete():
|
|
"""ClientSession.call_tool (non-MRTR) raises if server returns IncompleteResult."""
|
|
|
|
async def on_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.IncompleteResult:
|
|
return types.IncompleteResult(request_state="x")
|
|
|
|
server = Server("mrtr-test", on_call_tool=on_call_tool, on_list_tools=_mrtr_list_tools)
|
|
|
|
async with Client(server) as client:
|
|
with pytest.raises(RuntimeError, match="Use Client.call_tool"):
|
|
await client.session.call_tool("stuck", {})
|