elicitation.py

"""Elicitation utilities for MCP servers."""

from __future__ import annotations

from typing import Any, Generic, Literal, TypeVar

from mcp_types import RequestId

# Internal surface package; imported as the gate's source of truth for spec-valid property schemas.
from mcp_types._v2025_11_25 import PrimitiveSchemaDefinition
from pydantic import BaseModel, ValidationError
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
from pydantic_core import core_schema
from typing_extensions import TypeAliasType

from mcp.server.session import ServerSession

ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel)


class AcceptedElicitation(BaseModel, Generic[ElicitSchemaModelT]):
    """Result when user accepts the elicitation."""

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


class DeclinedElicitation(BaseModel):
    """Result when user declines the elicitation."""

    action: Literal["decline"] = "decline"


class CancelledElicitation(BaseModel):
    """Result when user cancels the elicitation."""

    action: Literal["cancel"] = "cancel"


ElicitationResult = TypeAliasType(
    "ElicitationResult",
    AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation,
    type_params=(ElicitSchemaModelT,),
)


class AcceptedUrlElicitation(BaseModel):
    """Result when user accepts a URL mode elicitation."""

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


UrlElicitationResult = AcceptedUrlElicitation | DeclinedElicitation | CancelledElicitation


class _ElicitationJsonSchema(GenerateJsonSchema):
    """JSON-Schema generator that flattens `T | None` to `T` and drops `None` defaults.

    The spec's `PrimitiveSchemaDefinition` admits no `anyOf` or null type; an
    optional field is expressed by leaving it out of `required`, which pydantic
    already does for any field with a default.
    """

    def nullable_schema(self, schema: core_schema.NullableSchema) -> JsonSchemaValue:
        return self.generate_inner(schema["schema"])

    def default_schema(self, schema: core_schema.WithDefaultSchema) -> JsonSchemaValue:
        result = super().default_schema(schema)
        if result.get("default") is None:
            result.pop("default", None)
        return result


def _validate_rendered_properties(json_schema: dict[str, Any]) -> None:
    """Reject any `properties` entry the spec's `PrimitiveSchemaDefinition` won't accept.

    Catches whatever the renderer let through that isn't spec-valid: bare
    `list[str]` (no enum), multi-primitive unions, nested models.
    """
    for field_name, prop in json_schema.get("properties", {}).items():
        try:
            PrimitiveSchemaDefinition.model_validate(prop)
        except ValidationError:
            raise TypeError(
                f"Elicitation schema field {field_name!r} rendered as {prop!r}, "
                f"which is not a valid PrimitiveSchemaDefinition"
            ) from None


def render_elicitation_schema(schema: type[BaseModel]) -> dict[str, Any]:
    """Render a model as the spec-valid `requested_schema` for an elicitation.

    Raises:
        TypeError: If a field renders as something the spec's
            `PrimitiveSchemaDefinition` does not accept.
    """
    json_schema = schema.model_json_schema(schema_generator=_ElicitationJsonSchema)
    _validate_rendered_properties(json_schema)
    return json_schema


async def elicit_with_validation(
    session: ServerSession,
    message: str,
    schema: type[ElicitSchemaModelT],
    related_request_id: RequestId | None = None,
) -> ElicitationResult[ElicitSchemaModelT]:
    """Elicit information from the client/user with schema validation (form mode).

    This method can be used to interactively ask for additional information from the
    client within a tool's execution. The client might display the message to the
    user and collect a response according to the provided schema. If the client
    is an agent, it might decide how to handle the elicitation -- either by asking
    the user or automatically generating a response.

    For sensitive data like credentials or OAuth flows, use elicit_url() instead.

    Raises:
        ValueError: If the client accepted the elicitation without supplying
            content, or with content that does not match the requested schema.
    """
    json_schema = render_elicitation_schema(schema)

    result = await session.elicit_form(
        message=message,
        requested_schema=json_schema,
        related_request_id=related_request_id,
    )

    if result.action == "accept":
        if result.content is None:
            raise ValueError("Received an accepted elicitation with no content")
        try:
            validated_data = schema.model_validate(result.content)
        except ValidationError as e:
            raise ValueError(
                "Received an accepted elicitation whose content does not match the requested schema"
            ) from e
        return AcceptedElicitation(data=validated_data)
    if result.action == "decline":
        return DeclinedElicitation()
    return CancelledElicitation()


async def elicit_url(
    session: ServerSession,
    message: str,
    url: str,
    elicitation_id: str,
    related_request_id: RequestId | None = None,
) -> UrlElicitationResult:
    """Elicit information from the user via out-of-band URL navigation (URL mode).

    This method directs the user to an external URL where sensitive interactions can
    occur without passing data through the MCP client. Use this for:
    - Collecting sensitive credentials (API keys, passwords)
    - OAuth authorization flows with third-party services
    - Payment and subscription flows
    - Any interaction where data should not pass through the LLM context

    The response indicates whether the user consented to navigate to the URL.
    The actual interaction happens out-of-band. When the elicitation completes,
    the server should send an ElicitCompleteNotification to notify the client.

    Args:
        session: The server session
        message: Human-readable explanation of why the interaction is needed
        url: The URL the user should navigate to
        elicitation_id: Unique identifier for tracking this elicitation
        related_request_id: Optional ID of the request that triggered this elicitation

    Returns:
        UrlElicitationResult indicating accept, decline, or cancel
    """
    result = await session.elicit_url(
        message=message,
        url=url,
        elicitation_id=elicitation_id,
        related_request_id=related_request_id,
    )

    if result.action == "accept":
        return AcceptedUrlElicitation()
    elif result.action == "decline":
        return DeclinedElicitation()
    elif result.action == "cancel":
        return CancelledElicitation()
    else:  # pragma: no cover
        # This should never happen, but handle it just in case
        raise ValueError(f"Unexpected elicitation action: {result.action}")