Implement SEP-1577 - Sampling With Tools (#1594)
Co-authored-by: Felix Weinberger <fweinberger@anthropic.com> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Olivier Chafik
GitHub
parent
c51936f61f
commit
71c475588f
1 changed files with 674 additions and 19 deletions
|
|
@ -9,6 +9,7 @@ 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.exceptions import McpError
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.types import (
|
||||
|
|
@ -288,6 +289,184 @@ async def test_ping_request_before_initialization():
|
|||
assert ping_response_id == 42
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_tool_result_validation():
|
||||
"""Test tool_use/tool_result validation in create_message."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
||||
|
||||
async with (
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server_to_client_receive,
|
||||
):
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test",
|
||||
server_version="0.1.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as session:
|
||||
# Set up client params with sampling.tools capability for the test
|
||||
session._client_params = types.InitializeRequestParams(
|
||||
protocolVersion=types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=types.ClientCapabilities(
|
||||
sampling=types.SamplingCapability(tools=types.SamplingToolsCapability())
|
||||
),
|
||||
clientInfo=types.Implementation(name="test", version="1.0"),
|
||||
)
|
||||
|
||||
tool = types.Tool(name="test_tool", inputSchema={"type": "object"})
|
||||
text = types.TextContent(type="text", text="hello")
|
||||
tool_use = types.ToolUseContent(type="tool_use", id="call_1", name="test_tool", input={})
|
||||
tool_result = types.ToolResultContent(type="tool_result", toolUseId="call_1", content=[])
|
||||
|
||||
# Case 1: tool_result mixed with other content
|
||||
with pytest.raises(ValueError, match="only tool_result content"):
|
||||
await session.create_message(
|
||||
messages=[
|
||||
types.SamplingMessage(role="user", content=text),
|
||||
types.SamplingMessage(role="assistant", content=tool_use),
|
||||
types.SamplingMessage(role="user", content=[tool_result, text]), # mixed!
|
||||
],
|
||||
max_tokens=100,
|
||||
tools=[tool],
|
||||
)
|
||||
|
||||
# Case 2: tool_result without previous message
|
||||
with pytest.raises(ValueError, match="requires a previous message"):
|
||||
await session.create_message(
|
||||
messages=[types.SamplingMessage(role="user", content=tool_result)],
|
||||
max_tokens=100,
|
||||
tools=[tool],
|
||||
)
|
||||
|
||||
# Case 3: tool_result without previous tool_use
|
||||
with pytest.raises(ValueError, match="do not match any tool_use"):
|
||||
await session.create_message(
|
||||
messages=[
|
||||
types.SamplingMessage(role="user", content=text),
|
||||
types.SamplingMessage(role="user", content=tool_result),
|
||||
],
|
||||
max_tokens=100,
|
||||
tools=[tool],
|
||||
)
|
||||
|
||||
# Case 4: mismatched tool IDs
|
||||
with pytest.raises(ValueError, match="ids of tool_result blocks and tool_use blocks"):
|
||||
await session.create_message(
|
||||
messages=[
|
||||
types.SamplingMessage(role="user", content=text),
|
||||
types.SamplingMessage(role="assistant", content=tool_use),
|
||||
types.SamplingMessage(
|
||||
role="user",
|
||||
content=types.ToolResultContent(type="tool_result", toolUseId="wrong_id", content=[]),
|
||||
),
|
||||
],
|
||||
max_tokens=100,
|
||||
tools=[tool],
|
||||
)
|
||||
|
||||
# Case 5: text-only message with tools (no tool_results) - passes validation
|
||||
# Covers has_tool_results=False branch.
|
||||
# We use move_on_after because validation happens synchronously before
|
||||
# send_request, which would block indefinitely waiting for a response.
|
||||
# The timeout lets validation pass, then cancels the blocked send.
|
||||
with anyio.move_on_after(0.01):
|
||||
await session.create_message(
|
||||
messages=[types.SamplingMessage(role="user", content=text)],
|
||||
max_tokens=100,
|
||||
tools=[tool],
|
||||
)
|
||||
|
||||
# Case 6: valid matching tool_result/tool_use IDs - passes validation
|
||||
# Covers tool_use_ids == tool_result_ids branch.
|
||||
# (see Case 5 comment for move_on_after explanation)
|
||||
with anyio.move_on_after(0.01):
|
||||
await session.create_message(
|
||||
messages=[
|
||||
types.SamplingMessage(role="user", content=text),
|
||||
types.SamplingMessage(role="assistant", content=tool_use),
|
||||
types.SamplingMessage(role="user", content=tool_result),
|
||||
],
|
||||
max_tokens=100,
|
||||
tools=[tool],
|
||||
)
|
||||
|
||||
# Case 7: validation runs even without `tools` parameter
|
||||
# (tool loop continuation may omit tools while containing tool_result)
|
||||
with pytest.raises(ValueError, match="do not match any tool_use"):
|
||||
await session.create_message(
|
||||
messages=[
|
||||
types.SamplingMessage(role="user", content=text),
|
||||
types.SamplingMessage(role="user", content=tool_result),
|
||||
],
|
||||
max_tokens=100,
|
||||
# Note: no tools parameter
|
||||
)
|
||||
|
||||
# Case 8: empty messages list - skips validation entirely
|
||||
# Covers the `if messages:` branch (line 280->302)
|
||||
with anyio.move_on_after(0.01):
|
||||
await session.create_message(
|
||||
messages=[],
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_without_tools_capability():
|
||||
"""Test that create_message raises McpError when tools are provided without capability."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
||||
|
||||
async with (
|
||||
client_to_server_send,
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server_to_client_receive,
|
||||
):
|
||||
async with ServerSession(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
InitializationOptions(
|
||||
server_name="test",
|
||||
server_version="0.1.0",
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as session:
|
||||
# Set up client params WITHOUT sampling.tools capability
|
||||
session._client_params = types.InitializeRequestParams(
|
||||
protocolVersion=types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=types.ClientCapabilities(sampling=types.SamplingCapability()),
|
||||
clientInfo=types.Implementation(name="test", version="1.0"),
|
||||
)
|
||||
|
||||
tool = types.Tool(name="test_tool", inputSchema={"type": "object"})
|
||||
text = types.TextContent(type="text", text="hello")
|
||||
|
||||
# Should raise McpError when tools are provided but client lacks capability
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await session.create_message(
|
||||
messages=[types.SamplingMessage(role="user", content=text)],
|
||||
max_tokens=100,
|
||||
tools=[tool],
|
||||
)
|
||||
assert "does not support sampling tools capability" in exc_info.value.error.message
|
||||
|
||||
# Should also raise McpError when tool_choice is provided
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await session.create_message(
|
||||
messages=[types.SamplingMessage(role="user", content=text)],
|
||||
max_tokens=100,
|
||||
tool_choice=types.ToolChoice(mode="auto"),
|
||||
)
|
||||
assert "does not support sampling tools capability" in exc_info.value.error.message
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_other_requests_blocked_before_initialization():
|
||||
"""Test that non-ping requests are still blocked before initialization."""
|
||||
|
|
|
|||
Loading…
Reference in a new issue