mcp_utils/client/
connection_attempt_manager.rs

1use super::connection::McpConnectAttempt;
2use std::collections::HashMap;
3use tokio::task::{AbortHandle, JoinError, JoinSet};
4
5#[derive(Default)]
6pub struct McpConnectionAttemptManager {
7    set: JoinSet<McpConnectAttempt>,
8    by_server: HashMap<String, AbortHandle>,
9}
10
11impl McpConnectionAttemptManager {
12    /// Spawn a fresh auth task for `server`, aborting any in-flight task for
13    /// the same server so the user can retry a stuck OAuth flow.
14    pub fn spawn(&mut self, server: String, fut: impl Future<Output = McpConnectAttempt> + Send + 'static) {
15        if let Some(prior) = self.by_server.remove(&server) {
16            prior.abort();
17        }
18        let handle = self.set.spawn(fut);
19        self.by_server.insert(server, handle);
20    }
21
22    pub fn is_empty(&self) -> bool {
23        self.set.is_empty()
24    }
25
26    pub async fn join_next(&mut self) -> Option<Result<McpConnectAttempt, JoinError>> {
27        let joined = self.set.join_next().await;
28        if let Some(Ok(attempt)) = &joined {
29            self.by_server.remove(&attempt.name);
30        }
31        joined
32    }
33
34    pub async fn shutdown(&mut self) {
35        self.set.abort_all();
36        while self.set.join_next().await.is_some() {}
37        self.by_server.clear();
38    }
39}
mcp_utils/client/connection_attempt_manager.rs

mcp_utils/client/
connection_attempt_manager.rs
