lash_sansio/tool_contract/
schema_validation.rs

1use std::collections::HashMap;
2use std::io::{self, Write};
3use std::sync::{Arc, Mutex, OnceLock};
4
5use serde_json::Value;
6use sha2::{Digest, Sha256};
7
8use crate::tool_contract::ToolContract;
9
10#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
11pub struct LashSchema {
12    pub schema: Value,
13}
14
15impl LashSchema {
16    pub fn new(schema: Value) -> Self {
17        Self { schema }
18    }
19
20    pub fn any() -> Self {
21        Self::new(serde_json::json!({}))
22    }
23
24    pub fn object(properties: serde_json::Map<String, Value>, required: Vec<String>) -> Self {
25        let mut schema = serde_json::Map::new();
26        schema.insert("type".to_string(), Value::String("object".to_string()));
27        schema.insert("properties".to_string(), Value::Object(properties));
28        if !required.is_empty() {
29            schema.insert(
30                "required".to_string(),
31                Value::Array(required.into_iter().map(Value::String).collect()),
32            );
33        }
34        schema.insert("additionalProperties".to_string(), Value::Bool(true));
35        Self::new(Value::Object(schema))
36    }
37
38    pub fn validate(&self, value: &Value) -> Result<(), String> {
39        validate_schema(&self.schema, value)
40    }
41}
42
43const COMPILED_SCHEMA_CACHE_CAPACITY: usize = 1_024;
44const COMPILED_SCHEMA_CACHE_SCHEMA_BYTES: usize = 16 * 1024 * 1024;
45
46struct CachedSchema {
47    schema: Value,
48    compiled: Result<Arc<jsonschema::JSONSchema>, String>,
49}
50
51#[derive(Default)]
52struct CompiledSchemaCache {
53    entries: HashMap<[u8; 32], Vec<CachedSchema>>,
54    entry_count: usize,
55    schema_bytes: usize,
56}
57
58impl CompiledSchemaCache {
59    fn find_compiled(
60        &self,
61        hash: &[u8; 32],
62        schema: &Value,
63    ) -> Option<Result<Arc<jsonschema::JSONSchema>, String>> {
64        self.entries
65            .get(hash)
66            .and_then(|entries| entries.iter().find(|entry| entry.schema == *schema))
67            .map(|entry| entry.compiled.clone())
68    }
69
70    fn insert(
71        &mut self,
72        hash: [u8; 32],
73        schema: &Value,
74        serialized_bytes: usize,
75        compiled: Result<Arc<jsonschema::JSONSchema>, String>,
76    ) {
77        // The byte cap accounts for serialized schema input rather than the
78        // validator's opaque heap use; the entry cap is a second backstop.
79        if self.entry_count >= COMPILED_SCHEMA_CACHE_CAPACITY
80            || serialized_bytes
81                > COMPILED_SCHEMA_CACHE_SCHEMA_BYTES.saturating_sub(self.schema_bytes)
82        {
83            return;
84        }
85        self.entries.entry(hash).or_default().push(CachedSchema {
86            schema: schema.clone(),
87            compiled,
88        });
89        self.entry_count += 1;
90        self.schema_bytes += serialized_bytes;
91    }
92}
93
94fn compiled_schema_cache() -> &'static Mutex<CompiledSchemaCache> {
95    static CACHE: OnceLock<Mutex<CompiledSchemaCache>> = OnceLock::new();
96    CACHE.get_or_init(|| Mutex::new(CompiledSchemaCache::default()))
97}
98
99fn schema_content_fingerprint(schema: &Value) -> Result<([u8; 32], usize), String> {
100    struct DigestWriter<'a> {
101        digest: &'a mut Sha256,
102        bytes_written: usize,
103    }
104
105    impl Write for DigestWriter<'_> {
106        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
107            self.digest.update(bytes);
108            self.bytes_written = self.bytes_written.saturating_add(bytes.len());
109            Ok(bytes.len())
110        }
111
112        fn flush(&mut self) -> io::Result<()> {
113            Ok(())
114        }
115    }
116
117    let mut digest = Sha256::new();
118    let serialized_bytes = {
119        let mut writer = DigestWriter {
120            digest: &mut digest,
121            bytes_written: 0,
122        };
123        serde_json::to_writer(&mut writer, schema).map_err(|error| error.to_string())?;
124        writer.bytes_written
125    };
126    Ok((digest.finalize().into(), serialized_bytes))
127}
128
129fn compiled_schema(schema: &Value) -> Result<Arc<jsonschema::JSONSchema>, String> {
130    let (hash, serialized_bytes) = schema_content_fingerprint(schema)?;
131    if let Some(cached) = compiled_schema_cache()
132        .lock()
133        .expect("compiled schema cache lock")
134        .find_compiled(&hash, schema)
135    {
136        return cached;
137    }
138
139    let compiled = reject_non_local_references(schema).and_then(|()| {
140        jsonschema::JSONSchema::compile(schema)
141            .map(Arc::new)
142            .map_err(|error| error.to_string())
143    });
144
145    let mut cache = compiled_schema_cache()
146        .lock()
147        .expect("compiled schema cache lock");
148    if let Some(existing) = cache.find_compiled(&hash, schema) {
149        return existing;
150    }
151    cache.insert(hash, schema, serialized_bytes, compiled.clone());
152    compiled
153}
154
155fn validate_schema(schema: &Value, value: &Value) -> Result<(), String> {
156    let compiled = compiled_schema(schema)?;
157    if compiled.is_valid(value) {
158        return Ok(());
159    }
160    compiled.validate(value).map_err(|mut errors| {
161        format_validation_error(errors.next().expect("validation failure contains an error"))
162    })
163}
164
165pub fn validate_tool_input(contract: &ToolContract, args: &Value) -> Result<(), String> {
166    validate_schema(contract.input_schema.canonical(), args)
167}
168
169fn reject_non_local_references(schema: &Value) -> Result<(), String> {
170    match schema {
171        Value::Array(values) => {
172            for value in values {
173                reject_non_local_references(value)?;
174            }
175        }
176        Value::Object(object) => {
177            for (keyword, value) in object {
178                if matches!(keyword.as_str(), "$ref" | "$dynamicRef")
179                    && !value
180                        .as_str()
181                        .is_some_and(|reference| reference.starts_with('#'))
182                {
183                    return Err(format!(
184                        "non-local schema reference rejected: `{keyword}` must start with `#`, got {value}"
185                    ));
186                }
187                reject_non_local_references(value)?;
188            }
189        }
190        _ => {}
191    }
192
193    Ok(())
194}
195
196fn format_validation_error(error: jsonschema::ValidationError<'_>) -> String {
197    let instance_path = error.instance_path.to_string();
198    if instance_path.is_empty() {
199        error.to_string()
200    } else {
201        format!("{instance_path}: {error}")
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208    use crate::ToolDefinition;
209    use std::time::{Duration, Instant};
210
211    #[test]
212    fn repeated_schema_compilation_reuses_the_cached_validator() {
213        let schema = serde_json::json!({
214            "type": "object",
215            "properties": { "cache_probe_20260723": { "type": "string" } }
216        });
217
218        let first = compiled_schema(&schema).expect("compile schema once");
219        let second = compiled_schema(&schema).expect("reuse compiled schema");
220
221        assert!(Arc::ptr_eq(&first, &second));
222    }
223
224    #[test]
225    fn cache_hits_require_structural_equality_after_hash_match() {
226        let hash = [7; 32];
227        let first_schema = serde_json::json!({ "type": "string" });
228        let second_schema = serde_json::json!({ "type": "integer" });
229        let first = Arc::new(
230            jsonschema::JSONSchema::compile(&first_schema).expect("compile first validator"),
231        );
232        let second = Arc::new(
233            jsonschema::JSONSchema::compile(&second_schema).expect("compile second validator"),
234        );
235        let mut cache = CompiledSchemaCache::default();
236        cache.insert(hash, &first_schema, 17, Ok(first.clone()));
237        cache.insert(hash, &second_schema, 18, Ok(second.clone()));
238
239        let first_hit = cache
240            .find_compiled(&hash, &first_schema)
241            .expect("find first colliding schema")
242            .expect("first validator");
243        let second_hit = cache
244            .find_compiled(&hash, &second_schema)
245            .expect("find second colliding schema")
246            .expect("second validator");
247
248        assert!(Arc::ptr_eq(&first, &first_hit));
249        assert!(Arc::ptr_eq(&second, &second_hit));
250        assert!(!Arc::ptr_eq(&first_hit, &second_hit));
251    }
252
253    #[test]
254    fn validation_rejects_values_that_violate_local_refs() {
255        for definitions_key in ["$defs", "definitions"] {
256            let schema = serde_json::json!({
257                "type": "object",
258                "properties": {
259                    "item": { "$ref": format!("#/{definitions_key}/Item") }
260                },
261                "required": ["item"],
262                "additionalProperties": false,
263                (definitions_key): {
264                    "Item": {
265                        "type": "object",
266                        "properties": {
267                            "name": { "type": "string" }
268                        },
269                        "required": ["name"],
270                        "additionalProperties": false
271                    }
272                }
273            });
274
275            let error = LashSchema::new(schema)
276                .validate(&serde_json::json!({ "item": { "name": 42 } }))
277                .unwrap_err();
278
279            assert_eq!(error, "/item/name: 42 is not of type \"string\"");
280        }
281    }
282
283    #[test]
284    fn validation_rejects_bad_value_through_all_of_wrapped_ref() {
285        let schema = LashSchema::new(serde_json::json!({
286            "definitions": {
287                "Inner": { "type": "string" }
288            },
289            "allOf": [
290                { "$ref": "#/definitions/Inner" }
291            ],
292            "description": "A documented field"
293        }));
294
295        let error = schema.validate(&serde_json::json!(42)).unwrap_err();
296
297        assert_eq!(error, "42 is not of type \"string\"");
298    }
299
300    #[test]
301    fn validation_of_chained_doubling_refs_is_bounded() {
302        // Keep the adversarial fanout large enough to catch eager expansion,
303        // but small enough that a regression cannot exhaust the test machine.
304        const DEPTH: usize = 16;
305        const TIME_LIMIT: Duration = Duration::from_secs(1);
306
307        let mut definitions = serde_json::Map::new();
308        for depth in 0..DEPTH {
309            definitions.insert(
310                format!("D{depth}"),
311                serde_json::json!({
312                    "type": "object",
313                    "properties": {
314                        "left": { "$ref": format!("#/definitions/D{}", depth + 1) },
315                        "right": { "$ref": format!("#/definitions/D{}", depth + 1) }
316                    }
317                }),
318            );
319        }
320        definitions.insert(format!("D{DEPTH}"), serde_json::json!({ "type": "string" }));
321        let schema = LashSchema::new(serde_json::json!({
322            "$ref": "#/definitions/D0",
323            "definitions": definitions
324        }));
325        let mut value = serde_json::json!(42);
326        for _ in 0..DEPTH {
327            value = serde_json::json!({ "left": value });
328        }
329
330        let started = Instant::now();
331        let error = schema.validate(&value).unwrap_err();
332        let elapsed = started.elapsed();
333
334        assert!(error.ends_with(": 42 is not of type \"string\""), "{error}");
335        assert!(error.contains("/left"), "{error}");
336        assert!(
337            elapsed < TIME_LIMIT,
338            "chained reference validation took {elapsed:?}, limit is {TIME_LIMIT:?}"
339        );
340    }
341
342    #[test]
343    fn validation_rejects_unresolvable_local_ref() {
344        let schema = LashSchema::new(serde_json::json!({
345            "$ref": "#/definitions/Missing"
346        }));
347
348        let error = schema.validate(&serde_json::json!(42)).unwrap_err();
349
350        assert!(error.starts_with("Invalid reference:"), "{error}");
351    }
352
353    #[test]
354    fn validation_rejects_external_ref_before_compilation() {
355        let schema = LashSchema::new(serde_json::json!({
356            "$ref": "https://example.com/schema.json"
357        }));
358
359        let error = schema.validate(&serde_json::json!(42)).unwrap_err();
360
361        assert_eq!(
362            error,
363            "non-local schema reference rejected: `$ref` must start with `#`, got \"https://example.com/schema.json\""
364        );
365    }
366
367    #[test]
368    fn validation_rejects_external_dynamic_ref_before_compilation() {
369        let schema = LashSchema::new(serde_json::json!({
370            "$dynamicRef": "https://example.com/schema.json#node"
371        }));
372
373        let error = schema.validate(&serde_json::json!(42)).unwrap_err();
374
375        assert!(
376            error.contains("`$dynamicRef` must start with `#`"),
377            "{error}"
378        );
379    }
380
381    #[test]
382    fn jsonschema_dependency_cannot_resolve_http_references() {
383        let schema = serde_json::json!({
384            "$ref": "https://example.com/schema.json"
385        });
386
387        let compiled = jsonschema::JSONSchema::compile(&schema).unwrap();
388        let error = compiled
389            .validate(&serde_json::json!(42))
390            .unwrap_err()
391            .next()
392            .expect("external reference resolution produces an error")
393            .to_string();
394
395        assert!(
396            error.contains("`resolve-http` feature or a custom resolver is required"),
397            "{error}"
398        );
399    }
400
401    #[test]
402    fn validation_rejects_deep_recursive_violation() {
403        let schema = LashSchema::new(serde_json::json!({
404            "$ref": "#/definitions/Node",
405            "definitions": {
406                "Node": {
407                    "type": "object",
408                    "properties": {
409                        "name": { "type": "string" },
410                        "child": { "$ref": "#/definitions/Node" }
411                    },
412                    "required": ["name"]
413                }
414            }
415        }));
416
417        let error = schema
418            .validate(&serde_json::json!({
419                "name": "root",
420                "child": {
421                    "name": "level 1",
422                    "child": {
423                        "name": "level 2",
424                        "child": { "name": 123 }
425                    }
426                }
427            }))
428            .unwrap_err();
429
430        assert_eq!(
431            error,
432            "/child/child/child/name: 123 is not of type \"string\""
433        );
434    }
435
436    #[test]
437    fn validation_combines_ref_target_and_sibling_constraints() {
438        let schema = LashSchema::new(serde_json::json!({
439            "$ref": "#/$defs/AtLeastFive",
440            "minimum": 1,
441            "$defs": {
442                "AtLeastFive": {
443                    "type": "number",
444                    "minimum": 5
445                }
446            }
447        }));
448
449        let error = schema.validate(&serde_json::json!(2)).unwrap_err();
450
451        assert_eq!(error, "2 is less than the minimum of 5");
452    }
453
454    #[test]
455    fn validation_reports_missing_required_property_by_path() {
456        let tool = ToolDefinition::raw(
457            "tool:spotify",
458            "spotify",
459            "",
460            serde_json::json!({
461                "type": "object",
462                "properties": {
463                    "access_token": { "type": "string" }
464                },
465                "required": ["access_token"],
466                "additionalProperties": false
467            }),
468            serde_json::json!({}),
469        );
470
471        let error = validate_tool_input(&tool.contract(), &serde_json::json!({})).unwrap_err();
472        assert_eq!(error, "\"access_token\" is a required property");
473    }
474
475    #[test]
476    fn validation_reports_numeric_limits_by_path() {
477        let tool = ToolDefinition::raw(
478            "tool:spotify",
479            "spotify",
480            "",
481            serde_json::json!({
482                "type": "object",
483                "properties": {
484                    "page_limit": { "type": "integer", "maximum": 20 }
485                },
486                "required": ["page_limit"],
487                "additionalProperties": false
488            }),
489            serde_json::json!({}),
490        );
491
492        let error =
493            validate_tool_input(&tool.contract(), &serde_json::json!({ "page_limit": 100 }))
494                .unwrap_err();
495        assert_eq!(error, "/page_limit: 100 is greater than the maximum of 20");
496    }
497
498    #[test]
499    fn validation_allows_unknown_property_when_additional_properties_is_omitted() {
500        let tool = ToolDefinition::raw(
501            "tool:mcp__appworld__venmo_show_transactions",
502            "mcp__appworld__venmo_show_transactions",
503            "",
504            serde_json::json!({
505                "type": "object",
506                "properties": {
507                    "min_created_at": { "type": "string" },
508                    "max_created_at": { "type": "string" },
509                    "limit": { "type": "integer", "maximum": 100 }
510                },
511                "required": ["limit"]
512            }),
513            serde_json::json!({}),
514        );
515
516        validate_tool_input(
517            &tool.contract(),
518            &serde_json::json!({
519                "min_datetime": "2024-01-01T00:00:00Z",
520                "limit": 20
521            }),
522        )
523        .unwrap();
524    }
525
526    #[test]
527    fn validation_allows_unknown_property_when_additional_properties_is_true() {
528        let tool = ToolDefinition::raw(
529            "tool:open",
530            "open",
531            "",
532            serde_json::json!({
533                "type": "object",
534                "properties": {
535                    "path": { "type": "string" }
536                },
537                "additionalProperties": true
538            }),
539            serde_json::json!({}),
540        );
541
542        validate_tool_input(
543            &tool.contract(),
544            &serde_json::json!({
545                "path": "README.md",
546                "unknown": "preserved"
547            }),
548        )
549        .unwrap();
550    }
551}
lash_sansio/tool_contract/schema_validation.rs

lash_sansio/tool_contract/
schema_validation.rs
