22
33import copy
44from collections import defaultdict
5+ from copy import deepcopy
6+ from typing import Any
7+
8+
9+ def _detect_self_reference (schema : dict ) -> bool :
10+ """
11+ Detect if the schema contains self-referencing definitions.
12+
13+ Args:
14+ schema: The JSON schema to check
15+
16+ Returns:
17+ True if self-referencing is detected
18+ """
19+ defs = schema .get ("$defs" , {})
20+
21+ def find_refs_in_value (value : Any , parent_def : str ) -> bool :
22+ """Check if a value contains a reference to its parent definition."""
23+ if isinstance (value , dict ):
24+ if "$ref" in value :
25+ ref_path = value ["$ref" ]
26+ # Check if this references the parent definition
27+ if ref_path == f"#/$defs/{ parent_def } " :
28+ return True
29+ # Check all values in the dict
30+ for v in value .values ():
31+ if find_refs_in_value (v , parent_def ):
32+ return True
33+ elif isinstance (value , list ):
34+ # Check all items in the list
35+ for item in value :
36+ if find_refs_in_value (item , parent_def ):
37+ return True
38+ return False
39+
40+ # Check each definition for self-reference
41+ for def_name , def_content in defs .items ():
42+ if find_refs_in_value (def_content , def_name ):
43+ # Self-reference detected, return original schema
44+ return True
45+
46+ return False
47+
48+
49+ def dereference_json_schema (schema : dict , max_depth : int = 5 ) -> dict :
50+ """
51+ Dereference a JSON schema by resolving $ref references while preserving $defs.
52+
53+ This function flattens schema properties by:
54+ 1. Check for self-reference - if found, return original schema
55+ 2. When encountering $refs in properties, resolve them on-demand
56+ 3. Track visited definitions globally to prevent circular expansion
57+ 4. Preserve original $defs in the final result
58+
59+ Args:
60+ schema: The JSON schema to flatten
61+ max_depth: Maximum depth for resolving references (default: 5)
62+
63+ Returns:
64+ Schema with references resolved in properties, keeping original $defs
65+ """
66+ # Step 1: Check for self-reference
67+ if _detect_self_reference (schema ):
68+ # Self-referencing detected, return original schema
69+ return schema
70+
71+ # Make a deep copy to work with
72+ result = deepcopy (schema )
73+
74+ # Keep original $defs for the final result
75+ defs = deepcopy (schema .get ("$defs" , {}))
76+
77+ # Step 2: Define resolution function that tracks visits globally
78+ def resolve_refs_in_value (value : Any , depth : int , visiting : set [str ]) -> Any :
79+ """
80+ Recursively resolve $refs in a value.
81+
82+ Args:
83+ value: The value to process
84+ depth: Current depth in resolution
85+ visiting: Set of definitions currently being resolved (for cycle detection)
86+
87+ Returns:
88+ Value with $refs resolved (or kept if max depth reached)
89+ """
90+ if depth >= max_depth :
91+ return value
92+
93+ if isinstance (value , dict ):
94+ if "$ref" in value :
95+ ref_path = value ["$ref" ]
96+
97+ # Only handle internal references to $defs
98+ if ref_path .startswith ("#/$defs/" ):
99+ def_name = ref_path .split ("/" )[- 1 ]
100+
101+ # Check for circular reference
102+ if def_name in visiting :
103+ # Circular reference detected, keep the $ref
104+ return value
105+
106+ if def_name in defs :
107+ # Add to visiting set
108+ visiting .add (def_name )
109+
110+ # Get the definition and resolve any refs within it
111+ resolved = resolve_refs_in_value (
112+ deepcopy (defs [def_name ]), depth + 1 , visiting
113+ )
114+
115+ # Remove from visiting set
116+ visiting .remove (def_name )
117+
118+ # Merge resolved definition with additional properties
119+ # Additional properties from the original object take precedence
120+ for key , val in value .items ():
121+ if key != "$ref" :
122+ resolved [key ] = val
123+
124+ return resolved
125+ else :
126+ # Definition not found, keep the $ref
127+ return value
128+ else :
129+ # External ref or other type - keep as is
130+ return value
131+ else :
132+ # Regular dict - process all values
133+ return {
134+ key : resolve_refs_in_value (val , depth , visiting )
135+ for key , val in value .items ()
136+ }
137+ elif isinstance (value , list ):
138+ # Process each item in the list
139+ return [resolve_refs_in_value (item , depth , visiting ) for item in value ]
140+ else :
141+ # Primitive value - return as is
142+ return value
143+
144+ # Step 3: Process main schema properties with shared visiting set
145+ for key , value in result .items ():
146+ if key != "$defs" :
147+ # Each top-level property gets its own visiting set
148+ # This allows the same definition to be used in different contexts
149+ result [key ] = resolve_refs_in_value (value , 0 , set ())
150+
151+ # Step 4: Preserve original $defs
152+ if "$defs" in result :
153+ result ["$defs" ] = defs
154+
155+ return result
5156
6157
7158def _prune_param (schema : dict , param : str ) -> dict :
@@ -144,6 +295,7 @@ def compress_schema(
144295 prune_defs : bool = True ,
145296 prune_additional_properties : bool = True ,
146297 prune_titles : bool = False ,
298+ dereference_refs : bool = False ,
147299) -> dict :
148300 """
149301 Remove the given parameters from the schema.
@@ -154,6 +306,7 @@ def compress_schema(
154306 prune_defs: Whether to remove unused definitions
155307 prune_additional_properties: Whether to remove additionalProperties: false
156308 prune_titles: Whether to remove title fields from the schema
309+ dereference_refs: Whether to completely flatten by inlining all $refs (fixes Claude Desktop crashes).
157310 """
158311 # Make a copy so we don't modify the original
159312 schema = copy .deepcopy (schema )
@@ -172,4 +325,8 @@ def compress_schema(
172325 if prune_defs :
173326 schema = _prune_unused_defs (schema )
174327
328+ # Dereference all $refs if requested
329+ if dereference_refs :
330+ schema = dereference_json_schema (schema )
331+
175332 return schema
0 commit comments