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

aether_core/mcp/
mcp_builder.rs
