1use crate::version::META_NS;
23use crate::wire::Implementation;
24use serde_json::{Value, json};
25
26pub fn inject_client_meta(
31 params: &mut Value,
32 protocol_version: &str,
33 client: &Implementation,
34 capabilities: &Value,
35) {
36 let Some(obj) = params.as_object_mut() else {
37 return;
38 };
39 let meta = obj
40 .entry("_meta")
41 .or_insert_with(|| Value::Object(Default::default()));
42 let Some(m) = meta.as_object_mut() else {
43 return;
44 };
45 m.insert(
46 format!("{META_NS}protocolVersion"),
47 Value::String(protocol_version.to_string()),
48 );
49 m.insert(
50 format!("{META_NS}clientInfo"),
51 json!({"name": client.name, "version": client.version}),
52 );
53 m.insert(format!("{META_NS}clientCapabilities"), capabilities.clone());
54}
55
56pub fn mcp_name(method: &str, params: &Value) -> Option<String> {
60 match method {
61 "tools/call" | "prompts/get" => {
62 params.get("name").and_then(Value::as_str).map(String::from)
63 }
64 "resources/read" => params.get("uri").and_then(Value::as_str).map(String::from),
65 _ => None,
66 }
67}
68
69pub fn routing_headers(method: &str, params: &Value) -> Vec<(&'static str, String)> {
73 let mut headers = vec![("Mcp-Method", method.to_string())];
74 if let Some(name) = mcp_name(method, params) {
75 headers.push(("Mcp-Name", header_value(&name)));
76 }
77 headers
78}
79
80pub fn header_value(raw: &str) -> String {
86 if is_header_safe(raw) {
87 raw.to_string()
88 } else {
89 format!("=?base64?{}?=", base64_encode(raw.as_bytes()))
90 }
91}
92
93fn is_header_safe(s: &str) -> bool {
97 !s.is_empty()
98 && s.bytes().all(|b| (0x21..=0x7e).contains(&b))
99 && !(s.starts_with("=?base64?") && s.ends_with("?="))
100}
101
102pub fn param_headers(input_schema: &Value, arguments: &Value) -> Vec<(String, String)> {
111 let mut out = Vec::new();
112 collect_param_headers(input_schema, arguments, &mut out);
113 out
114}
115
116fn collect_param_headers(schema: &Value, instance: &Value, out: &mut Vec<(String, String)>) {
117 let Some(props) = schema.get("properties").and_then(Value::as_object) else {
118 return;
119 };
120 for (key, sub) in props {
121 let value = instance.get(key);
122 if let Some(header_name) = sub.get("x-mcp-header").and_then(Value::as_str)
123 && let Some(v) = value
124 && let Some(s) = primitive_to_string(v)
125 {
126 out.push((format!("Mcp-Param-{header_name}"), header_value(&s)));
127 }
128 if let Some(v) = value
130 && sub.get("properties").is_some()
131 {
132 collect_param_headers(sub, v, out);
133 }
134 }
135}
136
137fn primitive_to_string(v: &Value) -> Option<String> {
142 match v {
143 Value::String(s) => Some(s.clone()),
144 Value::Bool(b) => Some(if *b { "true".into() } else { "false".into() }),
145 Value::Number(n) if n.is_i64() || n.is_u64() => Some(n.to_string()),
146 _ => None,
147 }
148}
149
150pub fn validate_x_mcp_headers(input_schema: &Value) -> Result<(), String> {
159 let mut seen = std::collections::HashSet::new();
160 validate_schema_node(input_schema, true, &mut seen)
161}
162
163fn validate_schema_node(
164 node: &Value,
165 reachable: bool,
166 seen: &mut std::collections::HashSet<String>,
167) -> Result<(), String> {
168 let Some(obj) = node.as_object() else {
169 return Ok(());
170 };
171 if let Some(h) = obj.get("x-mcp-header") {
172 let name = h.as_str().ok_or("x-mcp-header must be a string")?;
173 if !reachable {
174 return Err(format!("x-mcp-header '{name}' is not statically reachable"));
175 }
176 validate_header_name(name)?;
177 if !seen.insert(name.to_ascii_lowercase()) {
178 return Err(format!("duplicate x-mcp-header '{name}'"));
179 }
180 match obj.get("type").and_then(Value::as_str) {
181 Some("string") | Some("integer") | Some("boolean") => {}
182 Some("number") => return Err(format!("x-mcp-header '{name}' on a number type")),
183 _ => return Err(format!("x-mcp-header '{name}' on a non-primitive type")),
184 }
185 }
186 if let Some(props) = obj.get("properties").and_then(Value::as_object) {
189 for sub in props.values() {
190 validate_schema_node(sub, reachable, seen)?;
191 }
192 }
193 for key in ["items", "additionalProperties", "not", "if", "then", "else"] {
194 if let Some(sub) = obj.get(key) {
195 validate_schema_node(sub, false, seen)?;
196 }
197 }
198 for key in ["oneOf", "anyOf", "allOf", "prefixItems"] {
199 if let Some(arr) = obj.get(key).and_then(Value::as_array) {
200 for sub in arr {
201 validate_schema_node(sub, false, seen)?;
202 }
203 }
204 }
205 Ok(())
206}
207
208fn validate_header_name(name: &str) -> Result<(), String> {
211 if name.is_empty() {
212 return Err("empty x-mcp-header".into());
213 }
214 let is_tchar = |c: u8| c.is_ascii_alphanumeric() || b"!#$%&'*+-.^_`|~".contains(&c);
215 if !name.bytes().all(is_tchar) {
216 return Err(format!("x-mcp-header '{name}' is not a valid HTTP token"));
217 }
218 Ok(())
219}
220
221fn base64_encode(input: &[u8]) -> String {
224 const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
225 let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
226 for chunk in input.chunks(3) {
227 let b0 = chunk[0] as u32;
228 let b1 = *chunk.get(1).unwrap_or(&0) as u32;
229 let b2 = *chunk.get(2).unwrap_or(&0) as u32;
230 let n = (b0 << 16) | (b1 << 8) | b2;
231 out.push(ALPHABET[((n >> 18) & 63) as usize] as char);
232 out.push(ALPHABET[((n >> 12) & 63) as usize] as char);
233 out.push(if chunk.len() > 1 {
234 ALPHABET[((n >> 6) & 63) as usize] as char
235 } else {
236 '='
237 });
238 out.push(if chunk.len() > 2 {
239 ALPHABET[(n & 63) as usize] as char
240 } else {
241 '='
242 });
243 }
244 out
245}
246
247#[cfg(test)]
248mod tests {
249 use super::*;
250
251 fn client() -> Implementation {
252 Implementation {
253 name: "agentd".into(),
254 version: "1.0.0".into(),
255 title: None,
256 }
257 }
258
259 #[test]
260 fn injects_the_three_meta_fields() {
261 let mut params = json!({"name": "echo", "arguments": {"x": 1}});
262 let caps = json!({"extensions": {"io.modelcontextprotocol/tasks": {}}});
263 inject_client_meta(&mut params, "2026-07-28", &client(), &caps);
264 let meta = ¶ms["_meta"];
265 assert_eq!(
266 meta["io.modelcontextprotocol/protocolVersion"],
267 "2026-07-28"
268 );
269 assert_eq!(meta["io.modelcontextprotocol/clientInfo"]["name"], "agentd");
270 assert_eq!(
271 meta["io.modelcontextprotocol/clientInfo"]["version"],
272 "1.0.0"
273 );
274 assert_eq!(
276 meta["io.modelcontextprotocol/clientCapabilities"]["extensions"]["io.modelcontextprotocol/tasks"],
277 json!({})
278 );
279 assert_eq!(params["name"], "echo");
281 assert_eq!(params["arguments"]["x"], 1);
282 }
283
284 #[test]
285 fn routing_headers_carry_method_and_name() {
286 let p = json!({"name": "get_weather", "arguments": {}});
287 let h = routing_headers("tools/call", &p);
288 assert_eq!(h[0], ("Mcp-Method", "tools/call".to_string()));
289 assert_eq!(h[1], ("Mcp-Name", "get_weather".to_string()));
290
291 let p = json!({"uri": "file:///a.json"});
293 let h = routing_headers("resources/read", &p);
294 assert_eq!(h[1], ("Mcp-Name", "file:///a.json".to_string()));
295
296 let h = routing_headers("tools/list", &json!({}));
298 assert_eq!(h.len(), 1);
299 assert_eq!(h[0].0, "Mcp-Method");
300 }
301
302 #[test]
303 fn header_value_encodes_only_when_unsafe() {
304 assert_eq!(header_value("get_weather"), "get_weather");
305 assert_eq!(header_value("file:///a.json"), "file:///a.json");
306 assert_eq!(
308 header_value("Hello, 世界"),
309 "=?base64?SGVsbG8sIOS4lueVjA==?="
310 );
311 assert_eq!(
313 header_value("a b"),
314 format!("=?base64?{}?=", base64_encode(b"a b"))
315 );
316 assert!(header_value("=?base64?x?=").starts_with("=?base64?"));
318 }
319
320 #[test]
321 fn param_headers_extracts_annotated_values() {
322 let schema = json!({
323 "type": "object",
324 "properties": {
325 "region": {"type": "string", "x-mcp-header": "Region"},
326 "limit": {"type": "integer", "x-mcp-header": "Limit"},
327 "query": {"type": "string"}
328 }
329 });
330 let args = json!({"region": "us-west1", "limit": 42, "query": "SELECT 1"});
331 let mut h = param_headers(&schema, &args);
332 h.sort();
333 assert_eq!(
334 h,
335 vec![
336 ("Mcp-Param-Limit".to_string(), "42".to_string()),
337 ("Mcp-Param-Region".to_string(), "us-west1".to_string()),
338 ]
339 );
340 let h = param_headers(&schema, &json!({"query": "x"}));
342 assert!(h.is_empty());
343 }
344
345 #[test]
346 fn validate_accepts_valid_and_rejects_invalid() {
347 assert!(
349 validate_x_mcp_headers(&json!({
350 "type": "object",
351 "properties": {"r": {"type": "string", "x-mcp-header": "Region"}}
352 }))
353 .is_ok()
354 );
355 assert!(
357 validate_x_mcp_headers(&json!({
358 "properties": {"n": {"type": "number", "x-mcp-header": "N"}}
359 }))
360 .is_err()
361 );
362 assert!(
364 validate_x_mcp_headers(&json!({
365 "properties": {
366 "a": {"type": "string", "x-mcp-header": "Dup"},
367 "b": {"type": "string", "x-mcp-header": "dup"}
368 }
369 }))
370 .is_err()
371 );
372 assert!(
374 validate_x_mcp_headers(&json!({
375 "properties": {"list": {"type": "array",
376 "items": {"type": "object", "properties": {
377 "x": {"type": "string", "x-mcp-header": "X"}}}}}
378 }))
379 .is_err()
380 );
381 assert!(
383 validate_x_mcp_headers(&json!({
384 "properties": {"a": {"type": "string", "x-mcp-header": "bad name"}}
385 }))
386 .is_err()
387 );
388 }
389
390 #[test]
391 fn base64_matches_known_vectors() {
392 assert_eq!(base64_encode(b""), "");
393 assert_eq!(base64_encode(b"f"), "Zg==");
394 assert_eq!(base64_encode(b"fo"), "Zm8=");
395 assert_eq!(base64_encode(b"foo"), "Zm9v");
396 assert_eq!(base64_encode(b"foob"), "Zm9vYg==");
397 assert_eq!(base64_encode(b"foobar"), "Zm9vYmFy");
398 }
399}