fix: handle cyclic JSON schemas in MCP tool resolution

This commit is contained in:
Greyson LaLonde
2026-04-16 05:03:00 +08:00
committed by GitHub
parent 74d061e994
commit 29b5531f78
2 changed files with 245 additions and 35 deletions

View File

@@ -417,9 +417,18 @@ class MCPToolResolver:
args_schema = None
if tool_def.get("inputSchema"):
args_schema = self._json_schema_to_pydantic(
tool_name, tool_def["inputSchema"]
)
try:
args_schema = self._json_schema_to_pydantic(
tool_name, tool_def["inputSchema"]
)
except Exception as e:
self._logger.log(
"warning",
f"Failed to build args schema for MCP tool "
f"'{tool_name}': {e}. Registering tool without a "
"typed schema.",
)
args_schema = None
tool_schema = {
"description": tool_def.get("description", ""),

View File

@@ -19,7 +19,18 @@ from collections.abc import Callable
from copy import deepcopy
import datetime
import logging
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypedDict, Union, cast
from typing import (
TYPE_CHECKING,
Annotated,
Any,
Final,
ForwardRef,
Literal,
Optional,
TypedDict,
Union,
cast,
)
import uuid
import jsonref # type: ignore[import-untyped]
@@ -99,15 +110,22 @@ def resolve_refs(schema: dict[str, Any]) -> dict[str, Any]:
"""
defs = schema.get("$defs", {})
schema_copy = deepcopy(schema)
expanding: set[str] = set()
def _resolve(node: Any) -> Any:
if isinstance(node, dict):
ref = node.get("$ref")
if isinstance(ref, str) and ref.startswith("#/$defs/"):
def_name = ref.replace("#/$defs/", "")
if def_name in defs:
if def_name not in defs:
raise KeyError(f"Definition '{def_name}' not found in $defs.")
if def_name in expanding:
return {}
expanding.add(def_name)
try:
return _resolve(deepcopy(defs[def_name]))
raise KeyError(f"Definition '{def_name}' not found in $defs.")
finally:
expanding.discard(def_name)
return {k: _resolve(v) for k, v in node.items()}
if isinstance(node, list):
@@ -119,7 +137,11 @@ def resolve_refs(schema: dict[str, Any]) -> dict[str, Any]:
def add_key_in_dict_recursively(
d: dict[str, Any], key: str, value: Any, criteria: Callable[[dict[str, Any]], bool]
d: dict[str, Any],
key: str,
value: Any,
criteria: Callable[[dict[str, Any]], bool],
_seen: set[int] | None = None,
) -> dict[str, Any]:
"""Recursively adds a key/value pair to all nested dicts matching `criteria`.
@@ -128,22 +150,31 @@ def add_key_in_dict_recursively(
key: The key to add.
value: The value to add.
criteria: A function that returns True for dicts that should receive the key.
_seen: Internal set of visited ``id()``s, used to guard cyclic schemas.
Returns:
The modified dictionary.
"""
if _seen is None:
_seen = set()
if isinstance(d, dict):
if id(d) in _seen:
return d
_seen.add(id(d))
if criteria(d) and key not in d:
d[key] = value
for v in d.values():
add_key_in_dict_recursively(v, key, value, criteria)
add_key_in_dict_recursively(v, key, value, criteria, _seen)
elif isinstance(d, list):
if id(d) in _seen:
return d
_seen.add(id(d))
for i in d:
add_key_in_dict_recursively(i, key, value, criteria)
add_key_in_dict_recursively(i, key, value, criteria, _seen)
return d
def force_additional_properties_false(d: Any) -> Any:
def force_additional_properties_false(d: Any, _seen: set[int] | None = None) -> Any:
"""Force additionalProperties=false on all object-type dicts recursively.
OpenAI strict mode requires all objects to have additionalProperties=false.
@@ -154,11 +185,17 @@ def force_additional_properties_false(d: Any) -> Any:
Args:
d: The dictionary/list to modify.
_seen: Internal set of visited ``id()``s, used to guard cyclic schemas.
Returns:
The modified dictionary/list.
"""
if _seen is None:
_seen = set()
if isinstance(d, dict):
if id(d) in _seen:
return d
_seen.add(id(d))
if d.get("type") == "object":
d["additionalProperties"] = False
if "properties" not in d:
@@ -166,10 +203,13 @@ def force_additional_properties_false(d: Any) -> Any:
if "required" not in d:
d["required"] = []
for v in d.values():
force_additional_properties_false(v)
force_additional_properties_false(v, _seen)
elif isinstance(d, list):
if id(d) in _seen:
return d
_seen.add(id(d))
for i in d:
force_additional_properties_false(i)
force_additional_properties_false(i, _seen)
return d
@@ -183,7 +223,7 @@ OPENAI_SUPPORTED_FORMATS: Final[
}
def strip_unsupported_formats(d: Any) -> Any:
def strip_unsupported_formats(d: Any, _seen: set[int] | None = None) -> Any:
"""Remove format annotations that OpenAI strict mode doesn't support.
OpenAI only supports: date-time, date, time, duration.
@@ -191,11 +231,17 @@ def strip_unsupported_formats(d: Any) -> Any:
Args:
d: The dictionary/list to modify.
_seen: Internal set of visited ``id()``s, used to guard cyclic schemas.
Returns:
The modified dictionary/list.
"""
if _seen is None:
_seen = set()
if isinstance(d, dict):
if id(d) in _seen:
return d
_seen.add(id(d))
format_value = d.get("format")
if (
isinstance(format_value, str)
@@ -203,14 +249,17 @@ def strip_unsupported_formats(d: Any) -> Any:
):
del d["format"]
for v in d.values():
strip_unsupported_formats(v)
strip_unsupported_formats(v, _seen)
elif isinstance(d, list):
if id(d) in _seen:
return d
_seen.add(id(d))
for i in d:
strip_unsupported_formats(i)
strip_unsupported_formats(i, _seen)
return d
def ensure_type_in_schemas(d: Any) -> Any:
def ensure_type_in_schemas(d: Any, _seen: set[int] | None = None) -> Any:
"""Ensure all schema objects in anyOf/oneOf have a 'type' key.
OpenAI strict mode requires every schema to have a 'type' key.
@@ -218,11 +267,17 @@ def ensure_type_in_schemas(d: Any) -> Any:
Args:
d: The dictionary/list to modify.
_seen: Internal set of visited ``id()``s, used to guard cyclic schemas.
Returns:
The modified dictionary/list.
"""
if _seen is None:
_seen = set()
if isinstance(d, dict):
if id(d) in _seen:
return d
_seen.add(id(d))
for key in ("anyOf", "oneOf"):
if key in d:
schema_list = d[key]
@@ -230,12 +285,15 @@ def ensure_type_in_schemas(d: Any) -> Any:
if isinstance(schema, dict) and schema == {}:
schema_list[i] = {"type": "object"}
else:
ensure_type_in_schemas(schema)
ensure_type_in_schemas(schema, _seen)
for v in d.values():
ensure_type_in_schemas(v)
ensure_type_in_schemas(v, _seen)
elif isinstance(d, list):
if id(d) in _seen:
return d
_seen.add(id(d))
for item in d:
ensure_type_in_schemas(item)
ensure_type_in_schemas(item, _seen)
return d
@@ -318,7 +376,9 @@ def add_const_to_oneof_variants(schema: dict[str, Any]) -> dict[str, Any]:
return _process_oneof(deepcopy(schema))
def convert_oneof_to_anyof(schema: dict[str, Any]) -> dict[str, Any]:
def convert_oneof_to_anyof(
schema: dict[str, Any], _seen: set[int] | None = None
) -> dict[str, Any]:
"""Convert oneOf to anyOf for OpenAI compatibility.
OpenAI's Structured Outputs support anyOf better than oneOf.
@@ -326,26 +386,37 @@ def convert_oneof_to_anyof(schema: dict[str, Any]) -> dict[str, Any]:
Args:
schema: JSON schema dictionary.
_seen: Internal set of visited ``id()``s, used to guard cyclic schemas.
Returns:
Modified schema with anyOf instead of oneOf.
"""
if _seen is None:
_seen = set()
if isinstance(schema, dict):
if id(schema) in _seen:
return schema
_seen.add(id(schema))
if "oneOf" in schema:
schema["anyOf"] = schema.pop("oneOf")
for value in schema.values():
if isinstance(value, dict):
convert_oneof_to_anyof(value)
convert_oneof_to_anyof(value, _seen)
elif isinstance(value, list):
if id(value) in _seen:
continue
_seen.add(id(value))
for item in value:
if isinstance(item, dict):
convert_oneof_to_anyof(item)
convert_oneof_to_anyof(item, _seen)
return schema
def ensure_all_properties_required(schema: dict[str, Any]) -> dict[str, Any]:
def ensure_all_properties_required(
schema: dict[str, Any], _seen: set[int] | None = None
) -> dict[str, Any]:
"""Ensure all properties are in the required array for OpenAI strict mode.
OpenAI's strict structured outputs require all properties to be listed
@@ -354,11 +425,17 @@ def ensure_all_properties_required(schema: dict[str, Any]) -> dict[str, Any]:
Args:
schema: JSON schema dictionary.
_seen: Internal set of visited ``id()``s, used to guard cyclic schemas.
Returns:
Modified schema with all properties marked as required.
"""
if _seen is None:
_seen = set()
if isinstance(schema, dict):
if id(schema) in _seen:
return schema
_seen.add(id(schema))
if schema.get("type") == "object" and "properties" in schema:
properties = schema["properties"]
if properties:
@@ -366,16 +443,21 @@ def ensure_all_properties_required(schema: dict[str, Any]) -> dict[str, Any]:
for value in schema.values():
if isinstance(value, dict):
ensure_all_properties_required(value)
ensure_all_properties_required(value, _seen)
elif isinstance(value, list):
if id(value) in _seen:
continue
_seen.add(id(value))
for item in value:
if isinstance(item, dict):
ensure_all_properties_required(item)
ensure_all_properties_required(item, _seen)
return schema
def strip_null_from_types(schema: dict[str, Any]) -> dict[str, Any]:
def strip_null_from_types(
schema: dict[str, Any], _seen: set[int] | None = None
) -> dict[str, Any]:
"""Remove null type from anyOf/type arrays.
Pydantic generates `T | None` for optional fields, which creates schemas with
@@ -384,11 +466,17 @@ def strip_null_from_types(schema: dict[str, Any]) -> dict[str, Any]:
Args:
schema: JSON schema dictionary.
_seen: Internal set of visited ``id()``s, used to guard cyclic schemas.
Returns:
Modified schema with null types removed.
"""
if _seen is None:
_seen = set()
if isinstance(schema, dict):
if id(schema) in _seen:
return schema
_seen.add(id(schema))
if "anyOf" in schema:
any_of = schema["anyOf"]
non_null = [opt for opt in any_of if opt.get("type") != "null"]
@@ -408,11 +496,14 @@ def strip_null_from_types(schema: dict[str, Any]) -> dict[str, Any]:
for value in schema.values():
if isinstance(value, dict):
strip_null_from_types(value)
strip_null_from_types(value, _seen)
elif isinstance(value, list):
if id(value) in _seen:
continue
_seen.add(id(value))
for item in value:
if isinstance(item, dict):
strip_null_from_types(item)
strip_null_from_types(item, _seen)
return schema
@@ -451,16 +542,26 @@ _CLAUDE_STRICT_UNSUPPORTED: Final[tuple[str, ...]] = (
)
def _strip_keys_recursive(d: Any, keys: tuple[str, ...]) -> Any:
def _strip_keys_recursive(
d: Any, keys: tuple[str, ...], _seen: set[int] | None = None
) -> Any:
"""Recursively delete a fixed set of keys from a schema."""
if _seen is None:
_seen = set()
if isinstance(d, dict):
if id(d) in _seen:
return d
_seen.add(id(d))
for key in keys:
d.pop(key, None)
for v in d.values():
_strip_keys_recursive(v, keys)
_strip_keys_recursive(v, keys, _seen)
elif isinstance(d, list):
if id(d) in _seen:
return d
_seen.add(id(d))
for i in d:
_strip_keys_recursive(i, keys)
_strip_keys_recursive(i, keys, _seen)
return d
@@ -719,12 +820,70 @@ def create_model_from_schema( # type: ignore[no-any-unimported]
json_schema = force_additional_properties_false(json_schema)
effective_root = force_additional_properties_false(effective_root)
in_progress: dict[int, Any] = {}
model = _build_model_from_schema(
json_schema,
effective_root,
model_name=model_name,
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
__config__=__config__,
__base__=__base__,
__module__=__module__,
__validators__=__validators__,
__cls_kwargs__=__cls_kwargs__,
)
types_namespace: dict[str, Any] = {
entry.__name__: entry
for entry in in_progress.values()
if isinstance(entry, type) and issubclass(entry, BaseModel)
}
for entry in in_progress.values():
if (
isinstance(entry, type)
and issubclass(entry, BaseModel)
and not getattr(entry, "__pydantic_complete__", True)
):
try:
entry.model_rebuild(_types_namespace=types_namespace)
except Exception as e:
logger.debug("model_rebuild failed for %s: %s", entry.__name__, e)
return model
def _build_model_from_schema( # type: ignore[no-any-unimported]
json_schema: dict[str, Any],
effective_root: dict[str, Any],
*,
model_name: str | None,
enrich_descriptions: bool,
in_progress: dict[int, Any],
__config__: ConfigDict | None = None,
__base__: type[BaseModel] | None = None,
__module__: str = __name__,
__validators__: dict[str, AnyClassMethod] | None = None,
__cls_kwargs__: dict[str, Any] | None = None,
) -> type[BaseModel]:
"""Inner builder shared by the public entry point and recursive nested-object creation.
Preprocessing via ``jsonref.replace_refs`` and the sanitization walkers is
run once by the public entry; this helper walks the already-normalized
schema and emits Pydantic models. ``in_progress`` maps ``id(schema)`` to
the model being built for that schema, so a cyclic ``$ref`` graph
degrades to a ``ForwardRef`` back-edge instead of blowing the stack.
"""
original_id = id(json_schema)
if "allOf" in json_schema:
json_schema = _merge_all_of_schemas(json_schema["allOf"], effective_root)
if "title" not in json_schema and "title" in (root_schema or {}):
json_schema["title"] = (root_schema or {}).get("title")
effective_name = model_name or json_schema.get("title") or "DynamicModel"
schema_id = id(json_schema)
in_progress[original_id] = effective_name
if schema_id != original_id:
in_progress[schema_id] = effective_name
field_definitions = {
name: _json_schema_to_pydantic_field(
name,
@@ -732,13 +891,14 @@ def create_model_from_schema( # type: ignore[no-any-unimported]
json_schema.get("required", []),
effective_root,
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
for name, prop in (json_schema.get("properties", {}) or {}).items()
}
effective_config = __config__ or ConfigDict(extra="forbid")
return create_model_base(
model = create_model_base(
effective_name,
__config__=effective_config,
__base__=__base__,
@@ -747,6 +907,10 @@ def create_model_from_schema( # type: ignore[no-any-unimported]
__cls_kwargs__=__cls_kwargs__,
**field_definitions,
)
in_progress[original_id] = model
if schema_id != original_id:
in_progress[schema_id] = model
return model
def _json_schema_to_pydantic_field(
@@ -756,6 +920,7 @@ def _json_schema_to_pydantic_field(
root_schema: dict[str, Any],
*,
enrich_descriptions: bool = False,
in_progress: dict[int, Any] | None = None,
) -> Any:
"""Convert a JSON schema property to a Pydantic field definition.
@@ -774,6 +939,7 @@ def _json_schema_to_pydantic_field(
root_schema,
name_=name.title(),
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
is_required = name in required
@@ -833,7 +999,7 @@ def _json_schema_to_pydantic_field(
field_params["pattern"] = json_schema["pattern"]
if not is_required:
type_ = type_ | None
type_ = Optional[type_] # noqa: UP045 - ForwardRef does not support `|`
if schema_extra:
field_params["json_schema_extra"] = schema_extra
@@ -906,6 +1072,7 @@ def _json_schema_to_pydantic_type(
*,
name_: str | None = None,
enrich_descriptions: bool = False,
in_progress: dict[int, Any] | None = None,
) -> Any:
"""Convert a JSON schema to a Python/Pydantic type.
@@ -914,10 +1081,23 @@ def _json_schema_to_pydantic_type(
root_schema: The root schema for resolving $ref.
name_: Optional name for nested models.
enrich_descriptions: Propagated to nested model creation.
in_progress: Map of ``id(schema_dict)`` to the Pydantic model
currently being built for that schema, or to a placeholder name
as a plain ``str`` while the model is still being constructed.
Populated by :func:`_build_model_from_schema`. Enables cycle
detection so a self-referential ``$ref`` graph resolves to a
:class:`ForwardRef` back-edge rather than recursing forever.
Returns:
A Python type corresponding to the JSON schema.
"""
if in_progress is not None:
cached = in_progress.get(id(json_schema))
if isinstance(cached, str):
return ForwardRef(cached)
if cached is not None:
return cached
ref = json_schema.get("$ref")
if ref:
ref_schema = _resolve_ref(ref, root_schema)
@@ -926,6 +1106,7 @@ def _json_schema_to_pydantic_type(
root_schema,
name_=name_,
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
enum_values = json_schema.get("enum")
@@ -945,6 +1126,7 @@ def _json_schema_to_pydantic_type(
root_schema,
name_=f"{name_ or 'Union'}Option{i}",
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
for i, schema in enumerate(any_of_schemas)
]
@@ -958,6 +1140,15 @@ def _json_schema_to_pydantic_type(
root_schema,
name_=name_,
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
if in_progress is not None:
return _build_model_from_schema(
json_schema,
root_schema,
model_name=name_,
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
merged = _merge_all_of_schemas(all_of_schemas, root_schema)
return _json_schema_to_pydantic_type(
@@ -965,6 +1156,7 @@ def _json_schema_to_pydantic_type(
root_schema,
name_=name_,
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
type_ = json_schema.get("type")
@@ -985,12 +1177,21 @@ def _json_schema_to_pydantic_type(
root_schema,
name_=name_,
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
return list[item_type] # type: ignore[valid-type]
return list
if type_ == "object":
properties = json_schema.get("properties")
if properties:
if in_progress is not None:
return _build_model_from_schema(
json_schema,
root_schema,
model_name=name_,
enrich_descriptions=enrich_descriptions,
in_progress=in_progress,
)
json_schema_ = json_schema.copy()
if json_schema_.get("title") is None:
json_schema_["title"] = name_ or "DynamicModel"