aether_core/mcp/
config.rs

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

aether_core/mcp/
config.rs
