mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 21:44:18 +02:00
Fix elicitation enums (#1632)
This commit is contained in:
parent
33545ab7d5
commit
953d3eb64b
2 changed files with 138 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue