1use mcp_utils::client::{ClientOptions, ElicitationSink, Transport};
2use mcp_utils::gateway::{McpGateway, ServerSpec};
3use mcp_utils::model::Implementation;
4use mcp_utils::server::{McpServer, ServeError, UnixSocket};
5use utils::{SettingsStore, variables::Vars};
6
7use super::config::ServerSource;
8use super::{McpCatalog, McpConfig, McpServerConfig, ParseError, ToolExposure, ToolFilter};
9use crate::core::AgentDeps;
10
11use std::collections::HashMap;
12use std::path::{Path, PathBuf};
13use thiserror::Error;
14
15pub fn mcp(root_dir: impl AsRef<Path>) -> McpBuilder {
16 McpBuilder::new(root_dir)
17}
18
19#[derive(Clone)]
20pub struct RuntimeServices {
21 pub root_dir: PathBuf,
22 pub agent_deps: AgentDeps,
23 pub deferred_tools_socket: Option<PathBuf>,
24}
25
26pub struct McpBuilder {
27 servers: Vec<ConfiguredServer>,
28 factories: HashMap<String, ServerFactory>,
29 root_dir: PathBuf,
30 agent_deps: AgentDeps,
31 vars: Vars,
32 tool_filter: ToolFilter,
33 elicitation: Option<ElicitationSink>,
34}
35
36pub type ServerFactoryError = Box<dyn std::error::Error + Send + Sync>;
37
38#[derive(Debug, Error)]
39pub enum McpSpawnError {
40 #[error(transparent)]
41 Serve(#[from] ServeError),
42 #[error("No factory is registered for in-memory MCP server '{0}'")]
43 InMemoryFactoryNotFound(String),
44 #[error("Failed to create in-memory MCP server '{name}': {source}")]
45 ServerFactory { name: String, source: ServerFactoryError },
46}
47
48impl McpBuilder {
49 pub fn new(root_dir: impl AsRef<Path>) -> Self {
50 let mut vars = Vars::new().with("WORKSPACE", root_dir.as_ref().to_string_lossy().into_owned());
51
52 if let Some(store) = SettingsStore::new("AETHER_HOME", ".aether") {
53 vars.insert("AETHER_HOME", store.home().to_string_lossy().into_owned());
54 }
55
56 Self {
57 servers: Vec::new(),
58 factories: HashMap::new(),
59 root_dir: root_dir.as_ref().to_path_buf(),
60 agent_deps: AgentDeps::default(),
61 vars,
62 tool_filter: ToolFilter::default(),
63 elicitation: None,
64 }
65 }
66
67 pub fn with_config(mut self, config: McpConfig) -> Result<Self, ParseError> {
68 for (name, server) in config.servers {
69 let server = ConfiguredServer::from_config(name, server, &self.vars, &self.root_dir)?;
70 self.insert(server);
71 }
72 Ok(self)
73 }
74
75 pub fn with_servers(self, servers: Vec<ServerSpec>) -> Self {
76 servers.into_iter().fold(self, |builder, spec| builder.with_server(spec, ToolExposure::ModelVisible))
77 }
78
79 pub fn with_server(mut self, spec: ServerSpec, exposure: ToolExposure) -> Self {
80 self.insert(ConfiguredServer { name: spec.name, exposure, source: ServerSource::Ready(spec.transport) });
81 self
82 }
83
84 pub fn with_tool_filter(mut self, filter: ToolFilter) -> Self {
85 self.tool_filter = filter;
86 self
87 }
88
89 pub fn with_elicitation(mut self, sink: ElicitationSink) -> Self {
90 self.elicitation = Some(sink);
91 self
92 }
93
94 pub fn register_in_memory_server(
95 mut self,
96 name: impl Into<String>,
97 factory: impl Fn(Vec<String>, &RuntimeServices) -> Result<McpServer, ServerFactoryError> + Send + Sync + 'static,
98 ) -> Self {
99 self.factories.insert(name.into(), Box::new(factory));
100 self
101 }
102
103 pub fn with_agent_deps(mut self, deps: AgentDeps) -> Self {
104 self.agent_deps = deps;
105 self
106 }
107
108 pub fn spawn(self) -> Result<McpCatalog, McpSpawnError> {
109 let exposures: HashMap<String, ToolExposure> =
110 self.servers.iter().map(|server| (server.name.clone(), server.exposure.clone())).collect();
111 let defers_tools = exposures.values().any(ToolExposure::defers_tools);
112 let socket = defers_tools.then(UnixSocket::bind).transpose()?;
113 let options = self.client_options();
114 let services = RuntimeServices {
115 root_dir: self.root_dir,
116 agent_deps: self.agent_deps,
117 deferred_tools_socket: socket.as_ref().map(|socket| socket.path().to_path_buf()),
118 };
119 let servers = self
120 .servers
121 .into_iter()
122 .map(|server| server.build(&self.factories, &services))
123 .collect::<Result<Vec<_>, _>>()?;
124
125 Ok(McpCatalog::new(McpGateway::new(servers, &options), self.tool_filter, exposures, socket))
126 }
127
128 fn insert(&mut self, server: ConfiguredServer) {
129 match self.servers.iter_mut().find(|existing| existing.name == server.name) {
130 Some(existing) => *existing = server,
131 None => self.servers.push(server),
132 }
133 }
134
135 fn client_options(&self) -> ClientOptions {
136 let mut options =
137 ClientOptions::default().implementation(Implementation::new("aether", env!("CARGO_PKG_VERSION")));
138 if let Some(store) = self.agent_deps.oauth_credential_store.clone() {
139 options = options.oauth_store(store);
140 }
141 if let Some(sink) = self.elicitation.clone() {
142 options = options.elicitation(sink);
143 }
144 options
145 }
146}
147
148type ServerFactory = Box<dyn Fn(Vec<String>, &RuntimeServices) -> Result<McpServer, ServerFactoryError> + Send + Sync>;
149
150struct ConfiguredServer {
151 name: String,
152 exposure: ToolExposure,
153 source: ServerSource,
154}
155
156impl ConfiguredServer {
157 fn from_config(name: String, config: McpServerConfig, vars: &Vars, root_dir: &Path) -> Result<Self, ParseError> {
158 let exposure = config.defer_tools().clone();
159 Ok(Self { name, exposure, source: config.resolve(vars, root_dir)? })
160 }
161
162 fn build(
163 self,
164 factories: &HashMap<String, ServerFactory>,
165 services: &RuntimeServices,
166 ) -> Result<ServerSpec, McpSpawnError> {
167 let transport = match self.source {
168 ServerSource::Ready(transport) => transport,
169 ServerSource::InMemory(args) => {
170 let factory = factories
171 .get(&self.name)
172 .ok_or_else(|| McpSpawnError::InMemoryFactoryNotFound(self.name.clone()))?;
173 let server = factory(args, services)
174 .map_err(|source| McpSpawnError::ServerFactory { name: self.name.clone(), source })?;
175 Transport::InProcess(server)
176 }
177 };
178 Ok(ServerSpec { name: self.name, transport })
179 }
180}
181
182#[cfg(test)]
183mod tests {
184 use super::*;
185 use aether_auth::{FakeOAuthCredentialStore, OAuthCredentialStorage};
186 use mcp_utils::testing::FakeMcpServer;
187 use std::sync::atomic::{AtomicUsize, Ordering};
188 use std::sync::{Arc, Mutex};
189 use utils::mcp_status::McpServerStatus;
190
191 fn config(json: &str) -> McpConfig {
192 McpConfig::from_json(json).unwrap()
193 }
194
195 #[tokio::test]
196 async fn in_memory_factory_runs_once_at_spawn_with_runtime_services() {
197 let calls = Arc::new(AtomicUsize::new(0));
198 let received = Arc::new(Mutex::new(None::<RuntimeServices>));
199 let factory_calls = Arc::clone(&calls);
200 let factory_received = Arc::clone(&received);
201 let oauth_store: Arc<dyn OAuthCredentialStorage> = Arc::new(FakeOAuthCredentialStore::new());
202 let deps = AgentDeps::new(Arc::clone(&oauth_store), None);
203
204 let builder = McpBuilder::new("/workspace")
205 .with_agent_deps(deps)
206 .register_in_memory_server("test", move |args, services| {
207 factory_calls.fetch_add(1, Ordering::SeqCst);
208 assert_eq!(args, ["--root", "/workspace/tools"]);
209 *factory_received.lock().unwrap() = Some(services.clone());
210 Ok(FakeMcpServer::new().into())
211 })
212 .with_config(config(r#"{"servers":{"test":{"type":"in-memory","args":["--root","${WORKSPACE}/tools"]}}}"#))
213 .unwrap();
214 assert_eq!(calls.load(Ordering::SeqCst), 0);
215
216 let _spawned = builder.spawn().unwrap();
217 assert_eq!(calls.load(Ordering::SeqCst), 1);
218 let services = received.lock().unwrap().clone().expect("factory received runtime services");
219 assert_eq!(services.root_dir, PathBuf::from("/workspace"));
220 assert!(Arc::ptr_eq(
221 services.agent_deps.oauth_credential_store.as_ref().expect("factory received agent dependencies"),
222 &oauth_store,
223 ));
224 assert!(services.deferred_tools_socket.is_none());
225 }
226
227 #[tokio::test]
228 async fn deferred_gateway_is_bound_before_in_memory_factories_run() {
229 let received = Arc::new(Mutex::new(None::<RuntimeServices>));
230 let factory_received = Arc::clone(&received);
231 let mcp = McpBuilder::new("/workspace")
232 .register_in_memory_server("test", move |_, services| {
233 *factory_received.lock().unwrap() = Some(services.clone());
234 Ok(FakeMcpServer::new().into())
235 })
236 .with_config(config(r#"{"servers":{"test":{"type":"in-memory","deferTools":true}}}"#))
237 .unwrap()
238 .spawn()
239 .unwrap();
240
241 let services = received.lock().unwrap().clone().expect("factory received runtime services");
242 let socket = services.deferred_tools_socket.expect("factory receives the deferred tools socket");
243 assert_eq!(socket, mcp.deferred_tools_socket().expect("the session serves deferred tools"));
244 assert!(socket.exists());
245 }
246
247 #[tokio::test]
248 async fn shutting_down_a_session_stops_the_deferred_socket_and_the_servers_for_every_clone() {
249 let lazy = ServerSpec::in_process("lazy", FakeMcpServer::new());
250 let mcp = McpBuilder::new("/workspace").with_server(lazy, ToolExposure::deferred_all()).spawn().unwrap();
251 mcp.ready().await;
252 let held = mcp.clone();
253 let socket = mcp.deferred_tools_socket().expect("the session serves deferred tools").to_path_buf();
254
255 mcp.shutdown().await;
256
257 assert!(!socket.exists());
258 assert_eq!(held.statuses().len(), 0);
259 }
260
261 #[tokio::test]
262 async fn spawned_in_memory_servers_connect_in_the_background() {
263 let mcp = McpBuilder::new("/workspace")
264 .register_in_memory_server("test", |_, _| Ok(FakeMcpServer::math().into()))
265 .with_config(config(r#"{"servers":{"test":{"type":"in-memory"}}}"#))
266 .unwrap()
267 .spawn()
268 .unwrap();
269
270 mcp.ready().await;
271
272 assert_eq!(mcp.tool_definitions()[0].name, "test__add_numbers");
273 }
274
275 #[tokio::test]
276 async fn ready_servers_and_config_servers_are_added_in_order() {
277 let mcp = McpBuilder::new("/workspace")
278 .register_in_memory_server("configured", |_, _| Ok(FakeMcpServer::new().into()))
279 .with_servers(vec![ServerSpec::in_process("ready", FakeMcpServer::new())])
280 .with_config(config(r#"{"servers":{"configured":{"type":"in-memory"}}}"#))
281 .unwrap()
282 .spawn()
283 .unwrap();
284
285 mcp.ready().await;
286
287 assert_eq!(names(&mcp), ["ready", "configured"]);
288 }
289
290 #[tokio::test]
291 async fn a_server_replaces_an_earlier_one_with_the_same_name_in_its_position() {
292 let replaced = FakeMcpServer::new();
293 let alive = replaced.alive();
294 let replacement = FakeMcpServer::new().with_tool(mcp_utils::testing::FakeTool::new("replacement"));
295 let mcp = McpBuilder::new("/workspace")
296 .with_servers(vec![
297 ServerSpec::in_process("math", replaced),
298 ServerSpec::in_process("docs", FakeMcpServer::new()),
299 ])
300 .with_server(ServerSpec::in_process("math", replacement), ToolExposure::deferred_all())
301 .spawn()
302 .unwrap();
303
304 mcp.ready().await;
305
306 assert!(alive.upgrade().is_none());
307 assert_eq!(names(&mcp), ["math", "docs"]);
308 assert!(mcp.statuses()[0].deferred_tools);
309 assert_eq!(mcp.statuses()[0].status, McpServerStatus::Connected { tool_count: 1 });
310 }
311
312 #[tokio::test]
313 async fn missing_in_memory_factory_fails_at_spawn_with_server_name() {
314 let builder = McpBuilder::new("/workspace")
315 .with_config(config(r#"{"servers":{"custom":{"type":"in-memory"}}}"#))
316 .unwrap();
317
318 let Err(error) = builder.spawn() else {
319 panic!("spawn should reject an unregistered factory");
320 };
321 assert!(matches!(error, McpSpawnError::InMemoryFactoryNotFound(ref server) if server == "custom"));
322 }
323
324 #[tokio::test]
325 async fn rejected_in_memory_args_fail_spawn_with_server_name() {
326 let builder = McpBuilder::new("/workspace")
327 .register_in_memory_server("custom", |args, _| Err(format!("unexpected argument '{}'", args[0]).into()))
328 .with_config(config(r#"{"servers":{"custom":{"type":"in-memory","args":["--bad"]}}}"#))
329 .unwrap();
330
331 let Err(error) = builder.spawn() else {
332 panic!("spawn should surface the factory's error");
333 };
334 assert!(matches!(error, McpSpawnError::ServerFactory { ref name, .. } if name == "custom"));
335 assert!(error.to_string().contains("unexpected argument '--bad'"), "unexpected error: {error}");
336 }
337
338 #[tokio::test]
339 async fn config_args_expand_the_aether_home_var() {
340 let home = SettingsStore::new("AETHER_HOME", ".aether").expect("Aether home resolves").home().to_path_buf();
341
342 let args = spawned_args(
343 "/work",
344 r#"{"servers":{"notes":{"type":"in-memory","args":["--dir","${AETHER_HOME}/skills"]}}}"#,
345 );
346
347 assert_eq!(args, ["--dir".to_string(), home.join("skills").to_string_lossy().into_owned()]);
348 }
349
350 fn names(mcp: &McpCatalog) -> Vec<String> {
351 mcp.statuses().into_iter().map(|status| status.name).collect()
352 }
353
354 fn spawned_args(root_dir: &str, json: &str) -> Vec<String> {
355 let received = Arc::new(Mutex::new(Vec::new()));
356 let factory_received = Arc::clone(&received);
357 McpBuilder::new(root_dir)
358 .register_in_memory_server("notes", move |args, _| {
359 *factory_received.lock().unwrap() = args;
360 Ok(FakeMcpServer::new().into())
361 })
362 .with_config(config(json))
363 .unwrap()
364 .spawn()
365 .unwrap();
366 received.lock().unwrap().clone()
367 }
368}