Skip to main content

mcpzip/
types.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4pub const NAME_SEPARATOR: &str = "__";
5
6/// A cached tool from an upstream MCP server.
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ToolEntry {
9    pub name: String,
10    pub server_name: String,
11    pub original_name: String,
12    pub description: String,
13    pub input_schema: serde_json::Value,
14    pub compact_params: String,
15}
16
17/// Result from search engine.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct SearchResult {
20    pub name: String,
21    pub description: String,
22    pub compact_params: String,
23}
24
25/// How to connect to an upstream MCP server.
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct ServerConfig {
28    #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
29    pub server_type: Option<String>,
30    #[serde(default, skip_serializing_if = "Option::is_none")]
31    pub command: Option<String>,
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub args: Option<Vec<String>>,
34    #[serde(default, skip_serializing_if = "Option::is_none")]
35    pub env: Option<HashMap<String, String>>,
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub url: Option<String>,
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub headers: Option<HashMap<String, String>>,
40}
41
42impl ServerConfig {
43    pub fn effective_type(&self) -> &str {
44        match self.server_type.as_deref() {
45            Some(t) if !t.is_empty() => t,
46            _ => "stdio",
47        }
48    }
49}
50
51/// Search engine settings.
52#[derive(Debug, Clone, Default, Serialize, Deserialize)]
53pub struct SearchConfig {
54    #[serde(default, skip_serializing_if = "Option::is_none")]
55    pub default_limit: Option<usize>,
56    #[serde(default, skip_serializing_if = "Option::is_none")]
57    pub model: Option<String>,
58}
59
60/// Full proxy configuration.
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct ProxyConfig {
63    #[serde(default, skip_serializing_if = "Option::is_none")]
64    pub gemini_api_key: Option<String>,
65    #[serde(default)]
66    pub search: SearchConfig,
67    #[serde(default, skip_serializing_if = "Option::is_none")]
68    pub idle_timeout_minutes: Option<u64>,
69    #[serde(default, skip_serializing_if = "Option::is_none")]
70    pub call_timeout_seconds: Option<u64>,
71    #[serde(rename = "mcpServers")]
72    pub mcp_servers: HashMap<String, ServerConfig>,
73}
74
75/// Health info for an upstream server.
76#[derive(Debug, Clone, Serialize, Deserialize)]
77pub struct ServerStatus {
78    pub name: String,
79    pub connected: bool,
80    pub tool_count: usize,
81    pub last_refresh: String,
82    #[serde(default, skip_serializing_if = "Option::is_none")]
83    pub error: Option<String>,
84}
85
86/// Returns "server__tool".
87pub fn prefixed_name(server: &str, tool: &str) -> String {
88    format!("{}{}{}", server, NAME_SEPARATOR, tool)
89}
90
91/// Splits "server__tool" into (server, tool). Splits on first occurrence of "__".
92pub fn parse_prefixed_name(name: &str) -> Result<(&str, &str), crate::error::McpzipError> {
93    match name.find(NAME_SEPARATOR) {
94        Some(idx) => Ok((&name[..idx], &name[idx + NAME_SEPARATOR.len()..])),
95        None => Err(crate::error::McpzipError::Protocol(format!(
96            "invalid prefixed name {:?}: missing separator {:?}",
97            name, NAME_SEPARATOR
98        ))),
99    }
100}
101
102/// Generate compact parameter summary from a JSON Schema.
103/// Format: "param1:type*, param2:type" where * marks required params.
104pub fn compact_params_from_schema(schema: &serde_json::Value) -> String {
105    let obj = match schema.as_object() {
106        Some(o) => o,
107        None => return String::new(),
108    };
109
110    let properties = match obj.get("properties").and_then(|v| v.as_object()) {
111        Some(p) => p,
112        None => return String::new(),
113    };
114
115    if properties.is_empty() {
116        return String::new();
117    }
118
119    let required: std::collections::HashSet<&str> = obj
120        .get("required")
121        .and_then(|v| v.as_array())
122        .map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
123        .unwrap_or_default();
124
125    let mut names: Vec<&str> = properties.keys().map(|s| s.as_str()).collect();
126    names.sort();
127
128    let parts: Vec<String> = names
129        .iter()
130        .map(|name| {
131            let typ = extract_type(&properties[*name]);
132            if required.contains(name) {
133                format!("{}:{}*", name, typ)
134            } else {
135                format!("{}:{}", name, typ)
136            }
137        })
138        .collect();
139
140    parts.join(", ")
141}
142
143fn extract_type(value: &serde_json::Value) -> &str {
144    // Handle "type": "string"
145    if let Some(t) = value.get("type").and_then(|v| v.as_str()) {
146        return t;
147    }
148
149    // Handle "type": ["string", "null"]
150    if let Some(arr) = value.get("type").and_then(|v| v.as_array()) {
151        for item in arr {
152            if let Some(s) = item.as_str() {
153                if s != "null" {
154                    return s;
155                }
156            }
157        }
158        if let Some(first) = arr.first().and_then(|v| v.as_str()) {
159            return first;
160        }
161    }
162
163    // Handle anyOf
164    if let Some(any_of) = value.get("anyOf").and_then(|v| v.as_array()) {
165        for item in any_of {
166            if let Some(t) = item.get("type").and_then(|v| v.as_str()) {
167                if t != "null" {
168                    return t;
169                }
170            }
171        }
172        if let Some(first) = any_of.first() {
173            if let Some(t) = first.get("type").and_then(|v| v.as_str()) {
174                return t;
175            }
176        }
177    }
178
179    "any"
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185    use serde_json::json;
186
187    #[test]
188    fn test_prefixed_name() {
189        assert_eq!(
190            prefixed_name("slack", "send_message"),
191            "slack__send_message"
192        );
193    }
194
195    #[test]
196    fn test_parse_prefixed_name() {
197        let (server, tool) = parse_prefixed_name("slack__send_message").unwrap();
198        assert_eq!(server, "slack");
199        assert_eq!(tool, "send_message");
200    }
201
202    #[test]
203    fn test_parse_prefixed_name_first_occurrence() {
204        let (server, tool) = parse_prefixed_name("a__b__c").unwrap();
205        assert_eq!(server, "a");
206        assert_eq!(tool, "b__c");
207    }
208
209    #[test]
210    fn test_parse_prefixed_name_no_separator() {
211        assert!(parse_prefixed_name("no_separator").is_err());
212    }
213
214    #[test]
215    fn test_effective_type_default() {
216        let cfg = ServerConfig {
217            server_type: None,
218            command: Some("echo".into()),
219            args: None,
220            env: None,
221            url: None,
222            headers: None,
223        };
224        assert_eq!(cfg.effective_type(), "stdio");
225    }
226
227    #[test]
228    fn test_effective_type_http() {
229        let cfg = ServerConfig {
230            server_type: Some("http".into()),
231            command: None,
232            args: None,
233            env: None,
234            url: Some("https://example.com".into()),
235            headers: None,
236        };
237        assert_eq!(cfg.effective_type(), "http");
238    }
239
240    #[test]
241    fn test_effective_type_empty_string() {
242        let cfg = ServerConfig {
243            server_type: Some(String::new()),
244            command: Some("echo".into()),
245            args: None,
246            env: None,
247            url: None,
248            headers: None,
249        };
250        assert_eq!(cfg.effective_type(), "stdio");
251    }
252
253    #[test]
254    fn test_compact_params_basic() {
255        let schema = json!({
256            "type": "object",
257            "properties": {
258                "channel": {"type": "string"},
259                "message": {"type": "string"}
260            },
261            "required": ["channel"]
262        });
263        assert_eq!(
264            compact_params_from_schema(&schema),
265            "channel:string*, message:string"
266        );
267    }
268
269    #[test]
270    fn test_compact_params_nullable_type() {
271        let schema = json!({
272            "type": "object",
273            "properties": {
274                "name": {"type": ["string", "null"]}
275            }
276        });
277        assert_eq!(compact_params_from_schema(&schema), "name:string");
278    }
279
280    #[test]
281    fn test_compact_params_any_of() {
282        let schema = json!({
283            "type": "object",
284            "properties": {
285                "value": {"anyOf": [{"type": "integer"}, {"type": "null"}]}
286            }
287        });
288        assert_eq!(compact_params_from_schema(&schema), "value:integer");
289    }
290
291    #[test]
292    fn test_compact_params_empty() {
293        assert_eq!(compact_params_from_schema(&json!(null)), "");
294        assert_eq!(compact_params_from_schema(&json!({})), "");
295        assert_eq!(compact_params_from_schema(&json!({"properties": {}})), "");
296    }
297
298    #[test]
299    fn test_compact_params_no_required() {
300        let schema = json!({
301            "type": "object",
302            "properties": {
303                "a": {"type": "string"},
304                "b": {"type": "number"}
305            }
306        });
307        assert_eq!(compact_params_from_schema(&schema), "a:string, b:number");
308    }
309
310    #[test]
311    fn test_tool_entry_serde_roundtrip() {
312        let entry = ToolEntry {
313            name: "slack__send".into(),
314            server_name: "slack".into(),
315            original_name: "send".into(),
316            description: "Send a message".into(),
317            input_schema: json!({"type": "object"}),
318            compact_params: "msg:string*".into(),
319        };
320        let json_str = serde_json::to_string(&entry).unwrap();
321        let parsed: ToolEntry = serde_json::from_str(&json_str).unwrap();
322        assert_eq!(parsed.name, entry.name);
323        assert_eq!(parsed.server_name, entry.server_name);
324    }
325
326    #[test]
327    fn test_proxy_config_serde() {
328        let json_str = r#"{
329            "mcpServers": {
330                "slack": {
331                    "command": "slack-mcp",
332                    "args": ["--token", "xxx"]
333                },
334                "todoist": {
335                    "type": "http",
336                    "url": "https://todoist.com/mcp"
337                }
338            }
339        }"#;
340        let cfg: ProxyConfig = serde_json::from_str(json_str).unwrap();
341        assert_eq!(cfg.mcp_servers.len(), 2);
342        assert_eq!(cfg.mcp_servers["slack"].effective_type(), "stdio");
343        assert_eq!(cfg.mcp_servers["todoist"].effective_type(), "http");
344    }
345
346    #[test]
347    fn test_tool_entry_send_sync() {
348        fn assert_send_sync<T: Send + Sync>() {}
349        assert_send_sync::<ToolEntry>();
350        assert_send_sync::<SearchResult>();
351        assert_send_sync::<ProxyConfig>();
352    }
353
354    // --- New tests ---
355
356    #[test]
357    fn test_compact_params_deeply_nested_schema() {
358        let schema = json!({
359            "type": "object",
360            "properties": {
361                "config": {
362                    "type": "object",
363                    "properties": {
364                        "nested": {"type": "string"}
365                    }
366                }
367            },
368            "required": ["config"]
369        });
370        assert_eq!(compact_params_from_schema(&schema), "config:object*");
371    }
372
373    #[test]
374    fn test_compact_params_many_params() {
375        let schema = json!({
376            "type": "object",
377            "properties": {
378                "alpha": {"type": "string"},
379                "beta": {"type": "integer"},
380                "gamma": {"type": "boolean"},
381                "delta": {"type": "number"},
382                "epsilon": {"type": "array"}
383            },
384            "required": ["alpha", "beta"]
385        });
386        let result = compact_params_from_schema(&schema);
387        assert_eq!(
388            result,
389            "alpha:string*, beta:integer*, delta:number, epsilon:array, gamma:boolean"
390        );
391    }
392
393    #[test]
394    fn test_compact_params_all_required() {
395        let schema = json!({
396            "type": "object",
397            "properties": {
398                "a": {"type": "string"},
399                "b": {"type": "integer"}
400            },
401            "required": ["a", "b"]
402        });
403        assert_eq!(compact_params_from_schema(&schema), "a:string*, b:integer*");
404    }
405
406    #[test]
407    fn test_compact_params_no_type_returns_any() {
408        let schema = json!({
409            "type": "object",
410            "properties": {
411                "x": {}
412            }
413        });
414        assert_eq!(compact_params_from_schema(&schema), "x:any");
415    }
416
417    #[test]
418    fn test_compact_params_nullable_only() {
419        // type: ["null"] with no non-null type => returns "null"
420        let schema = json!({
421            "type": "object",
422            "properties": {
423                "v": {"type": ["null"]}
424            }
425        });
426        assert_eq!(compact_params_from_schema(&schema), "v:null");
427    }
428
429    #[test]
430    fn test_compact_params_any_of_null_only() {
431        let schema = json!({
432            "type": "object",
433            "properties": {
434                "v": {"anyOf": [{"type": "null"}]}
435            }
436        });
437        assert_eq!(compact_params_from_schema(&schema), "v:null");
438    }
439
440    #[test]
441    fn test_compact_params_no_properties_key() {
442        let schema = json!({"type": "object"});
443        assert_eq!(compact_params_from_schema(&schema), "");
444    }
445
446    #[test]
447    fn test_compact_params_non_object_schema() {
448        assert_eq!(compact_params_from_schema(&json!("string")), "");
449        assert_eq!(compact_params_from_schema(&json!(42)), "");
450        assert_eq!(compact_params_from_schema(&json!(true)), "");
451        assert_eq!(compact_params_from_schema(&json!([1, 2, 3])), "");
452    }
453
454    #[test]
455    fn test_parse_prefixed_name_empty_string() {
456        assert!(parse_prefixed_name("").is_err());
457    }
458
459    #[test]
460    fn test_parse_prefixed_name_separator_only() {
461        let (server, tool) = parse_prefixed_name("__").unwrap();
462        assert_eq!(server, "");
463        assert_eq!(tool, "");
464    }
465
466    #[test]
467    fn test_parse_prefixed_name_multiple_separators() {
468        let (server, tool) = parse_prefixed_name("a__b__c__d").unwrap();
469        assert_eq!(server, "a");
470        assert_eq!(tool, "b__c__d");
471    }
472
473    #[test]
474    fn test_effective_type_sse() {
475        let cfg = ServerConfig {
476            server_type: Some("sse".into()),
477            command: None,
478            args: None,
479            env: None,
480            url: Some("https://example.com/sse".into()),
481            headers: None,
482        };
483        assert_eq!(cfg.effective_type(), "sse");
484    }
485
486    #[test]
487    fn test_proxy_config_full_serde_roundtrip() {
488        let cfg = ProxyConfig {
489            gemini_api_key: Some("test-key-123".into()),
490            search: SearchConfig {
491                default_limit: Some(10),
492                model: Some("gemini-2.0-flash".into()),
493            },
494            idle_timeout_minutes: Some(5),
495            call_timeout_seconds: Some(120),
496            mcp_servers: {
497                let mut m = HashMap::new();
498                m.insert(
499                    "slack".into(),
500                    ServerConfig {
501                        server_type: None,
502                        command: Some("slack-mcp".into()),
503                        args: Some(vec!["--token".into(), "xxx".into()]),
504                        env: Some({
505                            let mut e = HashMap::new();
506                            e.insert("API_KEY".into(), "secret".into());
507                            e
508                        }),
509                        url: None,
510                        headers: None,
511                    },
512                );
513                m.insert(
514                    "todoist".into(),
515                    ServerConfig {
516                        server_type: Some("http".into()),
517                        command: None,
518                        args: None,
519                        env: None,
520                        url: Some("https://todoist.com/mcp".into()),
521                        headers: Some({
522                            let mut h = HashMap::new();
523                            h.insert("Authorization".into(), "Bearer token".into());
524                            h
525                        }),
526                    },
527                );
528                m
529            },
530        };
531
532        let json_str = serde_json::to_string(&cfg).unwrap();
533        let parsed: ProxyConfig = serde_json::from_str(&json_str).unwrap();
534
535        assert_eq!(parsed.gemini_api_key, Some("test-key-123".into()));
536        assert_eq!(parsed.search.default_limit, Some(10));
537        assert_eq!(parsed.search.model, Some("gemini-2.0-flash".into()));
538        assert_eq!(parsed.idle_timeout_minutes, Some(5));
539        assert_eq!(parsed.call_timeout_seconds, Some(120));
540        assert_eq!(parsed.mcp_servers.len(), 2);
541        assert_eq!(parsed.mcp_servers["slack"].effective_type(), "stdio");
542        assert_eq!(parsed.mcp_servers["todoist"].effective_type(), "http");
543        assert_eq!(
544            parsed.mcp_servers["todoist"].url,
545            Some("https://todoist.com/mcp".into())
546        );
547    }
548
549    #[test]
550    fn test_server_config_serialization_skip_none() {
551        let cfg = ServerConfig {
552            server_type: None,
553            command: Some("echo".into()),
554            args: None,
555            env: None,
556            url: None,
557            headers: None,
558        };
559        let json_str = serde_json::to_string(&cfg).unwrap();
560        assert!(!json_str.contains("type"));
561        assert!(!json_str.contains("args"));
562        assert!(!json_str.contains("env"));
563        assert!(!json_str.contains("url"));
564        assert!(!json_str.contains("headers"));
565        assert!(json_str.contains("command"));
566    }
567
568    #[test]
569    fn test_search_result_serde() {
570        let sr = SearchResult {
571            name: "slack__send".into(),
572            description: "Send a message".into(),
573            compact_params: "msg:string*".into(),
574        };
575        let json_str = serde_json::to_string(&sr).unwrap();
576        let parsed: SearchResult = serde_json::from_str(&json_str).unwrap();
577        assert_eq!(parsed.name, "slack__send");
578        assert_eq!(parsed.description, "Send a message");
579        assert_eq!(parsed.compact_params, "msg:string*");
580    }
581
582    #[test]
583    fn test_server_status_serde() {
584        let status = ServerStatus {
585            name: "slack".into(),
586            connected: true,
587            tool_count: 42,
588            last_refresh: "2024-01-01T00:00:00Z".into(),
589            error: None,
590        };
591        let json_str = serde_json::to_string(&status).unwrap();
592        let parsed: ServerStatus = serde_json::from_str(&json_str).unwrap();
593        assert_eq!(parsed.name, "slack");
594        assert!(parsed.connected);
595        assert_eq!(parsed.tool_count, 42);
596        assert!(parsed.error.is_none());
597        // Verify "error" is omitted from JSON when None
598        assert!(!json_str.contains("error"));
599    }
600
601    #[test]
602    fn test_server_status_serde_with_error() {
603        let status = ServerStatus {
604            name: "github".into(),
605            connected: false,
606            tool_count: 0,
607            last_refresh: "2024-01-01T00:00:00Z".into(),
608            error: Some("connection refused".into()),
609        };
610        let json_str = serde_json::to_string(&status).unwrap();
611        let parsed: ServerStatus = serde_json::from_str(&json_str).unwrap();
612        assert!(!parsed.connected);
613        assert_eq!(parsed.error, Some("connection refused".into()));
614    }
615
616    #[test]
617    fn test_prefixed_name_roundtrip() {
618        let server = "my_server";
619        let tool = "my_tool";
620        let prefixed = prefixed_name(server, tool);
621        let (parsed_server, parsed_tool) = parse_prefixed_name(&prefixed).unwrap();
622        assert_eq!(parsed_server, server);
623        assert_eq!(parsed_tool, tool);
624    }
625}