mcp_utils/tool_gateway/
transport.rs

1use rmcp::transport::async_rw::AsyncRwTransport;
2use rmcp::{RoleClient, RoleServer, ServiceExt};
3use std::env::{temp_dir, var_os};
4use std::fs::{Permissions, create_dir_all, remove_dir};
5use std::fs::{remove_file, set_permissions};
6use std::{
7    io,
8    os::unix::fs::PermissionsExt,
9    os::unix::net::UnixListener as StdUnixListener,
10    path::{Path, PathBuf},
11    sync::{Arc, Mutex},
12};
13use tokio::io::{ReadHalf, WriteHalf};
14use tokio::net::{UnixListener, UnixStream};
15use tokio::task::{JoinHandle, JoinSet};
16use tokio_util::sync::CancellationToken;
17use uuid::Uuid;
18
19#[derive(Debug, thiserror::Error)]
20pub enum UnixSocketTransportError {
21    #[error("failed to create MCP socket directory: {0}")]
22    CreateDirectory(#[source] io::Error),
23    #[error("failed to bind MCP socket: {0}")]
24    Bind(#[source] io::Error),
25    #[error("MCP socket path must be absolute")]
26    NotAbsolute,
27    #[error("MCP socket path is not valid UTF-8")]
28    InvalidPath,
29}
30
31/// A session endpoint allocated by this process, or inherited from the session environment.
32#[derive(Debug)]
33pub struct UnixSocketPath {
34    directory: PathBuf,
35    socket: PathBuf,
36    /// Only the allocating process owns endpoint removal; inherited paths never remove.
37    remove_on_drop: bool,
38}
39
40impl UnixSocketPath {
41    pub fn new() -> Result<Self, UnixSocketTransportError> {
42        let short_id = &Uuid::new_v4().simple().to_string()[..8];
43        let runtime_dir =
44            var_os("XDG_RUNTIME_DIR").filter(|path| Path::new(path).is_absolute()).map_or_else(temp_dir, PathBuf::from);
45        let socket_dir = runtime_dir.join("aether").join(format!("aether-{short_id}"));
46        let socket = socket_dir.join("ipc.sock");
47        create_dir_all(&socket_dir).map_err(UnixSocketTransportError::CreateDirectory)?;
48        set_permissions(&socket_dir, Permissions::from_mode(0o700))
49            .map_err(UnixSocketTransportError::CreateDirectory)?;
50        Ok(Self { directory: socket_dir, socket, remove_on_drop: true })
51    }
52
53    pub fn from_path(path: impl Into<PathBuf>) -> Result<Self, UnixSocketTransportError> {
54        let socket = path.into();
55        if !socket.is_absolute() {
56            return Err(UnixSocketTransportError::NotAbsolute);
57        }
58        let directory = socket.parent().ok_or(UnixSocketTransportError::InvalidPath)?.to_path_buf();
59        Ok(Self { directory, socket, remove_on_drop: false })
60    }
61
62    pub fn path(&self) -> &Path {
63        &self.socket
64    }
65
66    pub fn directory(&self) -> &Path {
67        &self.directory
68    }
69}
70
71impl Drop for UnixSocketPath {
72    fn drop(&mut self) {
73        if self.remove_on_drop {
74            let _ = remove_file(&self.socket);
75            let _ = remove_dir(&self.directory);
76        }
77    }
78}
79
80pub struct UnixSocketMcpTransport {
81    path: UnixSocketPath,
82    listener: UnixListener,
83}
84
85impl UnixSocketMcpTransport {
86    pub fn bind(path: UnixSocketPath) -> Result<Self, UnixSocketTransportError> {
87        let _ = remove_file(path.path());
88        let listener = StdUnixListener::bind(path.path()).map_err(UnixSocketTransportError::Bind)?;
89        listener.set_nonblocking(true).map_err(UnixSocketTransportError::Bind)?;
90        let listener = UnixListener::from_std(listener).map_err(UnixSocketTransportError::Bind)?;
91        Ok(Self { path, listener })
92    }
93
94    pub fn path(&self) -> &Path {
95        self.path.path()
96    }
97
98    /// Serve connections until the returned server is dropped.
99    pub fn spawn<T>(self, server: T) -> UnixSocketServer
100    where
101        T: Clone + ServiceExt<RoleServer> + Send + 'static,
102    {
103        let Self { path, listener } = self;
104        let cancellation = CancellationToken::new();
105        let accept_cancellation = cancellation.clone();
106        let connections = Arc::new(Mutex::new(JoinSet::new()));
107        let accept_connections = Arc::clone(&connections);
108        let task = tokio::spawn(async move {
109            loop {
110                let accepted = tokio::select! {
111                    () = accept_cancellation.cancelled() => break,
112                    result = listener.accept() => result,
113                };
114                let Ok((stream, _)) = accepted else { break };
115                let server = server.clone();
116                let connection_cancellation = accept_cancellation.clone();
117                accept_connections.lock().unwrap().spawn(async move {
118                    match server.serve(stream).await {
119                        Ok(running) => {
120                            let service_cancellation = running.cancellation_token();
121                            let mut waiting = Box::pin(running.waiting());
122                            tokio::select! {
123                                () = connection_cancellation.cancelled() => service_cancellation.cancel(),
124                                _ = &mut waiting => {}
125                            }
126                        }
127                        Err(error) => tracing::debug!(%error, "MCP Unix socket client ended during initialization"),
128                    }
129                });
130            }
131        });
132        UnixSocketServer { path, cancellation, task, connections }
133    }
134}
135
136/// Owns the accept task, its connection tasks, and the session endpoint.
137/// Dropping it cancels in-flight connections and removes the endpoint.
138pub struct UnixSocketServer {
139    path: UnixSocketPath,
140    cancellation: CancellationToken,
141    task: JoinHandle<()>,
142    connections: Arc<Mutex<JoinSet<()>>>,
143}
144
145impl UnixSocketServer {
146    pub fn path(&self) -> &Path {
147        self.path.path()
148    }
149}
150
151impl Drop for UnixSocketServer {
152    fn drop(&mut self) {
153        self.cancellation.cancel();
154        self.task.abort();
155        self.connections.lock().unwrap().detach_all();
156    }
157}
158
159/// Connect an rmcp client to an inherited session endpoint.
160pub async fn connect(
161    path: impl AsRef<Path>,
162) -> io::Result<AsyncRwTransport<RoleClient, ReadHalf<UnixStream>, WriteHalf<UnixStream>>> {
163    let stream = UnixStream::connect(path).await?;
164    let (read, write) = tokio::io::split(stream);
165    Ok(AsyncRwTransport::new_client(read, write))
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    use rmcp::ServerHandler;
172    use rmcp::handler::server::router::tool::ToolRouter;
173    use rmcp::model::{ServerCapabilities, ServerInfo};
174    use rmcp::{tool, tool_handler, tool_router};
175
176    #[derive(Clone)]
177    struct TestServer {
178        tool_router: ToolRouter<Self>,
179    }
180
181    #[tool_router]
182    impl TestServer {}
183
184    #[allow(clippy::unused_async_trait_impl)]
185    #[tool_handler(router = self.tool_router)]
186    impl ServerHandler for TestServer {
187        fn get_info(&self) -> ServerInfo {
188            ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
189        }
190    }
191
192    #[test]
193    fn allocated_endpoint_is_removed_on_drop_without_binding() {
194        let path = UnixSocketPath::new().unwrap();
195        let directory = path.directory().to_path_buf();
196        assert!(directory.exists());
197        drop(path);
198        assert!(!directory.exists());
199    }
200
201    #[test]
202    fn inherited_endpoint_is_not_removed_on_drop() {
203        let path = UnixSocketPath::new().unwrap();
204        let directory = path.directory().to_path_buf();
205        let inherited = UnixSocketPath::from_path(path.path()).unwrap();
206        drop(inherited);
207        assert!(directory.exists());
208        drop(path);
209        assert!(!directory.exists());
210    }
211
212    #[tokio::test]
213    async fn allocated_endpoint_is_private_and_removed_with_transport() {
214        let path = UnixSocketPath::new().unwrap();
215        let directory = path.directory().to_path_buf();
216        let transport = UnixSocketMcpTransport::bind(path).unwrap();
217        assert_eq!(std::fs::metadata(&directory).unwrap().permissions().mode() & 0o777, 0o700);
218        assert!(transport.path().exists());
219        drop(transport);
220        assert!(!directory.exists());
221    }
222
223    #[tokio::test]
224    async fn spawned_endpoint_is_removed_with_server() {
225        let path = UnixSocketPath::new().unwrap();
226        let directory = path.directory().to_path_buf();
227        let transport = UnixSocketMcpTransport::bind(path).unwrap();
228        let socket = transport.path().to_path_buf();
229        let server = transport.spawn(TestServer { tool_router: TestServer::tool_router() });
230        assert!(socket.exists());
231        drop(server);
232        assert!(!directory.exists());
233    }
234
235    #[tokio::test]
236    async fn connected_client_completes_initialization() {
237        let path = UnixSocketPath::new().unwrap();
238        let transport = UnixSocketMcpTransport::bind(path).unwrap();
239        let socket = transport.path().to_path_buf();
240        let _server = transport.spawn(TestServer { tool_router: TestServer::tool_router() });
241        let _client = ().serve(connect(&socket).await.unwrap()).await.unwrap();
242    }
243
244    #[tokio::test]
245    async fn dropping_server_cancels_in_flight_connections() {
246        use rmcp::model::CallToolRequestParams;
247        use tokio::sync::watch;
248
249        #[derive(Clone)]
250        struct SlowServer {
251            tool_router: ToolRouter<Self>,
252            started: watch::Sender<bool>,
253        }
254
255        #[tool_router]
256        impl SlowServer {
257            #[tool(description = "Blocks until the connection is aborted")]
258            async fn slow(&self) -> String {
259                let _ = self.started.send(true);
260                std::future::pending::<()>().await;
261                "done".to_string()
262            }
263        }
264
265        #[allow(clippy::unused_async_trait_impl)]
266        #[tool_handler(router = self.tool_router)]
267        impl ServerHandler for SlowServer {
268            fn get_info(&self) -> ServerInfo {
269                ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
270            }
271        }
272
273        let path = UnixSocketPath::new().unwrap();
274        let transport = UnixSocketMcpTransport::bind(path).unwrap();
275        let socket = transport.path().to_path_buf();
276        let (started_tx, mut started_rx) = watch::channel(false);
277        let server = transport.spawn(SlowServer { tool_router: SlowServer::tool_router(), started: started_tx });
278
279        let client = ().serve(connect(&socket).await.unwrap()).await.unwrap();
280        let call = tokio::spawn(async move {
281            let _ = client.call_tool_once(CallToolRequestParams::new("slow")).await;
282        });
283        started_rx.changed().await.unwrap();
284
285        drop(server);
286        call.await.unwrap();
287    }
288}
mcp_utils/tool_gateway/transport.rs

mcp_utils/tool_gateway/
transport.rs
