from __future__ import annotations import copy from collections import defaultdict from typing import Any from jsonref import JsonRefError, replace_refs def _defs_have_cycles(defs: dict[str, Any]) -> bool: """Check whether any definitions in ``$defs`` form a reference cycle. A cycle means a definition directly or transitively references itself (e.g. Node → children → Node, or A → B → A). ``jsonref.replace_refs`` silently produces Python-level object cycles for these, which Pydantic's serializer rejects. """ if not defs: return False # Build adjacency: def_name -> set of def_names it references. edges: dict[str, set[str]] = defaultdict(set) def _collect_refs(obj: Any, source: str) -> None: if isinstance(obj, dict): ref = obj.get("$ref") if isinstance(ref, str) and ref.startswith("#/$defs/"): edges[source].add(ref.split("/")[-1]) for v in obj.values(): _collect_refs(v, source) elif isinstance(obj, list): for item in obj: _collect_refs(item, source) for name, definition in defs.items(): _collect_refs(definition, name) # DFS cycle detection. UNVISITED, IN_STACK, DONE = 0, 1, 2 state: dict[str, int] = defaultdict(int) def _has_cycle(node: str) -> bool: state[node] = IN_STACK for neighbor in edges.get(node, ()): if neighbor not in defs: continue if state[neighbor] == IN_STACK: return True if state[neighbor] == UNVISITED and _has_cycle(neighbor): return True state[node] = DONE return False return any(state[name] == UNVISITED and _has_cycle(name) for name in defs) def _strip_remote_refs(obj: Any) -> Any: """Return a deep copy of *obj* with non-local ``$ref`` values removed. Local refs (starting with ``#``) are kept intact. Remote refs (``http://``, ``https://``, ``file://``, or any other URI scheme) are stripped so that ``jsonref.replace_refs`` never attempts to fetch an external resource. This prevents SSRF / LFI when proxying schemas from untrusted servers. """ if isinstance(obj, dict): ref = obj.get("$ref") if isinstance(ref, str) and not ref.startswith("#"): # Drop the remote $ref key; keep all other keys. return {k: _strip_remote_refs(v) for k, v in obj.items() if k != "$ref"} return {k: _strip_remote_refs(v) for k, v in obj.items()} if isinstance(obj, list): return [_strip_remote_refs(item) for item in obj] return obj def _strip_discriminator(obj: Any) -> Any: """Recursively remove OpenAPI ``discriminator`` keys from a schema. Pydantic emits ``discriminator.mapping`` with values like ``#/$defs/ClassName``. After ``$defs`` are inlined and removed by ``dereference_refs``, those mapping entries dangle. The keyword is an OpenAPI extension — the ``anyOf`` variants already carry ``const`` on the discriminant field, so the mapping is redundant. Only strips ``discriminator`` when it appears alongside ``anyOf`` or ``oneOf``, which is where the OpenAPI keyword lives. A property *named* ``discriminator`` (inside ``properties``) is left alone. """ if isinstance(obj, dict): skip = "discriminator" in obj and ("anyOf" in obj or "oneOf" in obj) if skip: obj = require_discriminator_property(obj) # Keys that hold instance data, not sub-schemas — don't recurse. _DATA_KEYS = {"default", "const", "examples", "enum"} return { k: (v if k in _DATA_KEYS else _strip_discriminator(v)) for k, v in obj.items() if not (k == "discriminator" and skip) } if isinstance(obj, list): return [_strip_discriminator(item) for item in obj] return obj def _require_property(schema: dict[str, Any], property_name: str) -> dict[str, Any]: """Return a copy of *schema* with *property_name* in ``required``.""" required = schema.get("required") if required is None: return {**schema, "required": [property_name]} if isinstance(required, list) and property_name not in required: return {**schema, "required": [*required, property_name]} return schema def require_discriminator_property(schema: dict[str, Any]) -> dict[str, Any]: """Keep an OpenAPI discriminator's tag mandatory after the keyword is dropped. Returns a copy of *schema* with ``discriminator.propertyName`` added to each ``anyOf``/``oneOf`` variant's ``required`` list. A Pydantic discriminated union whose tag has a default omits that tag from ``required``; without this, an untagged payload passes the generated schema but fails later in the source model with ``union_tag_not_found``. No-op if there is no string ``propertyName``. """ discriminator = schema.get("discriminator") if not isinstance(discriminator, dict): return schema property_name = discriminator.get("propertyName") if not isinstance(property_name, str): return schema result = schema.copy() for key in ("anyOf", "oneOf"): variants = result.get(key) if not isinstance(variants, list): continue result[key] = [ _require_property(variant, property_name) if isinstance(variant, dict) else variant for variant in variants ] return result def dereference_refs(schema: dict[str, Any]) -> dict[str, Any]: """Resolve all $ref references in a JSON schema by inlining definitions. This function resolves $ref references that point to $defs, replacing them with the actual definition content while preserving sibling keywords (like description, default, examples) that Pydantic places alongside $ref. This is necessary because some MCP clients (e.g., VS Code Copilot) don't properly handle $ref in tool input schemas. For self-referencing/circular schemas where full dereferencing is not possible, this function falls back to resolving only the root-level $ref while preserving $defs for nested references. Only local ``$ref`` values (those starting with ``#``) are resolved. Remote URIs (``http://``, ``file://``, etc.) are stripped before resolution to prevent SSRF / local-file-inclusion attacks when proxying schemas from untrusted servers. Args: schema: JSON schema dict that may contain $ref references Returns: A new schema dict with $ref resolved where possible and $defs removed when no longer needed Example: >>> schema = { ... "$defs": {"Category": {"enum": ["a", "b"], "type": "string"}}, ... "properties": {"cat": {"$ref": "#/$defs/Category", "default": "a"}} ... } >>> resolved = dereference_refs(schema) >>> # Result: {"properties": {"cat": {"enum": ["a", "b"], "type": "string", "default": "a"}}} """ # Strip any remote $ref values before processing to prevent SSRF / LFI. schema = _strip_remote_refs(schema) # Circular $defs can't be fully inlined — jsonref.replace_refs produces # Python dicts with object-identity cycles that Pydantic's model_dump # rejects with "Circular reference detected (id repeated)". # Detect cycles up front and fall back to root-only resolution. if _defs_have_cycles(schema.get("$defs", {})): return resolve_root_ref(schema) try: # Use jsonref to resolve all $ref references # proxies=False returns plain dicts (not proxy objects) # lazy_load=False resolves immediately dereferenced = replace_refs(schema, proxies=False, lazy_load=False) # Merge sibling keywords that were lost during dereferencing # Pydantic puts description, default, examples as siblings to $ref defs = schema.get("$defs", {}) merged = _merge_ref_siblings(schema, dereferenced, defs) # Type assertion: top-level schema is always a dict assert isinstance(merged, dict) dereferenced = merged # Remove $defs since all references have been resolved if "$defs" in dereferenced: dereferenced = {k: v for k, v in dereferenced.items() if k != "$defs"} # Strip `discriminator` keys — they contain `mapping` values that # point at `#/$defs/...` entries we just removed. `discriminator` # is an OpenAPI extension; after inlining, the `anyOf` variants # already carry `const` on the discriminant field, making the # mapping redundant. dereferenced = _strip_discriminator(dereferenced) return dereferenced except (JsonRefError, RecursionError): # Self-referencing/circular schemas can't be fully dereferenced. # RecursionError covers circular $ref using JSON Pointer paths # (e.g. "#/properties/nodes/items") that bypass $defs-based cycle # detection — common in schemas from .NET/System.Text.Json. # Fall back to resolving only root-level $ref (for MCP spec compliance) return resolve_root_ref(schema) def _merge_ref_siblings( original: Any, dereferenced: Any, defs: dict[str, Any], visited: set[str] | None = None, ) -> Any: """Merge sibling keywords from original $ref nodes into dereferenced schema. When jsonref resolves $ref, it replaces the entire node with the referenced definition, losing any sibling keywords like description, default, or examples. This function walks both trees in parallel and merges those siblings back. Args: original: The original schema with $ref and potential siblings dereferenced: The schema after jsonref processing defs: The $defs from the original schema, for looking up referenced definitions visited: Set of definition names already being processed (prevents cycles) Returns: The dereferenced schema with sibling keywords restored """ if visited is None: visited = set() if isinstance(original, dict) and isinstance(dereferenced, dict): # Check if original had a $ref if "$ref" in original: ref = original["$ref"] siblings = {k: v for k, v in original.items() if k not in ("$ref", "$defs")} # Look up the referenced definition to process its nested siblings if isinstance(ref, str) and ref.startswith("#/$defs/"): def_name = ref.split("/")[-1] # Prevent infinite recursion on circular references if def_name in defs and def_name not in visited: # Recursively process the definition's content for nested siblings dereferenced = _merge_ref_siblings( defs[def_name], dereferenced, defs, visited | {def_name} ) if siblings: # Merge local siblings, which take precedence merged = dict(dereferenced) merged.update(siblings) return merged return dereferenced # Recurse into nested structures result = {} for key, value in dereferenced.items(): if key in original: result[key] = _merge_ref_siblings(original[key], value, defs, visited) else: result[key] = value return result elif isinstance(original, list) and isinstance(dereferenced, list): # Process list items in parallel min_len = min(len(original), len(dereferenced)) return [ _merge_ref_siblings(o, d, defs, visited) for o, d in zip(original[:min_len], dereferenced[:min_len], strict=False) ] + dereferenced[min_len:] return dereferenced def resolve_root_ref(schema: dict[str, Any]) -> dict[str, Any]: """Resolve $ref at root level to meet MCP spec requirements. MCP specification requires outputSchema to have "type": "object" at the root level. When Pydantic generates schemas for self-referential models, it uses $ref at the root level pointing to $defs. This function resolves such references by inlining the referenced definition while preserving $defs for nested references. Args: schema: JSON schema dict that may have $ref at root level Returns: A new schema dict with root-level $ref resolved, or the original schema if no resolution is needed Example: >>> schema = { ... "$defs": {"Node": {"type": "object", "properties": {...}}}, ... "$ref": "#/$defs/Node" ... } >>> resolved = resolve_root_ref(schema) >>> # Result: {"type": "object", "properties": {...}, "$defs": {...}} """ # Only resolve if we have $ref at root level with $defs but no explicit type if "$ref" in schema and "$defs" in schema and "type" not in schema: ref = schema["$ref"] # Only handle local $defs references if isinstance(ref, str) and ref.startswith("#/$defs/"): def_name = ref.split("/")[-1] defs = schema["$defs"] if def_name in defs: # Create a new schema by copying the referenced definition resolved = dict(defs[def_name]) # Preserve root-level sibling metadata from the original schema. # Pydantic may put user-facing fields such as title, description, # default, or examples next to the root $ref. Those fields still # describe the root schema even when we can only resolve that # root reference for circular schemas. resolved.update( { key: value for key, value in schema.items() if key not in {"$ref", "$defs"} } ) # Preserve $defs for nested references (other fields may still use them) resolved["$defs"] = defs return resolved return schema def _prune_param(schema: dict[str, Any], param: str) -> dict[str, Any]: """Return a new schema with *param* removed from `properties`, `required`, and (if no longer referenced) `$defs`. """ schema = copy.deepcopy(schema) # ── 1. drop from properties/required ────────────────────────────── props = schema.get("properties", {}) removed = props.pop(param, None) if removed is None: # nothing to do return schema # Keep empty properties object rather than removing it entirely schema["properties"] = props if param in schema.get("required", []): schema["required"].remove(param) if not schema["required"]: schema.pop("required") return schema # JSON Schema structural keywords — a node containing any of these is a # schema, so a string "title" sibling is metadata we can safely drop. _SCHEMA_KEYWORDS = frozenset( { "type", "properties", "$ref", "items", "allOf", "oneOf", "anyOf", "required", } ) # Pure schema-metadata keys. A node containing only these (e.g. Pydantic's # `{"title": "X"}` for Any-typed fields) is also a schema, just one with no # structural keywords alongside — still safe to strip title from. _METADATA_KEYS = frozenset( { "title", "description", "deprecated", "readOnly", "writeOnly", } ) # Keywords whose values are literal user data, not sub-schemas. Skipping # recursion here prevents `default: {"title": "X"}` from losing the "title" # data value because it happens to look metadata-shaped. Includes both # `examples` (JSON Schema draft 7+) and `example` (OpenAPI/Swagger 2.0). _LITERAL_KEYWORDS = frozenset({"default", "const", "examples", "example", "enum"}) # Keys whose values are dicts of arbitrary-name -> sub-schema. When we see # these, we traverse into each sub-schema regardless of its name — the keys # are user property/definition names, not schema keywords, so a property # literally named "enum" or "default" must not be confused with the # schema keywords of the same name. # # `dependencies` is a draft-07 keyword whose values can be sub-schemas OR # lists of required property names; list values short-circuit in the list # branch, so including it here is safe for both shapes. _SUBSCHEMA_MAP_KEYS = frozenset( { "properties", "patternProperties", "$defs", "definitions", "dependentSchemas", "dependencies", } ) # Keys whose values are a single sub-schema (not a dict of sub-schemas). # We traverse into them and treat the result as a schema node. # `additionalItems` is the draft-07 predecessor of `unevaluatedItems`. # `contentSchema` is a 2019-09+ keyword for typed string payloads. _SUBSCHEMA_VALUE_KEYS = frozenset( { "items", "additionalItems", "additionalProperties", "contains", "contentSchema", "propertyNames", "unevaluatedItems", "unevaluatedProperties", "if", "then", "else", "not", } ) # Keys whose values are LISTS of sub-schemas. _SUBSCHEMA_LIST_KEYS = frozenset({"allOf", "anyOf", "oneOf", "prefixItems"}) def _single_pass_optimize( schema: dict[str, Any], prune_titles: bool = False, prune_additional_properties: bool = False, prune_defs: bool = True, ) -> dict[str, Any]: """ Optimize JSON schemas in a single traversal for better performance. This function combines three schema cleanup operations that would normally require separate tree traversals: 1. **Remove unused definitions** (prune_defs): Finds and removes `$defs` entries that aren't referenced anywhere in the schema, reducing schema size. 2. **Remove titles** (prune_titles): Strips `title` fields throughout the schema to reduce verbosity while preserving functional information. 3. **Remove restrictive additionalProperties** (prune_additional_properties): Removes `"additionalProperties": false` constraints to make schemas more flexible. **Performance Benefits:** - Single tree traversal instead of multiple passes (2-3x faster) - Immutable design prevents shared reference bugs - Early termination prevents runaway recursion on deeply nested schemas **Algorithm Overview:** 1. Traverse main schema, collecting $ref references and applying cleanups 2. Traverse $defs section to map inter-definition dependencies 3. Remove unused definitions based on reference analysis Args: schema: JSON schema dict to optimize (not modified) prune_titles: Remove title fields for cleaner output prune_additional_properties: Remove "additionalProperties": false constraints prune_defs: Remove unused $defs entries to reduce size Returns: A new optimized schema dict Example: >>> schema = { ... "type": "object", ... "title": "MySchema", ... "additionalProperties": False, ... "$defs": {"UnusedDef": {"type": "string"}} ... } >>> result = _single_pass_optimize(schema, prune_titles=True, prune_defs=True) >>> # Result: {"type": "object", "additionalProperties": False} """ if not (prune_defs or prune_titles or prune_additional_properties): return schema # Nothing to do # Phase 1: Collect references and apply simple cleanups # Track which $defs are referenced from the main schema and from other $defs root_refs: set[str] = set() # $defs referenced directly from main schema def_dependencies: defaultdict[str, list[str]] = defaultdict( list ) # def A references def B defs = schema.get("$defs") def traverse_and_clean( node: object, current_def_name: str | None = None, skip_defs_section: bool = False, depth: int = 0, in_schema: bool = True, ) -> None: """Traverse schema tree, collecting $ref info and applying cleanups. The `in_schema` flag tracks whether the current node is reached via a known JSON-Schema-valued position (root, `properties` value, `items`, `allOf` element, etc.). When False — e.g. we descended through a user extension key like `x-ui` whose payload is opaque to us — we still collect `$ref` references (they may point at `$defs` the user cares about) but we skip all cleanups so we don't mutate user data that happens to look metadata-shaped. """ if depth > 50: # Prevent infinite recursion return if isinstance(node, dict): # Collect $ref references for unused definition removal. We do # this regardless of `in_schema` — a $ref in a user extension # still pins the referenced $def as "used". if prune_defs: ref = node.get("$ref") if isinstance(ref, str) and ref.startswith("#/$defs/"): referenced_def = ref.split("/")[-1] if current_def_name: # We're inside a $def, so this is a def->def reference def_dependencies[referenced_def].append(current_def_name) else: # We're in the main schema, so this is a root reference root_refs.add(referenced_def) # Cleanups only run when we know this node is a schema, never on # user extension payloads (`json_schema_extra={"x-ui": {...}}`). if in_schema: # Only remove "title" when it's schema metadata. A schema # node is either (a) one containing a structural keyword or # (b) one containing only metadata keys — Pydantic emits # bare `{"title": "X"}` for Any-typed fields with no sibling # type/properties, and Gemini 2.5 Flash rejects those with # MALFORMED_FUNCTION_CALL. The `isinstance(str)` guard # protects against deleting a user property literally named # "title" (its value would be a dict, not a string). if ( prune_titles and "title" in node and isinstance(node["title"], str) # type: ignore and ( any(k in node for k in _SCHEMA_KEYWORDS) or all(k in _METADATA_KEYS for k in node) ) ): node.pop("title") # type: ignore if ( prune_additional_properties and node.get("additionalProperties") is False ): node.pop("additionalProperties") # type: ignore # Recursive traversal for key, value in node.items(): if skip_defs_section and key == "$defs": continue # Skip $defs during main schema traversal # If we're not in a schema context, keep $ref-collecting but # don't promote sub-values to schema context — user extension # payloads can contain anything and must not be interpreted # as schemas. if not in_schema: traverse_and_clean( value, current_def_name, depth=depth + 1, in_schema=False ) continue # Arbitrary-key dicts of sub-schemas. The keys are user names # (property/definition names), not schema keywords, so we # must NOT apply the literal-keyword skip to them — a user # property named "enum" or "default" still needs its # sub-schema traversed (e.g. to collect $ref references). if key in _SUBSCHEMA_MAP_KEYS and isinstance(value, dict): for sub_schema in value.values(): traverse_and_clean( sub_schema, current_def_name, depth=depth + 1, in_schema=True, ) continue # Don't descend into keywords that carry literal data, not # sub-schemas — `default: {"title": "X"}` is a user value, # not schema metadata, and stripping "title" there would # corrupt it. if key in _LITERAL_KEYWORDS: continue # Keywords whose values are sub-schemas (or lists thereof). if key in _SUBSCHEMA_LIST_KEYS and isinstance(value, list): for item in value: traverse_and_clean( item, current_def_name, depth=depth + 1, in_schema=True, ) continue if key in _SUBSCHEMA_VALUE_KEYS: traverse_and_clean( value, current_def_name, depth=depth + 1, in_schema=True, ) continue # Unknown keys (user extensions like `x-ui`, vendor # metadata, etc.) — descend for $ref collection but mark # in_schema=False so cleanups don't touch user payloads. traverse_and_clean( value, current_def_name, depth=depth + 1, in_schema=False ) elif isinstance(node, list): for item in node: traverse_and_clean( item, current_def_name, depth=depth + 1, in_schema=in_schema ) # Phase 2: Traverse main schema (excluding $defs section) traverse_and_clean(schema, skip_defs_section=True, in_schema=True) # Phase 3: Traverse $defs to find inter-definition references if prune_defs and defs: for def_name, def_schema in defs.items(): traverse_and_clean(def_schema, current_def_name=def_name, in_schema=True) # Phase 4: Remove unused definitions def is_def_used(def_name: str, visiting: set[str] | None = None) -> bool: """Check if a definition is used, handling circular references.""" if def_name in root_refs: return True # Used directly from main schema # Check if any definition that references this one is itself used referencing_defs = def_dependencies.get(def_name, []) if referencing_defs: if visiting is None: visiting = set() # Avoid infinite recursion on circular references if def_name in visiting: return False visiting = visiting | {def_name} # If any referencing def is used, then this def is used for referencing_def in referencing_defs: if referencing_def not in visiting and is_def_used( referencing_def, visiting ): return True return False # Remove unused definitions for def_name in list(defs.keys()): if not is_def_used(def_name): defs.pop(def_name) # Clean up empty $defs section if not defs: schema.pop("$defs", None) return schema def compress_schema( schema: dict[str, Any], prune_params: list[str] | None = None, prune_additional_properties: bool = False, prune_titles: bool = False, dereference: bool = False, ) -> dict[str, Any]: """ Compress and optimize a JSON schema for MCP compatibility. Args: schema: The schema to compress prune_params: List of parameter names to remove from properties prune_additional_properties: Whether to remove additionalProperties: false. Defaults to False to maintain MCP client compatibility, as some clients (e.g., Claude) require additionalProperties: false for strict validation. prune_titles: Whether to remove title fields from the schema dereference: Whether to dereference $ref by inlining definitions. Defaults to False; dereferencing is typically handled by middleware at serve-time instead. """ if dereference: schema = dereference_refs(schema) # Resolve root-level $ref for MCP spec compliance (requires type: object at root) schema = resolve_root_ref(schema) # Remove specific parameters if requested for param in prune_params or []: schema = _prune_param(schema, param=param) # Apply combined optimizations in a single tree traversal. # Always prune unused $defs to keep schemas clean after parameter removal. schema = _single_pass_optimize( schema, prune_titles=prune_titles, prune_additional_properties=prune_additional_properties, prune_defs=True, ) return schema