add validation for primitive types

Author: ihrpr

README.md

@@ -327,13 +327,13 @@ mcp = FastMCP("Booking System")
 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
+        message=f"Confirm booking for {party_size} on {date}?", schema=ConfirmBooking
     )
 
     if result.action == "accept" and result.data:

src/mcp/server/fastmcp/server.py

@@ -4,17 +4,19 @@
 
 import inspect
 import re
+import types
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
 from contextlib import (
     AbstractAsyncContextManager,
     asynccontextmanager,
 )
 from itertools import chain
-from typing import Any, Generic, Literal, TypeVar
+from typing import Any, Generic, Literal, TypeVar, Union, get_args, get_origin
 
 import anyio
 import pydantic_core
 from pydantic import BaseModel, Field, ValidationError
+from pydantic.fields import FieldInfo
 from pydantic.networks import AnyUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from starlette.applications import Starlette
@@ -70,13 +72,13 @@
 
 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."""
 
@@ -873,6 +875,43 @@ def _convert_to_content(
     return [TextContent(type="text", text=result)]
 
 
+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."
+            )
+
+
+# Primitive types allowed in elicitation schemas
+_ELICITATION_PRIMITIVE_TYPES = (str, int, float, bool)
+
+
+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 type(None):
+        return True
+
+    # Handle basic primitive types
+    if annotation in _ELICITATION_PRIMITIVE_TYPES:
+        return True
+
+    # Handle Union types (including Optional and Python 3.10+ union syntax)
+    origin = get_origin(annotation)
+    if origin is Union or (hasattr(types, 'UnionType') and isinstance(annotation, types.UnionType)):
+        args = get_args(annotation)
+        # All args must be primitive types or None
+        return all(arg is type(None) or arg in _ELICITATION_PRIMITIVE_TYPES for arg in args)
+
+    return False
+
+
 class Context(BaseModel, Generic[ServerSessionT, LifespanContextT, RequestT]):
     """Context object providing access to MCP capabilities.
 
@@ -996,6 +1035,9 @@ async def elicit(
             The result.data will only be populated if action is "accept" and validation succeeded.
         """
 
+        # Validate that schema only contains primitive types and fail loudly if not
+        _validate_elicitation_schema(schema)
+
         json_schema = schema.model_json_schema()
 
         result = await self.request_context.session.elicit(

tests/server/fastmcp/test_elicitation.py

@@ -10,92 +10,201 @@
 from mcp.types import ElicitResult, TextContent
 
 
-@pytest.mark.anyio
-async def test_stdio_elicitation():
-    """Test the elicitation feature using stdio transport."""
+# Shared schema for basic tests
+class AnswerSchema(BaseModel):
+    answer: str = Field(description="The user's answer to the question")
 
-    # Create a FastMCP server with a tool that uses elicitation
-    mcp = FastMCP(name="StdioElicitationServer")
+
+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:
-        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"
 
+    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):
-        # 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"
+    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")
-    
-    @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
+    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:
-        # 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"
+
+        # 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": "[email protected]", "subscribe": True},
+            "Name: John Doe, Age: 30, Email: [email protected], 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"],
+    )