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::client::{McpServerConfig, McpTransport, StdioServerConfig, StdioType, ToolExposure};
382    use mcp_utils::testing::FakeMcpServer;
383    use std::collections::{BTreeMap, HashMap};
384    use std::sync::atomic::{AtomicUsize, Ordering};
385    use std::sync::{Arc, Mutex};
386    use utils::mcp_status::McpServerStatus;
387
388    fn write_config_file(name: &str, json: &str) -> (tempfile::TempDir, PathBuf) {
389        let dir = tempfile::tempdir().unwrap();
390        let path = dir.path().join(name);
391        std::fs::write(&path, json).unwrap();
392        (dir, path)
393    }
394
395    fn json_source(json: &str) -> McpConfigSource {
396        McpConfigSource::Json(json.to_string())
397    }
398
399    fn builder_from_sources(sources: &[McpConfigSource]) -> McpBuilder {
400        McpBuilder::new("/workspace").from_mcp_config_sources(sources).unwrap()
401    }
402
403    #[tokio::test]
404    async fn in_memory_factory_runs_once_at_spawn_with_runtime_services() {
405        let calls = Arc::new(AtomicUsize::new(0));
406        let received = Arc::new(Mutex::new(None::<RuntimeServices>));
407        let factory_calls = Arc::clone(&calls);
408        let factory_received = Arc::clone(&received);
409        let oauth_store: Arc<dyn OAuthCredentialStorage> = Arc::new(FakeOAuthCredentialStore::new());
410        let deps = AgentDeps::new(Arc::clone(&oauth_store), None);
411        let factory: ServerFactory = Box::new(move |spec, services| {
412            factory_calls.fetch_add(1, Ordering::SeqCst);
413            assert_eq!(spec.args, ["--root", "/workspace/tools"]);
414            assert_eq!(spec.input, Some(serde_json::json!({"enabled": true})));
415            *factory_received.lock().unwrap() = Some(services);
416            async move { FakeMcpServer::new().into_dyn() }.boxed()
417        });
418
419        let builder = McpBuilder::new("/workspace")
420            .with_agent_deps(deps)
421            .register_in_memory_server("test", factory)
422            .from_mcp_config_sources(&[json_source(
423                r#"{"servers":{"test":{"type":"in-memory","args":["--root","${WORKSPACE}/tools"],"input":{"enabled":true}}}}"#,
424            )])
425            .unwrap();
426        assert_eq!(calls.load(Ordering::SeqCst), 0);
427
428        let spawn = builder.spawn().await.unwrap();
429        assert_eq!(calls.load(Ordering::SeqCst), 1);
430        let services = received.lock().unwrap().clone().expect("factory received runtime services");
431        assert_eq!(services.root_dir, PathBuf::from("/workspace"));
432        assert!(Arc::ptr_eq(&services.mcp.snapshot(), &spawn.handle().snapshot()));
433        assert!(Arc::ptr_eq(
434            services.agent_deps.oauth_credential_store.as_ref().expect("factory received agent dependencies"),
435            &oauth_store,
436        ));
437        assert!(services.shell_environment.is_empty());
438    }
439
440    #[tokio::test]
441    async fn deferred_gateway_is_bound_before_in_memory_factories_run() {
442        let received = Arc::new(Mutex::new(None::<RuntimeServices>));
443        let factory_received = Arc::clone(&received);
444        let factory: ServerFactory = Box::new(move |_, services| {
445            *factory_received.lock().unwrap() = Some(services);
446            async move { FakeMcpServer::new().into_dyn() }.boxed()
447        });
448        let spawn = McpBuilder::new("/workspace")
449            .register_in_memory_server("test", factory)
450            .from_mcp_config_sources(&[json_source(r#"{"servers":{"test":{"type":"in-memory","deferTools":true}}}"#)])
451            .unwrap()
452            .spawn()
453            .await
454            .unwrap();
455
456        let services = received.lock().unwrap().clone().expect("factory received runtime services");
457        let inherited =
458            services.shell_environment.get(AETHER_MCP_IPC_SOCKET).expect("factory receives gateway endpoint");
459        assert_eq!(Path::new(inherited), spawn.gateway_endpoint().expect("gateway endpoint exists"));
460        assert!(Path::new(inherited).exists());
461    }
462
463    #[tokio::test]
464    async fn snapshots_are_immutable_and_watch_observes_connection_changes() {
465        let factory: ServerFactory = Box::new(|_, _| async move { FakeMcpServer::new().into_dyn() }.boxed());
466        let mut spawn = McpBuilder::new("/workspace")
467            .register_in_memory_server("test", factory)
468            .from_mcp_config_sources(&[json_source(r#"{"servers":{"test":{"type":"in-memory"}}}"#)])
469            .unwrap()
470            .spawn()
471            .await
472            .unwrap();
473        let old = spawn.handle().snapshot();
474        let mut updates = spawn.handle().subscribe();
475
476        let ready = spawn.block_until_ready().await.expect("bootstrap completes");
477        updates.changed().await.expect("connection publishes a snapshot");
478        let observed = updates.borrow().clone();
479
480        assert!(old.tool_definitions().is_empty());
481        assert_eq!(ready.tool_definitions()[0].name, "test__add_numbers");
482        assert_eq!(observed.tool_definitions(), ready.tool_definitions());
483        assert!(!Arc::ptr_eq(&old, &ready));
484    }
485
486    #[tokio::test]
487    async fn missing_in_memory_factory_fails_at_spawn_with_server_and_factory() {
488        let builder = McpBuilder::new("/workspace")
489            .from_mcp_config_sources(&[json_source(r#"{"servers":{"custom":{"type":"in-memory"}}}"#)])
490            .unwrap();
491
492        let Err(error) = builder.spawn().await else {
493            panic!("spawn should reject an unregistered factory");
494        };
495        assert!(matches!(
496            error,
497            McpError::InMemoryFactoryNotFound { ref server, ref factory }
498                if server == "custom" && factory == "custom"
499        ));
500    }
501
502    #[tokio::test]
503    async fn mixed_direct_sources_preserve_last_wins_order() {
504        let (_dir, file_path) =
505            write_config_file("mcp.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#);
506        let inline = McpConfig::new(BTreeMap::from([(
507            "coding".to_string(),
508            McpServerConfig::Stdio(StdioServerConfig {
509                type_: StdioType::Stdio,
510                command: "from_inline".to_string(),
511                args: Vec::new(),
512                env: HashMap::new(),
513                defer_tools: ToolExposure::ModelVisible,
514            }),
515        )]));
516        let sources = vec![
517            McpConfigSource::model_visible(file_path),
518            json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json"}}}"#),
519            McpConfigSource::Inline(inline),
520        ];
521
522        let builder = builder_from_sources(&sources);
523
524        assert_eq!(command_for(&builder, "coding"), Some("from_inline"));
525        assert_eq!(deferred_tools_for(&builder, "coding"), Some(false));
526    }
527
528    #[tokio::test]
529    async fn file_sources_keep_their_position_relative_to_json_sources() {
530        let (_dir, file_path) =
531            write_config_file("mcp.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#);
532        let sources = vec![
533            json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json"}}}"#),
534            McpConfigSource::model_visible(file_path),
535        ];
536
537        let builder = builder_from_sources(&sources);
538
539        assert_eq!(command_for(&builder, "coding"), Some("from_file"));
540    }
541
542    #[tokio::test]
543    async fn file_source_defer_tools_marks_all_file_servers_deferred() {
544        let (_dir, file_path) = write_config_file(
545            "deferred.json",
546            r#"{"servers":{"github":{"type":"stdio","command":"g","deferTools":{"exclude":["status"]}},"browser":{"type":"stdio","command":"b"}}}"#,
547        );
548
549        let builder = McpBuilder::new("/workspace")
550            .from_mcp_config_sources(&[McpConfigSource::File { path: file_path, defer_tools: true }])
551            .unwrap();
552
553        assert_eq!(deferred_tools_for(&builder, "github"), Some(true));
554        assert_eq!(deferred_tools_for(&builder, "browser"), Some(true));
555        assert!(is_direct_tool(&builder, "github", "status"));
556    }
557
558    #[tokio::test]
559    async fn later_sources_override_defer_tools_flag() {
560        let (_dir, file_path) =
561            write_config_file("deferred.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#);
562        let sources = vec![
563            McpConfigSource::File { path: file_path, defer_tools: true },
564            json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json","deferTools":false}}}"#),
565        ];
566
567        let builder = builder_from_sources(&sources);
568
569        assert_eq!(command_for(&builder, "coding"), Some("from_json"));
570        assert_eq!(deferred_tools_for(&builder, "coding"), Some(false));
571    }
572
573    #[tokio::test]
574    async fn spawn_returns_immediately_and_emits_initial_connecting_status() {
575        let spawn = McpBuilder::new("/workspace")
576            .from_mcp_config_sources(&[json_source(
577                r#"{"servers":{"slow":{"type":"stdio","command":"sleep","args":["30"]}}}"#,
578            )])
579            .unwrap()
580            .spawn()
581            .await
582            .expect("spawn should succeed");
583
584        let (_runtime, mut event_rx) = spawn.split();
585        let event = event_rx.try_recv().expect("spawn() should buffer an initial ServerStatusesChanged");
586        let McpClientEvent::ServerStatusesChanged(statuses) = event else {
587            panic!("expected ServerStatusesChanged, got {event:?}");
588        };
589        assert!(matches!(statuses[0].status, McpServerStatus::Connecting));
590    }
591
592    #[tokio::test]
593    async fn from_mcp_config_sources_expands_workspace_var_in_stdio_args() {
594        let builder = McpBuilder::new("/work")
595            .from_mcp_config_sources(&[json_source(
596                r#"{"servers":{"notes":{"type":"stdio","command":"server","args":["--dir","${WORKSPACE}/notes"]}}}"#,
597            )])
598            .unwrap();
599
600        assert_eq!(args_for(&builder, "notes"), Some(vec!["--dir".to_string(), "/work/notes".to_string()]));
601    }
602
603    #[tokio::test]
604    async fn from_mcp_config_sources_expands_aether_home_var_in_stdio_args() {
605        let home = tempfile::tempdir().unwrap();
606
607        let builder = McpBuilder::new("/work")
608            .with_aether_home(home.path())
609            .from_mcp_config_sources(&[json_source(
610                r#"{"servers":{"skills":{"type":"stdio","command":"server","args":["--dir","${AETHER_HOME}/skills"]}}}"#,
611            )])
612            .unwrap();
613
614        assert_eq!(
615            args_for(&builder, "skills"),
616            Some(vec!["--dir".to_string(), home.path().join("skills").to_string_lossy().into_owned()])
617        );
618    }
619
620    #[tokio::test]
621    async fn reserved_progressive_discovery_server_is_rejected_when_gateway_is_enabled() {
622        let result = McpBuilder::new("/workspace")
623            .from_mcp_config_sources(&[json_source(
624                r#"{"servers":{"progressive-discovery":{"type":"stdio","command":"server"},"deferred":{"type":"stdio","command":"server","deferTools":true}}}"#,
625            )])
626            .unwrap()
627            .spawn()
628            .await;
629
630        assert!(matches!(result, Err(McpError::ReservedServerName(name)) if name == "progressive-discovery"));
631    }
632
633    #[test]
634    fn new_sets_root_directory_from_workspace_root() {
635        let builder = McpBuilder::new("/workspace");
636        assert_eq!(builder.root_dir, PathBuf::from("/workspace"));
637    }
638
639    fn command_for<'a>(builder: &'a McpBuilder, name: &str) -> Option<&'a str> {
640        builder.servers.iter().find_map(|server| match &server.transport {
641            McpTransport::Stdio { command, .. } if server.name == name => Some(command.as_str()),
642            _ => None,
643        })
644    }
645
646    fn args_for(builder: &McpBuilder, name: &str) -> Option<Vec<String>> {
647        builder.servers.iter().find_map(|server| match &server.transport {
648            McpTransport::Stdio { args, .. } if server.name == name => Some(args.clone()),
649            _ => None,
650        })
651    }
652
653    fn is_direct_tool(builder: &McpBuilder, server_name: &str, tool_name: &str) -> bool {
654        builder
655            .servers
656            .iter()
657            .find(|server| server.name == server_name)
658            .is_some_and(|server| server.tool_exposure.is_model_visible_tool(tool_name))
659    }
660
661    fn deferred_tools_for(builder: &McpBuilder, name: &str) -> Option<bool> {
662        builder.servers.iter().find(|server| server.name == name).map(mcp_utils::client::McpServer::has_deferred_tools)
663    }
664}
aether_core/mcp/mcp_builder.rs

aether_core/mcp/
mcp_builder.rs
