adjust types after the spec revision

Author: ihrpr

src/mcp/client/session.py

@@ -131,7 +131,9 @@ def __init__(
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
-        elicitation = types.ElicitationCapability()
+        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

src/mcp/server/fastmcp/server.py

@@ -10,11 +10,11 @@
     asynccontextmanager,
 )
 from itertools import chain
-from typing import Any, Generic, Literal
+from typing import Any, Generic, Literal, TypeVar
 
 import anyio
 import pydantic_core
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError
 from pydantic.networks import AnyUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from starlette.applications import Starlette
@@ -65,6 +65,8 @@
 
 logger = get_logger(__name__)
 
+ElicitedModelT = TypeVar("ElicitedModelT", bound=BaseModel)
+
 
 class Settings(BaseSettings, Generic[LifespanResultT]):
     """FastMCP server settings.
@@ -975,35 +977,48 @@ async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContent
     async def elicit(
         self,
         message: str,
-        requestedSchema: dict[str, Any],
-    ) -> dict[str, Any]:
+        schema: type[ElicitedModelT],
+    ) -> ElicitedModelT:
         """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.
+        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:
-            message: The message to present to the user
-            requestedSchema: JSON Schema defining the expected response structure
+            schema: A Pydantic model class defining the expected response structure
+            message: Optional message to present to the user. If not provided, will use
+                    a default message based on the schema
 
         Returns:
-            The user's response as a dict matching the request schema structure
+            An instance of the schema type with the user's response
 
         Raises:
-            ValueError: If elicitation is not supported by the client or fails
+            Exception: If the user declines or cancels the elicitation
+            ValidationError: If the response doesn't match the schema
         """
 
+        json_schema = schema.model_json_schema()
+
         result = await self.request_context.session.elicit(
             message=message,
-            requestedSchema=requestedSchema,
+            requestedSchema=json_schema,
             related_request_id=self.request_id,
         )
 
-        return result.content
+        if result.action == "accept" and result.content:
+            # Validate and parse the content using the schema
+            try:
+                return schema.model_validate(result.content)
+            except ValidationError as e:
+                raise ValueError(f"Invalid response: {e}")
+        elif result.action == "decline":
+            raise Exception("User declined to provide information")
+        else:
+            raise Exception("User cancelled the request")
 
     async def log(
         self,

src/mcp/server/session.py

@@ -258,14 +258,14 @@ async def list_roots(self) -> types.ListRootsResult:
     async def elicit(
         self,
         message: str,
-        requestedSchema: dict[str, Any],
+        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: JSON Schema defining the expected response structure
+            requestedSchema: Schema defining the expected response structure
 
         Returns:
             The client's response

src/mcp/types.py

@@ -1194,16 +1194,16 @@ class ClientNotification(
     pass
 
 
+# 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
-    """The message to present to the user."""
-
-    requestedSchema: dict[str, Any]
-    """
-    A JSON Schema object defining the expected structure of the response.
-    """
+    requestedSchema: ElicitRequestedSchema
     model_config = ConfigDict(extra="allow")
 
 
@@ -1215,10 +1215,21 @@ class ElicitRequest(Request[ElicitRequestParams, Literal["elicitation/create"]])
 
 
 class ElicitResult(Result):
-    """The client's response to an elicitation/create request from the server."""
+    """The client's response to an elicitation request."""
 
-    content: dict[str, Any]
-    """The response from the client, matching the structure of requestedSchema."""
+    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]):

tests/server/fastmcp/test_integration.py

@@ -14,7 +14,7 @@
 
 import pytest
 import uvicorn
-from pydantic import AnyUrl
+from pydantic import AnyUrl, BaseModel, Field
 from starlette.applications import Starlette
 from starlette.requests import Request
 
@@ -102,19 +102,15 @@ def echo(message: str) -> str:
     # Add a tool that uses elicitation
     @mcp.tool(description="A tool that uses elicitation")
     async def ask_user(prompt: str, ctx: Context) -> str:
-        schema = {
-            "type": "object",
-            "properties": {
-                "answer": {"type": "string"},
-            },
-            "required": ["answer"],
-        }
+        class AnswerSchema(BaseModel):
+            answer: str = Field(description="The user's answer to the question")
 
-        response = await ctx.elicit(
-            message=f"Tool wants to ask: {prompt}",
-            requestedSchema=schema,
-        )
-        return f"User answered: {response['answer']}"
+        try:
+            result = await ctx.elicit(message=f"Tool wants to ask: {prompt}", schema=AnswerSchema)
+            return f"User answered: {result.answer}"
+        except Exception as e:
+            # Handle cancellation or decline
+            return f"User cancelled or declined: {str(e)}"
 
     # Create the SSE app
     app = mcp.sse_app()
@@ -279,6 +275,47 @@ 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
+            try:
+                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.checkAlternative:
+                    alt_date = result.alternativeDate
+                    return f"✅ Booked table for {party_size} on {alt_date} at {time}"
+                else:
+                    return "❌ No booking made"
+            except Exception:
+                # User declined or cancelled
+                return "❌ Booking cancelled"
+        else:
+            # Available - book directly
+            return f"✅ Booked table for {party_size} on {date} at {time}"
+
     return mcp
 
 
@@ -670,6 +707,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.
@@ -765,6 +818,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()
@@ -905,8 +973,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
@@ -919,6 +985,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
@@ -951,6 +1018,7 @@ 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
@@ -965,7 +1033,7 @@ async def test_elicitation_feature(server: None, server_url: str) -> None:
     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"})
+            return ElicitResult(content={"answer": "Test User"}, action="accept")
         else:
             raise ValueError("Unexpected elicitation message")
 

tests/server/fastmcp/test_stdio_elicitation.py

@@ -3,6 +3,7 @@
 """
 
 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
@@ -18,25 +19,27 @@ async def test_stdio_elicitation():
 
     @mcp.tool(description="A tool that uses elicitation")
     async def ask_user(prompt: str, ctx: Context) -> str:
-        schema = {
-            "type": "object",
-            "properties": {
-                "answer": {"type": "string"},
-            },
-            "required": ["answer"],
-        }
-
-        response = await ctx.elicit(
-            message=f"Tool wants to ask: {prompt}",
-            requestedSchema=schema,
-        )
-        return f"User answered: {response['answer']}"
+        class AnswerSchema(BaseModel):
+            answer: str = Field(description="The user's answer to the question")
+
+        try:
+            result = await ctx.elicit(
+                message=f"Tool wants to ask: {prompt}",
+                schema=AnswerSchema,
+            )
+            return f"User answered: {result.answer}"
+        except Exception as e:
+            # Handle cancellation or decline
+            if "declined" in str(e):
+                return "User declined to answer"
+            else:
+                return "User cancelled"
 
     # 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"})
+            return ElicitResult(action="accept", content={"answer": "Test User"})
         else:
             raise ValueError(f"Unexpected elicitation message: {params.message}")
 
@@ -49,9 +52,7 @@ async def elicitation_callback(context, params):
         assert result.serverInfo.name == "StdioElicitationServer"
 
         # Call the tool that uses elicitation
-        tool_result = await client_session.call_tool(
-            "ask_user", {"prompt": "What is your name?"}
-        )
+        tool_result = await client_session.call_tool("ask_user", {"prompt": "What is your name?"})
 
         # Verify the result
         assert len(tool_result.content) == 1