Skip to main content

aether_core/mcp/
config.rs

1use super::ToolExposure;
2use aether_auth::{SigV4Signer, SigV4Signers};
3use http::header::{AUTHORIZATION, HeaderMap, HeaderName, HeaderValue, InvalidHeaderName, InvalidHeaderValue};
4use mcp_utils::client::{HttpAuth, OAuthClientRegistration, OAuthSettings, Transport};
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7use std::collections::{BTreeMap, HashMap};
8use std::num::NonZeroU16;
9use std::path::Path;
10use utils::variables::{VarError, Vars};
11
12#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, PartialEq)]
13pub struct McpConfig {
14    #[serde(alias = "mcpServers")]
15    pub servers: BTreeMap<String, McpServerConfig>,
16}
17
18#[doc = include_str!("../docs/mcp_server_config.md")]
19#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)]
20#[serde(untagged)]
21pub enum McpServerConfig {
22    Stdio(StdioServerConfig),
23    Remote(RemoteServerConfig),
24    InMemory(InMemoryServerConfig),
25}
26
27#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)]
28#[serde(deny_unknown_fields)]
29pub struct StdioServerConfig {
30    #[serde(rename = "type", default)]
31    pub type_: StdioType,
32
33    pub command: String,
34
35    #[serde(default)]
36    pub args: Vec<String>,
37
38    #[serde(default)]
39    pub env: HashMap<String, String>,
40
41    #[serde(rename = "deferTools", default, skip_serializing_if = "is_model_visible")]
42    pub defer_tools: ToolExposure,
43}
44
45#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, PartialEq)]
46#[serde(rename_all = "camelCase", deny_unknown_fields)]
47pub struct McpOAuthConfig {
48    #[serde(default, skip_serializing_if = "Option::is_none")]
49    pub client_id: Option<String>,
50    #[serde(default, skip_serializing_if = "Option::is_none")]
51    pub client_metadata_url: Option<String>,
52    #[serde(default, skip_serializing_if = "Option::is_none")]
53    pub callback_port: Option<NonZeroU16>,
54}
55
56#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)]
57#[serde(deny_unknown_fields)]
58pub struct McpSigV4Config {
59    pub service: String,
60    pub region: String,
61    #[serde(default, skip_serializing_if = "Option::is_none")]
62    pub profile: Option<String>,
63}
64
65#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)]
66#[serde(deny_unknown_fields)]
67pub struct RemoteServerConfig {
68    #[serde(rename = "type")]
69    pub type_: RemoteType,
70
71    pub url: String,
72
73    #[serde(default)]
74    pub headers: HashMap<String, String>,
75
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub oauth: Option<McpOAuthConfig>,
78
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub sigv4: Option<McpSigV4Config>,
81
82    #[serde(rename = "deferTools", default, skip_serializing_if = "is_model_visible")]
83    pub defer_tools: ToolExposure,
84}
85
86#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)]
87#[serde(deny_unknown_fields)]
88pub struct InMemoryServerConfig {
89    #[serde(rename = "type")]
90    pub type_: InMemoryType,
91
92    #[serde(default)]
93    pub args: Vec<String>,
94
95    #[serde(rename = "deferTools", default, skip_serializing_if = "is_model_visible")]
96    pub defer_tools: ToolExposure,
97}
98
99#[derive(Debug)]
100pub(super) enum ResolvedServer {
101    Transport(Transport),
102    InMemory(Vec<String>),
103}
104
105#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, JsonSchema, PartialEq)]
106pub enum StdioType {
107    #[default]
108    #[serde(rename = "stdio")]
109    Stdio,
110}
111
112#[derive(Debug, Clone, Copy, Deserialize, Serialize, JsonSchema, PartialEq)]
113pub enum RemoteType {
114    #[serde(rename = "http")]
115    Http,
116}
117
118#[derive(Debug, Clone, Copy, Deserialize, Serialize, JsonSchema, PartialEq)]
119pub enum InMemoryType {
120    #[serde(rename = "in-memory")]
121    InMemory,
122}
123
124#[derive(Debug, thiserror::Error)]
125pub enum ParseError {
126    #[error("Failed to read config file: {0}")]
127    IoError(#[from] std::io::Error),
128
129    #[error("Invalid JSON: {0}")]
130    JsonError(#[from] serde_json::Error),
131
132    #[error("Variable expansion failed: {0}")]
133    VarError(#[from] VarError),
134
135    #[error("Invalid HTTP header name {name:?}: {source}")]
136    InvalidHeaderName {
137        name: String,
138        #[source]
139        source: InvalidHeaderName,
140    },
141
142    #[error("Invalid HTTP header value for {name:?}: {source}")]
143    InvalidHeaderValue {
144        name: String,
145        #[source]
146        source: InvalidHeaderValue,
147    },
148
149    #[error("A server can set at most one of `oauth`, `sigv4` or an `Authorization` header")]
150    ConflictingAuth,
151}
152
153pub fn http_transport(url: String, headers: HeaderMap) -> Transport {
154    let auth = default_auth(&headers);
155    Transport::Http { url, headers, auth }
156}
157
158impl McpConfig {
159    pub fn new(servers: BTreeMap<String, McpServerConfig>) -> Self {
160        Self { servers }
161    }
162
163    pub fn from_json_file(path: impl AsRef<Path>) -> Result<Self, ParseError> {
164        let content = std::fs::read_to_string(path)?;
165        Self::from_json(&content)
166    }
167
168    pub fn from_json(json: &str) -> Result<Self, ParseError> {
169        Ok(serde_json::from_str(json)?)
170    }
171
172    pub fn is_empty(&self) -> bool {
173        self.servers.is_empty()
174    }
175
176    pub fn defer_all_tools(&mut self) {
177        for exposure in self.servers.values_mut().map(McpServerConfig::defer_tools_mut) {
178            if !exposure.defers_tools() {
179                *exposure = ToolExposure::deferred_all();
180            }
181        }
182    }
183}
184
185/// Merges configs in order; a server named again replaces the earlier definition.
186impl FromIterator<McpConfig> for McpConfig {
187    fn from_iter<T: IntoIterator<Item = McpConfig>>(configs: T) -> Self {
188        Self::new(configs.into_iter().flat_map(|config| config.servers).collect())
189    }
190}
191
192impl McpServerConfig {
193    pub fn defer_tools(&self) -> &ToolExposure {
194        match self {
195            McpServerConfig::Stdio(config) => &config.defer_tools,
196            McpServerConfig::Remote(config) => &config.defer_tools,
197            McpServerConfig::InMemory(config) => &config.defer_tools,
198        }
199    }
200
201    pub(super) fn resolve(self, vars: &Vars, cwd: &Path, signers: &SigV4Signers) -> Result<ResolvedServer, ParseError> {
202        Ok(match self {
203            McpServerConfig::Stdio(config) => ResolvedServer::Transport(config.into_transport(vars, cwd)?),
204            McpServerConfig::Remote(config) => ResolvedServer::Transport(config.into_transport(vars, signers)?),
205            McpServerConfig::InMemory(config) => ResolvedServer::InMemory(expand_all(vars, &config.args)?),
206        })
207    }
208
209    fn defer_tools_mut(&mut self) -> &mut ToolExposure {
210        match self {
211            McpServerConfig::Stdio(config) => &mut config.defer_tools,
212            McpServerConfig::Remote(config) => &mut config.defer_tools,
213            McpServerConfig::InMemory(config) => &mut config.defer_tools,
214        }
215    }
216}
217
218impl StdioServerConfig {
219    pub fn into_transport(self, vars: &Vars, cwd: &Path) -> Result<Transport, ParseError> {
220        let env = self.env.into_iter().map(|(name, value)| Ok((name, vars.expand(&value)?)));
221        Ok(Transport::Stdio {
222            command: vars.expand(&self.command)?,
223            args: expand_all(vars, &self.args)?,
224            env: env.collect::<Result<_, VarError>>()?,
225            cwd: Some(cwd.to_path_buf()),
226        })
227    }
228}
229
230impl McpOAuthConfig {
231    fn into_settings(self) -> OAuthSettings {
232        let callback_port = self.callback_port.map_or(AETHER_OAUTH_CALLBACK_PORT, NonZeroU16::get);
233        let registration = match self {
234            Self { client_id: Some(client_id), .. } => OAuthClientRegistration::PreRegistered(client_id),
235            Self { client_metadata_url: Some(url), .. } => OAuthClientRegistration::ClientMetadata(url),
236            Self { .. } => OAuthClientRegistration::ClientMetadata(AETHER_OAUTH_CLIENT_METADATA_URL.to_string()),
237        };
238        OAuthSettings { registration, callback_port }
239    }
240}
241
242impl McpSigV4Config {
243    fn into_signer(self, vars: &Vars, signers: &SigV4Signers) -> Result<SigV4Signer, VarError> {
244        let profile = self.profile.map(|profile| vars.expand(&profile)).transpose()?;
245        Ok(signers.default_chain(vars.expand(&self.service)?, vars.expand(&self.region)?, profile))
246    }
247}
248
249impl RemoteServerConfig {
250    pub fn into_transport(self, vars: &Vars, signers: &SigV4Signers) -> Result<Transport, ParseError> {
251        let mut headers = HeaderMap::with_capacity(self.headers.len());
252        for (name, value) in self.headers {
253            let header_name = name
254                .parse::<HeaderName>()
255                .map_err(|source| ParseError::InvalidHeaderName { name: name.clone(), source })?;
256            let mut header_value = vars
257                .expand(&value)?
258                .parse::<HeaderValue>()
259                .map_err(|source| ParseError::InvalidHeaderValue { name, source })?;
260            header_value.set_sensitive(true);
261            headers.insert(header_name, header_value);
262        }
263        let has_authorization = headers.contains_key(AUTHORIZATION);
264        let auth = match (self.oauth, self.sigv4) {
265            (Some(oauth), None) if !has_authorization => HttpAuth::OAuth(expand_oauth(vars, oauth)?.into_settings()),
266            (None, Some(sigv4)) if !has_authorization => HttpAuth::SigV4(sigv4.into_signer(vars, signers)?),
267            (None, None) => default_auth(&headers),
268            _ => return Err(ParseError::ConflictingAuth),
269        };
270        Ok(Transport::Http { url: vars.expand(&self.url)?, headers, auth })
271    }
272}
273
274const AETHER_OAUTH_CLIENT_METADATA_URL: &str = "https://aether-agent.io/oauth/client-metadata.json";
275const AETHER_OAUTH_CALLBACK_PORT: u16 = 3118;
276
277fn is_model_visible(exposure: &ToolExposure) -> bool {
278    !exposure.defers_tools()
279}
280
281fn expand_all(vars: &Vars, values: &[String]) -> Result<Vec<String>, VarError> {
282    values.iter().map(|value| vars.expand(value)).collect()
283}
284
285fn default_auth(headers: &HeaderMap) -> HttpAuth {
286    if headers.contains_key(AUTHORIZATION) {
287        HttpAuth::None
288    } else {
289        HttpAuth::OAuth(McpOAuthConfig::default().into_settings())
290    }
291}
292
293fn expand_oauth(vars: &Vars, oauth: McpOAuthConfig) -> Result<McpOAuthConfig, VarError> {
294    let expand = |value: Option<String>| value.map(|value| vars.expand(&value)).transpose();
295    Ok(McpOAuthConfig {
296        client_id: expand(oauth.client_id)?,
297        client_metadata_url: expand(oauth.client_metadata_url)?,
298        callback_port: oauth.callback_port,
299    })
300}