agy_bridge/config/
mcp.rs

1//! MCP (Model Context Protocol) server configuration types.
2
3use serde::{Deserialize, Serialize};
4
5use super::{default_mcp_sse_read_timeout, default_mcp_timeout, default_true};
6
7/// Configuration for an MCP server connected via stdio.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct McpStdioServer {
10    /// The command to run to start the server.
11    pub command: String,
12    /// Arguments to pass to the command.
13    #[serde(default)]
14    pub args: Vec<String>,
15}
16
17/// Configuration for an MCP server connected via SSE.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct McpSseServer {
20    /// The URL of the SSE endpoint.
21    pub url: String,
22    /// Optional headers to send with the connection request.
23    #[serde(default)]
24    pub headers: Option<std::collections::HashMap<String, String>>,
25}
26
27/// Configuration for an MCP server connected via Streamable HTTP.
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct McpStreamableHttpServer {
30    /// The URL of the HTTP endpoint.
31    pub url: String,
32    /// Optional headers to send with the connection request.
33    #[serde(default)]
34    pub headers: Option<std::collections::HashMap<String, String>>,
35    /// Connection timeout in seconds.
36    #[serde(default = "default_mcp_timeout")]
37    pub timeout: f64,
38    /// SSE read timeout in seconds.
39    #[serde(default = "default_mcp_sse_read_timeout")]
40    pub sse_read_timeout: f64,
41    /// Whether to terminate the connection on close.
42    #[serde(default = "default_true")]
43    pub terminate_on_close: bool,
44}
45
46/// An MCP server, identified by its transport.
47///
48/// All MCP transports speak JSON-RPC 2.0; the variants describe *how* the
49/// client connects to the server process.
50///
51/// Use the convenience constructors [`McpServer::stdio`], [`McpServer::sse`],
52/// and [`McpServer::http`] to avoid importing the inner transport types.
53#[non_exhaustive]
54#[derive(Debug, Clone, Serialize, Deserialize)]
55#[serde(tag = "type")]
56pub enum McpServer {
57    #[serde(rename = "stdio")]
58    Stdio(McpStdioServer),
59    #[serde(rename = "sse")]
60    Sse(McpSseServer),
61    #[serde(rename = "http")]
62    Http(McpStreamableHttpServer),
63}
64
65impl McpServer {
66    /// Create a stdio-transport MCP server that spawns `command` as a child process.
67    #[must_use]
68    pub fn stdio(command: impl Into<String>) -> McpStdioServer {
69        McpStdioServer::new(command)
70    }
71
72    /// Create an SSE-transport MCP server at the given `url`.
73    #[must_use]
74    pub fn sse(url: impl Into<String>) -> McpSseServer {
75        McpSseServer::new(url)
76    }
77
78    /// Create a Streamable-HTTP-transport MCP server at the given `url`.
79    #[must_use]
80    pub fn http(url: impl Into<String>) -> McpStreamableHttpServer {
81        McpStreamableHttpServer::new(url)
82    }
83}
84
85// ─── MCP Server Builders ───────────────────────────────────────────────────
86
87impl From<McpStdioServer> for McpServer {
88    fn from(val: McpStdioServer) -> Self {
89        Self::Stdio(val)
90    }
91}
92
93impl McpStdioServer {
94    /// Create a new Stdio MCP Server configuration.
95    #[must_use]
96    pub fn new(command: impl Into<String>) -> Self {
97        Self {
98            command: command.into(),
99            args: Vec::new(),
100        }
101    }
102
103    /// Add an argument to the command.
104    #[must_use]
105    pub fn arg(mut self, arg: impl Into<String>) -> Self {
106        self.args.push(arg.into());
107        self
108    }
109
110    /// Add multiple arguments to the command at once.
111    #[must_use]
112    pub fn args<I, S>(mut self, args: I) -> Self
113    where
114        I: IntoIterator<Item = S>,
115        S: Into<String>,
116    {
117        self.args.extend(args.into_iter().map(Into::into));
118        self
119    }
120
121    /// Build this stdio configuration into an [`McpServer`].
122    #[must_use]
123    pub fn build(self) -> McpServer {
124        McpServer::Stdio(self)
125    }
126}
127
128impl From<McpSseServer> for McpServer {
129    fn from(val: McpSseServer) -> Self {
130        Self::Sse(val)
131    }
132}
133
134impl McpSseServer {
135    /// Create a new SSE MCP Server configuration.
136    #[must_use]
137    pub fn new(url: impl Into<String>) -> Self {
138        Self {
139            url: url.into(),
140            headers: None,
141        }
142    }
143
144    /// Add a header to the SSE connection.
145    #[must_use]
146    pub fn header(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
147        self.headers
148            .get_or_insert_with(std::collections::HashMap::new)
149            .insert(k.into(), v.into());
150        self
151    }
152
153    /// Build this SSE configuration into an [`McpServer`].
154    #[must_use]
155    pub fn build(self) -> McpServer {
156        McpServer::Sse(self)
157    }
158}
159
160impl From<McpStreamableHttpServer> for McpServer {
161    fn from(val: McpStreamableHttpServer) -> Self {
162        Self::Http(val)
163    }
164}
165
166impl McpStreamableHttpServer {
167    /// Create a new Streamable HTTP MCP Server configuration.
168    #[must_use]
169    pub fn new(url: impl Into<String>) -> Self {
170        Self {
171            url: url.into(),
172            headers: None,
173            timeout: default_mcp_timeout(),
174            sse_read_timeout: default_mcp_sse_read_timeout(),
175            terminate_on_close: true,
176        }
177    }
178
179    /// Add a header to the HTTP connection.
180    #[must_use]
181    pub fn header(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
182        self.headers
183            .get_or_insert_with(std::collections::HashMap::new)
184            .insert(k.into(), v.into());
185        self
186    }
187
188    /// Set the HTTP connection/request timeout in seconds.
189    #[must_use]
190    pub const fn timeout(mut self, timeout: f64) -> Self {
191        self.timeout = timeout;
192        self
193    }
194
195    /// Set the streaming read timeout in seconds.
196    #[must_use]
197    pub const fn sse_read_timeout(mut self, timeout: f64) -> Self {
198        self.sse_read_timeout = timeout;
199        self
200    }
201
202    /// Build this HTTP configuration into an [`McpServer`].
203    #[must_use]
204    pub fn build(self) -> McpServer {
205        McpServer::Http(self)
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use pyo3::types::PyAnyMethods;
212
213    use super::{
214        super::{DEFAULT_MCP_SSE_READ_TIMEOUT_SECS, DEFAULT_MCP_TIMEOUT_SECS},
215        *,
216    };
217
218    fn py_pydantic_field_default(module: &str, class: &str, field: &str) -> f64 {
219        pyo3::Python::initialize();
220        pyo3::Python::attach(|py| {
221            crate::runtime::venv::configure_python_sys_path(py)
222                .unwrap_or_else(|e| panic!("Failed to configure python sys.path: {e}"));
223            let m = crate::runtime::py_scripts::import_serialized(py, module)
224                .unwrap_or_else(|e| panic!("Failed to import {module}: {e}"));
225            let cls = m
226                .getattr(class)
227                .unwrap_or_else(|e| panic!("Failed to get {module}.{class}: {e}"));
228            let fields = cls
229                .getattr("model_fields")
230                .unwrap_or_else(|e| panic!("Failed to get {module}.{class}.model_fields: {e}"));
231            let field_info = fields.get_item(field).unwrap_or_else(|e| {
232                panic!("Failed to get field '{field}' from {module}.{class}.model_fields: {e}")
233            });
234            field_info
235                .getattr("default")
236                .unwrap_or_else(|e| {
237                    panic!("Failed to get default for {module}.{class}.{field}: {e}")
238                })
239                .extract::<f64>()
240                .unwrap_or_else(|e| {
241                    panic!("Failed to extract {module}.{class}.{field} default as f64: {e}")
242                })
243        })
244    }
245
246    #[test]
247    fn mcp_server_config_stdio_roundtrip() {
248        let config = McpServer::Stdio(McpStdioServer {
249            command: "npx".to_string(),
250            args: vec![
251                "-y".to_string(),
252                "@modelcontextprotocol/server-filesystem".to_string(),
253            ],
254        });
255        let json = serde_json::to_string(&config).unwrap();
256        let parsed: McpServer = serde_json::from_str(&json).unwrap();
257        match parsed {
258            McpServer::Stdio(s) => {
259                assert_eq!(s.command, "npx");
260                assert_eq!(
261                    s.args,
262                    vec!["-y", "@modelcontextprotocol/server-filesystem"]
263                );
264            }
265            other => panic!("Expected Stdio, got {other:?}"),
266        }
267        // Verify the JSON contains the "type" tag from serde.
268        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
269        assert_eq!(value["type"], "stdio");
270    }
271
272    #[test]
273    fn mcp_server_config_sse_roundtrip() {
274        let config = McpServer::Sse(McpSseServer {
275            url: "http://localhost:8080/sse".to_string(),
276            headers: Some(std::collections::HashMap::from([(
277                "Authorization".to_string(),
278                "Bearer token123".to_string(),
279            )])),
280        });
281        let json = serde_json::to_string(&config).unwrap();
282        let parsed: McpServer = serde_json::from_str(&json).unwrap();
283        match parsed {
284            McpServer::Sse(s) => {
285                assert_eq!(s.url, "http://localhost:8080/sse");
286                assert_eq!(
287                    s.headers.as_ref().unwrap()["Authorization"],
288                    "Bearer token123"
289                );
290            }
291            other => panic!("Expected Sse, got {other:?}"),
292        }
293        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
294        assert_eq!(value["type"], "sse");
295    }
296
297    #[test]
298    fn mcp_server_config_http_roundtrip() {
299        let config = McpServer::Http(McpStreamableHttpServer {
300            url: "http://localhost:9090/mcp".to_string(),
301            headers: None,
302            timeout: 60.0,
303            sse_read_timeout: 120.0,
304            terminate_on_close: false,
305        });
306        let json = serde_json::to_string(&config).unwrap();
307        let parsed: McpServer = serde_json::from_str(&json).unwrap();
308        match parsed {
309            McpServer::Http(s) => {
310                assert_eq!(s.url, "http://localhost:9090/mcp");
311                assert!(s.headers.is_none());
312                assert!((s.timeout - 60.0).abs() < f64::EPSILON);
313                assert!((s.sse_read_timeout - 120.0).abs() < f64::EPSILON);
314                assert!(!s.terminate_on_close);
315            }
316            other => panic!("Expected Http, got {other:?}"),
317        }
318        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
319        assert_eq!(value["type"], "http");
320    }
321
322    #[test]
323    fn mcp_server_config_http_defaults_roundtrip() {
324        // Deserialize with only required fields to verify defaults.
325        let json = r#"{"type":"http","url":"http://example.com/mcp"}"#;
326        let parsed: McpServer = serde_json::from_str(json).unwrap();
327        match parsed {
328            McpServer::Http(s) => {
329                assert_eq!(s.url, "http://example.com/mcp");
330                assert!(s.headers.is_none());
331                assert!((s.timeout - 30.0).abs() < f64::EPSILON);
332                assert!((s.sse_read_timeout - 300.0).abs() < f64::EPSILON);
333                assert!(s.terminate_on_close);
334            }
335            other => panic!("Expected Http, got {other:?}"),
336        }
337    }
338
339    #[test]
340    fn mcp_timeout_matches_python_sdk() {
341        let py_val = py_pydantic_field_default(
342            "google.antigravity.types",
343            "McpStreamableHttpServer",
344            "timeout",
345        );
346        assert!(
347            (DEFAULT_MCP_TIMEOUT_SECS - py_val).abs() < f64::EPSILON,
348            "Rust DEFAULT_MCP_TIMEOUT_SECS ({DEFAULT_MCP_TIMEOUT_SECS}) != Python SDK ({py_val})"
349        );
350    }
351
352    #[test]
353    fn mcp_sse_read_timeout_matches_python_sdk() {
354        let py_val = py_pydantic_field_default(
355            "google.antigravity.types",
356            "McpStreamableHttpServer",
357            "sse_read_timeout",
358        );
359        assert!(
360            (DEFAULT_MCP_SSE_READ_TIMEOUT_SECS - py_val).abs() < f64::EPSILON,
361            "Rust DEFAULT_MCP_SSE_READ_TIMEOUT_SECS ({DEFAULT_MCP_SSE_READ_TIMEOUT_SECS}) != Python SDK ({py_val})"
362        );
363    }
364
365    #[test]
366    fn test_mcp_server_builders() {
367        let stdio = McpServer::stdio("npx")
368            .args(["-y", "@modelcontextprotocol/server-postgres"])
369            .build();
370        match stdio {
371            McpServer::Stdio(s) => {
372                assert_eq!(s.command, "npx");
373                assert_eq!(s.args, vec!["-y", "@modelcontextprotocol/server-postgres"]);
374            }
375            _ => panic!("Expected Stdio"),
376        }
377
378        let sse = McpServer::sse("http://example.com/sse")
379            .header("Auth", "token")
380            .build();
381        match sse {
382            McpServer::Sse(s) => {
383                assert_eq!(s.url, "http://example.com/sse");
384                assert_eq!(s.headers.as_ref().unwrap()["Auth"], "token");
385            }
386            _ => panic!("Expected Sse"),
387        }
388
389        let http = McpServer::http("http://example.com/http")
390            .header("Auth", "token")
391            .timeout(10.0)
392            .build();
393        match http {
394            McpServer::Http(s) => {
395                assert_eq!(s.url, "http://example.com/http");
396                assert_eq!(s.headers.as_ref().unwrap()["Auth"], "token");
397                assert!((s.timeout - 10.0).abs() < f64::EPSILON);
398            }
399            _ => panic!("Expected Http"),
400        }
401    }
402}
agy_bridge/config/mcp.rs

agy_bridge/config/
mcp.rs
