diff --git a/README.md b/README.md index 38154c878..1bdae1d20 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,8 @@ - [Images](#images) - [Context](#context) - [Completions](#completions) + - [Elicitation](#elicitation) + - [Authentication](#authentication) - [Running Your Server](#running-your-server) - [Development Mode](#development-mode) - [Claude Desktop Integration](#claude-desktop-integration) @@ -74,7 +76,7 @@ The Model Context Protocol allows applications to provide context for LLMs in a ### Adding MCP to your python project -We recommend using [uv](https://docs.astral.sh/uv/) to manage your Python projects. +We recommend using [uv](https://docs.astral.sh/uv/) to manage your Python projects. If you haven't created a uv-managed project yet, create one: @@ -372,6 +374,50 @@ async def handle_completion( return Completion(values=filtered) return None ``` +### Elicitation + +Request additional information from users during tool execution: + +```python +from mcp.server.fastmcp import FastMCP, Context +from mcp.server.elicitation import ( + AcceptedElicitation, + DeclinedElicitation, + CancelledElicitation, +) +from pydantic import BaseModel, Field + +mcp = FastMCP("Booking System") + + +@mcp.tool() +async def book_table(date: str, party_size: int, ctx: Context) -> str: + """Book a table with confirmation""" + + # Schema must only contain primitive types (str, int, float, bool) + class ConfirmBooking(BaseModel): + confirm: bool = Field(description="Confirm booking?") + notes: str = Field(default="", description="Special requests") + + result = await ctx.elicit( + message=f"Confirm booking for {party_size} on {date}?", schema=ConfirmBooking + ) + + match result: + case AcceptedElicitation(data=data): + if data.confirm: + return f"Booked! Notes: {data.notes or 'None'}" + return "Booking cancelled" + case DeclinedElicitation(): + return "Booking declined" + case CancelledElicitation(): + return "Booking cancelled" +``` + +The `elicit()` method returns an `ElicitationResult` with: +- `action`: "accept", "decline", or "cancel" +- `data`: The validated response (only when accepted) +- `validation_error`: Any validation error message ### Authentication diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 8b819ad6d..948817140 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -22,6 +22,14 @@ async def __call__( ) -> types.CreateMessageResult | types.ErrorData: ... +class ElicitationFnT(Protocol): + async def __call__( + self, + context: RequestContext["ClientSession", Any], + params: types.ElicitRequestParams, + ) -> types.ElicitResult | types.ErrorData: ... + + class ListRootsFnT(Protocol): async def __call__( self, context: RequestContext["ClientSession", Any] @@ -58,6 +66,16 @@ async def _default_sampling_callback( ) +async def _default_elicitation_callback( + context: RequestContext["ClientSession", Any], + params: types.ElicitRequestParams, +) -> types.ElicitResult | types.ErrorData: + return types.ErrorData( + code=types.INVALID_REQUEST, + message="Elicitation not supported", + ) + + async def _default_list_roots_callback( context: RequestContext["ClientSession", Any], ) -> types.ListRootsResult | types.ErrorData: @@ -91,6 +109,7 @@ def __init__( write_stream: MemoryObjectSendStream[SessionMessage], read_timeout_seconds: timedelta | None = None, sampling_callback: SamplingFnT | None = None, + elicitation_callback: ElicitationFnT | None = None, list_roots_callback: ListRootsFnT | None = None, logging_callback: LoggingFnT | None = None, message_handler: MessageHandlerFnT | None = None, @@ -105,12 +124,16 @@ def __init__( ) self._client_info = client_info or DEFAULT_CLIENT_INFO self._sampling_callback = sampling_callback or _default_sampling_callback + self._elicitation_callback = elicitation_callback or _default_elicitation_callback self._list_roots_callback = list_roots_callback or _default_list_roots_callback self._logging_callback = logging_callback or _default_logging_callback self._message_handler = message_handler or _default_message_handler async def initialize(self) -> types.InitializeResult: sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None + elicitation = ( + types.ElicitationCapability() if self._elicitation_callback is not _default_elicitation_callback else None + ) roots = ( # TODO: Should this be based on whether we # _will_ send notifications, or only whether @@ -128,6 +151,7 @@ async def initialize(self) -> types.InitializeResult: protocolVersion=types.LATEST_PROTOCOL_VERSION, capabilities=types.ClientCapabilities( sampling=sampling, + elicitation=elicitation, experimental=None, roots=roots, ), @@ -362,6 +386,12 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques client_response = ClientResponse.validate_python(response) await responder.respond(client_response) + case types.ElicitRequest(params=params): + with responder: + response = await self._elicitation_callback(ctx, params) + client_response = ClientResponse.validate_python(response) + await responder.respond(client_response) + case types.ListRootsRequest(): with responder: response = await self._list_roots_callback(ctx) diff --git a/src/mcp/server/elicitation.py b/src/mcp/server/elicitation.py new file mode 100644 index 000000000..1e48738c8 --- /dev/null +++ b/src/mcp/server/elicitation.py @@ -0,0 +1,111 @@ +"""Elicitation utilities for MCP servers.""" + +from __future__ import annotations + +import types +from typing import Generic, Literal, TypeVar, Union, get_args, get_origin + +from pydantic import BaseModel +from pydantic.fields import FieldInfo + +from mcp.server.session import ServerSession +from mcp.types import RequestId + +ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel) + + +class AcceptedElicitation(BaseModel, Generic[ElicitSchemaModelT]): + """Result when user accepts the elicitation.""" + + action: Literal["accept"] = "accept" + data: ElicitSchemaModelT + + +class DeclinedElicitation(BaseModel): + """Result when user declines the elicitation.""" + + action: Literal["decline"] = "decline" + + +class CancelledElicitation(BaseModel): + """Result when user cancels the elicitation.""" + + action: Literal["cancel"] = "cancel" + + +ElicitationResult = AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation + + +# Primitive types allowed in elicitation schemas +_ELICITATION_PRIMITIVE_TYPES = (str, int, float, bool) + + +def _validate_elicitation_schema(schema: type[BaseModel]) -> None: + """Validate that a Pydantic model only contains primitive field types.""" + for field_name, field_info in schema.model_fields.items(): + if not _is_primitive_field(field_info): + raise TypeError( + f"Elicitation schema field '{field_name}' must be a primitive type " + f"{_ELICITATION_PRIMITIVE_TYPES} or Optional of these types. " + f"Complex types like lists, dicts, or nested models are not allowed." + ) + + +def _is_primitive_field(field_info: FieldInfo) -> bool: + """Check if a field is a primitive type allowed in elicitation schemas.""" + annotation = field_info.annotation + + # Handle None type + if annotation is types.NoneType: + return True + + # Handle basic primitive types + if annotation in _ELICITATION_PRIMITIVE_TYPES: + return True + + # Handle Union types + origin = get_origin(annotation) + if origin is Union or origin is types.UnionType: + args = get_args(annotation) + # All args must be primitive types or None + return all(arg is types.NoneType or arg in _ELICITATION_PRIMITIVE_TYPES for arg in args) + + return False + + +async def elicit_with_validation( + session: ServerSession, + message: str, + schema: type[ElicitSchemaModelT], + related_request_id: RequestId | None = None, +) -> ElicitationResult[ElicitSchemaModelT]: + """Elicit information from the client/user with schema validation. + + This method can be used to interactively ask for additional information from the + client within a tool's execution. The client might display the message to the + user and collect a response according to the provided schema. Or in case a + client is an agent, it might decide how to handle the elicitation -- either by asking + the user or automatically generating a response. + """ + # Validate that schema only contains primitive types and fail loudly if not + _validate_elicitation_schema(schema) + + json_schema = schema.model_json_schema() + + result = await session.elicit( + message=message, + requestedSchema=json_schema, + related_request_id=related_request_id, + ) + + if result.action == "accept" and result.content: + # Validate and parse the content using the schema + validated_data = schema.model_validate(result.content) + return AcceptedElicitation(data=validated_data) + elif result.action == "decline": + return DeclinedElicitation() + elif result.action == "cancel": + return CancelledElicitation() + else: + # This should never happen, but handle it just in case + raise ValueError(f"Unexpected elicitation action: {result.action}") diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index a62974bc9..a85c7117c 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -34,6 +34,7 @@ from mcp.server.auth.settings import ( AuthSettings, ) +from mcp.server.elicitation import ElicitationResult, ElicitSchemaModelT, elicit_with_validation from mcp.server.fastmcp.exceptions import ResourceError from mcp.server.fastmcp.prompts import Prompt, PromptManager from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager @@ -972,6 +973,37 @@ async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContent assert self._fastmcp is not None, "Context is not available outside of a request" return await self._fastmcp.read_resource(uri) + async def elicit( + self, + message: str, + schema: type[ElicitSchemaModelT], + ) -> ElicitationResult[ElicitSchemaModelT]: + """Elicit information from the client/user. + + This method can be used to interactively ask for additional information from the + client within a tool's execution. The client might display the message to the + user and collect a response according to the provided schema. Or in case a + client is an agent, it might decide how to handle the elicitation -- either by asking + the user or automatically generating a response. + + Args: + schema: A Pydantic model class defining the expected response structure, according to the specification, + only primive types are allowed. + message: Optional message to present to the user. If not provided, will use + a default message based on the schema + + Returns: + An ElicitationResult containing the action taken and the data if accepted + + Note: + Check the result.action to determine if the user accepted, declined, or cancelled. + The result.data will only be populated if action is "accept" and validation succeeded. + """ + + return await elicit_with_validation( + session=self.request_context.session, message=message, schema=schema, related_request_id=self.request_id + ) + async def log( self, level: Literal["debug", "info", "warning", "error"], diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py index e6611b0d4..5c696b136 100644 --- a/src/mcp/server/session.py +++ b/src/mcp/server/session.py @@ -121,6 +121,10 @@ def check_client_capability(self, capability: types.ClientCapabilities) -> bool: if client_caps.sampling is None: return False + if capability.elicitation is not None: + if client_caps.elicitation is None: + return False + if capability.experimental is not None: if client_caps.experimental is None: return False @@ -251,6 +255,35 @@ async def list_roots(self) -> types.ListRootsResult: types.ListRootsResult, ) + async def elicit( + self, + message: str, + requestedSchema: types.ElicitRequestedSchema, + related_request_id: types.RequestId | None = None, + ) -> types.ElicitResult: + """Send an elicitation/create request. + + Args: + message: The message to present to the user + requestedSchema: Schema defining the expected response structure + + Returns: + The client's response + """ + return await self.send_request( + types.ServerRequest( + types.ElicitRequest( + method="elicitation/create", + params=types.ElicitRequestParams( + message=message, + requestedSchema=requestedSchema, + ), + ) + ), + types.ElicitResult, + metadata=ServerMessageMetadata(related_request_id=related_request_id), + ) + async def send_ping(self) -> types.EmptyResult: """Send a ping request.""" return await self.send_request( diff --git a/src/mcp/shared/memory.py b/src/mcp/shared/memory.py index f088d3f8b..c94e5e6ac 100644 --- a/src/mcp/shared/memory.py +++ b/src/mcp/shared/memory.py @@ -11,7 +11,14 @@ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream import mcp.types as types -from mcp.client.session import ClientSession, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT +from mcp.client.session import ( + ClientSession, + ElicitationFnT, + ListRootsFnT, + LoggingFnT, + MessageHandlerFnT, + SamplingFnT, +) from mcp.server import Server from mcp.shared.message import SessionMessage @@ -53,6 +60,7 @@ async def create_connected_server_and_client_session( message_handler: MessageHandlerFnT | None = None, client_info: types.Implementation | None = None, raise_exceptions: bool = False, + elicitation_callback: ElicitationFnT | None = None, ) -> AsyncGenerator[ClientSession, None]: """Creates a ClientSession that is connected to a running MCP server.""" async with create_client_server_memory_streams() as ( @@ -83,6 +91,7 @@ async def create_connected_server_and_client_session( logging_callback=logging_callback, message_handler=message_handler, client_info=client_info, + elicitation_callback=elicitation_callback, ) as client_session: await client_session.initialize() yield client_session diff --git a/src/mcp/types.py b/src/mcp/types.py index be8d7326f..a678f101d 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -216,6 +216,12 @@ class SamplingCapability(BaseModel): model_config = ConfigDict(extra="allow") +class ElicitationCapability(BaseModel): + """Capability for elicitation operations.""" + + model_config = ConfigDict(extra="allow") + + class ClientCapabilities(BaseModel): """Capabilities a client may support.""" @@ -223,6 +229,8 @@ class ClientCapabilities(BaseModel): """Experimental, non-standard capabilities that the client supports.""" sampling: SamplingCapability | None = None """Present if the client supports sampling from an LLM.""" + elicitation: ElicitationCapability | None = None + """Present if the client supports elicitation from the user.""" roots: RootsCapability | None = None """Present if the client supports listing roots.""" model_config = ConfigDict(extra="allow") @@ -1186,11 +1194,49 @@ class ClientNotification( pass -class ClientResult(RootModel[EmptyResult | CreateMessageResult | ListRootsResult]): +# Type for elicitation schema - a JSON Schema dict +ElicitRequestedSchema: TypeAlias = dict[str, Any] +"""Schema for elicitation requests.""" + + +class ElicitRequestParams(RequestParams): + """Parameters for elicitation requests.""" + + message: str + requestedSchema: ElicitRequestedSchema + model_config = ConfigDict(extra="allow") + + +class ElicitRequest(Request[ElicitRequestParams, Literal["elicitation/create"]]): + """A request from the server to elicit information from the client.""" + + method: Literal["elicitation/create"] + params: ElicitRequestParams + + +class ElicitResult(Result): + """The client's response to an elicitation request.""" + + action: Literal["accept", "decline", "cancel"] + """ + The user action in response to the elicitation. + - "accept": User submitted the form/confirmed the action + - "decline": User explicitly declined the action + - "cancel": User dismissed without making an explicit choice + """ + + content: dict[str, str | int | float | bool | None] | None = None + """ + The submitted form data, only present when action is "accept". + Contains values matching the requested schema. + """ + + +class ClientResult(RootModel[EmptyResult | CreateMessageResult | ListRootsResult | ElicitResult]): pass -class ServerRequest(RootModel[PingRequest | CreateMessageRequest | ListRootsRequest]): +class ServerRequest(RootModel[PingRequest | CreateMessageRequest | ListRootsRequest | ElicitRequest]): pass diff --git a/tests/server/fastmcp/test_elicitation.py b/tests/server/fastmcp/test_elicitation.py new file mode 100644 index 000000000..20937d91d --- /dev/null +++ b/tests/server/fastmcp/test_elicitation.py @@ -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"], + ) diff --git a/tests/server/fastmcp/test_integration.py b/tests/server/fastmcp/test_integration.py index 3caf994e0..3eb013946 100644 --- a/tests/server/fastmcp/test_integration.py +++ b/tests/server/fastmcp/test_integration.py @@ -14,14 +14,14 @@ 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 @@ 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 echo_context(custom_request_id: str, ctx: Context[Any, Any, Request]) -> str 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 @@ async def handle_generic_notification(self, message) -> None: 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 progress_callback(progress: float, total: float | None, message: str | 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 message_handler(message): 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 message_handler(message): 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"