feat: add $ref support to JSON Schema parser (#22)

Author: hashemix

src/lib.rs

@@ -9,6 +9,7 @@ mod templates;
 mod types;
 mod utils;
 
+use serde_json::{to_value, Map, Value};
 pub use templates::OutputTemplate;
 pub use types::{
     DiscoveryCommand, LogLevel, McpCapabilities, McpServerInfo, McpToolMeta, ParamTypes,
@@ -328,7 +329,9 @@ impl McpDiscovery {
         let mut tools: Vec<_> = tools_result
             .iter()
             .map(|tool| {
-                let params = tool_params(&tool.input_schema.properties);
+                let root_schema: serde_json::Value =
+                    to_value(&tool.input_schema).unwrap_or_else(|_| Value::Object(Map::new()));
+                let params = tool_params(&tool.input_schema.properties, &root_schema);
 
                 Ok::<McpToolMeta, DiscoveryError>(McpToolMeta {
                     name: tool.name.to_owned(),

src/schema.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use serde_json::{Map, Value};
 
@@ -7,8 +7,71 @@ use crate::{
     types::{McpToolSParams, ParamTypes},
 };
 
+/// Resolves a $ref path to its target value in the schema.
+fn resolve_ref<'a>(
+    ref_path: &str,
+    root_schema: &'a Value,
+    visited: &mut HashSet<String>,
+) -> DiscoveryResult<&'a Value> {
+    if !ref_path.starts_with("#/") {
+        return Err(DiscoveryError::InvalidSchema(format!(
+            "$ref '{}' must start with '#/'",
+            ref_path
+        )));
+    }
+
+    if !visited.insert(ref_path.to_string()) {
+        return Err(DiscoveryError::InvalidSchema(format!(
+            "Cycle detected in $ref path '{}'",
+            ref_path
+        )));
+    }
+
+    let path = ref_path.trim_start_matches("#/").split('/');
+    let mut current = root_schema;
+
+    for segment in path {
+        if segment.is_empty() {
+            return Err(DiscoveryError::InvalidSchema(format!(
+                "Invalid $ref path '{}': empty segment",
+                ref_path
+            )));
+        }
+        current = match current {
+            Value::Object(obj) => obj.get(segment).ok_or_else(|| {
+                DiscoveryError::InvalidSchema(format!(
+                    "Invalid $ref path '{}': segment '{}' not found",
+                    ref_path, segment
+                ))
+            })?,
+            Value::Array(arr) => segment
+                .parse::<usize>()
+                .ok()
+                .and_then(|i| arr.get(i))
+                .ok_or_else(|| {
+                    DiscoveryError::InvalidSchema(format!(
+                        "Invalid $ref path '{}': segment '{}' not found in array",
+                        ref_path, segment
+                    ))
+                })?,
+            _ => {
+                return Err(DiscoveryError::InvalidSchema(format!(
+                    "Invalid $ref path '{}': cannot traverse into non-object/array",
+                    ref_path
+                )))
+            }
+        };
+    }
+
+    Ok(current)
+}
+
 /// Parses an object schema into a vector of `McpToolSParams`.
-pub fn param_object(object_map: &Map<String, Value>) -> DiscoveryResult<Vec<McpToolSParams>> {
+pub fn param_object(
+    object_map: &Map<String, Value>,
+    root_schema: &Value,
+    visited: &mut HashSet<String>,
+) -> DiscoveryResult<Vec<McpToolSParams>> {
     let properties = object_map
         .get("properties")
         .and_then(|v| v.as_object())
@@ -31,7 +94,7 @@ pub fn param_object(object_map: &Map<String, Value>) -> DiscoveryResult<Vec<McpT
                     "Property '{}' is not an object",
                     param_name
                 )))?;
-            let param_type = param_type(param_value)?;
+            let param_type = param_type(param_value, root_schema, visited)?;
             let param_description = object_map
                 .get("description")
                 .and_then(|v| v.as_str())
@@ -50,7 +113,26 @@ pub fn param_object(object_map: &Map<String, Value>) -> DiscoveryResult<Vec<McpT
 }
 
 /// Determines the parameter type from a schema definition.
-pub fn param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTypes> {
+pub fn param_type(
+    type_info: &Map<String, Value>,
+    root_schema: &Value,
+    visited: &mut HashSet<String>,
+) -> DiscoveryResult<ParamTypes> {
+    // Handle $ref
+    if let Some(ref_path) = type_info.get("$ref") {
+        let ref_path_str = ref_path.as_str().ok_or(DiscoveryError::InvalidSchema(
+            "$ref must be a string".to_string(),
+        ))?;
+        let ref_value = resolve_ref(ref_path_str, root_schema, visited)?;
+        let ref_map = ref_value
+            .as_object()
+            .ok_or(DiscoveryError::InvalidSchema(format!(
+                "$ref '{}' does not point to an object",
+                ref_path_str
+            )))?;
+        return param_type(ref_map, root_schema, visited);
+    }
+
     // Check for 'enum' keyword
     if let Some(enum_values) = type_info.get("enum") {
         let values = enum_values.as_array().ok_or(DiscoveryError::InvalidSchema(
@@ -90,6 +172,7 @@ pub fn param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTypes>
         ));
     }
 
+    // Check for 'anyOf'
     if let Some(any_of) = type_info.get("anyOf") {
         let any_of_array = any_of.as_array().ok_or(DiscoveryError::InvalidSchema(
             "'anyOf' field must be an array".to_string(),
@@ -104,7 +187,7 @@ pub fn param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTypes>
             let item_map = item.as_object().ok_or(DiscoveryError::InvalidSchema(
                 "Items in 'anyOf' must be objects".to_string(),
             ))?;
-            enum_types.push(param_type(item_map)?);
+            enum_types.push(param_type(item_map, root_schema, visited)?);
         }
         return Ok(ParamTypes::Anyof(enum_types));
     }
@@ -124,7 +207,7 @@ pub fn param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTypes>
             let item_map = item.as_object().ok_or(DiscoveryError::InvalidSchema(
                 "Items in 'oneOf' must be objects".to_string(),
             ))?;
-            one_of_types.push(param_type(item_map)?);
+            one_of_types.push(param_type(item_map, root_schema, visited)?);
         }
         return Ok(ParamTypes::OneOf(one_of_types));
     }
@@ -144,12 +227,12 @@ pub fn param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTypes>
             let item_map = item.as_object().ok_or(DiscoveryError::InvalidSchema(
                 "Items in 'allOf' must be objects".to_string(),
             ))?;
-            all_of_types.push(param_type(item_map)?);
+            all_of_types.push(param_type(item_map, root_schema, visited)?);
         }
         return Ok(ParamTypes::AllOf(all_of_types));
     }
 
-    // other types
+    // Other types
     let type_name =
         type_info
             .get("type")
@@ -166,22 +249,34 @@ pub fn param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTypes>
                     "Missing or invalid 'items' field in array type".to_string(),
                 ),
             )?;
-            Ok(ParamTypes::Array(vec![param_type(items_map)?]))
+            Ok(ParamTypes::Array(vec![param_type(
+                items_map,
+                root_schema,
+                visited,
+            )?]))
         }
-        "object" => Ok(ParamTypes::Object(param_object(type_info)?)),
+        "object" => Ok(ParamTypes::Object(param_object(
+            type_info,
+            root_schema,
+            visited,
+        )?)),
         _ => Ok(ParamTypes::Primitive(type_name.to_string())),
     }
 }
 
+/// Processes tool parameters with a given properties map and root schema.
 pub fn tool_params(
     properties: &Option<HashMap<String, Map<String, Value>>>,
+    root_schema: &Value,
 ) -> Vec<McpToolSParams> {
+    let mut visited = HashSet::new();
     let result = properties.clone().map(|props| {
         let mut params: Vec<_> = props
             .iter()
             .map(|(prop_name, prop_map)| {
                 let param_name = prop_name.to_owned();
-                let prop_type = param_type(prop_map).unwrap();
+                let prop_type = param_type(prop_map, root_schema, &mut visited)
+                    .unwrap_or_else(|_| ParamTypes::Primitive("unknown".to_string()));
                 let prop_description = prop_map
                     .get("description")
                     .and_then(|v| v.as_str())

templates/html/html_tools.hbs

@@ -21,7 +21,7 @@
             <td>
                 <ul>
                     {{#each this.params}}
-                    <li style="white-space: nowrap;"> <code>{{{this.param_name}}}</code> : {{{tool_param_type
+                    <li style=""> <code>{{{this.param_name}}}</code> : {{{tool_param_type
                         this.param_type}}}<br /></li>
                     {{/each}}
                 </ul>
@@ -31,4 +31,4 @@
     </tbody>
 </table>
 
-{{/if}}
+{{/if}}

templates/markdown/md_tools.hbs

@@ -21,7 +21,7 @@
             <td>
                 <ul>
                     {{#each this.params}}
-                    <li style="white-space: nowrap;"> <code>{{{this.param_name}}}</code> : {{{tool_param_type
+                    <li> <code>{{{this.param_name}}}</code> : {{{tool_param_type
                         this.param_type}}}<br /></li>
                     {{/each}}
                 </ul>