Fix elicitation enums (#1632)

Author: jlowin

src/fastmcp/server/elicitation.py

@@ -8,6 +8,8 @@
     DeclinedElicitation,
 )
 from pydantic import BaseModel
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
+from pydantic_core import core_schema
 
 from fastmcp.utilities.json_schema import compress_schema
 from fastmcp.utilities.logging import get_logger
@@ -26,6 +28,60 @@
 T = TypeVar("T")
 
 
+class ElicitationJsonSchema(GenerateJsonSchema):
+    """Custom JSON schema generator for MCP elicitation that always inlines enums.
+
+    MCP elicitation requires inline enum schemas without $ref/$defs references.
+    This generator ensures enums are always generated inline for compatibility.
+    Optionally adds enumNames for better UI display when available.
+    """
+
+    def generate_inner(self, schema: core_schema.CoreSchema) -> JsonSchemaValue:
+        """Override to prevent ref generation for enums."""
+        # For enum schemas, bypass the ref mechanism entirely
+        if schema["type"] == "enum":
+            # Directly call our custom enum_schema without going through handler
+            # This prevents the ref/defs mechanism from being invoked
+            return self.enum_schema(schema)
+        # For all other types, use the default implementation
+        return super().generate_inner(schema)
+
+    def enum_schema(self, schema: core_schema.EnumSchema) -> JsonSchemaValue:
+        """Generate inline enum schema with optional enumNames for better UI.
+
+        If enum members have a _display_name_ attribute or custom __str__,
+        we'll include enumNames for better UI representation.
+        """
+        # Get the base schema from parent
+        result = super().enum_schema(schema)
+
+        # Try to add enumNames if the enum has display-friendly names
+        enum_cls = schema.get("cls")
+        if enum_cls:
+            members = schema.get("members", [])
+            enum_names = []
+            has_custom_names = False
+
+            for member in members:
+                # Check if member has a custom display name attribute
+                if hasattr(member, "_display_name_"):
+                    enum_names.append(member._display_name_)
+                    has_custom_names = True
+                # Or use the member name with better formatting
+                else:
+                    # Convert SNAKE_CASE to Title Case for display
+                    display_name = member.name.replace("_", " ").title()
+                    enum_names.append(display_name)
+                    if display_name != member.value:
+                        has_custom_names = True
+
+            # Only add enumNames if they differ from the values
+            if has_custom_names:
+                result["enumNames"] = enum_names
+
+        return result
+
+
 # we can't use the low-level AcceptedElicitation because it only works with BaseModels
 class AcceptedElicitation(BaseModel, Generic[T]):
     """Result when user accepts the elicitation."""
@@ -46,7 +102,10 @@ def get_elicitation_schema(response_type: type[T]) -> dict[str, Any]:
         response_type: The type of the response
     """
 
-    schema = get_cached_typeadapter(response_type).json_schema()
+    # Use custom schema generator that inlines enums for MCP compatibility
+    schema = get_cached_typeadapter(response_type).json_schema(
+        schema_generator=ElicitationJsonSchema
+    )
     schema = compress_schema(schema)
 
     # Validate the schema to ensure it follows MCP elicitation requirements

tests/client/test_elicitation.py

@@ -15,6 +15,7 @@
     AcceptedElicitation,
     CancelledElicitation,
     DeclinedElicitation,
+    get_elicitation_schema,
     validate_elicitation_json_schema,
 )
 from fastmcp.utilities.types import TypeAdapter
@@ -639,3 +640,80 @@ async def elicitation_handler(message, response_type, params, ctx):
             match="Elicitation responses must be serializable as a JSON object",
         ):
             await client.call_tool("ask_for_name")
+
+
+def test_enum_elicitation_schema_inline():
+    """Test that enum schemas are generated inline without $ref/$defs for MCP compatibility."""
+
+    class Priority(Enum):
+        LOW = "low"
+        MEDIUM = "medium"
+        HIGH = "high"
+
+    @dataclass
+    class TaskRequest:
+        title: str
+        priority: Priority
+
+    # Generate elicitation schema
+    schema = get_elicitation_schema(TaskRequest)
+
+    # Verify no $defs section exists (enums should be inlined)
+    assert "$defs" not in schema, (
+        "Schema should not contain $defs - enums must be inline"
+    )
+
+    # Verify no $ref in properties
+    for prop_name, prop_schema in schema.get("properties", {}).items():
+        assert "$ref" not in prop_schema, (
+            f"Property {prop_name} contains $ref - should be inline"
+        )
+
+    # Verify the priority field has inline enum values
+    priority_schema = schema["properties"]["priority"]
+    assert "enum" in priority_schema, "Priority should have enum values inline"
+    assert priority_schema["enum"] == ["low", "medium", "high"]
+    assert priority_schema.get("type") == "string"
+
+    # Verify title field is a simple string
+    assert schema["properties"]["title"]["type"] == "string"
+
+
+def test_enum_elicitation_schema_with_enum_names():
+    """Test that enum schemas can include enumNames for better UI display."""
+
+    class TaskStatus(Enum):
+        NOT_STARTED = "not_started"
+        IN_PROGRESS = "in_progress"
+        COMPLETED = "completed"
+        ON_HOLD = "on_hold"
+
+    @dataclass
+    class TaskUpdate:
+        task_id: str
+        status: TaskStatus
+
+    # Generate elicitation schema
+    schema = get_elicitation_schema(TaskUpdate)
+
+    # Verify enum is inline
+    assert "$defs" not in schema
+    assert "$ref" not in str(schema)
+
+    status_schema = schema["properties"]["status"]
+    assert "enum" in status_schema
+    assert status_schema["enum"] == [
+        "not_started",
+        "in_progress",
+        "completed",
+        "on_hold",
+    ]
+
+    # Check if enumNames were added for display
+    assert "enumNames" in status_schema
+    assert status_schema["enumNames"] == [
+        "Not Started",
+        "In Progress",
+        "Completed",
+        "On Hold",
+    ]