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