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