agentic_core/tool/mcp/
pool.rs

1use std::collections::HashMap;
2use std::sync::{Arc, OnceLock};
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
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(untagged)]
18pub enum McpServerEntry {
19    Http {
20        url: String,
21        #[serde(default, skip_serializing_if = "Option::is_none")]
22        headers: Option<HashMap<String, String>>,
23    },
24    Stdio {
25        command: String,
26        #[serde(default)]
27        args: Vec<String>,
28        #[serde(default, skip_serializing_if = "Option::is_none")]
29        env: Option<HashMap<String, String>>,
30        #[serde(default, skip_serializing_if = "Option::is_none")]
31        cwd: Option<String>,
32    },
33}
34
35#[derive(Default)]
36pub struct McpClientPool {
37    clients: HashMap<String, Arc<McpClient>>,
38    connection_errors: HashMap<String, String>,
39}
40
41impl McpClientPool {
42    pub async fn from_params(params: &[McpToolParam]) -> Self {
43        let servers: HashMap<String, McpServerEntry> = params.iter().filter_map(server_entry_from_param).collect();
44        Self::from_config(servers).await
45    }
46
47    pub async fn from_config(servers: HashMap<String, McpServerEntry>) -> Self {
48        let mut clients = HashMap::with_capacity(servers.len());
49        let mut connection_errors = HashMap::new();
50
51        for (server_label, entry) in servers {
52            let result = match entry {
53                McpServerEntry::Http { url, headers } => McpClient::connect(&url, headers).await,
54                McpServerEntry::Stdio {
55                    command,
56                    args,
57                    env,
58                    cwd,
59                } => McpClient::connect_stdio(&command, &args, env.as_ref(), cwd.as_deref()).await,
60            };
61
62            match result {
63                Ok(client) => {
64                    clients.insert(server_label, Arc::new(client));
65                }
66                Err(error) => {
67                    let error_message = error.to_string();
68                    tracing::warn!(
69                        server_label = %server_label,
70                        error = %error_message,
71                        "failed to connect MCP server from config"
72                    );
73                    connection_errors.insert(server_label, error_message);
74                }
75            }
76        }
77
78        Self {
79            clients,
80            connection_errors,
81        }
82    }
83
84    #[must_use]
85    pub fn get(&self, server_label: &str) -> Option<&Arc<McpClient>> {
86        self.clients.get(server_label)
87    }
88
89    #[must_use]
90    pub fn connection_error(&self, server_label: &str) -> Option<&str> {
91        self.connection_errors.get(server_label).map(String::as_str)
92    }
93}
94
95fn server_entry_from_param(param: &McpToolParam) -> Option<(String, McpServerEntry)> {
96    let Some(server_label) = clean_string(Some(&param.server_label)) else {
97        tracing::debug!("MCP tool param has no server_label");
98        return None;
99    };
100
101    if let Some(url) = clean_string(param.server_url.as_deref()) {
102        let url = match validate_request_server_url(&url) {
103            Ok(url) => url,
104            Err(reason) => {
105                tracing::warn!(server_label, url, reason, "MCP tool param server_url rejected");
106                return None;
107            }
108        };
109
110        return Some((
111            server_label,
112            McpServerEntry::Http {
113                url,
114                headers: request_headers(param),
115            },
116        ));
117    }
118
119    tracing::warn!(server_label, "MCP tool param has no server_url");
120    None
121}
122
123fn request_headers(param: &McpToolParam) -> Option<HashMap<String, String>> {
124    let mut headers = param.headers.clone().unwrap_or_default();
125    if let Some(authorization) = clean_string(param.authorization.as_deref()) {
126        headers.insert("Authorization".to_owned(), format!("Bearer {authorization}"));
127    }
128    (!headers.is_empty()).then_some(headers)
129}
130
131fn validate_request_server_url(value: &str) -> Result<String, String> {
132    let url = Url::parse(value).map_err(|error| format!("invalid URL: {error}"))?;
133    match url.scheme() {
134        "http" | "https" => {}
135        _ => return Err("URL scheme must be http or https".to_owned()),
136    }
137
138    if !url.username().is_empty() || url.password().is_some() {
139        return Err("URL must not include credentials".to_owned());
140    }
141
142    let host = url.host().ok_or_else(|| "URL must include a host".to_owned())?;
143    if is_allowed_request_host(&host) {
144        return Ok(value.to_owned());
145    }
146
147    Err(format!(
148        "MCP server_url host is not allowed; set {MCP_ALLOWED_HOSTS_ENV} to allow it"
149    ))
150}
151
152fn is_allowed_request_host(host: &url::Host<&str>) -> bool {
153    match host {
154        url::Host::Domain(host) => host.eq_ignore_ascii_case("localhost") || host_allowed_by_env(host),
155        url::Host::Ipv4(address) => address.is_loopback() || host_allowed_by_env(&address.to_string()),
156        url::Host::Ipv6(address) => address.is_loopback() || host_allowed_by_env(&address.to_string()),
157    }
158}
159
160fn host_allowed_by_env(host: &str) -> bool {
161    allowed_hosts()
162        .iter()
163        .any(|allowed_host| allowed_host.eq_ignore_ascii_case(host))
164}
165
166fn allowed_hosts() -> &'static [String] {
167    static ALLOWED_HOSTS: OnceLock<Vec<String>> = OnceLock::new();
168    ALLOWED_HOSTS.get_or_init(|| parse_allowed_hosts(&std::env::var(MCP_ALLOWED_HOSTS_ENV).unwrap_or_default()))
169}
170
171fn parse_allowed_hosts(value: &str) -> Vec<String> {
172    value
173        .split(',')
174        .map(str::trim)
175        .filter(|host| !host.is_empty())
176        .map(str::to_owned)
177        .collect()
178}
179
180fn clean_string(value: Option<&str>) -> Option<String> {
181    value
182        .map(str::trim)
183        .filter(|value| !value.is_empty())
184        .map(str::to_owned)
185}
186
187#[cfg(test)]
188mod tests {
189    use super::{McpServerEntry, parse_allowed_hosts, server_entry_from_param, validate_request_server_url};
190    use crate::types::tools::McpToolParam;
191
192    #[test]
193    fn mcp_server_entry_deserializes_http_config() {
194        let entry = serde_json::from_value::<McpServerEntry>(serde_json::json!({
195            "url": "http://localhost:9000",
196            "headers": {"Authorization": "Bearer token"}
197        }))
198        .unwrap();
199
200        match entry {
201            McpServerEntry::Http { url, headers } => {
202                assert_eq!(url, "http://localhost:9000");
203                assert_eq!(headers.unwrap()["Authorization"], "Bearer token");
204            }
205            McpServerEntry::Stdio { .. } => panic!("expected HTTP MCP config"),
206        }
207    }
208
209    #[test]
210    fn mcp_server_entry_deserializes_stdio_config() {
211        let entry = serde_json::from_value::<McpServerEntry>(serde_json::json!({
212            "command": "python3",
213            "args": ["/tmp/server.py"],
214            "env": {"TOKEN": "secret"},
215            "cwd": "/tmp"
216        }))
217        .unwrap();
218
219        match entry {
220            McpServerEntry::Stdio {
221                command,
222                args,
223                env,
224                cwd,
225            } => {
226                assert_eq!(command, "python3");
227                assert_eq!(args, vec!["/tmp/server.py".to_owned()]);
228                assert_eq!(env.unwrap()["TOKEN"], "secret");
229                assert_eq!(cwd.as_deref(), Some("/tmp"));
230            }
231            McpServerEntry::Http { .. } => panic!("expected stdio MCP config"),
232        }
233    }
234
235    #[test]
236    fn request_server_url_allows_loopback_http() {
237        let url = validate_request_server_url("http://127.0.0.1:8000/mcp").unwrap();
238        assert_eq!(url, "http://127.0.0.1:8000/mcp");
239    }
240
241    #[test]
242    fn request_server_url_allows_ipv6_loopback_http() {
243        let url = validate_request_server_url("http://[::1]:8000/mcp").unwrap();
244        assert_eq!(url, "http://[::1]:8000/mcp");
245    }
246
247    #[test]
248    fn request_server_url_rejects_unallowlisted_host() {
249        let error = validate_request_server_url("http://169.254.169.254/mcp").unwrap_err();
250        assert!(error.contains("not allowed"));
251    }
252
253    #[test]
254    fn request_params_ignore_stdio_fields_without_configuring_transport() {
255        let param = serde_json::from_value::<McpToolParam>(serde_json::json!({
256            "server_label": "repo",
257            "command": "python3",
258            "args": ["/tmp/server.py"]
259        }))
260        .unwrap();
261
262        assert!(server_entry_from_param(&param).is_none());
263    }
264
265    #[test]
266    fn allowed_host_parser_trims_and_discards_empty_entries() {
267        assert_eq!(
268            parse_allowed_hosts(" Example.COM, ,api.test "),
269            vec!["Example.COM", "api.test"]
270        );
271    }
272}
agentic_core/tool/mcp/pool.rs

agentic_core/tool/mcp/
pool.rs
