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
185impl 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}