turul_mcp_protocol_2025_06_18/
schema.rs

1//! JSON Schema Support for MCP
2//!
3//! This module provides JSON Schema types used throughout the MCP protocol.
4
5use std::collections::HashMap;
6use serde::{Deserialize, Serialize};
7
8/// Trait for generating JSON schemas from Rust types
9pub trait JsonSchemaGenerator {
10    /// Generate a ToolSchema for this type
11    fn json_schema() -> crate::tools::ToolSchema;
12}
13
14/// A JSON Schema definition
15#[derive(Debug, Clone, Serialize, Deserialize)]
16#[serde(tag = "type", rename_all = "lowercase")]
17pub enum JsonSchema {
18    /// String type
19    String {
20        #[serde(skip_serializing_if = "Option::is_none")]
21        description: Option<String>,
22        #[serde(skip_serializing_if = "Option::is_none")]
23        pattern: Option<String>,
24        #[serde(skip_serializing_if = "Option::is_none")]
25        min_length: Option<u64>,
26        #[serde(skip_serializing_if = "Option::is_none")]
27        max_length: Option<u64>,
28        #[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
29        enum_values: Option<Vec<String>>,
30    },
31    /// Number type
32    Number {
33        #[serde(skip_serializing_if = "Option::is_none")]
34        description: Option<String>,
35        #[serde(skip_serializing_if = "Option::is_none")]
36        minimum: Option<f64>,
37        #[serde(skip_serializing_if = "Option::is_none")]
38        maximum: Option<f64>,
39    },
40    /// Integer type
41    Integer {
42        #[serde(skip_serializing_if = "Option::is_none")]
43        description: Option<String>,
44        #[serde(skip_serializing_if = "Option::is_none")]
45        minimum: Option<i64>,
46        #[serde(skip_serializing_if = "Option::is_none")]
47        maximum: Option<i64>,
48    },
49    /// Boolean type
50    Boolean {
51        #[serde(skip_serializing_if = "Option::is_none")]
52        description: Option<String>,
53    },
54    /// Array type
55    Array {
56        #[serde(skip_serializing_if = "Option::is_none")]
57        description: Option<String>,
58        #[serde(skip_serializing_if = "Option::is_none")]
59        items: Option<Box<JsonSchema>>,
60        #[serde(skip_serializing_if = "Option::is_none")]
61        min_items: Option<u64>,
62        #[serde(skip_serializing_if = "Option::is_none")]
63        max_items: Option<u64>,
64    },
65    /// Object type
66    Object {
67        #[serde(skip_serializing_if = "Option::is_none")]
68        description: Option<String>,
69        #[serde(skip_serializing_if = "Option::is_none")]
70        properties: Option<HashMap<String, JsonSchema>>,
71        #[serde(skip_serializing_if = "Option::is_none")]
72        required: Option<Vec<String>>,
73        #[serde(skip_serializing_if = "Option::is_none")]
74        additional_properties: Option<bool>,
75    },
76}
77
78impl JsonSchema {
79    /// Create a string schema
80    pub fn string() -> Self {
81        Self::String {
82            description: None,
83            pattern: None,
84            min_length: None,
85            max_length: None,
86            enum_values: None,
87        }
88    }
89
90    /// Create a string schema with description
91    pub fn string_with_description(description: impl Into<String>) -> Self {
92        Self::String {
93            description: Some(description.into()),
94            pattern: None,
95            min_length: None,
96            max_length: None,
97            enum_values: None,
98        }
99    }
100
101    /// Create a string enum schema
102    pub fn string_enum(values: Vec<String>) -> Self {
103        Self::String {
104            description: None,
105            pattern: None,
106            min_length: None,
107            max_length: None,
108            enum_values: Some(values),
109        }
110    }
111
112    /// Create a number schema
113    pub fn number() -> Self {
114        Self::Number {
115            description: None,
116            minimum: None,
117            maximum: None,
118        }
119    }
120
121    /// Create a number schema with description
122    pub fn number_with_description(description: impl Into<String>) -> Self {
123        Self::Number {
124            description: Some(description.into()),
125            minimum: None,
126            maximum: None,
127        }
128    }
129
130    /// Create an integer schema
131    pub fn integer() -> Self {
132        Self::Integer {
133            description: None,
134            minimum: None,
135            maximum: None,
136        }
137    }
138
139    /// Create an integer schema with description
140    pub fn integer_with_description(description: impl Into<String>) -> Self {
141        Self::Integer {
142            description: Some(description.into()),
143            minimum: None,
144            maximum: None,
145        }
146    }
147
148    /// Create a boolean schema
149    pub fn boolean() -> Self {
150        Self::Boolean {
151            description: None,
152        }
153    }
154
155    /// Create a boolean schema with description
156    pub fn boolean_with_description(description: impl Into<String>) -> Self {
157        Self::Boolean {
158            description: Some(description.into()),
159        }
160    }
161
162    /// Create an array schema
163    pub fn array(items: JsonSchema) -> Self {
164        Self::Array {
165            description: None,
166            items: Some(Box::new(items)),
167            min_items: None,
168            max_items: None,
169        }
170    }
171
172    /// Create an array schema with description
173    pub fn array_with_description(items: JsonSchema, description: impl Into<String>) -> Self {
174        Self::Array {
175            description: Some(description.into()),
176            items: Some(Box::new(items)),
177            min_items: None,
178            max_items: None,
179        }
180    }
181
182    /// Create an object schema
183    pub fn object() -> Self {
184        Self::Object {
185            description: None,
186            properties: None,
187            required: None,
188            additional_properties: None,
189        }
190    }
191
192    /// Create an object schema with properties
193    pub fn object_with_properties(properties: HashMap<String, JsonSchema>) -> Self {
194        Self::Object {
195            description: None,
196            properties: Some(properties),
197            required: None,
198            additional_properties: None,
199        }
200    }
201
202    /// Create an object schema with properties and required fields
203    pub fn object_with_required(
204        properties: HashMap<String, JsonSchema>,
205        required: Vec<String>,
206    ) -> Self {
207        Self::Object {
208            description: None,
209            properties: Some(properties),
210            required: Some(required),
211            additional_properties: None,
212        }
213    }
214
215    /// Add description to any schema
216    pub fn with_description(mut self, description: impl Into<String>) -> Self {
217        match &mut self {
218            JsonSchema::String { description: d, .. } => *d = Some(description.into()),
219            JsonSchema::Number { description: d, .. } => *d = Some(description.into()),
220            JsonSchema::Integer { description: d, .. } => *d = Some(description.into()),
221            JsonSchema::Boolean { description: d, .. } => *d = Some(description.into()),
222            JsonSchema::Array { description: d, .. } => *d = Some(description.into()),
223            JsonSchema::Object { description: d, .. } => *d = Some(description.into()),
224        }
225        self
226    }
227
228    /// Add minimum constraint to number schema
229    pub fn with_minimum(mut self, minimum: f64) -> Self {
230        match &mut self {
231            JsonSchema::Number { minimum: m, .. } => *m = Some(minimum),
232            JsonSchema::Integer { minimum: m, .. } => *m = Some(minimum as i64),
233            _ => {}, // Ignore for non-numeric types
234        }
235        self
236    }
237
238    /// Add maximum constraint to number schema
239    pub fn with_maximum(mut self, maximum: f64) -> Self {
240        match &mut self {
241            JsonSchema::Number { maximum: m, .. } => *m = Some(maximum),
242            JsonSchema::Integer { maximum: m, .. } => *m = Some(maximum as i64),
243            _ => {}, // Ignore for non-numeric types
244        }
245        self
246    }
247
248    /// Add properties to object schema
249    pub fn with_properties(mut self, properties: HashMap<String, JsonSchema>) -> Self {
250        match &mut self {
251            JsonSchema::Object { properties: p, .. } => *p = Some(properties),
252            _ => {}, // Ignore for non-object types
253        }
254        self
255    }
256
257    /// Add required fields to object schema
258    pub fn with_required(mut self, required: Vec<String>) -> Self {
259        match &mut self {
260            JsonSchema::Object { required: r, .. } => *r = Some(required),
261            _ => {}, // Ignore for non-object types
262        }
263        self
264    }
265}
266
267/// Converts common Rust types to JsonSchema
268pub trait ToJsonSchema {
269    fn to_json_schema() -> JsonSchema;
270}
271
272impl ToJsonSchema for String {
273    fn to_json_schema() -> JsonSchema {
274        JsonSchema::string()
275    }
276}
277
278impl ToJsonSchema for &str {
279    fn to_json_schema() -> JsonSchema {
280        JsonSchema::string()
281    }
282}
283
284impl ToJsonSchema for i32 {
285    fn to_json_schema() -> JsonSchema {
286        JsonSchema::integer()
287    }
288}
289
290impl ToJsonSchema for i64 {
291    fn to_json_schema() -> JsonSchema {
292        JsonSchema::integer()
293    }
294}
295
296impl ToJsonSchema for f32 {
297    fn to_json_schema() -> JsonSchema {
298        JsonSchema::number()
299    }
300}
301
302impl ToJsonSchema for f64 {
303    fn to_json_schema() -> JsonSchema {
304        JsonSchema::number()
305    }
306}
307
308impl ToJsonSchema for bool {
309    fn to_json_schema() -> JsonSchema {
310        JsonSchema::boolean()
311    }
312}
313
314impl<T: ToJsonSchema> ToJsonSchema for Vec<T> {
315    fn to_json_schema() -> JsonSchema {
316        JsonSchema::array(T::to_json_schema())
317    }
318}
319
320impl<T: ToJsonSchema> ToJsonSchema for Option<T> {
321    fn to_json_schema() -> JsonSchema {
322        T::to_json_schema()
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329
330    #[test]
331    fn test_string_schema() {
332        let schema = JsonSchema::string_with_description("A test string");
333        let json = serde_json::to_string(&schema).unwrap();
334        assert!(json.contains("string"));
335        assert!(json.contains("A test string"));
336    }
337
338    #[test]
339    fn test_object_schema() {
340        let mut properties = HashMap::new();
341        properties.insert("name".to_string(), JsonSchema::string());
342        properties.insert("age".to_string(), JsonSchema::integer());
343
344        let schema = JsonSchema::object_with_required(
345            properties,
346            vec!["name".to_string()],
347        );
348
349        let json = serde_json::to_string(&schema).unwrap();
350        assert!(json.contains("object"));
351        assert!(json.contains("name"));
352        assert!(json.contains("age"));
353    }
354
355    #[test]
356    fn test_array_schema() {
357        let schema = JsonSchema::array(JsonSchema::string());
358        let json = serde_json::to_string(&schema).unwrap();
359        assert!(json.contains("array"));
360    }
361
362    #[test]
363    fn test_enum_schema() {
364        let schema = JsonSchema::string_enum(vec![
365            "option1".to_string(),
366            "option2".to_string(),
367        ]);
368        let json = serde_json::to_string(&schema).unwrap();
369        assert!(json.contains("option1"));
370        assert!(json.contains("option2"));
371    }
372
373    #[test]
374    fn test_to_json_schema_trait() {
375        assert!(matches!(String::to_json_schema(), JsonSchema::String { .. }));
376        assert!(matches!(i32::to_json_schema(), JsonSchema::Integer { .. }));
377        assert!(matches!(f64::to_json_schema(), JsonSchema::Number { .. }));
378        assert!(matches!(bool::to_json_schema(), JsonSchema::Boolean { .. }));
379    }
380}