"""Regenerate the per-version wire-shape surface packages from vendored schemas.

Runs `datamodel-code-generator` over each `schema/PINNED.json` entry and

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, 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

import re

import subprocess

import sys

import tempfile

from pathlib import Path

from typing import Any

REPO_ROOT = Path(__file__).resolve().parent.parent

SCHEMA_DIR = REPO_ROOT / "schema"

TYPES_DIR = REPO_ROOT / "src" / "mcp-types" / "mcp_types"

# The result-meta serverInfo stamp: every `$defs` entry carrying this property

# gets its typed `$ref` stripped by `make_server_info_opaque` below.

SERVER_INFO_META_PROPERTY = "io.modelcontextprotocol/serverInfo"

# schema.ts -> schema.json renders TypeScript `number` as JSON Schema

# `integer` at these sites; patch the JSON before codegen so floats validate.

# Patched to `["integer", "number"]` (not bare `"number"`) so codegen emits

# `int | float` and pydantic's smart-union preserves ints on round-trip.

# TODO: drop once modelcontextprotocol/modelcontextprotocol fixes the schema.ts -> schema.json number rendering.

SCHEMA_PATCHES: dict[str, list[tuple[str, Any, Any]]] = {

"2025-11-25": [

("$defs/NumberSchema/properties/default/type", "integer", ["integer", "number"]),

("$defs/NumberSchema/properties/maximum/type", "integer", ["integer", "number"]),

("$defs/NumberSchema/properties/minimum/type", "integer", ["integer", "number"]),

# `null` arm is monolith superset leniency: hosts may answer optional form fields with null.

(

"$defs/ElicitResult/properties/content/additionalProperties/anyOf/1/type",

["string", "integer", "boolean"],

["string", "integer", "number", "boolean", "null"],

),

# Older python-sdk releases emit `anyOf` for Optional fields; the callback's

# own schema validation is the real gate, so accept any property shape inbound.

# PrimitiveSchemaDefinition becomes an orphan $def after this patch but

# datamodel-codegen still emits it; elicitation.py imports it as the gate type.

(

"$defs/ElicitRequestFormParams/properties/requestedSchema/properties/properties/additionalProperties",

{"$ref": "#/$defs/PrimitiveSchemaDefinition"},

{},

),

# JSON Schema 2020-12 allows a boolean wherever a sub-schema is expected, and

# 2026-07-28 already leaves these free-form; accept `true`/`false` property schemas.

(

"$defs/Tool/properties/inputSchema/properties/properties/additionalProperties",

{"additionalProperties": True, "properties": {}, "type": "object"},

{"anyOf": [{"additionalProperties": True, "properties": {}, "type": "object"}, {"type": "boolean"}]},

),

(

"$defs/Tool/properties/outputSchema/properties/properties/additionalProperties",

{"additionalProperties": True, "properties": {}, "type": "object"},

{"anyOf": [{"additionalProperties": True, "properties": {}, "type": "object"}, {"type": "boolean"}]},

),

],

"2026-07-28": [

("$defs/NumberSchema/properties/default/type", "number", ["integer", "number"]),

("$defs/NumberSchema/properties/maximum/type", "number", ["integer", "number"]),

("$defs/NumberSchema/properties/minimum/type", "number", ["integer", "number"]),

# `null` arm is monolith superset leniency: hosts may answer optional form fields with null.

(

"$defs/ElicitResult/properties/content/additionalProperties/anyOf/1/type",

["string", "integer", "boolean"],

["string", "integer", "number", "boolean", "null"],

),

# Spec `JSONValue` includes `number` and `null`; the ts->json render dropped both.

(

"$defs/JSONValue/anyOf/2/type",

["string", "integer", "boolean"],

["string", "integer", "number", "boolean", "null"],

),

# Older python-sdk releases emit `anyOf` for Optional fields; the callback's

# own schema validation is the real gate, so accept any property shape inbound.

(

"$defs/ElicitRequestFormParams/properties/requestedSchema/properties/properties/additionalProperties",

{"$ref": "#/$defs/PrimitiveSchemaDefinition"},

{},

),

],

}

# Classes the spec defines as open key-value bags: `_meta` content, the

# JSON-Schema-document fields on `Tool`, and the schemas with explicit

# `additionalProperties: {}`. These keep `extra="allow"` so the sieve preserves

# arbitrary keys; every other class ignores extras. Per-version because codegen

# reuses class names across versions for unrelated schemas (e.g. `Data`).

OPEN_CLASSES: dict[str, frozenset[str]] = {

"2025-11-25": frozenset({"Meta", "InputSchema", "OutputSchema", "Result", "GetTaskPayloadResult", "Data"}),

"2026-07-28": frozenset(

{

"MetaObject",

"NotificationMetaObject",

"RequestMetaObject",

"ResultMetaObject",

"SubscriptionsListenResultMeta",

"InputSchema",

"OutputSchema",

"Result",

}

),

}

# Hand-written union aliases the wire-method maps reference by value; the schema

# has no named definition for "everything tools/call may return", so name it here.

EPILOGUES: dict[str, str] = {

"2026-07-28": (

"AnyCallToolResult = CallToolResult | InputRequiredResult\n"

"AnyGetPromptResult = GetPromptResult | InputRequiredResult\n"

"AnyReadResourceResult = ReadResourceResult | InputRequiredResult\n"

),

}

HEADER = (

'"""Internal wire-shape models for protocol {version}. Generated; do not edit.\n'

"\n"

"Regenerate with `scripts/gen_surface_types.py` from `schema/{version}.json`\n"

'(sha256 `{sha}`)."""\n'

"# pyright: reportIncompatibleVariableOverride=false, reportGeneralTypeIssues=false\n"

)

def load_pinned() -> list[dict[str, str]]:

"""Read `schema/PINNED.json` and verify each vendored file's sha256."""

entries: list[dict[str, str]] = json.loads((SCHEMA_DIR / "PINNED.json").read_text(encoding="utf-8"))

for entry in entries:

path = SCHEMA_DIR / f"{entry['protocol_version']}.json"

actual = hashlib.sha256(path.read_bytes()).hexdigest()

if actual != entry["sha256"]:

raise SystemExit(f"sha256 mismatch for {path.name}: PINNED={entry['sha256']} disk={actual}")

return entries

def patch_schema(schema: dict[str, Any], patches: list[tuple[str, Any, Any]]) -> None:

"""Apply `(path, old, new)` JSON-pointer-ish patches in place, asserting the old value.

Path segments use JSON-pointer escaping (`~1` for `/`, `~0` for `~`) so keys

that themselves contain a slash (the reserved `io.modelcontextprotocol/*`

`_meta` keys) are addressable.

"""

for path, old, new in patches:

*parts, leaf = (part.replace("~1", "/").replace("~0", "~") for part in path.split("/"))

node: Any = schema

for part in parts:

node = node[int(part) if part.isdigit() else part]

if node[leaf] != old:

raise SystemExit(f"schema patch {path}: expected {old!r}, found {node[leaf]!r}")

node[leaf] = new

def make_server_info_opaque(schema: dict[str, Any]) -> None:

"""Strip the typed `$ref` from every result-meta serverInfo property.

The stamp is display-only: the spec forbids acting on it, so a malformed

value must never fail a whole response (clients validate every inbound

result against this surface). Walking every `$defs` entry keeps future

result-meta definitions lenient by construction instead of relying on an

enumerated list; the typed, lenient parse happens at the read edge

(`ClientSession.server_info`). typescript-sdk does the same with a

schema-level catch-to-undefined.

"""

for definition in schema.get("$defs", {}).values():

prop = definition.get("properties", {}).get(SERVER_INFO_META_PROPERTY)

if prop is not None and "$ref" in prop:

del prop["$ref"]

def run_codegen(schema_path: Path, output_path: Path) -> None:

"""Run datamodel-code-generator at the version pinned in the `codegen` dependency group."""

# fmt: off

result = subprocess.run(

[

"uv", "run", "--frozen", "--group", "codegen", "datamodel-codegen",

"--input", str(schema_path),

"--input-file-type", "jsonschema",

"--output", str(output_path),

"--output-model-type", "pydantic_v2.BaseModel",

"--target-python-version", "3.11",

"--base-class", "mcp_types._wire_base.WireModel",

"--snake-case-field", "--remove-special-field-name-prefix",

"--use-annotated", "--use-field-description", "--use-schema-description",

"--enum-field-as-literal", "all",

"--default-values", '{"ElicitRequestFormParams.mode": "form"}',

"--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.

"--type-mappings", "byte=string", "uri=string", "uri-template=string",

"--disable-timestamp",

],

capture_output=True, encoding="utf-8", errors="replace",

)

# fmt: on

if result.returncode != 0:

raise SystemExit(f"datamodel-codegen failed:\n{result.stderr}")

def allow_open_class_extras(source: str, open_classes: frozenset[str]) -> str:

"""Restore `extra="allow"` on `open_classes` only.

Every other class uses `extra="ignore"` so the surface acts as a sieve;

`open_classes` are the places the spec defines as open key-value bags.

"""

def patch(match: re.Match[str]) -> str:

if match.group(1) not in open_classes:

return match.group(0)

return match.group(0).replace('extra="ignore"', 'extra="allow"')

source = re.sub(

r'^class (\w+)\(WireModel\):\n(?: {4}.*\n|\n)*? {4}model_config = ConfigDict\(\n {8}extra="ignore",\n {4}\)\n',

patch,

source,

flags=re.MULTILINE,

)

# Drift guard: substitution count must match the allow-list.

assert source.count('extra="allow"') == len(open_classes), (source.count('extra="allow"'), open_classes)

return source

def build(entry: dict[str, str]) -> str:

"""Generate, post-process, and format one version's surface module text."""

version = entry["protocol_version"]

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"

patched.write_text(json.dumps(schema), encoding="utf-8")

raw = Path(tmp) / "raw.py"

run_codegen(patched, raw)

source = raw.read_text(encoding="utf-8")

source = re.sub(r"\A# generated by datamodel-codegen:\n#[^\n]*\n", "", source)

# Explicit __annotations__ hides declared fields under Python 3.14's deferred annotations.

source = source.replace(' __annotations__ = {\n "__pydantic_extra__": Dict[str, Any],\n }\n', "")

source = source.replace(", Dict,", ",")

# 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"')

# Schema descriptions link to spec-site pages with site-absolute paths; expand

# them to full URLs so they resolve from the rendered API docs and pass the

# strict mkdocs link validation.

source = source.replace("](/", "](https://modelcontextprotocol.io/")

source = allow_open_class_extras(source, OPEN_CLASSES[version])

if epilogue := EPILOGUES.get(version, ""):

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

source = HEADER.format(version=version, sha=entry["sha256"]) + source

staging = TYPES_DIR / f"_staging_{version}.py"

try:

staging.write_text(source, encoding="utf-8")

subprocess.run(

["uv", "run", "--frozen", "ruff", "format", "--no-cache", str(staging)],

cwd=REPO_ROOT, capture_output=True, check=True,

) # fmt: skip

return staging.read_text(encoding="utf-8")

finally:

staging.unlink(missing_ok=True)

def main(argv: list[str] | None = None) -> int:

"""CLI entry point: write each surface package, or diff under `--check`."""

parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument("--check", action="store_true", help="diff regenerated output against committed files")

args = parser.parse_args(argv)

drift = False

for entry in load_pinned():

target = TYPES_DIR / ("_v" + entry["protocol_version"].replace("-", "_")) / "__init__.py"

candidate = build(entry)

if not args.check:

target.parent.mkdir(parents=True, exist_ok=True)

target.write_text(candidate, encoding="utf-8")

print(f"{entry['protocol_version']}: wrote {target.relative_to(REPO_ROOT)} ({len(candidate)} bytes)")

continue

committed = target.read_text(encoding="utf-8") if target.is_file() else ""

if committed != candidate:

drift = True

sys.stderr.writelines(

difflib.unified_diff(

committed.splitlines(keepends=True),

candidate.splitlines(keepends=True),

fromfile=str(target.relative_to(REPO_ROOT)),

tofile="<regenerated>",

)

)

return 1 if drift else 0

if __name__ == "__main__":

sys.exit(main())