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

scripts/gen_surface_types.py

@@ -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
 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:]}"