Add support for Elicitation (#625)

Co-authored-by: David Soria Parra <davidsp@anthropic.com>
This commit is contained in:
Inna Harper 2025-06-17 11:54:11 +01:00 • committed by GitHub
commit 69e6572d48
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 643 additions and 8 deletions

View file

@ -0,0 +1,210 @@
"""
Test the elicitation feature using stdio transport.
"""
import pytest
from pydantic import BaseModel, Field
from mcp.server.fastmcp import Context, FastMCP
from mcp.shared.memory import create_connected_server_and_client_session
from mcp.types import ElicitResult, TextContent
# Shared schema for basic tests
class AnswerSchema(BaseModel):
answer: str = Field(description="The user's answer to the question")
def create_ask_user_tool(mcp: FastMCP):
"""Create a standard ask_user tool that handles all elicitation responses."""
@mcp.tool(description="A tool that uses elicitation")
async def ask_user(prompt: str, ctx: Context) -> str:
result = await ctx.elicit(
message=f"Tool wants to ask: {prompt}",
schema=AnswerSchema,
)
if result.action == "accept" and result.data:
return f"User answered: {result.data.answer}"
elif result.action == "decline":
return "User declined to answer"
else:
return "User cancelled"
return ask_user
async def call_tool_and_assert(
mcp: FastMCP,
elicitation_callback,
tool_name: str,
args: dict,
expected_text: str | None = None,
text_contains: list[str] | None = None,
):
"""Helper to create session, call tool, and assert result."""
async with create_connected_server_and_client_session(
mcp._mcp_server, elicitation_callback=elicitation_callback
) as client_session:
await client_session.initialize()
result = await client_session.call_tool(tool_name, args)
assert len(result.content) == 1
assert isinstance(result.content[0], TextContent)
if expected_text is not None:
assert result.content[0].text == expected_text
elif text_contains is not None:
for substring in text_contains:
assert substring in result.content[0].text
return result
@pytest.mark.anyio
async def test_stdio_elicitation():
"""Test the elicitation feature using stdio transport."""
mcp = FastMCP(name="StdioElicitationServer")
create_ask_user_tool(mcp)
# Create a custom handler for elicitation requests
async def elicitation_callback(context, params):
if params.message == "Tool wants to ask: What is your name?":
return ElicitResult(action="accept", content={"answer": "Test User"})
else:
raise ValueError(f"Unexpected elicitation message: {params.message}")
await call_tool_and_assert(
mcp, elicitation_callback, "ask_user", {"prompt": "What is your name?"}, "User answered: Test User"
)
@pytest.mark.anyio
async def test_stdio_elicitation_decline():
"""Test elicitation with user declining."""
mcp = FastMCP(name="StdioElicitationDeclineServer")
create_ask_user_tool(mcp)
async def elicitation_callback(context, params):
return ElicitResult(action="decline")
await call_tool_and_assert(
mcp, elicitation_callback, "ask_user", {"prompt": "What is your name?"}, "User declined to answer"
)
@pytest.mark.anyio
async def test_elicitation_schema_validation():
"""Test that elicitation schemas must only contain primitive types."""
mcp = FastMCP(name="ValidationTestServer")
def create_validation_tool(name: str, schema_class: type[BaseModel]):
@mcp.tool(name=name, description=f"Tool testing {name}")
async def tool(ctx: Context) -> str:
try:
await ctx.elicit(message="This should fail validation", schema=schema_class)
return "Should not reach here"
except TypeError as e:
return f"Validation failed as expected: {str(e)}"
return tool
# Test cases for invalid schemas
class InvalidListSchema(BaseModel):
names: list[str] = Field(description="List of names")
class NestedModel(BaseModel):
value: str
class InvalidNestedSchema(BaseModel):
nested: NestedModel = Field(description="Nested model")
create_validation_tool("invalid_list", InvalidListSchema)
create_validation_tool("nested_model", InvalidNestedSchema)
# Dummy callback (won't be called due to validation failure)
async def elicitation_callback(context, params):
return ElicitResult(action="accept", content={})
async with create_connected_server_and_client_session(
mcp._mcp_server, elicitation_callback=elicitation_callback
) as client_session:
await client_session.initialize()
# Test both invalid schemas
for tool_name, field_name in [("invalid_list", "names"), ("nested_model", "nested")]:
result = await client_session.call_tool(tool_name, {})
assert len(result.content) == 1
assert isinstance(result.content[0], TextContent)
assert "Validation failed as expected" in result.content[0].text
assert field_name in result.content[0].text
@pytest.mark.anyio
async def test_elicitation_with_optional_fields():
"""Test that Optional fields work correctly in elicitation schemas."""
mcp = FastMCP(name="OptionalFieldServer")
class OptionalSchema(BaseModel):
required_name: str = Field(description="Your name (required)")
optional_age: int | None = Field(default=None, description="Your age (optional)")
optional_email: str | None = Field(default=None, description="Your email (optional)")
subscribe: bool | None = Field(default=False, description="Subscribe to newsletter?")
@mcp.tool(description="Tool with optional fields")
async def optional_tool(ctx: Context) -> str:
result = await ctx.elicit(message="Please provide your information", schema=OptionalSchema)
if result.action == "accept" and result.data:
info = [f"Name: {result.data.required_name}"]
if result.data.optional_age is not None:
info.append(f"Age: {result.data.optional_age}")
if result.data.optional_email is not None:
info.append(f"Email: {result.data.optional_email}")
info.append(f"Subscribe: {result.data.subscribe}")
return ", ".join(info)
else:
return f"User {result.action}"
# Test cases with different field combinations
test_cases = [
(
# All fields provided
{"required_name": "John Doe", "optional_age": 30, "optional_email": "john@example.com", "subscribe": True},
"Name: John Doe, Age: 30, Email: john@example.com, Subscribe: True",
),
(
# Only required fields
{"required_name": "Jane Smith"},
"Name: Jane Smith, Subscribe: False",
),
]
for content, expected in test_cases:
async def callback(context, params):
return ElicitResult(action="accept", content=content)
await call_tool_and_assert(mcp, callback, "optional_tool", {}, expected)
# Test invalid optional field
class InvalidOptionalSchema(BaseModel):
name: str = Field(description="Name")
optional_list: list[str] | None = Field(default=None, description="Invalid optional list")
@mcp.tool(description="Tool with invalid optional field")
async def invalid_optional_tool(ctx: Context) -> str:
try:
await ctx.elicit(message="This should fail", schema=InvalidOptionalSchema)
return "Should not reach here"
except TypeError as e:
return f"Validation failed: {str(e)}"
await call_tool_and_assert(
mcp,
lambda c, p: ElicitResult(action="accept", content={}),
"invalid_optional_tool",
{},
text_contains=["Validation failed:", "optional_list"],
)

View file

@ -14,14 +14,14 @@ from typing import Any
import pytest
import uvicorn
from pydantic import AnyUrl
from pydantic import AnyUrl, BaseModel, Field
from starlette.applications import Starlette
from starlette.requests import Request
from mcp.client.session import ClientSession
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamablehttp_client
from mcp.server.fastmcp import FastMCP
from mcp.server.fastmcp import Context, FastMCP
from mcp.server.fastmcp.resources import FunctionResource
from mcp.shared.context import RequestContext
from mcp.types import (
@ -30,6 +30,7 @@ from mcp.types import (
CompletionContext,
CreateMessageRequestParams,
CreateMessageResult,
ElicitResult,
GetPromptResult,
InitializeResult,
LoggingMessageNotification,
@ -98,6 +99,20 @@ def make_fastmcp_app():
def echo(message: str) -> str:
return f"Echo: {message}"
# Add a tool that uses elicitation
@mcp.tool(description="A tool that uses elicitation")
async def ask_user(prompt: str, ctx: Context) -> str:
class AnswerSchema(BaseModel):
answer: str = Field(description="The user's answer to the question")
result = await ctx.elicit(message=f"Tool wants to ask: {prompt}", schema=AnswerSchema)
if result.action == "accept" and result.data:
return f"User answered: {result.data.answer}"
else:
# Handle cancellation or decline
return f"User cancelled or declined: {result.action}"
# Create the SSE app
app = mcp.sse_app()
@ -261,6 +276,49 @@ def make_everything_fastmcp() -> FastMCP:
context_data["path"] = request.url.path
return json.dumps(context_data)
# Restaurant booking tool with elicitation
@mcp.tool(description="Book a table at a restaurant with elicitation")
async def book_restaurant(
date: str,
time: str,
party_size: int,
ctx: Context,
) -> str:
"""Book a table - uses elicitation if requested date is unavailable."""
class AlternativeDateSchema(BaseModel):
checkAlternative: bool = Field(description="Would you like to try another date?")
alternativeDate: str = Field(
default="2024-12-26",
description="What date would you prefer? (YYYY-MM-DD)",
)
# For testing: assume dates starting with "2024-12-25" are unavailable
if date.startswith("2024-12-25"):
# Use elicitation to ask about alternatives
result = await ctx.elicit(
message=(
f"No tables available for {party_size} people on {date} "
f"at {time}. Would you like to check another date?"
),
schema=AlternativeDateSchema,
)
if result.action == "accept" and result.data:
if result.data.checkAlternative:
alt_date = result.data.alternativeDate
return f"✅ Booked table for {party_size} on {alt_date} at {time}"
else:
return "❌ No booking made"
elif result.action in ("decline", "cancel"):
return "❌ Booking cancelled"
else:
# Handle case where action is "accept" but data is None
return "❌ No booking data received"
else:
# Available - book directly
return f"✅ Booked table for {party_size} on {date} at {time}"
return mcp
@ -652,6 +710,22 @@ class NotificationCollector:
await self.handle_tool_list_changed(message.root.params)
async def create_test_elicitation_callback(context, params):
"""Shared elicitation callback for tests.
Handles elicitation requests for restaurant booking tests.
"""
# For restaurant booking test
if "No tables available" in params.message:
return ElicitResult(
action="accept",
content={"checkAlternative": True, "alternativeDate": "2024-12-26"},
)
else:
# Default response
return ElicitResult(action="decline")
async def call_all_mcp_features(session: ClientSession, collector: NotificationCollector) -> None:
"""
Test all MCP features using the provided session.
@ -747,6 +821,21 @@ async def call_all_mcp_features(session: ClientSession, collector: NotificationC
assert "info" in log_levels
assert "warning" in log_levels
# 5. Test elicitation tool
# Test restaurant booking with unavailable date (triggers elicitation)
booking_result = await session.call_tool(
"book_restaurant",
{
"date": "2024-12-25", # Unavailable date to trigger elicitation
"time": "19:00",
"party_size": 4,
},
)
assert len(booking_result.content) == 1
assert isinstance(booking_result.content[0], TextContent)
# Should have booked the alternative date from elicitation callback
assert "✅ Booked table for 4 on 2024-12-26" in booking_result.content[0].text
# Test resources
# 1. Static resource
resources = await session.list_resources()
@ -887,8 +976,6 @@ async def test_fastmcp_all_features_sse(everything_server: None, everything_serv
# Create notification collector
collector = NotificationCollector()
# Create a sampling callback that simulates an LLM
# Connect to the server with callbacks
async with sse_client(everything_server_url + "/sse") as streams:
# Set up message handler to capture notifications
@ -901,6 +988,7 @@ async def test_fastmcp_all_features_sse(everything_server: None, everything_serv
async with ClientSession(
*streams,
sampling_callback=sampling_callback,
elicitation_callback=create_test_elicitation_callback,
message_handler=message_handler,
) as session:
# Run the common test suite
@ -933,7 +1021,37 @@ async def test_fastmcp_all_features_streamable_http(
read_stream,
write_stream,
sampling_callback=sampling_callback,
elicitation_callback=create_test_elicitation_callback,
message_handler=message_handler,
) as session:
# Run the common test suite with HTTP-specific test suffix
await call_all_mcp_features(session, collector)
@pytest.mark.anyio
async def test_elicitation_feature(server: None, server_url: str) -> None:
"""Test the elicitation feature."""
# Create a custom handler for elicitation requests
async def elicitation_callback(context, params):
# Verify the elicitation parameters
if params.message == "Tool wants to ask: What is your name?":
return ElicitResult(content={"answer": "Test User"}, action="accept")
else:
raise ValueError("Unexpected elicitation message")
# Connect to the server with our custom elicitation handler
async with sse_client(server_url + "/sse") as streams:
async with ClientSession(*streams, elicitation_callback=elicitation_callback) as session:
# First initialize the session
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.serverInfo.name == "NoAuthServer"
# Call the tool that uses elicitation
tool_result = await session.call_tool("ask_user", {"prompt": "What is your name?"})
# Verify the result
assert len(tool_result.content) == 1
assert isinstance(tool_result.content[0], TextContent)
# # The test should only succeed with the successful elicitation response
assert tool_result.content[0].text == "User answered: Test User"