Replace RootModel wrappers with type aliases and TypeAdapter validation (#3470)
Some checks failed
CI / checks (push) Has been cancelled
CI / all-green (push) Has been cancelled

This commit is contained in:
Marcelo Trylesinski 2026-09-07 17:40:21 +02:00 • committed by GitHub
commit 9972c21aa4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 872 additions and 1127 deletions

View file

@ -169,6 +169,8 @@ result.structured_content # {"London": 16.2, "Reykjavik": 4.4}
The keys must be `str`. A `dict[int, float]` can't be a JSON object, so it falls back to the `{"result": ...}` wrapper. The keys must be `str`. A `dict[int, float]` can't be a JSON object, so it falls back to the `{"result": ...}` wrapper.
Dictionary results use Pydantic's `TypeAdapter` for validation and serialization. If you inspect a tool's `FuncMetadata.output_model`, it holds the dictionary type annotation with its schema title.
## Validation ## Validation
`output_schema` is not documentation. Whatever your function returns is **validated against it** before it leaves the server. `output_schema` is not documentation. Whatever your function returns is **validated against it** before it leaves the server.

View file

@ -219,6 +219,7 @@ ignore = ["PERF203"]
[tool.ruff.lint.flake8-tidy-imports.banned-api] [tool.ruff.lint.flake8-tidy-imports.banned-api]
"pydantic.RootModel".msg = "Use `pydantic.TypeAdapter` instead." "pydantic.RootModel".msg = "Use `pydantic.TypeAdapter` instead."
"pydantic.root_model.RootModel".msg = "Use `pydantic.TypeAdapter` instead."
[tool.ruff.lint.mccabe] [tool.ruff.lint.mccabe]
@ -228,8 +229,8 @@ max-complexity = 24 # Default is 10
"__init__.py" = ["F401"] "__init__.py" = ["F401"]
# The mcp.types package is an alias that mirrors mcp_types namespaces by design. # The mcp.types package is an alias that mirrors mcp_types namespaces by design.
"src/mcp/types/*.py" = ["F403"] "src/mcp/types/*.py" = ["F403"]
# Generated by scripts/gen_surface_types.py: raw datamodel-codegen output (TID251 lifts the repo-wide RootModel ban for these generated validators). # Generated by scripts/gen_surface_types.py.
"src/mcp-types/mcp_types/_v*/__init__.py" = ["D212", "E501", "I001", "TID251", "UP007", "UP037"] "src/mcp-types/mcp_types/_v*/__init__.py" = ["D212", "E501", "I001", "UP007", "UP037"]
"tests/server/mcpserver/test_func_metadata.py" = ["E501"] "tests/server/mcpserver/test_func_metadata.py" = ["E501"]
# Inline snapshots of the translation tool's output carry long status/prompt lines verbatim. # Inline snapshots of the translation tool's output carry long status/prompt lines verbatim.
"tests/docs/test_translations.py" = ["E501"] "tests/docs/test_translations.py" = ["E501"]

View file

@ -5,13 +5,14 @@ writes the result to `src/mcp-types/mcp_types/_v<version>/__init__.py` (the
underscore marks these as internal validators, not public API) with only underscore marks these as internal validators, not public API) with only
the fixes the raw output needs: a small JSON pre-patch for the known the fixes the raw output needs: a small JSON pre-patch for the known
`number`-as-`integer` schema.json defect, a header, full URLs for the spec's `number`-as-`integer` schema.json defect, a header, full URLs for the spec's
site-absolute doc links, and per-version epilogue aliases. Run with site-absolute doc links, plain type aliases, and per-version epilogue aliases. Run with
`uv run --frozen --group codegen python scripts/gen_surface_types.py [--check]`. `uv run --frozen --group codegen python scripts/gen_surface_types.py [--check]`.
""" """
from __future__ import annotations from __future__ import annotations
import argparse import argparse
import ast
import difflib import difflib
import hashlib import hashlib
import json import json
@ -195,6 +196,7 @@ def run_codegen(schema_path: Path, output_path: Path) -> None:
"--use-annotated", "--use-field-description", "--use-schema-description", "--use-annotated", "--use-field-description", "--use-schema-description",
"--enum-field-as-literal", "all", "--enum-field-as-literal", "all",
"--use-union-operator", "--use-double-quotes", "--use-union-operator", "--use-double-quotes",
"--use-type-alias", "--skip-root-model",
"--extra-fields", "ignore", "--extra-fields", "ignore",
# JSON Schema `format` is annotation-only; codegen's defaults # JSON Schema `format` is annotation-only; codegen's defaults
# (Base64Str, AnyUrl) over-assert and reject valid wire data. # (Base64Str, AnyUrl) over-assert and reject valid wire data.
@ -237,6 +239,10 @@ def build(entry: dict[str, str]) -> str:
schema = json.loads((SCHEMA_DIR / f"{version}.json").read_text(encoding="utf-8")) schema = json.loads((SCHEMA_DIR / f"{version}.json").read_text(encoding="utf-8"))
patch_schema(schema, SCHEMA_PATCHES.get(version, [])) patch_schema(schema, SCHEMA_PATCHES.get(version, []))
make_server_info_opaque(schema) make_server_info_opaque(schema)
if "JSONValue" in schema["$defs"]:
# A single recursive alias avoids mutually recursive alias evaluation in type checkers.
assert schema["$defs"]["JSONValue"]["anyOf"][0] == {"$ref": "#/$defs/JSONObject"}
schema["$defs"]["JSONValue"]["anyOf"][0] = schema["$defs"]["JSONObject"]
with tempfile.TemporaryDirectory() as tmp: with tempfile.TemporaryDirectory() as tmp:
patched = Path(tmp) / "schema.json" patched = Path(tmp) / "schema.json"
@ -246,7 +252,27 @@ def build(entry: dict[str, str]) -> str:
source = raw.read_text(encoding="utf-8") source = raw.read_text(encoding="utf-8")
source = re.sub(r"\A# generated by datamodel-codegen:\n#[^\n]*\n", "", source) source = re.sub(r"\A# generated by datamodel-codegen:\n#[^\n]*\n", "", source)
source = re.sub(r"^class Model\(RootModel\[Any\]\):\n {4}root: Any\n+", "", source, count=1, flags=re.MULTILINE) # Keep named aliases only for recursive types; other aliases remain ordinary Python types and unions.
for node in reversed(ast.parse(source).body):
if not (
isinstance(node, ast.Assign)
and isinstance(node.value, ast.Call)
and isinstance(node.value.func, ast.Name)
and node.value.func.id == "TypeAliasType"
):
continue
value = node.value.args[1]
if any(
isinstance(part, ast.Constant) and isinstance(part.value, str) and part.value in schema["$defs"]
for part in ast.walk(value)
):
continue
original = ast.get_source_segment(source, node.value)
replacement = ast.get_source_segment(source, value)
assert original is not None and replacement is not None
source = source.replace(original, f"({replacement})", 1)
if "= TypeAliasType(" not in source:
source = source.replace("from typing_extensions import TypeAliasType\n", "")
# Codegen appends `| None` to forward refs of nullable models, which is a # Codegen appends `| None` to forward refs of nullable models, which is a
# runtime TypeError on a string ref and redundant since `JSONValue` includes None. # runtime TypeError on a string ref and redundant since `JSONValue` includes None.
source = source.replace('"JSONValue" | None', '"JSONValue"') source = source.replace('"JSONValue" | None', '"JSONValue"')
@ -256,8 +282,7 @@ def build(entry: dict[str, str]) -> str:
source = source.replace("](/", "](https://modelcontextprotocol.io/") source = source.replace("](/", "](https://modelcontextprotocol.io/")
source = allow_open_class_extras(source, OPEN_CLASSES[version]) source = allow_open_class_extras(source, OPEN_CLASSES[version])
if epilogue := EPILOGUES.get(version, ""): if epilogue := EPILOGUES.get(version, ""):
# Insert before the trailing model_rebuild() block: pyright's evaluation # Resolve aliases before rebuilding models with forward references.
# order for the recursive RootModel block is sensitive to placement.
match = re.search(r"^\w+\.model_rebuild\(\)$", source, flags=re.MULTILINE) match = re.search(r"^\w+\.model_rebuild\(\)$", source, flags=re.MULTILINE)
cut = match.start() if match else len(source) cut = match.start() if match else len(source)
source = f"{source[:cut]}{epilogue}\n\n{source[cut:]}" source = f"{source[:cut]}{epilogue}\n\n{source[cut:]}"

View file

@ -9,7 +9,7 @@ from __future__ import annotations
from typing import Annotated, Any, Literal from typing import Annotated, Any, Literal
from mcp_types._wire_base import WireModel from mcp_types._wire_base import WireModel
from pydantic import ConfigDict, Field, RootModel from pydantic import ConfigDict, Field
class BaseMetadata(WireModel): class BaseMetadata(WireModel):
@ -285,11 +285,10 @@ class CompleteResult(WireModel):
completion: Completion completion: Completion
class Cursor(RootModel[str]): Cursor = str
root: str """
""" An opaque token used to represent a cursor for pagination.
An opaque token used to represent a cursor for pagination. """
"""
class RequestedSchema(WireModel): class RequestedSchema(WireModel):
@ -556,27 +555,13 @@ class LegacyTitledEnumSchema(WireModel):
type: Literal["string"] type: Literal["string"]
class LoggingLevel( LoggingLevel = Literal["alert", "critical", "debug", "emergency", "error", "info", "notice", "warning"]
RootModel[ """
Literal[ The severity of a log message.
"alert",
"critical",
"debug",
"emergency",
"error",
"info",
"notice",
"warning",
]
]
):
root: Literal["alert", "critical", "debug", "emergency", "error", "info", "notice", "warning"]
"""
The severity of a log message.
These map to syslog message severities, as specified in RFC-5424: These map to syslog message severities, as specified in RFC-5424:
https://datatracker.ietf.org/doc/html/rfc5424#section-6.2.1 https://datatracker.ietf.org/doc/html/rfc5424#section-6.2.1
""" """
class LoggingMessageNotificationParams(WireModel): class LoggingMessageNotificationParams(WireModel):
@ -723,11 +708,10 @@ class PaginatedResult(WireModel):
""" """
class ProgressToken(RootModel[str | int]): ProgressToken = str | int
root: str | int """
""" A progress token, used to associate progress notifications with the original request.
A progress token, used to associate progress notifications with the original request. """
"""
class PromptArgument(WireModel): class PromptArgument(WireModel):
@ -853,11 +837,10 @@ class Request(WireModel):
params: dict[str, Any] | None = None params: dict[str, Any] | None = None
class RequestId(RootModel[str | int]): RequestId = str | int
root: str | int """
""" A uniquely identifying ID for a request in JSON-RPC.
A uniquely identifying ID for a request in JSON-RPC. """
"""
class RequestParams(WireModel): class RequestParams(WireModel):
@ -970,11 +953,10 @@ class Result(WireModel):
""" """
class Role(RootModel[Literal["assistant", "user"]]): Role = Literal["assistant", "user"]
root: Literal["assistant", "user"] """
""" The sender or recipient of messages and data in a conversation.
The sender or recipient of messages and data in a conversation. """
"""
class Root(WireModel): class Root(WireModel):
@ -1216,11 +1198,10 @@ class TaskMetadata(WireModel):
""" """
class TaskStatus(RootModel[Literal["cancelled", "completed", "failed", "input_required", "working"]]): TaskStatus = Literal["cancelled", "completed", "failed", "input_required", "working"]
root: Literal["cancelled", "completed", "failed", "input_required", "working"] """
""" The status of a task.
The status of a task. """
"""
class TextResourceContents(WireModel): class TextResourceContents(WireModel):
@ -1867,26 +1848,16 @@ class EmbeddedResource(WireModel):
type: Literal["resource"] type: Literal["resource"]
class EmptyResult(RootModel[Result]): EmptyResult = Result
root: Result
class EnumSchema( EnumSchema = (
RootModel[ UntitledSingleSelectEnumSchema
UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema
| TitledSingleSelectEnumSchema | UntitledMultiSelectEnumSchema
| UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema
| TitledMultiSelectEnumSchema | LegacyTitledEnumSchema
| LegacyTitledEnumSchema )
]
):
root: (
UntitledSingleSelectEnumSchema
| TitledSingleSelectEnumSchema
| UntitledMultiSelectEnumSchema
| TitledMultiSelectEnumSchema
| LegacyTitledEnumSchema
)
class GetPromptRequestParams(WireModel): class GetPromptRequestParams(WireModel):
@ -2115,8 +2086,7 @@ class LoggingMessageNotification(WireModel):
params: LoggingMessageNotificationParams params: LoggingMessageNotificationParams
class MultiSelectEnumSchema(RootModel[UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema]): MultiSelectEnumSchema = UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema
root: UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema
class PaginatedRequestParams(WireModel): class PaginatedRequestParams(WireModel):
@ -2152,32 +2122,20 @@ class PingRequest(WireModel):
params: RequestParams | None = None params: RequestParams | None = None
class PrimitiveSchemaDefinition( PrimitiveSchemaDefinition = (
RootModel[ StringSchema
StringSchema | NumberSchema
| NumberSchema | BooleanSchema
| BooleanSchema | UntitledSingleSelectEnumSchema
| UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema
| TitledSingleSelectEnumSchema | UntitledMultiSelectEnumSchema
| UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema
| TitledMultiSelectEnumSchema | LegacyTitledEnumSchema
| LegacyTitledEnumSchema )
] """
): Restricted schema definitions that only allow primitive types
root: ( without nested objects or arrays.
StringSchema """
| NumberSchema
| BooleanSchema
| UntitledSingleSelectEnumSchema
| TitledSingleSelectEnumSchema
| UntitledMultiSelectEnumSchema
| TitledMultiSelectEnumSchema
| LegacyTitledEnumSchema
)
"""
Restricted schema definitions that only allow primitive types
without nested objects or arrays.
"""
class ProgressNotificationParams(WireModel): class ProgressNotificationParams(WireModel):
@ -2499,8 +2457,7 @@ class SetLevelRequest(WireModel):
params: SetLevelRequestParams params: SetLevelRequestParams
class SingleSelectEnumSchema(RootModel[UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema]): SingleSelectEnumSchema = UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema
root: UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema
class SubscribeRequest(WireModel): class SubscribeRequest(WireModel):
@ -2793,8 +2750,7 @@ class CompleteRequest(WireModel):
params: CompleteRequestParams params: CompleteRequestParams
class ContentBlock(RootModel[TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource]): ContentBlock = TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource
root: TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource
class CreateTaskResult(WireModel): class CreateTaskResult(WireModel):
@ -2812,11 +2768,10 @@ class CreateTaskResult(WireModel):
task: Task task: Task
class ElicitRequestParams(RootModel[ElicitRequestURLParams | ElicitRequestFormParams]): ElicitRequestParams = ElicitRequestURLParams | ElicitRequestFormParams
root: ElicitRequestURLParams | ElicitRequestFormParams """
""" The parameters for a request to elicit additional information from the user via the client.
The parameters for a request to elicit additional information from the user via the client. """
"""
class GetPromptRequest(WireModel): class GetPromptRequest(WireModel):
@ -2857,18 +2812,16 @@ class InitializeRequest(WireModel):
params: InitializeRequestParams params: InitializeRequestParams
class JSONRPCMessage(RootModel[JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse]): JSONRPCMessage = JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse
root: JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse """
""" Refers to any valid JSON-RPC object that can be decoded off the wire, or encoded to be sent.
Refers to any valid JSON-RPC object that can be decoded off the wire, or encoded to be sent. """
"""
class JSONRPCResponse(RootModel[JSONRPCResultResponse | JSONRPCErrorResponse]): JSONRPCResponse = JSONRPCResultResponse | JSONRPCErrorResponse
root: JSONRPCResultResponse | JSONRPCErrorResponse """
""" A response to a request, containing either the result or error.
A response to a request, containing either the result or error. """
"""
class ListPromptsRequest(WireModel): class ListPromptsRequest(WireModel):
@ -3173,64 +3126,34 @@ class CallToolResult(WireModel):
""" """
class ClientNotification( ClientNotification = (
RootModel[ CancelledNotification
CancelledNotification | InitializedNotification
| InitializedNotification | ProgressNotification
| ProgressNotification | TaskStatusNotification
| TaskStatusNotification | RootsListChangedNotification
| RootsListChangedNotification )
]
):
root: (
CancelledNotification
| InitializedNotification
| ProgressNotification
| TaskStatusNotification
| RootsListChangedNotification
)
class ClientRequest( ClientRequest = (
RootModel[ InitializeRequest
InitializeRequest | PingRequest
| PingRequest | ListResourcesRequest
| ListResourcesRequest | ListResourceTemplatesRequest
| ListResourceTemplatesRequest | ReadResourceRequest
| ReadResourceRequest | SubscribeRequest
| SubscribeRequest | UnsubscribeRequest
| UnsubscribeRequest | ListPromptsRequest
| ListPromptsRequest | GetPromptRequest
| GetPromptRequest | ListToolsRequest
| ListToolsRequest | CallToolRequest
| CallToolRequest | GetTaskRequest
| GetTaskRequest | GetTaskPayloadRequest
| GetTaskPayloadRequest | CancelTaskRequest
| CancelTaskRequest | ListTasksRequest
| ListTasksRequest | SetLevelRequest
| SetLevelRequest | CompleteRequest
| CompleteRequest )
]
):
root: (
InitializeRequest
| PingRequest
| ListResourcesRequest
| ListResourceTemplatesRequest
| ReadResourceRequest
| SubscribeRequest
| UnsubscribeRequest
| ListPromptsRequest
| GetPromptRequest
| ListToolsRequest
| CallToolRequest
| GetTaskRequest
| GetTaskPayloadRequest
| CancelTaskRequest
| ListTasksRequest
| SetLevelRequest
| CompleteRequest
)
class ElicitRequest(WireModel): class ElicitRequest(WireModel):
@ -3266,72 +3189,38 @@ class GetPromptResult(WireModel):
messages: list[PromptMessage] messages: list[PromptMessage]
class SamplingMessageContentBlock( SamplingMessageContentBlock = TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
RootModel[TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent]
):
root: TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
class ServerNotification( ServerNotification = (
RootModel[ CancelledNotification
CancelledNotification | ProgressNotification
| ProgressNotification | ResourceListChangedNotification
| ResourceListChangedNotification | ResourceUpdatedNotification
| ResourceUpdatedNotification | PromptListChangedNotification
| PromptListChangedNotification | ToolListChangedNotification
| ToolListChangedNotification | TaskStatusNotification
| TaskStatusNotification | LoggingMessageNotification
| LoggingMessageNotification | ElicitationCompleteNotification
| ElicitationCompleteNotification )
]
):
root: (
CancelledNotification
| ProgressNotification
| ResourceListChangedNotification
| ResourceUpdatedNotification
| PromptListChangedNotification
| ToolListChangedNotification
| TaskStatusNotification
| LoggingMessageNotification
| ElicitationCompleteNotification
)
class ServerResult( ServerResult = (
RootModel[ Result
Result | InitializeResult
| InitializeResult | ListResourcesResult
| ListResourcesResult | ListResourceTemplatesResult
| ListResourceTemplatesResult | ReadResourceResult
| ReadResourceResult | ListPromptsResult
| ListPromptsResult | GetPromptResult
| GetPromptResult | ListToolsResult
| ListToolsResult | CallToolResult
| CallToolResult | GetTaskResult
| GetTaskResult | GetTaskPayloadResult
| GetTaskPayloadResult | CancelTaskResult
| CancelTaskResult | ListTasksResult
| ListTasksResult | CompleteResult
| CompleteResult )
]
):
root: (
Result
| InitializeResult
| ListResourcesResult
| ListResourceTemplatesResult
| ReadResourceResult
| ListPromptsResult
| GetPromptResult
| ListToolsResult
| CallToolResult
| GetTaskResult
| GetTaskPayloadResult
| CancelTaskResult
| ListTasksResult
| CompleteResult
)
class CreateMessageResult(WireModel): class CreateMessageResult(WireModel):
@ -3398,28 +3287,16 @@ class SamplingMessage(WireModel):
role: Role role: Role
class ClientResult( ClientResult = (
RootModel[ Result
Result | GetTaskResult
| GetTaskResult | GetTaskPayloadResult
| GetTaskPayloadResult | CancelTaskResult
| CancelTaskResult | ListTasksResult
| ListTasksResult | CreateMessageResult
| CreateMessageResult | ListRootsResult
| ListRootsResult | ElicitResult
| ElicitResult )
]
):
root: (
Result
| GetTaskResult
| GetTaskPayloadResult
| CancelTaskResult
| ListTasksResult
| CreateMessageResult
| ListRootsResult
| ElicitResult
)
class CreateMessageRequestParams(WireModel): class CreateMessageRequestParams(WireModel):
@ -3502,25 +3379,13 @@ class CreateMessageRequest(WireModel):
params: CreateMessageRequestParams params: CreateMessageRequestParams
class ServerRequest( ServerRequest = (
RootModel[ PingRequest
PingRequest | GetTaskRequest
| GetTaskRequest | GetTaskPayloadRequest
| GetTaskPayloadRequest | CancelTaskRequest
| CancelTaskRequest | ListTasksRequest
| ListTasksRequest | CreateMessageRequest
| CreateMessageRequest | ListRootsRequest
| ListRootsRequest | ElicitRequest
| ElicitRequest )
]
):
root: (
PingRequest
| GetTaskRequest
| GetTaskPayloadRequest
| CancelTaskRequest
| ListTasksRequest
| CreateMessageRequest
| ListRootsRequest
| ElicitRequest
)

File diff suppressed because it is too large Load diff

View file

@ -80,19 +80,12 @@ def _clamp_inbound_ttl(raw: dict[str, Any]) -> None:
@cache @cache
def _wire_fields(target: type[BaseModel] | UnionType) -> frozenset[str]: def _wire_fields(target: type[BaseModel] | UnionType) -> frozenset[str]:
"""Top-level wire keys `target` declares (its members', for a union). """Top-level wire keys `target` declares (its members', for a union)."""
A `RootModel` row (e.g. an empty result carried as `RootModel[Result]`)
reports its wrapped type's keys, not the pydantic-internal `root`.
"""
members: tuple[Any, ...] = get_args(target) if isinstance(target, UnionType) else (target,) members: tuple[Any, ...] = get_args(target) if isinstance(target, UnionType) else (target,)
models = [m for m in members if isinstance(m, type) and issubclass(m, BaseModel)] models = [m for m in members if isinstance(m, type) and issubclass(m, BaseModel)]
fields: set[str] = set() fields: set[str] = set()
for model in models: for model in models:
if getattr(model, "__pydantic_root_model__", False): # a RootModel wrapper row fields.update(field.alias or name for name, field in model.model_fields.items())
fields |= _wire_fields(model.model_fields["root"].annotation)
else:
fields.update(field.alias or name for name, field in model.model_fields.items())
return frozenset(fields) return frozenset(fields)

View file

@ -2,8 +2,7 @@ import logging
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Literal from typing import Any, Literal
# TODO(Marcelo): We should drop the `RootModel`. from pydantic import AnyUrl, BaseModel, Field, TypeAdapter, ValidationError
from pydantic import AnyUrl, BaseModel, Field, RootModel, ValidationError # noqa: TID251
from starlette.datastructures import FormData, QueryParams from starlette.datastructures import FormData, QueryParams
from starlette.requests import Request from starlette.requests import Request
from starlette.responses import RedirectResponse, Response from starlette.responses import RedirectResponse, Response
@ -20,6 +19,7 @@ from mcp.server.auth.provider import (
from mcp.shared.auth import InvalidRedirectUriError, InvalidScopeError from mcp.shared.auth import InvalidRedirectUriError, InvalidScopeError
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_ANY_URL_ADAPTER = TypeAdapter(AnyUrl)
class AuthorizationRequest(BaseModel): class AuthorizationRequest(BaseModel):
@ -59,10 +59,6 @@ def best_effort_extract_string(key: str, params: None | FormData | QueryParams)
return None return None
class AnyUrlModel(RootModel[AnyUrl]):
root: AnyUrl
@dataclass @dataclass
class AuthorizationHandler: class AuthorizationHandler:
provider: OAuthAuthorizationServerProvider[Any, Any, Any] provider: OAuthAuthorizationServerProvider[Any, Any, Any]
@ -107,9 +103,9 @@ class AuthorizationHandler:
if params is not None and "redirect_uri" not in params: if params is not None and "redirect_uri" not in params:
raw_redirect_uri = None raw_redirect_uri = None
else: else:
raw_redirect_uri = AnyUrlModel.model_validate( raw_redirect_uri = _ANY_URL_ADAPTER.validate_python(
best_effort_extract_string("redirect_uri", params) best_effort_extract_string("redirect_uri", params)
).root )
redirect_uri = client.validate_redirect_uri(raw_redirect_uri) redirect_uri = client.validate_redirect_uri(raw_redirect_uri)
except (ValidationError, InvalidRedirectUriError): except (ValidationError, InvalidRedirectUriError):
# if the redirect URI is invalid, ignore it & just return the # if the redirect URI is invalid, ignore it & just return the

View file

@ -8,7 +8,7 @@ from mcp_types import RequestId
# Internal surface package; imported as the gate's source of truth for spec-valid property schemas. # 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 mcp_types._v2025_11_25 import PrimitiveSchemaDefinition
from pydantic import BaseModel, ValidationError from pydantic import BaseModel, TypeAdapter, ValidationError
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
from pydantic_core import core_schema from pydantic_core import core_schema
from typing_extensions import TypeAliasType from typing_extensions import TypeAliasType
@ -16,6 +16,7 @@ from typing_extensions import TypeAliasType
from mcp.server.session import ServerSession from mcp.server.session import ServerSession
ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel) ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel)
_PRIMITIVE_SCHEMA_ADAPTER = TypeAdapter[PrimitiveSchemaDefinition](PrimitiveSchemaDefinition)
class AcceptedElicitation(BaseModel, Generic[ElicitSchemaModelT]): class AcceptedElicitation(BaseModel, Generic[ElicitSchemaModelT]):
@ -79,7 +80,7 @@ def _validate_rendered_properties(json_schema: dict[str, Any]) -> None:
""" """
for field_name, prop in json_schema.get("properties", {}).items(): for field_name, prop in json_schema.get("properties", {}).items():
try: try:
PrimitiveSchemaDefinition.model_validate(prop) _PRIMITIVE_SCHEMA_ADAPTER.validate_python(prop)
except ValidationError: except ValidationError:
raise TypeError( raise TypeError(
f"Elicitation schema field {field_name!r} rendered as {prop!r}, " f"Elicitation schema field {field_name!r} rendered as {prop!r}, "

View file

@ -114,15 +114,15 @@ class ArgModelBase(BaseModel):
class FuncMetadata(BaseModel): class FuncMetadata(BaseModel):
"""A tool function's argument model plus, for structured output, the published `output_schema` and the """A tool function's argument model plus, for structured output, the published `output_schema` and the
`output_model` results are validated against. Constructing one with an `output_model` and no schema derives `output_model` type annotation results are validated against. Constructing one with an `output_model` and no
the schema (and raises if pydantic can't); the fields are read live, so clearing or reassigning them later schema derives the schema (and raises if pydantic can't); the fields are read live, so reassigning them later
takes effect on the next call.""" takes effect on the next call."""
arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)] arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)]
output_schema: dict[str, Any] | None = None output_schema: dict[str, Any] | None = None
output_model: Annotated[type[Any], WithJsonSchema(None)] | None = None output_model: Annotated[Any, WithJsonSchema(None)] = None
wrap_output: bool = False wrap_output: bool = False
_adapter: tuple[type[Any], TypeAdapter[Any]] | None = PrivateAttr(default=None) _adapter: tuple[Any, TypeAdapter[Any]] | None = PrivateAttr(default=None)
def model_post_init(self, context: Any, /) -> None: def model_post_init(self, context: Any, /) -> None:
if self.output_model is not None and self.output_schema is None: if self.output_model is not None and self.output_schema is None:
@ -130,7 +130,7 @@ class FuncMetadata(BaseModel):
schema = self._output_adapter(self.output_model).json_schema(schema_generator=StrictJsonSchema) schema = self._output_adapter(self.output_model).json_schema(schema_generator=StrictJsonSchema)
self.output_schema = _inline_root_ref(schema) self.output_schema = _inline_root_ref(schema)
def _output_adapter(self, output_model: type[Any]) -> TypeAdapter[Any]: def _output_adapter(self, output_model: Any) -> TypeAdapter[Any]:
"""The validator/serializer for `output_model`, built once and rebuilt only if the field is reassigned.""" """The validator/serializer for `output_model`, built once and rebuilt only if the field is reassigned."""
if self._adapter is None or self._adapter[0] is not output_model: if self._adapter is None or self._adapter[0] is not output_model:
self._adapter = (output_model, TypeAdapter(_pydantic_readable_typeddict(output_model))) self._adapter = (output_model, TypeAdapter(_pydantic_readable_typeddict(output_model)))
@ -477,7 +477,7 @@ def func_metadata(
return FuncMetadata(arg_model=arguments_model) return FuncMetadata(arg_model=arguments_model)
def _create_output_model(original_annotation: Any, type_expr: Any, func_name: str) -> tuple[type[Any] | None, bool]: def _create_output_model(original_annotation: Any, type_expr: Any, func_name: str) -> tuple[Any, bool]:
"""Pick the type structured output is validated against for the given return annotation. """Pick the type structured output is validated against for the given return annotation.
Args: Args:
@ -491,7 +491,7 @@ def _create_output_model(original_annotation: Any, type_expr: Any, func_name: st
Model is None if the type cannot carry structured output. Model is None if the type cannot carry structured output.
wrap_output is True if the result needs to be wrapped in {"result": ...} wrap_output is True if the result needs to be wrapped in {"result": ...}
""" """
model: type[Any] | None = None model: Any = None
wrap_output = False wrap_output = False
# First handle special case: None # First handle special case: None
@ -503,13 +503,12 @@ def _create_output_model(original_annotation: Any, type_expr: Any, func_name: st
elif isinstance(type_expr, GenericAlias): elif isinstance(type_expr, GenericAlias):
origin = get_origin(type_expr) origin = get_origin(type_expr)
# Special case: dict with string keys can use RootModel
if origin is dict: if origin is dict:
args = get_args(type_expr) args = get_args(type_expr)
if len(args) == 2 and args[0] is str: if len(args) == 2 and args[0] is str:
# TODO: should we use the original annotation? We are losing any potential `Annotated` # TODO: should we use the original annotation? We are losing any potential `Annotated`
# metadata for Pydantic here: # metadata for Pydantic here:
model = _create_dict_model(func_name, type_expr) model = Annotated[type_expr, Field(title=f"{func_name}DictOutput")]
else: else:
# dict with non-str keys needs wrapping # dict with non-str keys needs wrapping
model = _create_wrapped_model(func_name, original_annotation) model = _create_wrapped_model(func_name, original_annotation)
@ -622,21 +621,6 @@ def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]:
return create_model(model_name, result=annotation) return create_model(model_name, result=annotation)
def _create_dict_model(func_name: str, dict_annotation: Any) -> type[BaseModel]:
"""Create a RootModel for dict[str, T] types."""
# TODO(Marcelo): We should not rely on RootModel for this.
from pydantic import RootModel # noqa: TID251
class DictModel(RootModel[dict_annotation]):
pass
# Give it a meaningful name
DictModel.__name__ = f"{func_name}DictOutput"
DictModel.__qualname__ = f"{func_name}DictOutput"
return DictModel
def _convert_to_content(result: Any) -> list[ContentBlock]: def _convert_to_content(result: Any) -> list[ContentBlock]:
"""Convert a result to a sequence of content objects. """Convert a result to a sequence of content objects.

View file

@ -454,11 +454,20 @@ def test_elicit_result_surface_accepts_null_content_values_at_every_version_that
surface.model_validate({"action": "accept", "content": {"name": "x", "age": None}}) surface.model_validate({"action": "accept", "content": {"name": "x", "age": None}})
def test_server_capabilities_extensions_with_null_json_value_round_trips_at_2026(): def test_discovery_capabilities_preserve_nested_json_values() -> None:
"""Spec `JSONValue` includes `null`; the ts->json render dropped it from the vendored schema.""" """Spec `JSONValue` permits nested containers and every JSON scalar, including numbers and null."""
raw: dict[str, Any] = {"extensions": {"x": {"k": None}}} extensions = {"x": {"nested": [None, True, 1, 1.5, "text", {"child": [False]}]}}
parsed = v2026.ServerCapabilities.model_validate(raw) raw = {
assert parsed.model_dump(mode="json")["extensions"] == {"x": {"k": None}} "supportedVersions": ["2026-07-28"],
"capabilities": {"extensions": extensions},
"resultType": "complete",
"ttlMs": 0,
"cacheScope": "private",
}
parsed = methods.parse_server_result("server/discover", "2026-07-28", raw)
assert isinstance(parsed, types.DiscoverResult)
assert parsed.capabilities.extensions == extensions
assert methods.serialize_server_result("server/discover", "2026-07-28", raw) == raw
def test_elicit_request_surface_accepts_loose_property_schemas(): def test_elicit_request_surface_accepts_loose_property_schemas():

View file

@ -153,8 +153,6 @@ def _surface_classes(module: ModuleType) -> list[tuple[str, type[BaseModel]]]:
continue continue
if obj.__module__ != module.__name__ or obj.__name__ != name: if obj.__module__ != module.__name__ or obj.__name__ != name:
continue # re-export or alias to another model continue # re-export or alias to another model
if getattr(obj, "__pydantic_root_model__", False):
continue # RootModel alias wrapper; the field-subset property does not apply
out.append((f"{tail}.{name}", obj)) out.append((f"{tail}.{name}", obj))
return out return out