Replace RootModel wrappers with type aliases and TypeAdapter validation (#3470)

Author: Kludex

docs/servers/structured-output.md

@@ -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.
 
+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
 
 `output_schema` is not documentation. Whatever your function returns is **validated against it** before it leaves the server.

pyproject.toml

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

scripts/gen_surface_types.py

@@ -5,13 +5,14 @@
 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
 `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]`.
 """
 
 from __future__ import annotations
 
 import argparse
+import ast
 import difflib
 import hashlib
 import json
@@ -195,6 +196,7 @@ def run_codegen(schema_path: Path, output_path: Path) -> None:
             "--use-annotated", "--use-field-description", "--use-schema-description",
             "--enum-field-as-literal", "all",
             "--use-union-operator", "--use-double-quotes",
+            "--use-type-alias", "--skip-root-model",
             "--extra-fields", "ignore",
             # JSON Schema `format` is annotation-only; codegen's defaults
             # (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"))
     patch_schema(schema, SCHEMA_PATCHES.get(version, []))
     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:
         patched = Path(tmp) / "schema.json"
@@ -246,7 +252,27 @@ def build(entry: dict[str, str]) -> str:
         source = raw.read_text(encoding="utf-8")
 
     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
     # runtime TypeError on a string ref and redundant since `JSONValue` includes None.
     source = source.replace('"JSONValue" | None', '"JSONValue"')
@@ -256,8 +282,7 @@ def build(entry: dict[str, str]) -> str:
     source = source.replace("](/", "](https://modelcontextprotocol.io/")
     source = allow_open_class_extras(source, OPEN_CLASSES[version])
     if epilogue := EPILOGUES.get(version, ""):
-        # Insert before the trailing model_rebuild() block: pyright's evaluation
-        # order for the recursive RootModel block is sensitive to placement.
+        # Resolve aliases before rebuilding models with forward references.
         match = re.search(r"^\w+\.model_rebuild\(\)$", source, flags=re.MULTILINE)
         cut = match.start() if match else len(source)
         source = f"{source[:cut]}{epilogue}\n\n{source[cut:]}"