Fix elicitation enums (#1632)

This commit is contained in:
Jeremiah Lowin 2025-08-25 22:14:00 -04:00 committed by GitHub
commit 953d3eb64b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 138 additions and 1 deletions

View file

@ -8,6 +8,8 @@ from mcp.server.elicitation import (
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 @@ logger = get_logger(__name__)
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

View file

@ -15,6 +15,7 @@ from fastmcp.server.elicitation import (
AcceptedElicitation,
CancelledElicitation,
DeclinedElicitation,
get_elicitation_schema,
validate_elicitation_json_schema,
)
from fastmcp.utilities.types import TypeAdapter
@ -639,3 +640,80 @@ async def test_elicitation_implicit_acceptance_must_be_dict(fastmcp_server):
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",
]