add ElicitationResult to fastMCP

Author: ihrpr

src/mcp/server/fastmcp/server.py

@@ -65,7 +65,20 @@
 
 logger = get_logger(__name__)
 
-ElicitedModelT = TypeVar("ElicitedModelT", bound=BaseModel)
+ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel)
+
+
+class ElicitationResult(BaseModel, Generic[ElicitSchemaModelT]):
+    """Result of an elicitation request."""
+    
+    action: Literal["accept", "decline", "cancel"]
+    """The user's action in response to the elicitation."""
+    
+    data: ElicitSchemaModelT | None = None
+    """The validated data if action is 'accept', None otherwise."""
+    
+    validation_error: str | None = None
+    """Validation error message if data failed to validate."""
 
 
 class Settings(BaseSettings, Generic[LifespanResultT]):
@@ -977,28 +990,28 @@ async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContent
     async def elicit(
         self,
         message: str,
-        schema: type[ElicitedModelT],
-    ) -> ElicitedModelT:
+        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
+        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
+            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 instance of the schema type with the user's response
+            An ElicitationResult containing the action taken and the data if accepted
 
-        Raises:
-            Exception: If the user declines or cancels the elicitation
-            ValidationError: If the response doesn't match the schema
+        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.
         """
 
         json_schema = schema.model_json_schema()
@@ -1012,13 +1025,12 @@ async def elicit(
         if result.action == "accept" and result.content:
             # Validate and parse the content using the schema
             try:
-                return schema.model_validate(result.content)
+                validated_data = schema.model_validate(result.content)
+                return ElicitationResult(action="accept", data=validated_data)
             except ValidationError as e:
-                raise ValueError(f"Invalid response: {e}")
-        elif result.action == "decline":
-            raise Exception("User declined to provide information")
+                return ElicitationResult(action="accept", validation_error=str(e))
         else:
-            raise Exception("User cancelled the request")
+            return ElicitationResult(action=result.action)
 
     async def log(
         self,

tests/server/fastmcp/test_elicitation.py

@@ -0,0 +1,101 @@
+"""
+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
+
+
+@pytest.mark.anyio
+async def test_stdio_elicitation():
+    """Test the elicitation feature using stdio transport."""
+
+    # Create a FastMCP server with a tool that uses elicitation
+    mcp = FastMCP(name="StdioElicitationServer")
+
+    @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}"
+        elif result.action == "decline":
+            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(action="accept", content={"answer": "Test User"})
+        else:
+            raise ValueError(f"Unexpected elicitation message: {params.message}")
+
+    # Use memory-based session to test with stdio transport
+    async with create_connected_server_and_client_session(
+        mcp._mcp_server, elicitation_callback=elicitation_callback
+    ) as client_session:
+        # First initialize the session
+        result = await client_session.initialize()
+        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?"})
+
+        # Verify the result
+        assert len(tool_result.content) == 1
+        assert isinstance(tool_result.content[0], TextContent)
+        assert tool_result.content[0].text == "User answered: Test User"
+
+
+@pytest.mark.anyio
+async def test_stdio_elicitation_decline():
+    """Test elicitation with user declining."""
+    
+    mcp = FastMCP(name="StdioElicitationDeclineServer")
+    
+    @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}"
+        elif result.action == "decline":
+            return "User declined to answer"
+        else:
+            return "User cancelled"
+    
+    # Create a custom handler that declines
+    async def elicitation_callback(context, params):
+        return ElicitResult(action="decline")
+    
+    async with create_connected_server_and_client_session(
+        mcp._mcp_server, elicitation_callback=elicitation_callback
+    ) as client_session:
+        # Initialize
+        await client_session.initialize()
+        
+        # Call the tool
+        tool_result = await client_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)
+        assert tool_result.content[0].text == "User declined to answer"

tests/server/fastmcp/test_integration.py

@@ -105,12 +105,13 @@ async def ask_user(prompt: str, ctx: Context) -> str:
         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:
+        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: {str(e)}"
+            return f"User cancelled or declined: {result.action}"
 
     # Create the SSE app
     app = mcp.sse_app()
@@ -295,23 +296,25 @@ class AlternativeDateSchema(BaseModel):
         # 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,
-                )
+            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
+            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"
-            except Exception:
-                # User declined or cancelled
+            elif result.action in ("decline", "cancel"):
                 return "❌ Booking cancelled"
+            else:
+                # Validation error
+                return f"❌ Invalid input: {result.validation_error}"
         else:
             # Available - book directly
             return f"✅ Booked table for {party_size} on {date} at {time}"