agentic_core/tool/mcp/
pool.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use reqwest::Url;
5use serde::{Deserialize, Serialize};
6
7use super::client::McpClient;
8use crate::types::tools::McpToolParam;
9
10// Hostnames configured here are a trust boundary. The HTTP client resolves a
11// configured name once per connection and pins all returned addresses for that
12// transport. Only add names whose DNS records are controlled by a trusted
13// administrator.
14const MCP_ALLOWED_HOSTS_ENV: &str = "AGENTIC_MCP_ALLOWED_HOSTS";
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16#[serde(untagged, deny_unknown_fields)]
17pub enum McpServerEntry {
18    Http {
19        url: String,
20        #[serde(default, alias = "http_headers", skip_serializing_if = "Option::is_none")]
21        headers: Option<HashMap<String, String>>,
22        #[serde(default, skip_serializing_if = "Option::is_none")]
23        allowed_tools: Option<Vec<String>>,
24        #[serde(default, skip_serializing_if = "Option::is_none")]
25        require_approval: Option<String>,
26    },
27    Stdio {
28        command: String,
29        #[serde(default)]
30        args: Vec<String>,
31        #[serde(default, skip_serializing_if = "Option::is_none")]
32        env: Option<HashMap<String, String>>,
33        #[serde(default, skip_serializing_if = "Option::is_none")]
34        cwd: Option<String>,
35        #[serde(default, skip_serializing_if = "Option::is_none")]
36        allowed_tools: Option<Vec<String>>,
37        #[serde(default, skip_serializing_if = "Option::is_none")]
38        require_approval: Option<String>,
39    },
40}
41
42impl McpServerEntry {
43    #[must_use]
44    pub fn allowed_tools(&self) -> Option<&[String]> {
45        match self {
46            Self::Http { allowed_tools, .. } | Self::Stdio { allowed_tools, .. } => allowed_tools.as_deref(),
47        }
48    }
49
50    #[must_use]
51    pub fn require_approval(&self) -> Option<&str> {
52        match self {
53            Self::Http { require_approval, .. } | Self::Stdio { require_approval, .. } => require_approval.as_deref(),
54        }
55    }
56}
57
58#[derive(Default)]
59pub struct McpClientPool {
60    clients: HashMap<String, Arc<McpClient>>,
61    connection_errors: HashMap<String, String>,
62}
63
64impl McpClientPool {
65    pub async fn from_params(params: &[McpToolParam]) -> Self {
66        Self::from_params_with_allowed_hosts(params, &allowed_hosts_from_env()).await
67    }
68
69    pub async fn from_params_with_allowed_hosts(params: &[McpToolParam], allowed_hosts: &[String]) -> Self {
70        let servers: HashMap<String, McpServerEntry> = params
71            .iter()
72            .filter_map(|param| server_entry_from_param(param, allowed_hosts))
73            .collect();
74        Self::from_config(servers).await
75    }
76
77    pub async fn from_config(servers: HashMap<String, McpServerEntry>) -> Self {
78        let mut clients = HashMap::with_capacity(servers.len());
79        let mut connection_errors = HashMap::new();
80
81        for (server_label, entry) in servers {
82            let result = match entry {
83                McpServerEntry::Http { url, headers, .. } => McpClient::connect(&url, headers).await,
84                McpServerEntry::Stdio {
85                    command,
86                    args,
87                    env,
88                    cwd,
89                    ..
90                } => McpClient::connect_stdio(&command, &args, env.as_ref(), cwd.as_deref()).await,
91            };
92
93            match result {
94                Ok(client) => {
95                    clients.insert(server_label, Arc::new(client));
96                }
97                Err(error) => {
98                    let error_message = error.to_string();
99                    tracing::warn!(
100                        server_label = %server_label,
101                        error = %error_message,
102                        "failed to connect MCP server from config"
103                    );
104                    connection_errors.insert(server_label, error_message);
105                }
106            }
107        }
108
109        Self {
110            clients,
111            connection_errors,
112        }
113    }
114
115    #[must_use]
116    pub fn get(&self, server_label: &str) -> Option<&Arc<McpClient>> {
117        self.clients.get(server_label)
118    }
119
120    #[must_use]
121    pub fn connection_error(&self, server_label: &str) -> Option<&str> {
122        self.connection_errors.get(server_label).map(String::as_str)
123    }
124}
125
126fn server_entry_from_param(param: &McpToolParam, allowed_hosts: &[String]) -> Option<(String, McpServerEntry)> {
127    let Some(server_label) = clean_string(Some(&param.server_label)) else {
128        tracing::debug!("MCP tool param has no server_label");
129        return None;
130    };
131
132    if let Some(url) = clean_string(param.server_url.as_deref()) {
133        let url = match validate_request_server_url_with_allowed_hosts(&url, allowed_hosts) {
134            Ok(url) => url,
135            Err(reason) => {
136                tracing::warn!(server_label, url, reason, "MCP tool param server_url rejected");
137                return None;
138            }
139        };
140
141        return Some((
142            server_label,
143            McpServerEntry::Http {
144                url,
145                headers: request_headers(param),
146                allowed_tools: None,
147                require_approval: None,
148            },
149        ));
150    }
151
152    tracing::warn!(server_label, "MCP tool param has no server_url");
153    None
154}
155
156fn request_headers(param: &McpToolParam) -> Option<HashMap<String, String>> {
157    let mut headers = param.headers.clone().unwrap_or_default();
158    if let Some(authorization) = clean_string(param.authorization.as_deref()) {
159        headers.insert("Authorization".to_owned(), format!("Bearer {authorization}"));
160    }
161    (!headers.is_empty()).then_some(headers)
162}
163
164fn validate_request_server_url_with_allowed_hosts(value: &str, allowed_hosts: &[String]) -> Result<String, String> {
165    let url = Url::parse(value).map_err(|error| format!("invalid URL: {error}"))?;
166    match url.scheme() {
167        "http" | "https" => {}
168        _ => return Err("URL scheme must be http or https".to_owned()),
169    }
170
171    if !url.username().is_empty() || url.password().is_some() {
172        return Err("URL must not include credentials".to_owned());
173    }
174
175    let host = url.host().ok_or_else(|| "URL must include a host".to_owned())?;
176    if is_allowed_request_host(&host, allowed_hosts) {
177        return Ok(value.to_owned());
178    }
179
180    Err(format!(
181        "MCP server_url host is not allowed; set {MCP_ALLOWED_HOSTS_ENV} to allow it"
182    ))
183}
184
185fn is_allowed_request_host(host: &url::Host<&str>, allowed_hosts: &[String]) -> bool {
186    match host {
187        url::Host::Domain(host) => host.eq_ignore_ascii_case("localhost") || host_allowed(host, allowed_hosts),
188        url::Host::Ipv4(address) => address.is_loopback() || host_allowed(&address.to_string(), allowed_hosts),
189        url::Host::Ipv6(address) => address.is_loopback() || host_allowed(&address.to_string(), allowed_hosts),
190    }
191}
192
193fn host_allowed(host: &str, allowed_hosts: &[String]) -> bool {
194    allowed_hosts
195        .iter()
196        .any(|allowed_host| allowed_host.eq_ignore_ascii_case(host))
197}
198
199pub(crate) fn allowed_hosts_from_env() -> Vec<String> {
200    parse_allowed_hosts(&std::env::var(MCP_ALLOWED_HOSTS_ENV).unwrap_or_default())
201}
202
203fn parse_allowed_hosts(value: &str) -> Vec<String> {
204    value
205        .split(',')
206        .map(str::trim)
207        .filter(|host| !host.is_empty())
208        .map(str::to_owned)
209        .collect()
210}
211
212fn clean_string(value: Option<&str>) -> Option<String> {
213    value
214        .map(str::trim)
215        .filter(|value| !value.is_empty())
216        .map(str::to_owned)
217}
218
219#[cfg(test)]
220mod tests {
221    use super::{
222        McpServerEntry, parse_allowed_hosts, server_entry_from_param, validate_request_server_url_with_allowed_hosts,
223    };
224    use crate::types::tools::McpToolParam;
225
226    #[test]
227    fn mcp_server_entry_deserializes_http_config() {
228        let entry = serde_json::from_value::<McpServerEntry>(serde_json::json!({
229            "url": "http://localhost:9000",
230            "headers": {"Authorization": "Bearer token"},
231            "allowed_tools": ["say_hello", "sum"],
232            "require_approval": "never"
233        }))
234        .unwrap();
235
236        match entry {
237            McpServerEntry::Http {
238                url,
239                headers,
240                allowed_tools,
241                require_approval,
242            } => {
243                assert_eq!(url, "http://localhost:9000");
244                assert_eq!(headers.unwrap()["Authorization"], "Bearer token");
245                assert_eq!(allowed_tools.unwrap(), ["say_hello", "sum"]);
246                assert_eq!(require_approval.as_deref(), Some("never"));
247            }
248            McpServerEntry::Stdio { .. } => panic!("expected HTTP MCP config"),
249        }
250    }
251
252    #[test]
253    fn mcp_server_entry_deserializes_stdio_config() {
254        let entry = serde_json::from_value::<McpServerEntry>(serde_json::json!({
255            "command": "python3",
256            "args": ["/tmp/server.py"],
257            "env": {"TOKEN": "secret"},
258            "cwd": "/tmp"
259        }))
260        .unwrap();
261
262        match entry {
263            McpServerEntry::Stdio {
264                command,
265                args,
266                env,
267                cwd,
268                allowed_tools,
269                require_approval,
270            } => {
271                assert_eq!(command, "python3");
272                assert_eq!(args, vec!["/tmp/server.py".to_owned()]);
273                assert_eq!(env.unwrap()["TOKEN"], "secret");
274                assert_eq!(cwd.as_deref(), Some("/tmp"));
275                assert!(allowed_tools.is_none());
276                assert!(require_approval.is_none());
277            }
278            McpServerEntry::Http { .. } => panic!("expected stdio MCP config"),
279        }
280    }
281
282    #[test]
283    fn request_server_url_allows_loopback_http() {
284        let url = validate_request_server_url_with_allowed_hosts("http://127.0.0.1:8000/mcp", &[]).unwrap();
285        assert_eq!(url, "http://127.0.0.1:8000/mcp");
286    }
287
288    #[test]
289    fn request_server_url_allows_ipv6_loopback_http() {
290        let url = validate_request_server_url_with_allowed_hosts("http://[::1]:8000/mcp", &[]).unwrap();
291        assert_eq!(url, "http://[::1]:8000/mcp");
292    }
293
294    #[test]
295    fn request_server_url_rejects_unallowlisted_host() {
296        let error = validate_request_server_url_with_allowed_hosts("http://169.254.169.254/mcp", &[]).unwrap_err();
297        assert!(error.contains("not allowed"));
298    }
299
300    #[test]
301    fn request_server_url_uses_supplied_allowlist() {
302        assert!(
303            validate_request_server_url_with_allowed_hosts(
304                "https://mcp.example.com/mcp",
305                &["mcp.example.com".to_owned()]
306            )
307            .is_ok()
308        );
309        assert!(validate_request_server_url_with_allowed_hosts("https://mcp.example.com/mcp", &[]).is_err());
310    }
311
312    #[test]
313    fn request_params_ignore_stdio_fields_without_configuring_transport() {
314        let param = serde_json::from_value::<McpToolParam>(serde_json::json!({
315            "server_label": "repo",
316            "command": "python3",
317            "args": ["/tmp/server.py"]
318        }))
319        .unwrap();
320
321        assert!(server_entry_from_param(&param, &[]).is_none());
322    }
323
324    #[test]
325    fn allowed_host_parser_trims_and_discards_empty_entries() {
326        assert_eq!(
327            parse_allowed_hosts(" Example.COM, ,api.test "),
328            vec!["Example.COM", "api.test"]
329        );
330    }
331}
agentic_core/tool/mcp/pool.rs

agentic_core/tool/mcp/
pool.rs
