aether_core/mcp/
mcp_builder.rs

1use mcp_utils::client::{
2    InMemoryServerSpec, McpClientEvent, McpConfig, McpConnectionDetails, McpError, McpManager, McpServer, McpTransport,
3    OAuthHandlerFactory, PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME, ParseError, RuntimeMcpServer, RuntimeMcpTransport,
4    ToolFilter,
5};
6use mcp_utils::tool_gateway::{AETHER_MCP_IPC_SOCKET, UnixSocketMcpTransport, UnixSocketPath, UnixSocketServer};
7use utils::{SettingsStore, variables::Vars};
8
9use crate::agent_spec::McpConfigSource;
10use crate::core::AgentDeps;
11use crate::events::{AgentCommand, Command};
12
13use super::{
14    gateway_service::GatewayService,
15    mcp_handle::McpHandle,
16    run_mcp_task::{ManagerCommand, run_mcp_task},
17};
18use futures::future::BoxFuture;
19use rmcp::{RoleServer, service::DynService};
20use std::collections::{BTreeMap, BTreeSet, HashMap};
21use std::path::{Path, PathBuf};
22use std::sync::Arc;
23use tokio::{
24    sync::{
25        mpsc::{self, Receiver},
26        watch,
27    },
28    task::JoinHandle,
29};
30
31pub fn mcp(root_dir: impl AsRef<Path>) -> McpBuilder {
32    McpBuilder::new(root_dir)
33}
34
35#[derive(Clone)]
36pub struct RuntimeServices {
37    pub mcp: McpHandle,
38    pub root_dir: PathBuf,
39    pub agent_deps: AgentDeps,
40    pub shell_environment: BTreeMap<String, String>,
41}
42
43pub type ServerFactory = Box<
44    dyn Fn(InMemoryServerSpec, RuntimeServices) -> BoxFuture<'static, Box<dyn DynService<RoleServer>>> + Send + Sync,
45>;
46
47/// Owns a spawned MCP manager. Dropping this value aborts the manager task.
48pub struct McpRuntime {
49    mcp: McpHandle,
50    handle: JoinHandle<()>,
51    agent_sync_handle: Option<JoinHandle<()>>,
52    gateway: Option<UnixSocketServer>,
53}
54
55impl McpRuntime {
56    pub async fn shutdown(&mut self) {
57        self.handle.abort();
58        if let Some(handle) = self.agent_sync_handle.take() {
59            handle.abort();
60            let _ = handle.await;
61        }
62        let _ = (&mut self.handle).await;
63        self.gateway.take();
64    }
65
66    pub fn handle(&self) -> &McpHandle {
67        &self.mcp
68    }
69
70    pub fn gateway_endpoint(&self) -> Option<&Path> {
71        self.gateway.as_ref().map(UnixSocketServer::path)
72    }
73}
74
75impl Drop for McpRuntime {
76    fn drop(&mut self) {
77        self.handle.abort();
78        if let Some(handle) = &self.agent_sync_handle {
79            handle.abort();
80        }
81    }
82}
83
84/// A freshly spawned MCP manager paired with its event stream. Consumers
85/// receive incremental updates over the event stream (starting with an initial
86/// `ServerStatusesChanged` reflecting every configured server in `Connecting`)
87/// and can call [`split`](Self::split) to separate the stream from the
88/// [`McpRuntime`] that keeps the manager task alive.
89pub struct McpSession {
90    runtime: McpRuntime,
91    event_rx: Receiver<McpClientEvent>,
92}
93
94impl McpSession {
95    pub fn handle(&self) -> &McpHandle {
96        self.runtime.handle()
97    }
98
99    pub fn gateway_endpoint(&self) -> Option<&Path> {
100        self.runtime.gateway_endpoint()
101    }
102
103    /// Synchronize this session's current and future tools and instructions with
104    /// one agent. Initial state is sent before this method returns.
105    pub async fn connect_agent(mut self, agent_tx: mpsc::Sender<Command>) -> Self {
106        assert!(self.runtime.agent_sync_handle.is_none(), "an MCP session can only connect one agent");
107        let mut snapshots = self.runtime.handle().subscribe();
108        let initial = snapshots.borrow_and_update().clone();
109        let mut previous_tools = initial.tool_definitions();
110        let mut previous_instructions = initial.model_instructions();
111        if agent_tx.send(Command::agent(AgentCommand::UpdateTools(previous_tools.clone()))).await.is_err() {
112            return self;
113        }
114        for (server, body) in &previous_instructions {
115            if agent_tx
116                .send(Command::agent(AgentCommand::UpdateMcpInstructions {
117                    server: server.clone(),
118                    body: Some(body.clone()),
119                }))
120                .await
121                .is_err()
122            {
123                return self;
124            }
125        }
126
127        let agent_tx = agent_tx.downgrade();
128        self.runtime.agent_sync_handle = Some(tokio::spawn(async move {
129            while snapshots.changed().await.is_ok() {
130                let Some(agent_tx) = agent_tx.upgrade() else {
131                    break;
132                };
133                let snapshot = snapshots.borrow_and_update().clone();
134                let tools = snapshot.tool_definitions();
135                if tools != previous_tools {
136                    if agent_tx.send(Command::agent(AgentCommand::UpdateTools(tools.clone()))).await.is_err() {
137                        break;
138                    }
139                    previous_tools = tools;
140                }
141
142                let instructions = snapshot.model_instructions();
143                let servers = previous_instructions.keys().chain(instructions.keys()).cloned().collect::<BTreeSet<_>>();
144                for server in servers {
145                    let previous = previous_instructions.get(&server);
146                    let next = instructions.get(&server);
147                    if previous != next
148                        && agent_tx
149                            .send(Command::agent(AgentCommand::UpdateMcpInstructions { server, body: next.cloned() }))
150                            .await
151                            .is_err()
152                    {
153                        return;
154                    }
155                }
156                previous_instructions = instructions;
157            }
158        }));
159        self
160    }
161
162    /// Block until the manager finishes bootstrapping every initially-configured
163    /// server, then return the consolidated snapshot. Returns `None` if the
164    /// event channel closes before `ConnectionReady` is received.
165    pub async fn block_until_ready(&mut self) -> Option<McpConnectionDetails> {
166        while let Some(event) = self.event_rx.recv().await {
167            if let McpClientEvent::ConnectionReady(snapshot) = event {
168                return Some(snapshot);
169            }
170        }
171        None
172    }
173
174    pub fn split(self) -> (McpRuntime, Receiver<McpClientEvent>) {
175        (self.runtime, self.event_rx)
176    }
177}
178
179pub struct McpBuilder {
180    servers: Vec<McpServer>,
181    factories: HashMap<String, ServerFactory>,
182    mcp_channel_capacity: usize,
183    root_dir: PathBuf,
184    oauth_handler_factory: Option<OAuthHandlerFactory>,
185    agent_deps: AgentDeps,
186    aether_home: Option<PathBuf>,
187    vars: Vars,
188    tool_filter: ToolFilter,
189    progressive_discovery_instructions: Option<String>,
190}
191
192impl McpBuilder {
193    pub fn new(root_dir: impl AsRef<Path>) -> Self {
194        let mut vars = Vars::new().with("WORKSPACE", root_dir.as_ref().to_string_lossy().into_owned());
195
196        if let Some(store) = SettingsStore::new("AETHER_HOME", ".aether") {
197            vars.insert("AETHER_HOME", store.home().to_string_lossy().into_owned());
198        }
199
200        Self {
201            servers: Vec::new(),
202            factories: HashMap::new(),
203            mcp_channel_capacity: 1000,
204            root_dir: root_dir.as_ref().to_path_buf(),
205            oauth_handler_factory: None,
206            agent_deps: AgentDeps::default(),
207            aether_home: None,
208            vars,
209            tool_filter: ToolFilter::default(),
210            progressive_discovery_instructions: None,
211        }
212    }
213
214    pub fn with_servers(mut self, servers: Vec<McpServer>) -> Self {
215        self.servers.extend(servers);
216        self
217    }
218
219    pub fn with_tool_filter(mut self, filter: ToolFilter) -> Self {
220        self.tool_filter = filter;
221        self
222    }
223
224    pub fn with_progressive_discovery_instructions(mut self, instructions: impl Into<String>) -> Self {
225        self.progressive_discovery_instructions = Some(instructions.into());
226        self
227    }
228
229    pub fn register_in_memory_server(mut self, name: impl Into<String>, factory: ServerFactory) -> Self {
230        self.factories.insert(name.into(), factory);
231        self
232    }
233
234    pub fn root_dir(&self) -> &Path {
235        &self.root_dir
236    }
237
238    /// Cross-cutting dependencies handed to every agent spawned behind this
239    /// builder's in-memory servers.
240    pub fn agent_deps(&self) -> AgentDeps {
241        self.agent_deps.clone()
242    }
243
244    pub fn with_agent_deps(mut self, deps: AgentDeps) -> Self {
245        self.agent_deps = deps;
246        self
247    }
248
249    pub fn with_oauth_handler_factory(mut self, factory: OAuthHandlerFactory) -> Self {
250        self.oauth_handler_factory = Some(factory);
251        self
252    }
253
254    pub fn with_aether_home(mut self, aether_home: impl Into<PathBuf>) -> Self {
255        let aether_home = aether_home.into();
256        self.vars.insert("AETHER_HOME", aether_home.to_string_lossy().into_owned());
257        self.aether_home = Some(aether_home);
258        self
259    }
260
261    pub fn from_json_files<T: AsRef<Path>>(mut self, paths: &[T]) -> Result<Self, ParseError> {
262        if paths.is_empty() {
263            return Ok(self);
264        }
265        let raw = McpConfig::from_json_files(paths)?;
266        self.servers.extend(raw.into_servers(&self.vars)?);
267        Ok(self)
268    }
269
270    pub fn from_mcp_config_sources(mut self, sources: &[McpConfigSource]) -> Result<Self, ParseError> {
271        if sources.is_empty() {
272            return Ok(self);
273        }
274
275        let mut merged = McpConfig::default();
276        for source in sources {
277            let config = match source {
278                McpConfigSource::File { path, defer_tools } => {
279                    let mut config = McpConfig::from_json_file(path)?;
280                    if *defer_tools {
281                        config.defer_all_tools();
282                    }
283                    config
284                }
285                McpConfigSource::Json(json) => McpConfig::from_json(json)?,
286                McpConfigSource::Inline(config) => config.clone(),
287            };
288            merged.servers.extend(config.servers);
289        }
290
291        self.servers.extend(merged.into_servers(&self.vars)?);
292        Ok(self)
293    }
294
295    pub async fn spawn(self) -> Result<McpSession, McpError> {
296        let McpBuilder {
297            servers,
298            factories,
299            mcp_channel_capacity,
300            root_dir,
301            oauth_handler_factory,
302            agent_deps,
303            aether_home: _,
304            vars: _,
305            tool_filter,
306            progressive_discovery_instructions,
307        } = self;
308        if servers.iter().any(|server| server.tool_exposure.has_deferred_tools())
309            && servers.iter().any(|server| server.name == PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME)
310        {
311            return Err(McpError::ReservedServerName(PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME.to_string()));
312        }
313        let (manager_tx, manager_rx) = mpsc::channel::<ManagerCommand>(mcp_channel_capacity);
314        let (snapshot_tx, snapshot_rx) = watch::channel(Arc::new(mcp_utils::client::McpSnapshot::default()));
315        let (event_tx, event_rx) = mpsc::channel::<McpClientEvent>(mcp_channel_capacity);
316        let mcp = McpHandle::new(manager_tx, snapshot_rx);
317        let gateway_transport = if servers.iter().any(|server| server.tool_exposure.has_deferred_tools()) {
318            let path = UnixSocketPath::new().map_err(|error| McpError::TransportError(error.to_string()))?;
319            Some(UnixSocketMcpTransport::bind(path).map_err(|error| McpError::TransportError(error.to_string()))?)
320        } else {
321            None
322        };
323        let shell_environment = gateway_transport
324            .as_ref()
325            .map(|transport| {
326                BTreeMap::from([(AETHER_MCP_IPC_SOCKET.to_string(), transport.path().to_string_lossy().into_owned())])
327            })
328            .unwrap_or_default();
329        let services = RuntimeServices { mcp: mcp.clone(), root_dir: root_dir.clone(), agent_deps, shell_environment };
330        let servers = resolve_servers(servers, &factories, &services).await?;
331
332        let mut mcp_manager = McpManager::new(event_tx, oauth_handler_factory)
333            .with_tool_filter(tool_filter)
334            .with_snapshot_sender(snapshot_tx);
335        if let Some(capabilities) = services.agent_deps.mcp_client_capabilities.clone() {
336            mcp_manager = mcp_manager.with_client_capabilities(capabilities);
337        }
338        if let Some(instructions) = progressive_discovery_instructions {
339            mcp_manager = mcp_manager.with_progressive_discovery_instructions(instructions);
340        }
341        if let Some(store) = services.agent_deps.oauth_credential_store.clone() {
342            mcp_manager = mcp_manager.with_oauth_credential_store(store);
343        }
344        mcp_manager = mcp_manager.with_root_dir(root_dir);
345        let pending = mcp_manager.register_pending(servers).await?;
346        let task = tokio::spawn(run_mcp_task(mcp_manager, manager_rx, pending));
347        let gateway = gateway_transport.map(|transport| transport.spawn(GatewayService::new(mcp.clone())));
348
349        Ok(McpSession { runtime: McpRuntime { mcp, handle: task, agent_sync_handle: None, gateway }, event_rx })
350    }
351}
352
353async fn resolve_servers(
354    servers: Vec<McpServer>,
355    factories: &HashMap<String, ServerFactory>,
356    services: &RuntimeServices,
357) -> Result<Vec<RuntimeMcpServer>, McpError> {
358    let mut resolved = Vec::with_capacity(servers.len());
359    for McpServer { name, transport, tool_exposure } in servers {
360        let transport = match transport {
361            McpTransport::Stdio { command, args, env } => RuntimeMcpTransport::Stdio { command, args, env },
362            McpTransport::Http(config) => RuntimeMcpTransport::Http(config),
363            McpTransport::InMemory { spec } => {
364                let factory = factories.get(&spec.factory).ok_or_else(|| McpError::InMemoryFactoryNotFound {
365                    server: name.clone(),
366                    factory: spec.factory.clone(),
367                })?;
368                RuntimeMcpTransport::InMemory { server: factory(spec, services.clone()).await }
369            }
370        };
371        resolved.push(RuntimeMcpServer::new(name, transport, tool_exposure));
372    }
373    Ok(resolved)
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379    use aether_auth::{FakeOAuthCredentialStore, OAuthCredentialStorage};
380    use futures::FutureExt;
381    use mcp_utils::testing::FakeMcpServer;
382    use mcp_utils::{
383        client::{McpServerConfig, McpTransport, StdioServerConfig, StdioType, ToolExposure},
384        status::McpServerStatus,
385    };
386    use std::collections::{BTreeMap, HashMap};
387    use std::sync::atomic::{AtomicUsize, Ordering};
388    use std::sync::{Arc, Mutex};
389
390    fn write_config_file(name: &str, json: &str) -> (tempfile::TempDir, PathBuf) {
391        let dir = tempfile::tempdir().unwrap();
392        let path = dir.path().join(name);
393        std::fs::write(&path, json).unwrap();
394        (dir, path)
395    }
396
397    fn json_source(json: &str) -> McpConfigSource {
398        McpConfigSource::Json(json.to_string())
399    }
400
401    fn builder_from_sources(sources: &[McpConfigSource]) -> McpBuilder {
402        McpBuilder::new("/workspace").from_mcp_config_sources(sources).unwrap()
403    }
404
405    #[tokio::test]
406    async fn in_memory_factory_runs_once_at_spawn_with_runtime_services() {
407        let calls = Arc::new(AtomicUsize::new(0));
408        let received = Arc::new(Mutex::new(None::<RuntimeServices>));
409        let factory_calls = Arc::clone(&calls);
410        let factory_received = Arc::clone(&received);
411        let oauth_store: Arc<dyn OAuthCredentialStorage> = Arc::new(FakeOAuthCredentialStore::new());
412        let deps = AgentDeps::new(Arc::clone(&oauth_store), None);
413        let factory: ServerFactory = Box::new(move |spec, services| {
414            factory_calls.fetch_add(1, Ordering::SeqCst);
415            assert_eq!(spec.args, ["--root", "/workspace/tools"]);
416            assert_eq!(spec.input, Some(serde_json::json!({"enabled": true})));
417            *factory_received.lock().unwrap() = Some(services);
418            async move { FakeMcpServer::new().into_dyn() }.boxed()
419        });
420
421        let builder = McpBuilder::new("/workspace")
422            .with_agent_deps(deps)
423            .register_in_memory_server("test", factory)
424            .from_mcp_config_sources(&[json_source(
425                r#"{"servers":{"test":{"type":"in-memory","args":["--root","${WORKSPACE}/tools"],"input":{"enabled":true}}}}"#,
426            )])
427            .unwrap();
428        assert_eq!(calls.load(Ordering::SeqCst), 0);
429
430        let spawn = builder.spawn().await.unwrap();
431        assert_eq!(calls.load(Ordering::SeqCst), 1);
432        let services = received.lock().unwrap().clone().expect("factory received runtime services");
433        assert_eq!(services.root_dir, PathBuf::from("/workspace"));
434        assert!(Arc::ptr_eq(&services.mcp.snapshot(), &spawn.handle().snapshot()));
435        assert!(Arc::ptr_eq(
436            services.agent_deps.oauth_credential_store.as_ref().expect("factory received agent dependencies"),
437            &oauth_store,
438        ));
439        assert!(services.shell_environment.is_empty());
440    }
441
442    #[tokio::test]
443    async fn deferred_gateway_is_bound_before_in_memory_factories_run() {
444        let received = Arc::new(Mutex::new(None::<RuntimeServices>));
445        let factory_received = Arc::clone(&received);
446        let factory: ServerFactory = Box::new(move |_, services| {
447            *factory_received.lock().unwrap() = Some(services);
448            async move { FakeMcpServer::new().into_dyn() }.boxed()
449        });
450        let spawn = McpBuilder::new("/workspace")
451            .register_in_memory_server("test", factory)
452            .from_mcp_config_sources(&[json_source(r#"{"servers":{"test":{"type":"in-memory","deferTools":true}}}"#)])
453            .unwrap()
454            .spawn()
455            .await
456            .unwrap();
457
458        let services = received.lock().unwrap().clone().expect("factory received runtime services");
459        let inherited =
460            services.shell_environment.get(AETHER_MCP_IPC_SOCKET).expect("factory receives gateway endpoint");
461        assert_eq!(Path::new(inherited), spawn.gateway_endpoint().expect("gateway endpoint exists"));
462        assert!(Path::new(inherited).exists());
463    }
464
465    #[tokio::test]
466    async fn snapshots_are_immutable_and_watch_observes_connection_changes() {
467        let factory: ServerFactory = Box::new(|_, _| async move { FakeMcpServer::new().into_dyn() }.boxed());
468        let mut spawn = McpBuilder::new("/workspace")
469            .register_in_memory_server("test", factory)
470            .from_mcp_config_sources(&[json_source(r#"{"servers":{"test":{"type":"in-memory"}}}"#)])
471            .unwrap()
472            .spawn()
473            .await
474            .unwrap();
475        let old = spawn.handle().snapshot();
476        let mut updates = spawn.handle().subscribe();
477
478        let ready = spawn.block_until_ready().await.expect("bootstrap completes");
479        updates.changed().await.expect("connection publishes a snapshot");
480        let observed = updates.borrow().clone();
481
482        assert!(old.tool_definitions().is_empty());
483        assert_eq!(ready.tool_definitions()[0].name, "test__add_numbers");
484        assert_eq!(observed.tool_definitions(), ready.tool_definitions());
485        assert!(!Arc::ptr_eq(&old, &ready));
486    }
487
488    #[tokio::test]
489    async fn missing_in_memory_factory_fails_at_spawn_with_server_and_factory() {
490        let builder = McpBuilder::new("/workspace")
491            .from_mcp_config_sources(&[json_source(r#"{"servers":{"custom":{"type":"in-memory"}}}"#)])
492            .unwrap();
493
494        let Err(error) = builder.spawn().await else {
495            panic!("spawn should reject an unregistered factory");
496        };
497        assert!(matches!(
498            error,
499            McpError::InMemoryFactoryNotFound { ref server, ref factory }
500                if server == "custom" && factory == "custom"
501        ));
502    }
503
504    #[tokio::test]
505    async fn mixed_direct_sources_preserve_last_wins_order() {
506        let (_dir, file_path) =
507            write_config_file("mcp.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#);
508        let inline = McpConfig::new(BTreeMap::from([(
509            "coding".to_string(),
510            McpServerConfig::Stdio(StdioServerConfig {
511                type_: StdioType::Stdio,
512                command: "from_inline".to_string(),
513                args: Vec::new(),
514                env: HashMap::new(),
515                defer_tools: ToolExposure::ModelVisible,
516            }),
517        )]));
518        let sources = vec![
519            McpConfigSource::model_visible(file_path),
520            json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json"}}}"#),
521            McpConfigSource::Inline(inline),
522        ];
523
524        let builder = builder_from_sources(&sources);
525
526        assert_eq!(command_for(&builder, "coding"), Some("from_inline"));
527        assert_eq!(deferred_tools_for(&builder, "coding"), Some(false));
528    }
529
530    #[tokio::test]
531    async fn file_sources_keep_their_position_relative_to_json_sources() {
532        let (_dir, file_path) =
533            write_config_file("mcp.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#);
534        let sources = vec![
535            json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json"}}}"#),
536            McpConfigSource::model_visible(file_path),
537        ];
538
539        let builder = builder_from_sources(&sources);
540
541        assert_eq!(command_for(&builder, "coding"), Some("from_file"));
542    }
543
544    #[tokio::test]
545    async fn file_source_defer_tools_marks_all_file_servers_deferred() {
546        let (_dir, file_path) = write_config_file(
547            "deferred.json",
548            r#"{"servers":{"github":{"type":"stdio","command":"g","deferTools":{"exclude":["status"]}},"browser":{"type":"stdio","command":"b"}}}"#,
549        );
550
551        let builder = McpBuilder::new("/workspace")
552            .from_mcp_config_sources(&[McpConfigSource::File { path: file_path, defer_tools: true }])
553            .unwrap();
554
555        assert_eq!(deferred_tools_for(&builder, "github"), Some(true));
556        assert_eq!(deferred_tools_for(&builder, "browser"), Some(true));
557        assert!(is_direct_tool(&builder, "github", "status"));
558    }
559
560    #[tokio::test]
561    async fn later_sources_override_defer_tools_flag() {
562        let (_dir, file_path) =
563            write_config_file("deferred.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#);
564        let sources = vec![
565            McpConfigSource::File { path: file_path, defer_tools: true },
566            json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json","deferTools":false}}}"#),
567        ];
568
569        let builder = builder_from_sources(&sources);
570
571        assert_eq!(command_for(&builder, "coding"), Some("from_json"));
572        assert_eq!(deferred_tools_for(&builder, "coding"), Some(false));
573    }
574
575    #[tokio::test]
576    async fn spawn_returns_immediately_and_emits_initial_connecting_status() {
577        let spawn = McpBuilder::new("/workspace")
578            .from_mcp_config_sources(&[json_source(
579                r#"{"servers":{"slow":{"type":"stdio","command":"sleep","args":["30"]}}}"#,
580            )])
581            .unwrap()
582            .spawn()
583            .await
584            .expect("spawn should succeed");
585
586        let (_runtime, mut event_rx) = spawn.split();
587        let event = event_rx.try_recv().expect("spawn() should buffer an initial ServerStatusesChanged");
588        let McpClientEvent::ServerStatusesChanged(statuses) = event else {
589            panic!("expected ServerStatusesChanged, got {event:?}");
590        };
591        assert!(matches!(statuses[0].status, McpServerStatus::Connecting));
592    }
593
594    #[tokio::test]
595    async fn from_mcp_config_sources_expands_workspace_var_in_stdio_args() {
596        let builder = McpBuilder::new("/work")
597            .from_mcp_config_sources(&[json_source(
598                r#"{"servers":{"notes":{"type":"stdio","command":"server","args":["--dir","${WORKSPACE}/notes"]}}}"#,
599            )])
600            .unwrap();
601
602        assert_eq!(args_for(&builder, "notes"), Some(vec!["--dir".to_string(), "/work/notes".to_string()]));
603    }
604
605    #[tokio::test]
606    async fn from_mcp_config_sources_expands_aether_home_var_in_stdio_args() {
607        let home = tempfile::tempdir().unwrap();
608
609        let builder = McpBuilder::new("/work")
610            .with_aether_home(home.path())
611            .from_mcp_config_sources(&[json_source(
612                r#"{"servers":{"skills":{"type":"stdio","command":"server","args":["--dir","${AETHER_HOME}/skills"]}}}"#,
613            )])
614            .unwrap();
615
616        assert_eq!(
617            args_for(&builder, "skills"),
618            Some(vec!["--dir".to_string(), home.path().join("skills").to_string_lossy().into_owned()])
619        );
620    }
621
622    #[tokio::test]
623    async fn reserved_progressive_discovery_server_is_rejected_when_gateway_is_enabled() {
624        let result = McpBuilder::new("/workspace")
625            .from_mcp_config_sources(&[json_source(
626                r#"{"servers":{"progressive-discovery":{"type":"stdio","command":"server"},"deferred":{"type":"stdio","command":"server","deferTools":true}}}"#,
627            )])
628            .unwrap()
629            .spawn()
630            .await;
631
632        assert!(matches!(result, Err(McpError::ReservedServerName(name)) if name == "progressive-discovery"));
633    }
634
635    #[test]
636    fn new_sets_root_directory_from_workspace_root() {
637        let builder = McpBuilder::new("/workspace");
638        assert_eq!(builder.root_dir, PathBuf::from("/workspace"));
639    }
640
641    fn command_for<'a>(builder: &'a McpBuilder, name: &str) -> Option<&'a str> {
642        builder.servers.iter().find_map(|server| match &server.transport {
643            McpTransport::Stdio { command, .. } if server.name == name => Some(command.as_str()),
644            _ => None,
645        })
646    }
647
648    fn args_for(builder: &McpBuilder, name: &str) -> Option<Vec<String>> {
649        builder.servers.iter().find_map(|server| match &server.transport {
650            McpTransport::Stdio { args, .. } if server.name == name => Some(args.clone()),
651            _ => None,
652        })
653    }
654
655    fn is_direct_tool(builder: &McpBuilder, server_name: &str, tool_name: &str) -> bool {
656        builder
657            .servers
658            .iter()
659            .find(|server| server.name == server_name)
660            .is_some_and(|server| server.tool_exposure.is_model_visible_tool(tool_name))
661    }
662
663    fn deferred_tools_for(builder: &McpBuilder, name: &str) -> Option<bool> {
664        builder.servers.iter().find(|server| server.name == name).map(mcp_utils::client::McpServer::has_deferred_tools)
665    }
666}
aether_core/mcp/mcp_builder.rs

aether_core/mcp/
mcp_builder.rs
