fastmcp/fastmcp_slim/fastmcp/server/elicitation.py

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Any, Generic, Literal, cast, get_origin

from mcp.server.elicitation import (
    AcceptedUrlElicitation,
    CancelledElicitation,
    DeclinedElicitation,
    UrlElicitationResult,
)
from pydantic import BaseModel
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
from pydantic_core import core_schema
from typing_extensions import TypeVar

from fastmcp.utilities.json_schema import compress_schema
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.types import get_cached_typeadapter

__all__ = [
    "AcceptedElicitation",
    "AcceptedUrlElicitation",
    "CancelledElicitation",
    "DeclinedElicitation",
    "ElicitConfig",
    "ScalarElicitationType",
    "UrlElicitationResult",
    "get_elicitation_schema",
    "handle_elicit_accept",
    "parse_elicit_response_type",
]

logger = get_logger(__name__)

T = TypeVar("T", default=Any)


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:  # type: ignore[override]  # ty:ignore[invalid-method-override]
        """Override to prevent ref generation for enums and handle list schemas."""
        # 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 list schemas, check if items are enums
        if schema["type"] == "list":
            return self.list_schema(schema)
        # For all other types, use the default implementation
        return super().generate_inner(schema)

    def list_schema(self, schema: core_schema.ListSchema) -> JsonSchemaValue:
        """Generate schema for list types, detecting enum items for multi-select."""
        items_schema = schema.get("items_schema")

        # Check if items are enum/Literal
        if items_schema and items_schema.get("type") == "enum":
            # Generate array with enum items
            items = self.enum_schema(items_schema)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
            # If items have oneOf pattern, convert to anyOf for multi-select per SEP-1330
            if "oneOf" in items:
                items = {"anyOf": items["oneOf"]}
            return {
                "type": "array",
                "items": items,  # Will be {"enum": [...]} or {"anyOf": [...]}
            }

        # Check if items are Literal (which Pydantic represents differently)
        if items_schema:
            # Try to detect Literal patterns
            items_result = super().generate_inner(items_schema)
            # If it's a const pattern or enum-like, allow it
            if (
                "const" in items_result
                or "enum" in items_result
                or "oneOf" in items_result
            ):
                # Convert oneOf to anyOf for multi-select
                if "oneOf" in items_result:
                    items_result = {"anyOf": items_result["oneOf"]}
                return {
                    "type": "array",
                    "items": items_result,
                }

        # Default behavior for non-enum arrays
        return super().list_schema(schema)

    def enum_schema(self, schema: core_schema.EnumSchema) -> JsonSchemaValue:
        """Generate inline enum schema.

        Always generates enum pattern: `{"enum": [value, ...]}`
        Titled enums are handled separately via dict-based syntax in ctx.elicit().
        """
        # Get the base schema from parent - always use simple enum pattern
        return super().enum_schema(schema)


# 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."""

    action: Literal["accept"] = "accept"
    data: T


@dataclass
class ScalarElicitationType(Generic[T]):
    value: T


@dataclass
class ElicitConfig:
    """Configuration for an elicitation request.

    Attributes:
        schema: The JSON schema to send to the client
        response_type: The type to validate responses with (None for raw schemas)
        is_raw: True if schema was built directly (extract "value" from response)
    """

    schema: dict[str, Any]
    response_type: type | None
    is_raw: bool


def parse_elicit_response_type(
    response_type: Any,
    response_title: str | None = None,
    response_description: str | None = None,
) -> ElicitConfig:
    """Parse response_type into schema and handling configuration.

    Supports multiple syntaxes:
    - None: Empty object schema, expect empty response
    - dict: `{"low": {"title": "..."}}` -> single-select titled enum
    - list patterns:
        - `[["a", "b"]]` -> multi-select untitled
        - `[{"low": {...}}]` -> multi-select titled
        - `["a", "b"]` -> single-select untitled
    - `list[X]` type annotation: multi-select with type
    - Scalar types (bool, int, float, str, Literal, Enum): single value
    - Other types (dataclass, BaseModel): use directly

    The ``response_title`` and ``response_description`` arguments customize the
    label and description of the wrapped ``value`` property for the scalar/dict/list
    shorthand forms. They are only valid when FastMCP is wrapping the response
    type; passing them with a full BaseModel/dataclass (or ``None``) raises
    ``TypeError``, because in those cases the user already controls field
    metadata via ``Field(title=..., description=...)``.
    """
    has_response_metadata = (
        response_title is not None or response_description is not None
    )

    if response_type is None:
        if has_response_metadata:
            raise TypeError(
                "response_title and response_description are not supported when "
                "response_type is None, because the elicitation schema has no "
                "fields to label."
            )
        return ElicitConfig(
            schema={"type": "object", "properties": {}},
            response_type=None,
            is_raw=False,
        )

    if isinstance(response_type, dict):
        config = _parse_dict_syntax(response_type)
    elif isinstance(response_type, list):
        config = _parse_list_syntax(response_type)
    elif get_origin(response_type) is list:
        config = _parse_generic_list(response_type)
    elif _is_scalar_type(response_type):
        config = _parse_scalar_type(response_type)
    else:
        # Other types (dataclass, BaseModel, etc.) - use directly
        if has_response_metadata:
            raise TypeError(
                "response_title and response_description are only supported when "
                "response_type is a scalar, Literal, Enum, or the dict/list "
                "shorthand forms. For BaseModel or dataclass response types, use "
                "Field(title=..., description=...) on the individual fields."
            )
        return ElicitConfig(
            schema=get_elicitation_schema(response_type),
            response_type=response_type,
            is_raw=False,
        )

    if has_response_metadata:
        _apply_value_metadata(config.schema, response_title, response_description)
    return config


def _apply_value_metadata(
    schema: dict[str, Any],
    title: str | None,
    description: str | None,
) -> None:
    """Override title/description on the wrapped ``value`` property in-place."""
    value_schema = schema.get("properties", {}).get("value")
    if value_schema is None:
        return
    if title is not None:
        value_schema["title"] = title
    if description is not None:
        value_schema["description"] = description


def _is_scalar_type(response_type: Any) -> bool:
    """Check if response_type is a scalar type that needs wrapping."""
    return (
        response_type in {bool, int, float, str}
        or get_origin(response_type) is Literal
        or (isinstance(response_type, type) and issubclass(response_type, Enum))
    )


def _parse_dict_syntax(d: dict[str, Any]) -> ElicitConfig:
    """Parse dict syntax: {"low": {"title": "..."}} -> single-select titled."""
    if not d:
        raise ValueError("Dict response_type cannot be empty.")
    enum_schema = _dict_to_enum_schema(d, multi_select=False)
    return ElicitConfig(
        schema={
            "type": "object",
            "properties": {"value": enum_schema},
            "required": ["value"],
        },
        response_type=None,
        is_raw=True,
    )


def _parse_list_syntax(lst: list[Any]) -> ElicitConfig:
    """Parse list patterns: [[...]], [{...}], or [...]."""
    # [["a", "b", "c"]] -> multi-select untitled
    if (
        len(lst) == 1
        and isinstance(lst[0], list)
        and lst[0]
        and all(isinstance(item, str) for item in lst[0])
    ):
        return ElicitConfig(
            schema={
                "type": "object",
                "properties": {"value": {"type": "array", "items": {"enum": lst[0]}}},
                "required": ["value"],
            },
            response_type=None,
            is_raw=True,
        )

    # [{"low": {"title": "..."}}] -> multi-select titled
    if len(lst) == 1 and isinstance(lst[0], dict) and lst[0]:
        enum_schema = _dict_to_enum_schema(lst[0], multi_select=True)
        return ElicitConfig(
            schema={
                "type": "object",
                "properties": {"value": {"type": "array", "items": enum_schema}},
                "required": ["value"],
            },
            response_type=None,
            is_raw=True,
        )

    # ["a", "b", "c"] -> single-select untitled
    if lst and all(isinstance(item, str) for item in lst):
        # Construct Literal type from tuple - use cast since we can't construct Literal dynamically
        # but we know the values are all strings
        choice_literal: type[Any] = cast(type[Any], Literal[tuple(lst)])  # type: ignore[valid-type]  # ty:ignore[invalid-type-form]
        wrapped = ScalarElicitationType[choice_literal]  # type: ignore[valid-type]  # ty:ignore[invalid-type-form]
        return ElicitConfig(
            schema=get_elicitation_schema(wrapped),
            response_type=wrapped,
            is_raw=False,
        )

    raise ValueError(f"Invalid list response_type format. Received: {lst}")


def _parse_generic_list(response_type: Any) -> ElicitConfig:
    """Parse list[X] type annotation -> multi-select."""
    wrapped = ScalarElicitationType[response_type]
    return ElicitConfig(
        schema=get_elicitation_schema(wrapped),
        response_type=wrapped,
        is_raw=False,
    )


def _parse_scalar_type(response_type: Any) -> ElicitConfig:
    """Parse scalar types (bool, int, float, str, Literal, Enum)."""
    wrapped = ScalarElicitationType[response_type]
    return ElicitConfig(
        schema=get_elicitation_schema(wrapped),
        response_type=wrapped,
        is_raw=False,
    )


def handle_elicit_accept(
    config: ElicitConfig, content: Any
) -> AcceptedElicitation[Any]:
    """Handle an accepted elicitation response.

    Args:
        config: The elicitation configuration from parse_elicit_response_type
        content: The response content from the client

    Returns:
        AcceptedElicitation with the extracted/validated data
    """
    # For raw schemas (dict/nested-list syntax), extract value directly
    if config.is_raw:
        if not isinstance(content, dict) or "value" not in content:
            raise ValueError("Elicitation response missing required 'value' field.")
        return AcceptedElicitation[Any](data=content["value"])

    # For typed schemas, validate with Pydantic
    if config.response_type is not None:
        type_adapter = get_cached_typeadapter(config.response_type)
        validated_data = type_adapter.validate_python(content)
        if isinstance(validated_data, ScalarElicitationType):
            return AcceptedElicitation[Any](data=validated_data.value)
        return AcceptedElicitation[Any](data=validated_data)

    # For None response_type, expect empty response
    if content:
        raise ValueError(
            f"Elicitation expected an empty response, but received: {content}"
        )
    return AcceptedElicitation[dict[str, Any]](data={})


def _dict_to_enum_schema(
    enum_dict: dict[str, dict[str, str]], multi_select: bool = False
) -> dict[str, Any]:
    """Convert dict enum to SEP-1330 compliant schema pattern.

    Args:
        enum_dict: {"low": {"title": "Low Priority"}, "medium": {"title": "Medium Priority"}}
        multi_select: If True, use anyOf pattern; if False, use oneOf pattern

    Returns:
        {"type": "string", "oneOf": [...]} for single-select
        {"anyOf": [...]} for multi-select (used as array items)
    """
    pattern_key = "anyOf" if multi_select else "oneOf"
    pattern = []
    for value, metadata in enum_dict.items():
        title = metadata.get("title", value)
        pattern.append({"const": value, "title": title})

    result: dict[str, Any] = {pattern_key: pattern}
    if not multi_select:
        result["type"] = "string"
    return result


def get_elicitation_schema(response_type: type[T]) -> dict[str, Any]:
    """Get the schema for an elicitation response.

    Args:
        response_type: The type of the response
    """

    # 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
    validate_elicitation_json_schema(schema)

    return schema


def validate_elicitation_json_schema(schema: dict[str, Any]) -> None:
    """Validate that a JSON schema follows MCP elicitation requirements.

    This ensures the schema is compatible with MCP elicitation requirements:
    - Must be an object schema
    - Must only contain primitive field types (string, number, integer, boolean)
    - Must be flat (no nested objects or arrays of objects)
    - Allows const fields (for Literal types) and enum fields (for Enum types)
    - Only primitive types and their nullable variants are allowed

    Args:
        schema: The JSON schema to validate

    Raises:
        TypeError: If the schema doesn't meet MCP elicitation requirements
    """
    ALLOWED_TYPES = {"string", "number", "integer", "boolean"}

    # Check that the schema is an object
    if schema.get("type") != "object":
        raise TypeError(
            f"Elicitation schema must be an object schema, got type '{schema.get('type')}'. "
            "Elicitation schemas are limited to flat objects with primitive properties only."
        )

    properties = schema.get("properties", {})

    for prop_name, prop_schema in properties.items():
        prop_type = prop_schema.get("type")

        # Handle nullable types
        if isinstance(prop_type, list):
            if "null" in prop_type:
                prop_type = [t for t in prop_type if t != "null"]
                if len(prop_type) == 1:
                    prop_type = prop_type[0]
        elif prop_schema.get("nullable", False):
            continue  # Nullable with no other type is fine

        # Handle const fields (Literal types)
        if "const" in prop_schema:
            continue  # const fields are allowed regardless of type

        # Handle enum fields (Enum types)
        if "enum" in prop_schema:
            continue  # enum fields are allowed regardless of type

        # Handle references to definitions (like Enum types)
        if "$ref" in prop_schema:
            # Get the referenced definition
            ref_path = prop_schema["$ref"]
            if ref_path.startswith("#/$defs/"):
                def_name = ref_path[8:]  # Remove "#/$defs/" prefix
                ref_def = schema.get("$defs", {}).get(def_name, {})
                # If the referenced definition has an enum, it's allowed
                if "enum" in ref_def:
                    continue
                # If the referenced definition has a type that's allowed, it's allowed
                ref_type = ref_def.get("type")
                if ref_type in ALLOWED_TYPES:
                    continue
            # If we can't determine what the ref points to, reject it for safety
            raise TypeError(
                f"Elicitation schema field '{prop_name}' contains a reference '{ref_path}' "
                "that could not be validated. Only references to enum types or primitive types are allowed."
            )

        # Handle union types (oneOf/anyOf)
        if "oneOf" in prop_schema or "anyOf" in prop_schema:
            union_schemas = prop_schema.get("oneOf", []) + prop_schema.get("anyOf", [])
            for union_schema in union_schemas:
                # Allow const and enum in unions
                if "const" in union_schema or "enum" in union_schema:
                    continue
                union_type = union_schema.get("type")
                if union_type not in ALLOWED_TYPES:
                    raise TypeError(
                        f"Elicitation schema field '{prop_name}' has union type '{union_type}' which is not "
                        f"a primitive type. Only {ALLOWED_TYPES} are allowed in elicitation schemas."
                    )
            continue

        # Check for arrays before checking primitive types
        if prop_type == "array":
            items_schema = prop_schema.get("items", {})
            if items_schema.get("type") == "object":
                raise TypeError(
                    f"Elicitation schema field '{prop_name}' is an array of objects, but arrays of objects are not allowed. "
                    "Elicitation schemas must be flat objects with primitive properties only."
                )

            # Allow arrays with enum patterns (for multi-select)
            if "enum" in items_schema:
                continue  # Allowed: {"type": "array", "items": {"enum": [...]}}

            # Allow arrays with oneOf/anyOf const patterns (SEP-1330)
            if "oneOf" in items_schema or "anyOf" in items_schema:
                union_schemas = items_schema.get("oneOf", []) + items_schema.get(
                    "anyOf", []
                )
                if union_schemas and all("const" in s for s in union_schemas):
                    continue  # Allowed: {"type": "array", "items": {"anyOf": [{"const": ...}, ...]}}

            # Reject other array types (e.g., arrays of primitives without enum pattern)
            raise TypeError(
                f"Elicitation schema field '{prop_name}' is an array, but arrays are only allowed "
                "when items are enums (for multi-select). Only enum arrays are supported in elicitation schemas."
            )

        # Check for nested objects (not allowed)
        if prop_type == "object":
            raise TypeError(
                f"Elicitation schema field '{prop_name}' is an object, but nested objects are not allowed. "
                "Elicitation schemas must be flat objects with primitive properties only."
            )

        # Check if it's a primitive type
        if prop_type not in ALLOWED_TYPES:
            raise TypeError(
                f"Elicitation schema field '{prop_name}' has type '{prop_type}' which is not "
                f"a primitive type. Only {ALLOWED_TYPES} are allowed in elicitation schemas."
            )