Skip to main content

flymcprust_core/
mirror.rs

1//! Tool parameters mirrored into HTTP headers (`x-mcp-header`, SEP-2243).
2//!
3//! A property of a tool's `inputSchema` may carry `"x-mcp-header": "Name"`;
4//! Streamable HTTP clients then send its value as `Mcp-Param-Name`, and the
5//! server must check that header against the body. Annotations are validated
6//! when the server is built, so an invalid tool definition is never served.
7
8use serde_json::{Map, Value};
9
10/// JSON Schema primitive type of a mirrored parameter.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum ParamKind {
13    /// `"type": "string"`.
14    String,
15    /// `"type": "integer"`.
16    Integer,
17    /// `"type": "boolean"`.
18    Boolean,
19}
20
21/// A tool parameter mirrored into an `Mcp-Param-{name}` header.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct MirroredParam {
24    header: String,
25    path: Vec<String>,
26    kind: ParamKind,
27}
28
29impl MirroredParam {
30    /// The `{name}` part of the `Mcp-Param-{name}` header.
31    #[must_use]
32    pub fn header(&self) -> &str {
33        &self.header
34    }
35
36    /// The chain of property names leading to the parameter.
37    #[must_use]
38    pub fn path(&self) -> &[String] {
39        &self.path
40    }
41
42    /// The parameter's primitive type.
43    #[must_use]
44    pub fn kind(&self) -> ParamKind {
45        self.kind
46    }
47
48    /// The parameter's value in the call arguments, if present.
49    #[must_use]
50    pub fn value_in<'v>(&self, arguments: &'v Value) -> Option<&'v Value> {
51        self.path
52            .iter()
53            .try_fold(arguments, |current, key| current.get(key))
54    }
55}
56
57const ANNOTATION: &str = "x-mcp-header";
58
59/// Collects and validates every `x-mcp-header` annotation of a schema.
60///
61/// # Errors
62///
63/// Returns a description of the first violated constraint.
64pub(crate) fn collect(schema: &Map<String, Value>) -> Result<Vec<MirroredParam>, String> {
65    let mut found = Vec::new();
66    walk_properties(schema, &mut Vec::new(), &mut found)?;
67
68    // Annotations outside a pure `properties` chain (inside `items`,
69    // `anyOf`, `$defs`, …) make the definition invalid.
70    let total = count_annotations(&Value::Object(schema.clone()));
71    if total != found.len() {
72        return Err(format!(
73            "`{ANNOTATION}` may only annotate properties reachable through `properties` keys"
74        ));
75    }
76
77    let mut seen: Vec<String> = Vec::with_capacity(found.len());
78    for param in &found {
79        let lower = param.header.to_ascii_lowercase();
80        if seen.contains(&lower) {
81            return Err(format!(
82                "`{ANNOTATION}` value `{}` is used twice (names are case-insensitive)",
83                param.header
84            ));
85        }
86        seen.push(lower);
87    }
88    Ok(found)
89}
90
91fn walk_properties(
92    schema: &Map<String, Value>,
93    path: &mut Vec<String>,
94    found: &mut Vec<MirroredParam>,
95) -> Result<(), String> {
96    let Some(Value::Object(properties)) = schema.get("properties") else {
97        return Ok(());
98    };
99    for (name, property) in properties {
100        let Value::Object(property) = property else {
101            continue;
102        };
103        path.push(name.clone());
104        if let Some(annotation) = property.get(ANNOTATION) {
105            found.push(validate(annotation, property, path)?);
106        }
107        walk_properties(property, path, found)?;
108        path.pop();
109    }
110    Ok(())
111}
112
113fn validate(
114    annotation: &Value,
115    property: &Map<String, Value>,
116    path: &[String],
117) -> Result<MirroredParam, String> {
118    let at = path.join(".");
119    let header = annotation
120        .as_str()
121        .ok_or_else(|| format!("`{ANNOTATION}` on `{at}` must be a string"))?;
122    if header.is_empty() || !header.bytes().all(is_tchar) {
123        return Err(format!(
124            "`{ANNOTATION}` value `{header}` on `{at}` is not a valid HTTP field name token"
125        ));
126    }
127    let kind = match property.get("type").and_then(Value::as_str) {
128        Some("string") => ParamKind::String,
129        Some("integer") => ParamKind::Integer,
130        Some("boolean") => ParamKind::Boolean,
131        _ => {
132            return Err(format!(
133                "`{ANNOTATION}` on `{at}` requires `type` to be string, integer or boolean"
134            ));
135        }
136    };
137    Ok(MirroredParam {
138        header: header.to_owned(),
139        path: path.to_vec(),
140        kind,
141    })
142}
143
144fn count_annotations(value: &Value) -> usize {
145    match value {
146        Value::Object(map) => {
147            usize::from(map.contains_key(ANNOTATION))
148                + map.values().map(count_annotations).sum::<usize>()
149        }
150        Value::Array(items) => items.iter().map(count_annotations).sum(),
151        _ => 0,
152    }
153}
154
155/// `tchar` from RFC 9110 §5.6.2.
156fn is_tchar(b: u8) -> bool {
157    b.is_ascii_alphanumeric()
158        || matches!(
159            b,
160            b'!' | b'#'
161                | b'$'
162                | b'%'
163                | b'&'
164                | b'\''
165                | b'*'
166                | b'+'
167                | b'-'
168                | b'.'
169                | b'^'
170                | b'_'
171                | b'`'
172                | b'|'
173                | b'~'
174        )
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use serde_json::json;
181
182    fn schema(value: Value) -> Map<String, Value> {
183        match value {
184            Value::Object(map) => map,
185            _ => Map::new(),
186        }
187    }
188
189    #[test]
190    fn collects_nested_annotations() {
191        let s = schema(json!({
192            "type": "object",
193            "properties": {
194                "region": { "type": "string", "x-mcp-header": "Region" },
195                "query": { "type": "string" },
196                "opts": {
197                    "type": "object",
198                    "properties": { "dry": { "type": "boolean", "x-mcp-header": "Dry-Run" } }
199                }
200            }
201        }));
202        let params = collect(&s).unwrap();
203        assert_eq!(params.len(), 2);
204        // Property order depends on serde_json's `preserve_order` feature.
205        let by_header = |h: &str| params.iter().find(|p| p.header() == h).unwrap();
206        let args = json!({ "region": "us-west1", "opts": { "dry": true } });
207        assert_eq!(by_header("Dry-Run").value_in(&args), Some(&json!(true)));
208        assert_eq!(by_header("Dry-Run").path(), ["opts", "dry"]);
209        assert_eq!(
210            by_header("Region").value_in(&args),
211            Some(&json!("us-west1"))
212        );
213    }
214
215    #[test]
216    fn rejects_invalid_annotations() {
217        let bad_type = schema(json!({ "properties": {
218            "n": { "type": "number", "x-mcp-header": "N" } } }));
219        assert!(collect(&bad_type).is_err());
220
221        let bad_token = schema(json!({ "properties": {
222            "n": { "type": "string", "x-mcp-header": "a b" } } }));
223        assert!(collect(&bad_token).is_err());
224
225        let duplicate = schema(json!({ "properties": {
226            "a": { "type": "string", "x-mcp-header": "Name" },
227            "b": { "type": "string", "x-mcp-header": "name" } } }));
228        assert!(collect(&duplicate).is_err());
229
230        let in_items = schema(json!({ "properties": {
231            "list": { "type": "array", "items": { "type": "string", "x-mcp-header": "Item" } } } }));
232        assert!(collect(&in_items).is_err());
233    }
234}