Extend resolver DI to sampling and roots requests (#3049)
This commit is contained in:
Max
GitHub
parent
53117cb3a9
commit
d287c9868f
2 changed files with 1134 additions and 200 deletions
|
|
@ -7,22 +7,40 @@ from typing import Annotated, Any, Literal, TypeVar, cast
|
|||
|
||||
import anyio
|
||||
import pytest
|
||||
from inline_snapshot import snapshot
|
||||
from mcp_types import (
|
||||
MISSING_REQUIRED_CLIENT_CAPABILITY,
|
||||
CallToolResult,
|
||||
CreateMessageRequest,
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
CreateMessageResultWithTools,
|
||||
ElicitRequest,
|
||||
ElicitRequestFormParams,
|
||||
ElicitRequestParams,
|
||||
ElicitResult,
|
||||
InputRequiredResult,
|
||||
InputResponses,
|
||||
JSONRPCError,
|
||||
JSONRPCNotification,
|
||||
JSONRPCRequest,
|
||||
ListRootsResult,
|
||||
Root,
|
||||
SamplingCapability,
|
||||
SamplingMessage,
|
||||
SamplingToolsCapability,
|
||||
TextContent,
|
||||
ToolChoice,
|
||||
)
|
||||
from pydantic import BaseModel, Field, ValidationError, create_model
|
||||
from mcp_types import (
|
||||
Tool as SamplingTool,
|
||||
)
|
||||
from pydantic import BaseModel, Field, FileUrl, ValidationError, create_model
|
||||
from typing_extensions import TypeAliasType
|
||||
|
||||
from mcp import Client, InputRequiredRoundsExceededError
|
||||
from mcp.client import ClientRequestContext
|
||||
from mcp.client._memory import InMemoryTransport
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.server.mcpserver import (
|
||||
AcceptedElicitation,
|
||||
|
|
@ -32,18 +50,20 @@ from mcp.server.mcpserver import (
|
|||
DeclinedElicitation,
|
||||
Elicit,
|
||||
ElicitationResult,
|
||||
ListRoots,
|
||||
MCPServer,
|
||||
RequestStateBoundary,
|
||||
RequestStateSecurity,
|
||||
Resolve,
|
||||
Sample,
|
||||
)
|
||||
from mcp.server.mcpserver.exceptions import InvalidSignature
|
||||
from mcp.server.mcpserver.resolve import (
|
||||
_check_elicit_return,
|
||||
_decode_state,
|
||||
_elicit_request,
|
||||
_encode_state,
|
||||
_outcome_from_state,
|
||||
_render_request,
|
||||
_request_digest,
|
||||
_resolver_key,
|
||||
_state_key,
|
||||
|
|
@ -54,11 +74,12 @@ from mcp.server.mcpserver.resolve import (
|
|||
)
|
||||
from mcp.server.mcpserver.tools.base import Tool
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.shared.message import SessionMessage
|
||||
|
||||
|
||||
def _question_digest(elicit: Elicit[Any]) -> str:
|
||||
"""The digest `_elicit` pins: the rendered request the client would be shown."""
|
||||
return _request_digest(_elicit_request(elicit))
|
||||
"""The digest `_fulfil` pins: the rendered request the client would be shown."""
|
||||
return _request_digest(_render_request(elicit))
|
||||
|
||||
|
||||
class Login(BaseModel):
|
||||
|
|
@ -177,42 +198,6 @@ def _wire_key(fn: Callable[..., Any]) -> str:
|
|||
return f"{fn.__module__}:{fn.__qualname__}"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_resolver_returns_value_directly_without_eliciting():
|
||||
mcp = MCPServer(name="Direct", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def login(ctx: Context) -> Login | Elicit[Login]:
|
||||
username = (ctx.headers or {}).get("x-github-user")
|
||||
if username: # pragma: no cover - no headers on in-memory transport
|
||||
return Login(username=username)
|
||||
return Login(username="from-resolver")
|
||||
|
||||
@mcp.tool()
|
||||
async def whoami(login: Annotated[Login, Resolve(login)]) -> str:
|
||||
return login.username
|
||||
|
||||
async def never(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult: # pragma: no cover
|
||||
raise AssertionError("should not elicit")
|
||||
|
||||
async with Client(mcp, mode="legacy", elicitation_callback=never) as client:
|
||||
assert await _text(client, "whoami", {}) == "from-resolver"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_resolver_elicits_and_injects_unwrapped_model_on_accept():
|
||||
mcp = MCPServer(name="Accept", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def login(ctx: Context) -> Login | Elicit[Login]:
|
||||
return Elicit("GitHub username?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def whoami(login: Annotated[Login, Resolve(login)]) -> str:
|
||||
return login.username
|
||||
|
||||
async with Client(mcp, mode="legacy", elicitation_callback=_accept({"username": "octocat"})) as client:
|
||||
assert await _text(client, "whoami", {}) == "octocat"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_consumer_receives_result_union_and_branches():
|
||||
mcp = MCPServer(name="Union", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
|
@ -232,43 +217,6 @@ async def test_consumer_receives_result_union_and_branches():
|
|||
assert await _text(client, "whoami", {}) == "hi octocat"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_decline_reaches_union_consumer_without_aborting():
|
||||
mcp = MCPServer(name="UnionDecline", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def login(ctx: Context) -> Login | Elicit[Login]:
|
||||
return Elicit("GitHub username?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def whoami(
|
||||
login: Annotated[AcceptedElicitation[Login] | DeclinedElicitation | CancelledElicitation, Resolve(login)],
|
||||
) -> str:
|
||||
if isinstance(login, DeclinedElicitation):
|
||||
return "declined gracefully"
|
||||
raise NotImplementedError
|
||||
|
||||
async with Client(mcp, mode="legacy", elicitation_callback=_decline) as client:
|
||||
assert await _text(client, "whoami", {}) == "declined gracefully"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_decline_aborts_when_consumer_wants_unwrapped():
|
||||
mcp = MCPServer(name="UnwrappedDecline", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def login(ctx: Context) -> Login | Elicit[Login]:
|
||||
return Elicit("GitHub username?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def whoami(login: Annotated[Login, Resolve(login)]) -> str:
|
||||
raise NotImplementedError # pragma: no cover - never reached
|
||||
|
||||
async with Client(mcp, mode="legacy", elicitation_callback=_decline) as client:
|
||||
result = await client.call_tool("whoami", {})
|
||||
assert result.is_error
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert "decline" in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_nested_resolver_sees_dependency_and_tool_args():
|
||||
mcp = MCPServer(name="Nested", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
|
@ -348,22 +296,6 @@ async def test_sync_resolver():
|
|||
assert await _text(client, "whoami", {}) == "sync-user"
|
||||
|
||||
|
||||
def test_resolved_params_absent_from_input_schema():
|
||||
async def login(ctx: Context) -> Login:
|
||||
return Login(username="x") # pragma: no cover - only the schema is inspected
|
||||
|
||||
async def tool(
|
||||
repo: Annotated[str, Field(description="repo name")],
|
||||
login: Annotated[Login, Resolve(login)],
|
||||
) -> str:
|
||||
return repo # pragma: no cover - only the schema is inspected
|
||||
|
||||
built = Tool.from_function(tool)
|
||||
properties = built.parameters["properties"]
|
||||
assert "repo" in properties
|
||||
assert "login" not in properties
|
||||
|
||||
|
||||
def test_cycle_detection_raises_at_registration():
|
||||
async def a(dep: Login) -> Login:
|
||||
return dep # pragma: no cover
|
||||
|
|
@ -425,7 +357,7 @@ def test_multiple_elicit_arms_raise_at_registration():
|
|||
async def tool(login: Annotated[Login, Resolve(ambiguous)]) -> str:
|
||||
return login.username # pragma: no cover
|
||||
|
||||
with pytest.raises(InvalidSignature, match="multiple Elicit arms"):
|
||||
with pytest.raises(InvalidSignature, match="multiple Elicit/Sample/ListRoots arms"):
|
||||
Tool.from_function(tool)
|
||||
|
||||
|
||||
|
|
@ -938,15 +870,16 @@ def test_state_round_trips_accept_decline_cancel():
|
|||
assert decoded == entries # encode-restore is the identity on the stored entries
|
||||
assert state.asked == {"e": "asked-digest"}
|
||||
|
||||
accepted = _outcome_from_state(decoded["a"], Login)
|
||||
ask = Elicit("q", Login)
|
||||
accepted = _outcome_from_state(decoded["a"], ask)
|
||||
assert isinstance(accepted, AcceptedElicitation) and accepted.data == Login(username="octocat")
|
||||
# Decline/cancel entries carry no data; the schema is not consulted for them.
|
||||
assert isinstance(_outcome_from_state(decoded["b"], Login), DeclinedElicitation)
|
||||
assert isinstance(_outcome_from_state(decoded["c"], Login), CancelledElicitation)
|
||||
assert isinstance(_outcome_from_state(decoded["b"], ask), DeclinedElicitation)
|
||||
assert isinstance(_outcome_from_state(decoded["c"], ask), CancelledElicitation)
|
||||
# An accepted restore always validates against the question's live schema -
|
||||
# data that doesn't fit is rejected, never passed through raw.
|
||||
with pytest.raises(ValidationError):
|
||||
_outcome_from_state(decoded["d"], Login)
|
||||
_outcome_from_state(decoded["d"], ask)
|
||||
|
||||
|
||||
def test_check_elicit_return_allows_one_arm_and_rejects_two():
|
||||
|
|
@ -955,7 +888,7 @@ def test_check_elicit_return_allows_one_arm_and_rejects_two():
|
|||
_check_elicit_return(Login, "r") # no Elicit arm
|
||||
_check_elicit_return(None, "r") # unannotated
|
||||
# A resolver asks one question: two distinct Elicit arms mean it should be split.
|
||||
with pytest.raises(InvalidSignature, match="'r' return annotation has multiple Elicit arms"):
|
||||
with pytest.raises(InvalidSignature, match="'r' return annotation has multiple Elicit/Sample/ListRoots arms"):
|
||||
_check_elicit_return(Elicit[Login] | Elicit[Confirm], "r")
|
||||
|
||||
|
||||
|
|
@ -1118,28 +1051,6 @@ async def test_accept_with_no_content_is_an_error_not_a_cancel(mode: Literal["le
|
|||
assert "no content" in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_eliciting_tool_without_client_capability_is_a_protocol_error():
|
||||
# The server must not send an `input_requests` entry the client has not declared
|
||||
# capability for: with no `elicitation` declared (no callback), the call fails as
|
||||
# a -32021 protocol error, not a CallToolResult execution failure.
|
||||
mcp = MCPServer(name="NoElicitationCapability", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def ask(ctx: Context) -> Elicit[Login]:
|
||||
return Elicit("user?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def tool(login: Annotated[Login, Resolve(ask)]) -> str:
|
||||
return login.username # pragma: no cover
|
||||
|
||||
async with Client(mcp) as client:
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await client.session.call_tool("tool", {}, allow_input_required=True)
|
||||
assert exc_info.value.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
||||
assert exc_info.value.error.data is not None
|
||||
assert "elicitation" in exc_info.value.error.data["requiredCapabilities"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_independent_nested_deps_batch_into_one_round():
|
||||
mcp = MCPServer(name="NestedBatch", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
|
@ -2365,3 +2276,593 @@ async def test_resolver_elicitation_seals_and_completes_on_a_fully_default_serve
|
|||
assert isinstance(final, CallToolResult)
|
||||
assert isinstance(final.content[0], TextContent)
|
||||
assert final.content[0].text == "went:True"
|
||||
|
||||
|
||||
# --- Sample / ListRoots markers ---
|
||||
|
||||
|
||||
async def _sample_never( # pragma: no cover - declares the capability; never invoked
|
||||
context: ClientRequestContext, params: CreateMessageRequestParams
|
||||
) -> CreateMessageResult:
|
||||
raise AssertionError("should not be called")
|
||||
|
||||
|
||||
async def _roots_never(context: ClientRequestContext) -> ListRootsResult: # pragma: no cover - see _sample_never
|
||||
raise AssertionError("should not be called")
|
||||
|
||||
|
||||
def _sample_capital(ctx: Context) -> Sample:
|
||||
return Sample(
|
||||
[SamplingMessage(role="user", content=TextContent(type="text", text="Capital of France?"))],
|
||||
max_tokens=16,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.filterwarnings("error::mcp.MCPDeprecationWarning")
|
||||
@pytest.mark.parametrize("mode", ["legacy", "auto"])
|
||||
async def test_sample_resolver_injects_result(mode: Literal["legacy", "auto"]):
|
||||
# The marker form is the 2026-blessed carrier: no SEP-2577 deprecation warning on either mode.
|
||||
mcp = MCPServer(name="Sampler", request_state_security=RequestStateSecurity.ephemeral())
|
||||
prompts: list[str] = []
|
||||
|
||||
async def sampler(context: ClientRequestContext, params: CreateMessageRequestParams) -> CreateMessageResult:
|
||||
content = params.messages[0].content
|
||||
assert isinstance(content, TextContent)
|
||||
prompts.append(content.text)
|
||||
return CreateMessageResult(role="assistant", content=TextContent(type="text", text="Paris"), model="m")
|
||||
|
||||
@mcp.tool()
|
||||
async def capital(answer: Annotated[CreateMessageResult, Resolve(_sample_capital)]) -> str:
|
||||
assert isinstance(answer.content, TextContent)
|
||||
return answer.content.text
|
||||
|
||||
async with Client(mcp, mode=mode, sampling_callback=sampler) as client:
|
||||
assert await _text(client, "capital", {}) == "Paris"
|
||||
assert prompts == ["Capital of France?"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.filterwarnings("error::mcp.MCPDeprecationWarning")
|
||||
@pytest.mark.parametrize("mode", ["legacy", "auto"])
|
||||
async def test_list_roots_resolver_injects_result(mode: Literal["legacy", "auto"]):
|
||||
mcp = MCPServer(name="Rooted", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def client_roots(context: ClientRequestContext) -> ListRootsResult:
|
||||
return ListRootsResult(roots=[Root(uri=FileUrl("file:///workspace"))])
|
||||
|
||||
def fetch_roots(ctx: Context) -> ListRoots:
|
||||
return ListRoots()
|
||||
|
||||
@mcp.tool()
|
||||
async def workspace(roots: Annotated[ListRootsResult, Resolve(fetch_roots)]) -> str:
|
||||
return str(len(roots.roots))
|
||||
|
||||
async with Client(mcp, mode=mode, list_roots_callback=client_roots) as client:
|
||||
assert await _text(client, "workspace", {}) == "1"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_mixed_kinds_batch_into_one_round():
|
||||
mcp = MCPServer(name="Mixed", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def ask_name(ctx: Context) -> Elicit[Login]:
|
||||
return Elicit("user?", Login)
|
||||
|
||||
async def fetch_roots(ctx: Context) -> ListRoots:
|
||||
return ListRoots()
|
||||
|
||||
@mcp.tool()
|
||||
async def combo(
|
||||
login: Annotated[Login, Resolve(ask_name)],
|
||||
answer: Annotated[CreateMessageResult, Resolve(_sample_capital)],
|
||||
roots: Annotated[ListRootsResult, Resolve(fetch_roots)],
|
||||
) -> str:
|
||||
assert isinstance(answer.content, TextContent)
|
||||
return f"{login.username}/{answer.content.text}/{len(roots.roots)}"
|
||||
|
||||
async with Client(
|
||||
mcp, elicitation_callback=_never, sampling_callback=_sample_never, list_roots_callback=_roots_never
|
||||
) as client:
|
||||
first = await client.session.call_tool("combo", {}, allow_input_required=True)
|
||||
assert isinstance(first, InputRequiredResult)
|
||||
assert first.input_requests is not None
|
||||
kinds = sorted(type(req).__name__ for req in first.input_requests.values())
|
||||
assert kinds == ["CreateMessageRequest", "ElicitRequest", "ListRootsRequest"]
|
||||
responses: InputResponses = {}
|
||||
for key, req in first.input_requests.items():
|
||||
if isinstance(req, ElicitRequest):
|
||||
responses[key] = ElicitResult(action="accept", content={"username": "octocat"})
|
||||
elif isinstance(req, CreateMessageRequest):
|
||||
responses[key] = CreateMessageResult(
|
||||
role="assistant", content=TextContent(type="text", text="hey"), model="m"
|
||||
)
|
||||
else:
|
||||
responses[key] = ListRootsResult(roots=[])
|
||||
final = await client.session.call_tool(
|
||||
"combo", {}, input_responses=responses, request_state=first.request_state, allow_input_required=True
|
||||
)
|
||||
assert isinstance(final, CallToolResult)
|
||||
assert isinstance(final.content[0], TextContent)
|
||||
assert final.content[0].text == "octocat/hey/0"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sampling_tool_without_client_capability_is_a_protocol_error():
|
||||
mcp = MCPServer(name="NoSamplingCapability", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
@mcp.tool()
|
||||
async def capital(answer: Annotated[CreateMessageResult, Resolve(_sample_capital)]) -> str:
|
||||
return "unreachable" # pragma: no cover
|
||||
|
||||
async with Client(mcp) as client:
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await client.session.call_tool("capital", {}, allow_input_required=True)
|
||||
assert exc_info.value.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
||||
assert exc_info.value.error.data is not None
|
||||
assert "sampling" in exc_info.value.error.data["requiredCapabilities"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_roots_tool_without_client_capability_is_a_protocol_error():
|
||||
mcp = MCPServer(name="NoRootsCapability", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def fetch_roots(ctx: Context) -> ListRoots:
|
||||
return ListRoots()
|
||||
|
||||
@mcp.tool()
|
||||
async def workspace(roots: Annotated[ListRootsResult, Resolve(fetch_roots)]) -> str:
|
||||
return "unreachable" # pragma: no cover
|
||||
|
||||
async with Client(mcp) as client:
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await client.session.call_tool("workspace", {}, allow_input_required=True)
|
||||
assert exc_info.value.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
||||
assert exc_info.value.error.data is not None
|
||||
assert "roots" in exc_info.value.error.data["requiredCapabilities"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_legacy_eliciting_tool_without_capability_is_a_protocol_error():
|
||||
# Same egress gate as the input_requests leg; the session stays usable after the refusal.
|
||||
mcp = MCPServer(name="LegacyGate", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def ask(ctx: Context) -> Elicit[Login]:
|
||||
return Elicit("user?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def tool(login: Annotated[Login, Resolve(ask)]) -> str:
|
||||
return login.username # pragma: no cover
|
||||
|
||||
@mcp.tool()
|
||||
def plain() -> str:
|
||||
return "ok"
|
||||
|
||||
async with Client(mcp, mode="legacy") as client:
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await client.call_tool("tool", {})
|
||||
assert exc_info.value.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
||||
assert await _text(client, "plain", {}) == "ok"
|
||||
|
||||
|
||||
def _ask_with_tools(ctx: Context) -> Sample:
|
||||
return Sample(
|
||||
[SamplingMessage(role="user", content=TextContent(type="text", text="2+2?"))],
|
||||
max_tokens=16,
|
||||
tools=[SamplingTool(name="calc", input_schema={"type": "object"})],
|
||||
)
|
||||
|
||||
|
||||
def _ask_with_tool_choice(ctx: Context) -> Sample:
|
||||
return Sample(
|
||||
[SamplingMessage(role="user", content=TextContent(type="text", text="2+2?"))],
|
||||
max_tokens=16,
|
||||
tool_choice=ToolChoice(mode="none"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("ask", [_ask_with_tools, _ask_with_tool_choice])
|
||||
async def test_sample_tools_require_the_tools_subcapability(ask: Callable[[Context], Sample]):
|
||||
mcp = MCPServer(name="NoToolsSubcapability", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
@mcp.tool()
|
||||
async def calc(answer: Annotated[CreateMessageResultWithTools, Resolve(ask)]) -> str:
|
||||
return "unreachable" # pragma: no cover
|
||||
|
||||
# The callback declares base `sampling` but not `sampling.tools`.
|
||||
async with Client(mcp, sampling_callback=_sample_never) as client:
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await client.session.call_tool("calc", {}, allow_input_required=True)
|
||||
assert exc_info.value.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
||||
assert exc_info.value.error.data is not None
|
||||
assert exc_info.value.error.data["requiredCapabilities"] == {"sampling": {"tools": {}}}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sample_with_tools_round_trips_with_declared_subcapability():
|
||||
mcp = MCPServer(name="ToolsSampling", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def sampler(
|
||||
context: ClientRequestContext, params: CreateMessageRequestParams
|
||||
) -> CreateMessageResultWithTools:
|
||||
assert params.tools is not None and params.tools[0].name == "calc"
|
||||
return CreateMessageResultWithTools(role="assistant", content=[TextContent(type="text", text="4")], model="m")
|
||||
|
||||
@mcp.tool()
|
||||
async def calc(answer: Annotated[CreateMessageResultWithTools, Resolve(_ask_with_tools)]) -> str:
|
||||
assert isinstance(answer.content, list) and isinstance(answer.content[0], TextContent)
|
||||
return answer.content[0].text
|
||||
|
||||
async with Client(
|
||||
mcp,
|
||||
sampling_callback=sampler,
|
||||
sampling_capabilities=SamplingCapability(tools=SamplingToolsCapability()),
|
||||
) as client:
|
||||
assert await _text(client, "calc", {}) == "4"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_no_tool_use_answer_to_a_tools_request_is_accepted():
|
||||
# The answer parses off the wire as plain CreateMessageResult but must inject as CreateMessageResultWithTools.
|
||||
mcp = MCPServer(name="NoToolUse", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
@mcp.tool()
|
||||
async def calc(answer: Annotated[CreateMessageResultWithTools, Resolve(_ask_with_tools)]) -> str:
|
||||
assert isinstance(answer, CreateMessageResultWithTools)
|
||||
assert isinstance(answer.content, TextContent)
|
||||
return answer.content.text
|
||||
|
||||
async with Client(
|
||||
mcp,
|
||||
sampling_callback=_sample_never,
|
||||
sampling_capabilities=SamplingCapability(tools=SamplingToolsCapability()),
|
||||
) as client:
|
||||
first = await client.session.call_tool("calc", {}, allow_input_required=True)
|
||||
assert isinstance(first, InputRequiredResult)
|
||||
assert first.input_requests is not None
|
||||
(key,) = first.input_requests
|
||||
final = await client.session.call_tool(
|
||||
"calc",
|
||||
{},
|
||||
input_responses={
|
||||
key: CreateMessageResult(role="assistant", content=TextContent(type="text", text="4"), model="m")
|
||||
},
|
||||
request_state=first.request_state,
|
||||
allow_input_required=True,
|
||||
)
|
||||
assert isinstance(final, CallToolResult)
|
||||
assert not final.is_error
|
||||
assert isinstance(final.content[0], TextContent)
|
||||
assert final.content[0].text == "4"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sample_outcome_persists_across_rounds():
|
||||
# The confirm arm depends on the sample, forcing extra rounds that restore the result instead of re-sampling.
|
||||
mcp = MCPServer(name="Chain", request_state_security=RequestStateSecurity.ephemeral())
|
||||
samples = 0
|
||||
|
||||
async def sampler(context: ClientRequestContext, params: CreateMessageRequestParams) -> CreateMessageResult:
|
||||
nonlocal samples
|
||||
samples += 1
|
||||
return CreateMessageResult(role="assistant", content=TextContent(type="text", text="Paris"), model="m")
|
||||
|
||||
async def confirm(
|
||||
answer: Annotated[CreateMessageResult, Resolve(_sample_capital)], ctx: Context
|
||||
) -> Elicit[Confirm]:
|
||||
return Elicit("Accept the model's answer?", Confirm)
|
||||
|
||||
@mcp.tool()
|
||||
async def tool(
|
||||
ok: Annotated[Confirm, Resolve(confirm)],
|
||||
answer: Annotated[CreateMessageResult, Resolve(_sample_capital)],
|
||||
) -> str:
|
||||
assert isinstance(answer.content, TextContent)
|
||||
return f"{answer.content.text}:{ok.ok}"
|
||||
|
||||
async with Client(mcp, sampling_callback=sampler, elicitation_callback=_accept({"ok": True})) as client:
|
||||
assert await _text(client, "tool", {}) == "Paris:True"
|
||||
assert samples == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wrong_kind_response_for_sample_raises():
|
||||
mcp = MCPServer(name="WrongKind", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
@mcp.tool()
|
||||
async def capital(answer: Annotated[CreateMessageResult, Resolve(_sample_capital)]) -> str:
|
||||
return "unreachable" # pragma: no cover
|
||||
|
||||
async with Client(mcp, sampling_callback=_sample_never) as client:
|
||||
first = await client.session.call_tool("capital", {}, allow_input_required=True)
|
||||
assert isinstance(first, InputRequiredResult)
|
||||
assert first.input_requests is not None
|
||||
(key,) = first.input_requests
|
||||
final = await client.session.call_tool(
|
||||
"capital",
|
||||
{},
|
||||
input_responses={key: ElicitResult(action="accept", content={"x": "y"})},
|
||||
request_state=first.request_state,
|
||||
allow_input_required=True,
|
||||
)
|
||||
assert isinstance(final, CallToolResult)
|
||||
assert final.is_error
|
||||
assert isinstance(final.content[0], TextContent)
|
||||
assert "wrong kind" in final.content[0].text
|
||||
|
||||
|
||||
def test_mixed_marker_arms_raise_at_registration():
|
||||
async def ambiguous(ctx: Context) -> Sample | Elicit[Login]:
|
||||
raise NotImplementedError # pragma: no cover
|
||||
|
||||
async def tool(login: Annotated[Login, Resolve(ambiguous)]) -> str:
|
||||
return login.username # pragma: no cover
|
||||
|
||||
with pytest.raises(InvalidSignature, match="multiple Elicit/Sample/ListRoots arms"):
|
||||
Tool.from_function(tool)
|
||||
|
||||
|
||||
def test_marker_union_with_generic_alias_member_registers():
|
||||
# dict[str, Any] passes isinstance(c, type) on Python 3.10; the arm filter must not feed it to issubclass.
|
||||
async def maybe_ask(ctx: Context) -> Sample | dict[str, Any]:
|
||||
raise NotImplementedError # pragma: no cover
|
||||
|
||||
async def tool(answer: Annotated[CreateMessageResult, Resolve(maybe_ask)]) -> str:
|
||||
return "ok" # pragma: no cover
|
||||
|
||||
Tool.from_function(tool)
|
||||
|
||||
|
||||
def test_decline_entry_for_a_sample_marker_is_invalid():
|
||||
# Decline outcomes exist only for elicitations; for a Sample the entry's None data fails validation.
|
||||
with pytest.raises(ValidationError):
|
||||
_outcome_from_state(_StateEntry(action="decline"), _sample_capital(cast(Context, None)))
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_bare_initialized_session_is_still_gated():
|
||||
# notifications/initialized alone commits the handshake: a live back-channel, no declared capabilities.
|
||||
mcp = MCPServer(name="BareInit", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def ask(ctx: Context) -> Elicit[Login]:
|
||||
return Elicit("user?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def tool(login: Annotated[Login, Resolve(ask)]) -> str:
|
||||
return login.username # pragma: no cover
|
||||
|
||||
async with InMemoryTransport(mcp) as (read, write):
|
||||
await write.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized")))
|
||||
await write.send(
|
||||
SessionMessage(
|
||||
JSONRPCRequest(jsonrpc="2.0", id=1, method="tools/call", params={"name": "tool", "arguments": {}})
|
||||
)
|
||||
)
|
||||
with anyio.fail_after(5):
|
||||
message = await read.receive()
|
||||
assert isinstance(message, SessionMessage)
|
||||
assert isinstance(message.message, JSONRPCError)
|
||||
assert message.message.error.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("mode", ["legacy", "auto"])
|
||||
async def test_tool_choice_only_sample_validates_as_tools_mode(mode: Literal["legacy", "auto"]):
|
||||
# Gate and answer model share one predicate: tool_choice alone is tools-mode,
|
||||
# so a single-content answer still validates (WithTools accepts both shapes).
|
||||
mcp = MCPServer(name="ToolChoiceOnly", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
async def sampler(context: ClientRequestContext, params: CreateMessageRequestParams) -> CreateMessageResult:
|
||||
assert params.tool_choice is not None and params.tools is None
|
||||
return CreateMessageResult(role="assistant", content=TextContent(type="text", text="4"), model="m")
|
||||
|
||||
@mcp.tool()
|
||||
async def calc(answer: Annotated[CreateMessageResultWithTools, Resolve(_ask_with_tool_choice)]) -> str:
|
||||
assert isinstance(answer, CreateMessageResultWithTools)
|
||||
assert isinstance(answer.content, TextContent)
|
||||
return answer.content.text
|
||||
|
||||
async with Client(
|
||||
mcp,
|
||||
mode=mode,
|
||||
sampling_callback=sampler,
|
||||
sampling_capabilities=SamplingCapability(tools=SamplingToolsCapability()),
|
||||
) as client:
|
||||
assert await _text(client, "calc", {}) == "4"
|
||||
|
||||
|
||||
# --- Feature walk: interaction-level tests (public API only, one behaviour per test) ---
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_resolver_fills_its_parameter_without_any_client_interaction():
|
||||
"""A resolver that computes its value server-side injects it with no client
|
||||
callbacks involved. SDK-defined resolver contract."""
|
||||
mcp = MCPServer(name="Walk", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def price_of(title: str) -> int:
|
||||
return 42 if title == "Dune" else 0
|
||||
|
||||
@mcp.tool()
|
||||
async def quote(title: str, price: Annotated[int, Resolve(price_of)]) -> str:
|
||||
return f"{title}: {price}"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("quote", {"title": "Dune"})
|
||||
|
||||
assert not result.is_error
|
||||
assert result.content == [TextContent(type="text", text="Dune: 42")]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_resolver_may_depend_on_another_resolvers_value():
|
||||
"""A resolver declares its own Resolve dependency and receives that resolver's
|
||||
value before the tool body runs. SDK-defined resolver contract."""
|
||||
mcp = MCPServer(name="Walk", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def base_price(title: str) -> int:
|
||||
return 10 if title == "Dune" else 1
|
||||
|
||||
def with_tax(price: Annotated[int, Resolve(base_price)]) -> int:
|
||||
return price * 2
|
||||
|
||||
@mcp.tool()
|
||||
async def quote(title: str, total: Annotated[int, Resolve(with_tax)]) -> str:
|
||||
return f"{title} costs {total}"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("quote", {"title": "Dune"})
|
||||
|
||||
assert not result.is_error
|
||||
assert result.content == [TextContent(type="text", text="Dune costs 20")]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_supplied_value_for_a_resolved_parameter_is_discarded():
|
||||
"""A value the client sends under a resolved parameter's name never reaches the
|
||||
tool; the resolver's value wins. SDK-defined: resolved parameters are server-side only."""
|
||||
mcp = MCPServer(name="Walk", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def price_of(title: str) -> int:
|
||||
return 42
|
||||
|
||||
@mcp.tool()
|
||||
async def quote(title: str, price: Annotated[int, Resolve(price_of)]) -> str:
|
||||
return f"{title}: {price}"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("quote", {"title": "Dune", "price": 999})
|
||||
|
||||
assert not result.is_error
|
||||
assert result.content == [TextContent(type="text", text="Dune: 42")]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_resolved_parameters_are_absent_from_the_advertised_tool_schema():
|
||||
"""tools/list advertises only the model-facing parameters; resolved ones are
|
||||
invisible to the client. SDK-defined: resolvers are server-side dependency injection."""
|
||||
mcp = MCPServer(name="Walk", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def price_of(title: str) -> int:
|
||||
raise NotImplementedError # pragma: no cover - only the schema is inspected
|
||||
|
||||
@mcp.tool()
|
||||
async def quote(title: str, price: Annotated[int, Resolve(price_of)]) -> str:
|
||||
raise NotImplementedError # pragma: no cover - only the schema is inspected
|
||||
|
||||
async with Client(mcp) as client:
|
||||
(advertised,) = (await client.list_tools()).tools
|
||||
|
||||
assert advertised.input_schema == snapshot(
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"title": {"title": "Title", "type": "string"}},
|
||||
"required": ["title"],
|
||||
"title": "quoteArguments",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("mode", ["legacy", "auto"])
|
||||
async def test_an_elicit_marker_asks_the_client_and_injects_the_accepted_answer(mode: Literal["legacy", "auto"]):
|
||||
"""A resolver returning Elicit puts its question to the client's elicitation
|
||||
callback and the tool receives the validated model, identically over the 2025
|
||||
back-channel and the 2026 multi-round-trip flow. SDK-defined injection contract."""
|
||||
mcp = MCPServer(name="Walk", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def ask(ctx: Context) -> Elicit[Login]:
|
||||
return Elicit("GitHub username?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def whoami(login: Annotated[Login, Resolve(ask)]) -> str:
|
||||
return login.username
|
||||
|
||||
asked: list[str] = []
|
||||
|
||||
async def on_elicit(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:
|
||||
asked.append(params.message)
|
||||
return ElicitResult(action="accept", content={"username": "octocat"})
|
||||
|
||||
async with Client(mcp, mode=mode, elicitation_callback=on_elicit) as client:
|
||||
# The parametrize is only a real era matrix while the arms negotiate different revisions.
|
||||
assert client.protocol_version == ("2025-11-25" if mode == "legacy" else "2026-07-28")
|
||||
result = await client.call_tool("whoami", {})
|
||||
|
||||
assert not result.is_error
|
||||
assert result.content == [TextContent(type="text", text="octocat")]
|
||||
assert asked == ["GitHub username?"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("mode", ["legacy", "auto"])
|
||||
async def test_a_declined_elicitation_fails_the_call_for_a_plain_model_consumer(mode: Literal["legacy", "auto"]):
|
||||
"""Declining the question aborts a tool whose consumer asked for the unwrapped
|
||||
model, on either protocol era. Decline semantics are spec-defined; the abort is
|
||||
the SDK's contract for plain-model consumers."""
|
||||
mcp = MCPServer(name="Walk", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def ask(ctx: Context) -> Elicit[Login]:
|
||||
return Elicit("GitHub username?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def whoami(login: Annotated[Login, Resolve(ask)]) -> str:
|
||||
raise NotImplementedError # pragma: no cover - the decline aborts before the body
|
||||
|
||||
async with Client(mcp, mode=mode, elicitation_callback=_decline) as client:
|
||||
result = await client.call_tool("whoami", {})
|
||||
|
||||
assert result.is_error
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == snapshot(
|
||||
"Error executing tool whoami: Resolver for parameter 'login' could not resolve: elicitation was decline"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("mode", ["legacy", "auto"])
|
||||
async def test_a_declined_elicitation_reaches_a_consumer_that_asked_for_the_outcome_union(
|
||||
mode: Literal["legacy", "auto"],
|
||||
):
|
||||
"""Annotating the outcome union hands the tool the decline to branch on instead
|
||||
of aborting the call. SDK-defined consumer-annotation contract."""
|
||||
mcp = MCPServer(name="Walk", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def ask(ctx: Context) -> Elicit[Login]:
|
||||
return Elicit("GitHub username?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def whoami(login: Annotated[ElicitationResult[Login], Resolve(ask)]) -> str:
|
||||
if isinstance(login, AcceptedElicitation):
|
||||
return login.data.username # pragma: no cover - declined in this test
|
||||
return "anonymous"
|
||||
|
||||
async with Client(mcp, mode=mode, elicitation_callback=_decline) as client:
|
||||
result = await client.call_tool("whoami", {})
|
||||
|
||||
assert not result.is_error
|
||||
assert result.content == [TextContent(type="text", text="anonymous")]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("mode", ["legacy", "auto"])
|
||||
async def test_an_undeclared_capability_refuses_the_call_instead_of_asking(mode: Literal["legacy", "auto"]):
|
||||
"""A client that never declared the elicitation capability is refused with the
|
||||
missing-capability protocol error rather than sent a question it cannot handle.
|
||||
The egress rule is spec-mandated; the SDK applies it on both eras."""
|
||||
mcp = MCPServer(name="Walk", request_state_security=RequestStateSecurity.ephemeral())
|
||||
|
||||
def ask(ctx: Context) -> Elicit[Login]:
|
||||
return Elicit("GitHub username?", Login)
|
||||
|
||||
@mcp.tool()
|
||||
async def whoami(login: Annotated[Login, Resolve(ask)]) -> str:
|
||||
raise NotImplementedError # pragma: no cover - the gate refuses before the body
|
||||
|
||||
async with Client(mcp, mode=mode) as client:
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await client.call_tool("whoami", {})
|
||||
|
||||
assert exc_info.value.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
||||
assert exc_info.value.error.data == {"requiredCapabilities": {"elicitation": {"form": {}}}}
|
||||
|
|
|
|||
|
|
@ -225,6 +225,21 @@ async def test_create_message_with_tools_returns_with_tools_result():
|
|||
assert params is not None and params["tools"][0]["name"] == "t"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_with_tool_choice_only_returns_with_tools_result():
|
||||
# tool_choice alone is tools-mode: the answer may carry array content.
|
||||
outbound = StubOutbound(result={"role": "assistant", "content": [{"type": "text", "text": "ok"}], "model": "m"})
|
||||
session = _make_session(
|
||||
outbound, capabilities=ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability()))
|
||||
)
|
||||
result = await session.create_message( # pyright: ignore[reportDeprecated]
|
||||
messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hi"))],
|
||||
max_tokens=10,
|
||||
tool_choice=types.ToolChoice(mode="none"),
|
||||
)
|
||||
assert isinstance(result, types.CreateMessageResultWithTools)
|
||||
|
||||
|
||||
def test_check_client_capability_delegates_to_connection():
|
||||
outbound = StubOutbound()
|
||||
session = _make_session(outbound, capabilities=ClientCapabilities(sampling=SamplingCapability()))
|
||||
|
|
|
|||
Loading…
Reference in a new issue