rustlavel_mcp/
schema.rs

1//! The input schema a tool declares, and the validator that guards it.
2//!
3//! Two jobs, one definition. An agent needs a JSON Schema to know how to call
4//! a tool, and the server needs to check what actually arrives — a tool's
5//! arguments come from a language model, which is to say from the open
6//! internet. Deriving both from the same `Schema` means the description an
7//! agent reads can never drift from the rule the server enforces.
8//!
9//! Nobody writes schema JSON by hand:
10//!
11//! ```ignore
12//! Schema::object()
13//!     .string("city", "The city to look up")
14//!     .field(Field::integer("days").describe("How many days ahead").optional())
15//! ```
16
17use rustlavel_core::Json;
18
19/// The kinds of value a field may hold.
20#[derive(Debug, Clone, PartialEq)]
21pub enum FieldType {
22    String,
23    /// Any JSON number.
24    Number,
25    /// A number with no fractional part — checked, not just declared.
26    Integer,
27    Boolean,
28    Array(Box<FieldType>),
29    /// A nested object with its own fields.
30    Object(Box<Schema>),
31    /// Anything at all, for a tool that genuinely takes free-form JSON.
32    Any,
33}
34
35impl FieldType {
36    /// The JSON Schema `type` keyword for this field.
37    fn keyword(&self) -> Option<&'static str> {
38        match self {
39            FieldType::String => Some("string"),
40            FieldType::Number => Some("number"),
41            FieldType::Integer => Some("integer"),
42            FieldType::Boolean => Some("boolean"),
43            FieldType::Array(_) => Some("array"),
44            FieldType::Object(_) => Some("object"),
45            FieldType::Any => None,
46        }
47    }
48
49    /// What a value of the wrong type should be called in an error message.
50    fn describe(value: &Json) -> &'static str {
51        match value {
52            Json::Null => "null",
53            Json::Bool(_) => "a boolean",
54            Json::Number(_) => "a number",
55            Json::String(_) => "a string",
56            Json::Array(_) => "an array",
57            Json::Object(_) => "an object",
58        }
59    }
60
61    fn to_json(&self) -> Json {
62        match self {
63            FieldType::Array(item) => {
64                Json::object([("type", Json::from("array")), ("items", item.to_json())])
65            }
66            FieldType::Object(schema) => schema.to_json(),
67            FieldType::Any => Json::Object(Default::default()),
68            other => Json::object([("type", Json::from(other.keyword().unwrap_or("string")))]),
69        }
70    }
71
72    /// Check one value, appending a human sentence per problem found.
73    fn check(&self, path: &str, value: &Json, errors: &mut Vec<String>) {
74        let mismatch = |errors: &mut Vec<String>, expected: &str| {
75            errors.push(format!(
76                "`{path}` must be {expected}, found {}",
77                FieldType::describe(value)
78            ));
79        };
80
81        match self {
82            FieldType::Any => {}
83            FieldType::String if !matches!(value, Json::String(_)) => mismatch(errors, "a string"),
84            FieldType::Boolean if !matches!(value, Json::Bool(_)) => mismatch(errors, "a boolean"),
85            FieldType::Number if !matches!(value, Json::Number(_)) => mismatch(errors, "a number"),
86            FieldType::Integer => match value {
87                Json::Number(n) if n.fract() == 0.0 => {}
88                Json::Number(_) => errors
89                    .push(format!("`{path}` must be an integer, found a fractional number")),
90                _ => mismatch(errors, "an integer"),
91            },
92            FieldType::Array(item) => match value {
93                Json::Array(values) => {
94                    for (index, entry) in values.iter().enumerate() {
95                        item.check(&format!("{path}[{index}]"), entry, errors);
96                    }
97                }
98                _ => mismatch(errors, "an array"),
99            },
100            FieldType::Object(schema) => match value {
101                Json::Object(_) => schema.check(path, value, errors),
102                _ => mismatch(errors, "an object"),
103            },
104            _ => {}
105        }
106    }
107}
108
109/// One property of an object schema.
110#[derive(Debug, Clone, PartialEq)]
111pub struct Field {
112    pub name: String,
113    pub description: Option<String>,
114    pub kind: FieldType,
115    pub required: bool,
116    /// When set, the value must be one of these — a closed set an agent can see.
117    pub allowed: Vec<String>,
118}
119
120impl Field {
121    pub fn new(name: impl Into<String>, kind: FieldType) -> Self {
122        Field {
123            name: name.into(),
124            description: None,
125            kind,
126            // Required by default: the common case, and the safe one.
127            required: true,
128            allowed: Vec::new(),
129        }
130    }
131
132    pub fn string(name: impl Into<String>) -> Self {
133        Field::new(name, FieldType::String)
134    }
135
136    pub fn number(name: impl Into<String>) -> Self {
137        Field::new(name, FieldType::Number)
138    }
139
140    pub fn integer(name: impl Into<String>) -> Self {
141        Field::new(name, FieldType::Integer)
142    }
143
144    pub fn boolean(name: impl Into<String>) -> Self {
145        Field::new(name, FieldType::Boolean)
146    }
147
148    pub fn array_of(name: impl Into<String>, item: FieldType) -> Self {
149        Field::new(name, FieldType::Array(Box::new(item)))
150    }
151
152    pub fn object(name: impl Into<String>, schema: Schema) -> Self {
153        Field::new(name, FieldType::Object(Box::new(schema)))
154    }
155
156    pub fn any(name: impl Into<String>) -> Self {
157        Field::new(name, FieldType::Any)
158    }
159
160    /// The sentence the calling model reads to decide what to put here.
161    pub fn describe(mut self, description: impl Into<String>) -> Self {
162        self.description = Some(description.into());
163        self
164    }
165
166    pub fn optional(mut self) -> Self {
167        self.required = false;
168        self
169    }
170
171    /// Restrict to a fixed set of string values.
172    pub fn one_of<I, S>(mut self, values: I) -> Self
173    where
174        I: IntoIterator<Item = S>,
175        S: Into<String>,
176    {
177        self.allowed = values.into_iter().map(Into::into).collect();
178        self
179    }
180
181    fn to_json(&self) -> Json {
182        let mut value = self.kind.to_json();
183        if let Json::Object(map) = &mut value {
184            if let Some(description) = &self.description {
185                map.insert("description".into(), Json::from(description.clone()));
186            }
187            if !self.allowed.is_empty() {
188                let allowed: Vec<Json> = self.allowed.iter().cloned().map(Json::from).collect();
189                map.insert("enum".into(), Json::Array(allowed));
190            }
191        }
192        value
193    }
194}
195
196/// An object schema: what a tool accepts.
197#[derive(Debug, Clone, PartialEq, Default)]
198pub struct Schema {
199    fields: Vec<Field>,
200}
201
202impl Schema {
203    pub fn object() -> Self {
204        Schema::default()
205    }
206
207    /// Add a fully-specified field.
208    pub fn field(mut self, field: Field) -> Self {
209        self.fields.push(field);
210        self
211    }
212
213    /// Shorthand for the common case: a required, described property.
214    pub fn string(self, name: impl Into<String>, description: impl Into<String>) -> Self {
215        self.field(Field::string(name).describe(description))
216    }
217
218    pub fn number(self, name: impl Into<String>, description: impl Into<String>) -> Self {
219        self.field(Field::number(name).describe(description))
220    }
221
222    pub fn integer(self, name: impl Into<String>, description: impl Into<String>) -> Self {
223        self.field(Field::integer(name).describe(description))
224    }
225
226    pub fn boolean(self, name: impl Into<String>, description: impl Into<String>) -> Self {
227        self.field(Field::boolean(name).describe(description))
228    }
229
230    pub fn fields(&self) -> &[Field] {
231        &self.fields
232    }
233
234    pub fn is_empty(&self) -> bool {
235        self.fields.is_empty()
236    }
237
238    /// The JSON Schema an agent sees as a tool's `inputSchema`.
239    pub fn to_json(&self) -> Json {
240        let properties: Vec<(String, Json)> =
241            self.fields.iter().map(|f| (f.name.clone(), f.to_json())).collect();
242        let required: Vec<Json> = self
243            .fields
244            .iter()
245            .filter(|f| f.required)
246            .map(|f| Json::from(f.name.clone()))
247            .collect();
248
249        let mut pairs = vec![
250            ("type".to_string(), Json::from("object")),
251            ("properties".to_string(), Json::object(properties)),
252        ];
253        // An empty `required` array is noise; the specification allows omitting it.
254        if !required.is_empty() {
255            pairs.push(("required".to_string(), Json::Array(required)));
256        }
257        Json::object(pairs)
258    }
259
260    /// Check arguments against the schema.
261    ///
262    /// Returns every problem at once rather than the first, because the caller
263    /// is a language model that will otherwise fix one field per round trip.
264    pub fn validate(&self, arguments: &Json) -> Result<(), Vec<String>> {
265        let mut errors = Vec::new();
266        // A tool called with no arguments at all is the same as `{}`; whether
267        // that is acceptable is then decided by the required fields.
268        let arguments = if arguments.is_null() {
269            Json::Object(Default::default())
270        } else {
271            arguments.clone()
272        };
273
274        if arguments.as_object().is_none() {
275            return Err(vec![format!(
276                "arguments must be an object, found {}",
277                FieldType::describe(&arguments)
278            )]);
279        }
280
281        self.check("arguments", &arguments, &mut errors);
282        if errors.is_empty() { Ok(()) } else { Err(errors) }
283    }
284
285    fn check(&self, path: &str, value: &Json, errors: &mut Vec<String>) {
286        let Some(map) = value.as_object() else { return };
287
288        for field in &self.fields {
289            let child = if path == "arguments" {
290                field.name.clone()
291            } else {
292                format!("{path}.{}", field.name)
293            };
294
295            match map.get(&field.name) {
296                None | Some(Json::Null) if field.required => {
297                    errors.push(format!("`{child}` is required"));
298                }
299                None | Some(Json::Null) => {}
300                Some(found) => {
301                    field.kind.check(&child, found, errors);
302                    if !field.allowed.is_empty() {
303                        let text = found.as_str().unwrap_or_default();
304                        if !field.allowed.iter().any(|allowed| allowed == text) {
305                            errors.push(format!(
306                                "`{child}` must be one of {}",
307                                field.allowed.join(", ")
308                            ));
309                        }
310                    }
311                }
312            }
313        }
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320
321    fn weather_schema() -> Schema {
322        Schema::object()
323            .string("city", "The city to look up")
324            .field(Field::integer("days").describe("How many days ahead").optional())
325    }
326
327    #[test]
328    fn a_schema_renders_the_json_an_agent_reads() {
329        let schema = weather_schema();
330
331        assert_eq!(
332            schema.to_json().to_string(),
333            r#"{"properties":{"city":{"description":"The city to look up","type":"string"},"days":{"description":"How many days ahead","type":"integer"}},"required":["city"],"type":"object"}"#
334        );
335    }
336
337    #[test]
338    fn a_schema_with_no_required_fields_omits_the_required_array() {
339        let schema = Schema::object().field(Field::string("note").optional());
340        assert!(schema.to_json().get("required").is_none());
341    }
342
343    #[test]
344    fn valid_arguments_pass() {
345        let arguments =
346            Json::object([("city", Json::from("Jakarta")), ("days", Json::from(3))]);
347        assert!(weather_schema().validate(&arguments).is_ok());
348    }
349
350    #[test]
351    fn an_optional_field_may_be_left_out() {
352        let arguments = Json::object([("city", Json::from("Jakarta"))]);
353        assert!(weather_schema().validate(&arguments).is_ok());
354    }
355
356    #[test]
357    fn a_missing_required_field_is_reported() {
358        let errors = weather_schema().validate(&Json::object([("days", Json::from(1))])).unwrap_err();
359        assert_eq!(errors, ["`city` is required"]);
360    }
361
362    #[test]
363    fn a_wrong_type_is_rejected_with_the_type_it_saw() {
364        let arguments = Json::object([("city", Json::from(42))]);
365        let errors = weather_schema().validate(&arguments).unwrap_err();
366
367        assert_eq!(errors, ["`city` must be a string, found a number"]);
368    }
369
370    #[test]
371    fn a_fractional_number_is_not_an_integer() {
372        let arguments =
373            Json::object([("city", Json::from("Bandung")), ("days", Json::from(1.5))]);
374        let errors = weather_schema().validate(&arguments).unwrap_err();
375
376        assert_eq!(errors, ["`days` must be an integer, found a fractional number"]);
377    }
378
379    #[test]
380    fn every_problem_is_reported_at_once() {
381        let arguments = Json::object([("days", Json::from("soon"))]);
382        let errors = weather_schema().validate(&arguments).unwrap_err();
383
384        assert_eq!(errors.len(), 2);
385        assert!(errors.iter().any(|e| e.contains("`city` is required")));
386        assert!(errors.iter().any(|e| e.contains("`days` must be an integer")));
387    }
388
389    #[test]
390    fn arguments_that_are_not_an_object_are_rejected() {
391        let errors = weather_schema().validate(&Json::from("just a string")).unwrap_err();
392        assert_eq!(errors, ["arguments must be an object, found a string"]);
393    }
394
395    #[test]
396    fn absent_arguments_are_treated_as_an_empty_object() {
397        assert!(Schema::object().validate(&Json::Null).is_ok());
398        assert!(weather_schema().validate(&Json::Null).is_err());
399    }
400
401    #[test]
402    fn arrays_check_every_element() {
403        let schema = Schema::object().field(Field::array_of("tags", FieldType::String));
404        let good = Json::object([("tags", Json::Array(vec![Json::from("a")]))]);
405        let bad = Json::object([("tags", Json::Array(vec![Json::from("a"), Json::from(2)]))]);
406
407        assert!(schema.validate(&good).is_ok());
408        assert_eq!(schema.validate(&bad).unwrap_err(), ["`tags[1]` must be a string, found a number"]);
409    }
410
411    #[test]
412    fn nested_objects_are_validated_through() {
413        let schema = Schema::object()
414            .field(Field::object("filter", Schema::object().string("status", "Which status")));
415        let bad = Json::object([(
416            "filter",
417            Json::object([("status", Json::from(true))]),
418        )]);
419
420        assert_eq!(
421            schema.validate(&bad).unwrap_err(),
422            ["`filter.status` must be a string, found a boolean"]
423        );
424    }
425
426    #[test]
427    fn an_enum_restricts_the_accepted_values() {
428        let schema = Schema::object().field(Field::string("unit").one_of(["celsius", "fahrenheit"]));
429
430        assert!(schema.validate(&Json::object([("unit", Json::from("celsius"))])).is_ok());
431        assert_eq!(
432            schema.validate(&Json::object([("unit", Json::from("kelvin"))])).unwrap_err(),
433            ["`unit` must be one of celsius, fahrenheit"]
434        );
435        assert_eq!(
436            schema.to_json().get("properties.unit.enum").unwrap().to_string(),
437            r#"["celsius","fahrenheit"]"#
438        );
439    }
440
441    #[test]
442    fn an_any_field_accepts_whatever_arrives() {
443        let schema = Schema::object().field(Field::any("payload"));
444
445        assert!(schema.validate(&Json::object([("payload", Json::from(1))])).is_ok());
446        assert!(schema.validate(&Json::object([("payload", Json::from("x"))])).is_ok());
447    }
448}
rustlavel_mcp/schema.rs

rustlavel_mcp/
schema.rs
