| 1 | //! Schema sanitizer for tool `input_schema` before sending to DeepSeek. |
| 2 | //! |
| 3 | //! DeepSeek's `/beta/chat/completions` strict tool mode is harsh. MCP tool |
| 4 | //! schemas frequently arrive with Pydantic-style `anyOf:[{type:"string"}, |
| 5 | //! {type:"null"}]` unions, bare `{type:"object"}` with no `properties`, or |
| 6 | //! `required` entries that don't appear in `properties`. These dirty schemas |
| 7 | //! cause silent 400s that users can't diagnose. |
| 8 | //! |
| 9 | //! The sanitizer runs in-place on every schema returned by |
| 10 | //! `ToolRegistry::tools_for_api()` before the registry hands them off. |
| 11 | //! Output is cached so the per-tool overhead is paid once per registration. |
| 12 | |
| 13 | use serde_json::{Map, Value}; |
| 14 | |
| 15 | /// Sanitize a JSON Schema in-place for DeepSeek strict-tool compatibility. |
| 16 | /// |
| 17 | /// Applies a sequence of normalisations chosen to be semantics-preserving: |
| 18 | /// - Collapse `{"anyOf":[X, {"type":"null"}]}` → `X ∪ {"nullable": true}` |
| 19 | /// - Inject `"properties": {}` on bare-object schemas |
| 20 | /// - Prune dangling `required` entries |
| 21 | /// - Collapse single-element `oneOf` / `allOf` |
| 22 | /// - Walk recursively through all subschemas |
| 23 | pub fn sanitize(schema: &mut Value) { |
| 24 | collapse_nullable_unions(schema); |
| 25 | inject_properties_on_bare_objects(schema); |
| 26 | prune_dangling_required(schema); |
| 27 | collapse_single_element_unions(schema); |
| 28 | // Recurse into all sub-schemas |
| 29 | if let Some(obj) = schema.as_object_mut() { |
| 30 | for (_, v) in obj.iter_mut() { |
| 31 | sanitize(v); |
| 32 | } |
| 33 | } else if let Some(arr) = schema.as_array_mut() { |
| 34 | for v in arr.iter_mut() { |
| 35 | sanitize(v); |
| 36 | } |
| 37 | } |
| 38 | } |
| 39 | |
| 40 | /// Collapse `{"anyOf":[X, {"type":"null"}]}` → `X ∪ {"nullable": true}`. |
| 41 | /// |
| 42 | /// Same treatment for `oneOf`. Only collapses when exactly one non-null |
| 43 | /// member and exactly one null-type member are present. |
| 44 | fn collapse_nullable_unions(schema: &mut Value) { |
| 45 | let Some(obj) = schema.as_object_mut() else { |
| 46 | return; |
| 47 | }; |
| 48 | for key in ["anyOf", "oneOf"] { |
| 49 | let members: Vec<Value> = match obj.get(key).and_then(|v| v.as_array()) { |
| 50 | Some(arr) => arr.clone(), |
| 51 | None => continue, |
| 52 | }; |
| 53 | let (nulls, nons): (Vec<_>, Vec<_>) = members.into_iter().partition(is_null_type); |
| 54 | if nulls.len() == 1 && nons.len() == 1 { |
| 55 | obj.remove(key); |
| 56 | if let Value::Object(non_obj) = nons.into_iter().next().unwrap() { |
| 57 | for (k, v) in non_obj { |
| 58 | if k != "type" || v != "null" { |
| 59 | obj.insert(k, v); |
| 60 | } |
| 61 | } |
| 62 | } |
| 63 | obj.insert("nullable".into(), Value::Bool(true)); |
| 64 | } |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | fn is_null_type(v: &Value) -> bool { |
| 69 | v.as_object() |
| 70 | .and_then(|o| o.get("type")) |
| 71 | .and_then(|t| t.as_str()) |
| 72 | == Some("null") |
| 73 | } |
| 74 | |
| 75 | /// Bare `{"type": "object"}` (no `properties`, no `additionalProperties`) |
| 76 | /// → inject `"properties": {}` so DeepSeek's strict validator doesn't 400. |
| 77 | fn inject_properties_on_bare_objects(schema: &mut Value) { |
| 78 | let Some(obj) = schema.as_object_mut() else { |
| 79 | return; |
| 80 | }; |
| 81 | if obj.get("type").and_then(|t| t.as_str()) != Some("object") { |
| 82 | return; |
| 83 | } |
| 84 | if obj.contains_key("properties") || obj.contains_key("additionalProperties") { |
| 85 | return; |
| 86 | } |
| 87 | obj.insert("properties".into(), Value::Object(Map::new())); |
| 88 | } |
| 89 | |
| 90 | /// Remove entries from `required` that aren't keys in `properties`. |
| 91 | fn prune_dangling_required(schema: &mut Value) { |
| 92 | let Some(obj) = schema.as_object_mut() else { |
| 93 | return; |
| 94 | }; |
| 95 | // Collect known property names first (immutable borrow), then prune. |
| 96 | let known_keys: Vec<String> = obj |
| 97 | .get("properties") |
| 98 | .and_then(|v| v.as_object()) |
| 99 | .map(|props| props.keys().cloned().collect()) |
| 100 | .unwrap_or_default(); |
| 101 | let Some(required) = obj.get_mut("required").and_then(|v| v.as_array_mut()) else { |
| 102 | return; |
| 103 | }; |
| 104 | required.retain(|entry| { |
| 105 | entry |
| 106 | .as_str() |
| 107 | .is_some_and(|k| known_keys.iter().any(|known| known == k)) |
| 108 | }); |
| 109 | if required.is_empty() { |
| 110 | obj.remove("required"); |
| 111 | } |
| 112 | } |
| 113 | |
| 114 | /// Collapse `{"oneOf": [X]}` → X, same for `allOf`. |
| 115 | /// |
| 116 | /// Single-element unions are semantically equivalent to the element itself; |
| 117 | /// DeepSeek's strict validator doesn't always flatten them. |
| 118 | fn collapse_single_element_unions(schema: &mut Value) { |
| 119 | let Some(obj) = schema.as_object_mut() else { |
| 120 | return; |
| 121 | }; |
| 122 | for key in ["oneOf", "allOf", "anyOf"] { |
| 123 | let single = match obj.get(key).and_then(|v| v.as_array()) { |
| 124 | Some(arr) if arr.len() == 1 => arr[0].clone(), |
| 125 | _ => continue, |
| 126 | }; |
| 127 | obj.remove(key); |
| 128 | if let Value::Object(inner) = single { |
| 129 | for (k, v) in inner { |
| 130 | if !obj.contains_key(&k) { |
| 131 | obj.insert(k, v); |
| 132 | } |
| 133 | } |
| 134 | } |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | #[cfg(test)] |
| 139 | mod tests { |
| 140 | use super::*; |
| 141 | use serde_json::json; |
| 142 | |
| 143 | #[test] |
| 144 | fn collapses_nullable_anyof() { |
| 145 | let mut schema = json!({ |
| 146 | "anyOf": [ |
| 147 | {"type": "string"}, |
| 148 | {"type": "null"} |
| 149 | ] |
| 150 | }); |
| 151 | sanitize(&mut schema); |
| 152 | assert_eq!(schema["type"], "string"); |
| 153 | assert_eq!(schema["nullable"], true); |
| 154 | assert!(schema.get("anyOf").is_none()); |
| 155 | } |
| 156 | |
| 157 | #[test] |
| 158 | fn collapses_nullable_oneof() { |
| 159 | let mut schema = json!({ |
| 160 | "oneOf": [ |
| 161 | {"type": "null"}, |
| 162 | {"type": "integer", "minimum": 0} |
| 163 | ] |
| 164 | }); |
| 165 | sanitize(&mut schema); |
| 166 | assert_eq!(schema["type"], "integer"); |
| 167 | assert_eq!(schema["minimum"], 0); |
| 168 | assert_eq!(schema["nullable"], true); |
| 169 | } |
| 170 | |
| 171 | #[test] |
| 172 | fn preserves_non_null_anyof() { |
| 173 | let original = json!({ |
| 174 | "anyOf": [ |
| 175 | {"type": "string"}, |
| 176 | {"type": "integer"} |
| 177 | ] |
| 178 | }); |
| 179 | let mut schema = original.clone(); |
| 180 | sanitize(&mut schema); |
| 181 | // Multi-typed anyOf should collapse to single element after |
| 182 | // recursive walk — but here neither is null so the collapse |
| 183 | // doesn't trigger. The anyOf array itself remains. |
| 184 | assert!(schema.get("anyOf").is_some()); |
| 185 | } |
| 186 | |
| 187 | #[test] |
| 188 | fn injects_properties_on_bare_object() { |
| 189 | let mut schema = json!({"type": "object"}); |
| 190 | sanitize(&mut schema); |
| 191 | assert!(schema.get("properties").is_some()); |
| 192 | assert_eq!(schema["properties"], json!({})); |
| 193 | } |
| 194 | |
| 195 | #[test] |
| 196 | fn does_not_inject_properties_when_present() { |
| 197 | let mut schema = json!({ |
| 198 | "type": "object", |
| 199 | "properties": {"name": {"type": "string"}} |
| 200 | }); |
| 201 | let expected = schema.clone(); |
| 202 | sanitize(&mut schema); |
| 203 | assert_eq!(schema, expected); |
| 204 | } |
| 205 | |
| 206 | #[test] |
| 207 | fn prunes_dangling_required() { |
| 208 | let mut schema = json!({ |
| 209 | "type": "object", |
| 210 | "properties": {"name": {"type": "string"}}, |
| 211 | "required": ["name", "email"] |
| 212 | }); |
| 213 | sanitize(&mut schema); |
| 214 | let required = schema["required"].as_array().unwrap(); |
| 215 | assert_eq!(required.len(), 1); |
| 216 | assert_eq!(required[0], "name"); |
| 217 | } |
| 218 | |
| 219 | #[test] |
| 220 | fn removes_required_when_all_pruned() { |
| 221 | let mut schema = json!({ |
| 222 | "type": "object", |
| 223 | "properties": {}, |
| 224 | "required": ["ghost"] |
| 225 | }); |
| 226 | sanitize(&mut schema); |
| 227 | assert!(schema.get("required").is_none()); |
| 228 | } |
| 229 | |
| 230 | #[test] |
| 231 | fn collapses_single_element_oneof() { |
| 232 | let mut schema = json!({ |
| 233 | "oneOf": [{"type": "string", "minLength": 1}] |
| 234 | }); |
| 235 | sanitize(&mut schema); |
| 236 | assert!(schema.get("oneOf").is_none()); |
| 237 | assert_eq!(schema["type"], "string"); |
| 238 | assert_eq!(schema["minLength"], 1); |
| 239 | } |
| 240 | |
| 241 | #[test] |
| 242 | fn collapses_single_element_anyof() { |
| 243 | let mut schema = json!({ |
| 244 | "anyOf": [{"type": "boolean"}] |
| 245 | }); |
| 246 | sanitize(&mut schema); |
| 247 | assert!(schema.get("anyOf").is_none()); |
| 248 | assert_eq!(schema["type"], "boolean"); |
| 249 | } |
| 250 | |
| 251 | #[test] |
| 252 | fn recursive_walk_into_properties() { |
| 253 | let mut schema = json!({ |
| 254 | "type": "object", |
| 255 | "properties": { |
| 256 | "opt_name": { |
| 257 | "anyOf": [ |
| 258 | {"type": "string"}, |
| 259 | {"type": "null"} |
| 260 | ] |
| 261 | } |
| 262 | } |
| 263 | }); |
| 264 | sanitize(&mut schema); |
| 265 | let prop = &schema["properties"]["opt_name"]; |
| 266 | assert_eq!(prop["type"], "string"); |
| 267 | assert_eq!(prop["nullable"], true); |
| 268 | } |
| 269 | |
| 270 | #[test] |
| 271 | fn recursive_walk_into_items() { |
| 272 | let mut schema = json!({ |
| 273 | "type": "array", |
| 274 | "items": { |
| 275 | "anyOf": [ |
| 276 | {"type": "integer"}, |
| 277 | {"type": "null"} |
| 278 | ] |
| 279 | } |
| 280 | }); |
| 281 | sanitize(&mut schema); |
| 282 | let items = &schema["items"]; |
| 283 | assert_eq!(items["type"], "integer"); |
| 284 | assert_eq!(items["nullable"], true); |
| 285 | } |
| 286 | |
| 287 | #[test] |
| 288 | fn nested_anyof_in_anyof_collapses() { |
| 289 | // Pydantic can nest unions: Optional[Union[str, int]]. |
| 290 | let mut schema = json!({ |
| 291 | "anyOf": [ |
| 292 | { |
| 293 | "anyOf": [ |
| 294 | {"type": "string"}, |
| 295 | {"type": "integer"} |
| 296 | ] |
| 297 | }, |
| 298 | {"type": "null"} |
| 299 | ] |
| 300 | }); |
| 301 | sanitize(&mut schema); |
| 302 | // Outer anyOf is single non-null → collapsed. Inner anyOf is |
| 303 | // multi-typed → preserved, but the outer null is handled. |
| 304 | assert_eq!(schema["nullable"], true); |
| 305 | assert!(schema.get("anyOf").is_some()); |
| 306 | } |
| 307 | |
| 308 | #[test] |
| 309 | fn idempotent() { |
| 310 | let mut schema = json!({ |
| 311 | "type": "object", |
| 312 | "properties": { |
| 313 | "name": {"type": "string"}, |
| 314 | "maybe": { |
| 315 | "anyOf": [{"type": "integer"}, {"type": "null"}] |
| 316 | } |
| 317 | }, |
| 318 | "required": ["name", "missing_field"] |
| 319 | }); |
| 320 | sanitize(&mut schema); |
| 321 | let after_first = schema.clone(); |
| 322 | sanitize(&mut schema); |
| 323 | assert_eq!(schema, after_first, "sanitize must be idempotent"); |
| 324 | } |
| 325 | } |
| 326 |