fastmcp 2.12.5__py3-none-any.whl → 2.14.0__py3-none-any.whl

fastmcp/server/elicitation.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Generic, Literal
+from enum import Enum
+from typing import Any, Generic, Literal, get_origin
 
 from mcp.server.elicitation import (
     CancelledElicitation,
@@ -20,8 +21,11 @@ __all__ = [
     "AcceptedElicitation",
     "CancelledElicitation",
     "DeclinedElicitation",
-    "get_elicitation_schema",
+    "ElicitConfig",
     "ScalarElicitationType",
+    "get_elicitation_schema",
+    "handle_elicit_accept",
+    "parse_elicit_response_type",
 ]
 
 logger = get_logger(__name__)
@@ -37,50 +41,64 @@ class ElicitationJsonSchema(GenerateJsonSchema):
     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."""
+    def generate_inner(self, schema: core_schema.CoreSchema) -> JsonSchemaValue:  # type: ignore[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)
+            return self.enum_schema(schema)  # type: ignore[arg-type]
+        # For list schemas, check if items are enums
+        if schema["type"] == "list":
+            return self.list_schema(schema)  # type: ignore[arg-type]
         # 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]
+            # 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 with optional enumNames for better UI.
+        """Generate inline enum schema.
 
-        If enum members have a _display_name_ attribute or custom __str__,
-        we'll include enumNames for better UI representation.
+        Always generates enum pattern: {"enum": [value, ...]}
+        Titled enums are handled separately via dict-based syntax in ctx.elicit().
         """
-        # 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
+        # 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
@@ -96,6 +114,207 @@ 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) -> 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
+    """
+    if response_type is None:
+        return ElicitConfig(
+            schema={"type": "object", "properties": {}},
+            response_type=None,
+            is_raw=False,
+        )
+
+    if isinstance(response_type, dict):
+        return _parse_dict_syntax(response_type)
+
+    if isinstance(response_type, list):
+        return _parse_list_syntax(response_type)
+
+    if get_origin(response_type) is list:
+        return _parse_generic_list(response_type)
+
+    if _is_scalar_type(response_type):
+        return _parse_scalar_type(response_type)
+
+    # Other types (dataclass, BaseModel, etc.) - use directly
+    return ElicitConfig(
+        schema=get_elicitation_schema(response_type),
+        response_type=response_type,
+        is_raw=False,
+    )
+
+
+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):
+        choice_literal = Literal[tuple(lst)]  # type: ignore[valid-type]
+        wrapped = ScalarElicitationType[choice_literal]  # type: ignore[valid-type]
+        return ElicitConfig(
+            schema=get_elicitation_schema(wrapped),  # type: ignore[arg-type]
+            response_type=wrapped,  # type: ignore[assignment]
+            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]  # type: ignore[valid-type]
+    return ElicitConfig(
+        schema=get_elicitation_schema(wrapped),  # type: ignore[arg-type]
+        response_type=wrapped,  # type: ignore[assignment]
+        is_raw=False,
+    )
+
+
+def _parse_scalar_type(response_type: Any) -> ElicitConfig:
+    """Parse scalar types (bool, int, float, str, Literal, Enum)."""
+    wrapped = ScalarElicitationType[response_type]  # type: ignore[valid-type]
+    return ElicitConfig(
+        schema=get_elicitation_schema(wrapped),  # type: ignore[arg-type]
+        response_type=wrapped,  # type: ignore[assignment]
+        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:
+        {"oneOf": [{"const": "low", "title": "Low Priority"}, ...]} for single-select
+        {"anyOf": [{"const": "low", "title": "Low Priority"}, ...]} for multi-select
+    """
+    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})
+    return {pattern_key: pattern}
+
+
 def get_elicitation_schema(response_type: type[T]) -> dict[str, Any]:
     """Get the schema for an elicitation response.
 
@@ -197,24 +416,43 @@ def validate_elicitation_json_schema(schema: dict[str, Any]) -> None:
                     )
             continue
 
-        # Check if it's a primitive type
-        if prop_type not in ALLOWED_TYPES:
+        # 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}' has type '{prop_type}' which is not "
-                f"a primitive type. Only {ALLOWED_TYPES} are allowed in elicitation schemas."
+                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 or arrays of objects (not allowed)
+        # 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."
             )
 
-        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."
-                )
+        # 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."
+            )

fastmcp/server/event_store.py

@@ -0,0 +1,177 @@
+"""EventStore implementation backed by AsyncKeyValue.
+
+This module provides an EventStore implementation that enables SSE polling/resumability
+for Streamable HTTP transports. Events are stored using the key_value package's
+AsyncKeyValue protocol, allowing users to configure any compatible backend
+(in-memory, Redis, etc.) following the same pattern as ResponseCachingMiddleware.
+"""
+
+from __future__ import annotations
+
+from uuid import uuid4
+
+from key_value.aio.adapters.pydantic import PydanticAdapter
+from key_value.aio.protocols import AsyncKeyValue
+from key_value.aio.stores.memory import MemoryStore
+from mcp.server.streamable_http import EventCallback, EventId, EventMessage, StreamId
+from mcp.server.streamable_http import EventStore as SDKEventStore
+from mcp.types import JSONRPCMessage
+from pydantic import BaseModel
+
+from fastmcp.utilities.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class EventEntry(BaseModel):
+    """Stored event entry."""
+
+    event_id: str
+    stream_id: str
+    message: dict | None  # JSONRPCMessage serialized to dict
+
+
+class StreamEventList(BaseModel):
+    """List of event IDs for a stream."""
+
+    event_ids: list[str]
+
+
+class EventStore(SDKEventStore):
+    """EventStore implementation backed by AsyncKeyValue.
+
+    Enables SSE polling/resumability by storing events that can be replayed
+    when clients reconnect. Works with any AsyncKeyValue backend (memory, Redis, etc.)
+    following the same pattern as ResponseCachingMiddleware and OAuthProxy.
+
+    Example:
+        ```python
+        from fastmcp import FastMCP
+        from fastmcp.server.event_store import EventStore
+
+        # Default in-memory storage
+        event_store = EventStore()
+
+        # Or with a custom backend
+        from key_value.aio.stores.redis import RedisStore
+        redis_backend = RedisStore(url="redis://localhost")
+        event_store = EventStore(storage=redis_backend)
+
+        mcp = FastMCP("MyServer")
+        app = mcp.http_app(event_store=event_store, retry_interval=2000)
+        ```
+
+    Args:
+        storage: AsyncKeyValue backend. Defaults to MemoryStore.
+        max_events_per_stream: Maximum events to retain per stream. Default 100.
+        ttl: Event TTL in seconds. Default 3600 (1 hour). Set to None for no expiration.
+    """
+
+    def __init__(
+        self,
+        storage: AsyncKeyValue | None = None,
+        max_events_per_stream: int = 100,
+        ttl: int | None = 3600,
+    ):
+        self._storage: AsyncKeyValue = storage or MemoryStore()
+        self._max_events_per_stream = max_events_per_stream
+        self._ttl = ttl
+
+        # PydanticAdapter for type-safe storage (following OAuth proxy pattern)
+        self._event_store: PydanticAdapter[EventEntry] = PydanticAdapter[EventEntry](
+            key_value=self._storage,
+            pydantic_model=EventEntry,
+            default_collection="fastmcp_events",
+        )
+        self._stream_store: PydanticAdapter[StreamEventList] = PydanticAdapter[
+            StreamEventList
+        ](
+            key_value=self._storage,
+            pydantic_model=StreamEventList,
+            default_collection="fastmcp_streams",
+        )
+
+    async def store_event(
+        self, stream_id: StreamId, message: JSONRPCMessage | None
+    ) -> EventId:
+        """Store an event and return its ID.
+
+        Args:
+            stream_id: ID of the stream the event belongs to
+            message: The JSON-RPC message to store, or None for priming events
+
+        Returns:
+            The generated event ID for the stored event
+        """
+        event_id = str(uuid4())
+
+        # Store the event entry
+        entry = EventEntry(
+            event_id=event_id,
+            stream_id=stream_id,
+            message=message.model_dump(mode="json") if message else None,
+        )
+        await self._event_store.put(key=event_id, value=entry, ttl=self._ttl)
+
+        # Update stream's event list
+        stream_data = await self._stream_store.get(key=stream_id)
+        event_ids = stream_data.event_ids if stream_data else []
+        event_ids.append(event_id)
+
+        # Trim to max events (delete old events)
+        if len(event_ids) > self._max_events_per_stream:
+            for old_id in event_ids[: -self._max_events_per_stream]:
+                await self._event_store.delete(key=old_id)
+            event_ids = event_ids[-self._max_events_per_stream :]
+
+        await self._stream_store.put(
+            key=stream_id,
+            value=StreamEventList(event_ids=event_ids),
+            ttl=self._ttl,
+        )
+
+        return event_id
+
+    async def replay_events_after(
+        self,
+        last_event_id: EventId,
+        send_callback: EventCallback,
+    ) -> StreamId | None:
+        """Replay events that occurred after the specified event ID.
+
+        Args:
+            last_event_id: The ID of the last event the client received
+            send_callback: A callback function to send events to the client
+
+        Returns:
+            The stream ID of the replayed events, or None if the event ID was not found
+        """
+        # Look up the event to find its stream
+        entry = await self._event_store.get(key=last_event_id)
+        if not entry:
+            logger.warning(f"Event ID {last_event_id} not found in store")
+            return None
+
+        stream_id = entry.stream_id
+        stream_data = await self._stream_store.get(key=stream_id)
+        if not stream_data:
+            logger.warning(f"Stream {stream_id} not found in store")
+            return None
+
+        event_ids = stream_data.event_ids
+
+        # Find events after last_event_id
+        try:
+            start_idx = event_ids.index(last_event_id) + 1
+        except ValueError:
+            logger.warning(f"Event ID {last_event_id} not found in stream {stream_id}")
+            return None
+
+        # Replay events after the last one
+        for event_id in event_ids[start_idx:]:
+            event = await self._event_store.get(key=event_id)
+            if event and event.message:
+                msg = JSONRPCMessage.model_validate(event.message)
+                await send_callback(EventMessage(msg, event.event_id))
+
+        return stream_id