mcp_utils/client/
connection.rs

1use super::{
2    McpClientEvent, McpError, OAuthHandlerFactory, Result,
3    config::McpHttpConfig,
4    manager::{RuntimeMcpServer, RuntimeMcpTransport, ToolListChangedRequest},
5    mcp_client::McpClient,
6};
7use crate::{client::OAuthHandlerContext, protocol::client_lifecycle_mode, transport::create_in_memory_transport};
8use aether_auth::{OAuthCredentialStorage, create_auth_manager_from_store, perform_oauth_flow};
9use llm::ToolAnnotations;
10use rmcp::{
11    RoleClient, RoleServer, ServiceExt,
12    model::{ClientInfo, Tool as RmcpTool},
13    serve_client_with_lifecycle,
14    service::{DynService, RunningService},
15    transport::{
16        StreamableHttpClientTransport, TokioChildProcess, auth::AuthClient,
17        streamable_http_client::StreamableHttpClientTransportConfig,
18    },
19};
20use serde_json::Value;
21use std::collections::HashMap;
22use std::path::PathBuf;
23use std::process::Stdio;
24use std::sync::{
25    Arc,
26    atomic::{AtomicU64, Ordering},
27};
28use tokio::{
29    io::{AsyncBufReadExt, BufReader},
30    process::{ChildStderr, Command},
31    sync::mpsc,
32    task::JoinHandle,
33};
34
35#[derive(Debug, Clone)]
36pub struct Tool {
37    pub name: String,
38    pub description: String,
39    pub parameters: Value,
40    pub annotations: Option<ToolAnnotations>,
41}
42
43pub(crate) fn convert_tool_annotations(annotations: &rmcp::model::ToolAnnotations) -> ToolAnnotations {
44    ToolAnnotations {
45        title: annotations.title.clone(),
46        read_only_hint: annotations.read_only_hint,
47        destructive_hint: annotations.destructive_hint,
48        idempotent_hint: annotations.idempotent_hint,
49        open_world_hint: annotations.open_world_hint,
50    }
51}
52
53impl From<RmcpTool> for Tool {
54    fn from(tool: RmcpTool) -> Self {
55        Self::from(&tool)
56    }
57}
58
59impl From<&RmcpTool> for Tool {
60    fn from(tool: &RmcpTool) -> Self {
61        Self {
62            name: tool.name.to_string(),
63            description: tool.description.clone().unwrap_or_default().to_string(),
64            parameters: serde_json::Value::Object((*tool.input_schema).clone()),
65            annotations: tool.annotations.as_ref().map(convert_tool_annotations),
66        }
67    }
68}
69
70pub(super) struct ConnectConfig {
71    pub client_info: ClientInfo,
72    pub event_sender: mpsc::Sender<McpClientEvent>,
73    pub tool_refresh_sender: mpsc::Sender<ToolListChangedRequest>,
74    pub next_connection_generation: Arc<AtomicU64>,
75    pub root_dir: PathBuf,
76    pub oauth_handler_factory: Option<OAuthHandlerFactory>,
77    pub oauth_credential_store: Option<Arc<dyn OAuthCredentialStorage>>,
78}
79
80/// The result of attempting to connect (or authenticate) to an MCP server.
81pub struct McpConnectAttempt {
82    pub name: String,
83    pub outcome: McpConnectOutcome,
84}
85
86pub enum McpConnectOutcome {
87    Connected { conn: McpServerConnection, reauth_config: Option<McpHttpConfig> },
88    NeedsOAuth { config: McpHttpConfig, challenge: Option<String>, error: McpError },
89    Failed { error: McpError },
90}
91
92impl McpConnectAttempt {
93    pub fn failed(name: impl Into<String>, error: McpError) -> Self {
94        Self { name: name.into(), outcome: McpConnectOutcome::Failed { error } }
95    }
96}
97
98pub struct McpServerConnection {
99    pub(super) client: Arc<RunningService<RoleClient, McpClient>>,
100    pub(super) server_task: Option<JoinHandle<()>>,
101    pub(super) instructions: Option<String>,
102    generation: u64,
103}
104
105impl McpServerConnection {
106    pub(super) async fn reconnect_with_auth(
107        name: &str,
108        config: StreamableHttpClientTransportConfig,
109        auth_client: AuthClient<reqwest::Client>,
110        mcp_client: McpClient,
111        generation: u64,
112    ) -> Result<Self> {
113        let transport = StreamableHttpClientTransport::with_client(auth_client, config);
114        let client = serve_client_with_lifecycle(mcp_client, transport, client_lifecycle_mode())
115            .await
116            .map_err(|e| McpError::ConnectionFailed(format!("reconnect failed for '{name}': {e}")))?;
117        Ok(Self::from_parts(client, None, generation))
118    }
119
120    pub(super) async fn list_tools(&self) -> Result<Vec<RmcpTool>> {
121        self.client
122            .list_all_tools()
123            .await
124            .map_err(|e| McpError::ToolDiscoveryFailed(format!("Failed to list tools: {e}")))
125    }
126
127    fn from_parts(
128        client: RunningService<RoleClient, McpClient>,
129        server_task: Option<JoinHandle<()>>,
130        generation: u64,
131    ) -> Self {
132        let instructions = client.peer_info().and_then(|info| info.instructions.clone()).filter(|s| !s.is_empty());
133        Self { client: Arc::new(client), server_task, instructions, generation }
134    }
135
136    pub(super) fn generation(&self) -> u64 {
137        self.generation
138    }
139}
140
141pub(super) async fn connect_server(server: RuntimeMcpServer, ctx: &ConnectConfig) -> McpConnectAttempt {
142    let RuntimeMcpServer { name, transport, tool_exposure: _ } = server;
143    let reauth_config = reauth_config_for(&transport, ctx.oauth_handler_factory.as_ref());
144    let generation = ctx.next_connection_generation.fetch_add(1, Ordering::Relaxed);
145    let mcp_client = McpClient::new(ctx.client_info.clone(), name.clone(), ctx.event_sender.clone())
146        .with_tool_refresh(ctx.tool_refresh_sender.clone(), generation);
147
148    let outcome = match transport {
149        RuntimeMcpTransport::Stdio { command, args, env } => {
150            connect_stdio(&name, command, args, env, mcp_client, ctx.root_dir.clone(), generation).await
151        }
152        RuntimeMcpTransport::InMemory { server } => connect_in_memory(&name, server, mcp_client, generation).await,
153        RuntimeMcpTransport::Http(config) => {
154            connect_http(
155                &name,
156                config,
157                mcp_client,
158                ctx.oauth_handler_factory.as_ref(),
159                ctx.oauth_credential_store.as_ref(),
160                generation,
161            )
162            .await
163        }
164    };
165
166    McpConnectAttempt { name, outcome: outcome.with_reauth(reauth_config) }
167}
168
169pub async fn authenticate_http(
170    name: String,
171    config: McpHttpConfig,
172    challenge: Option<String>,
173    ctx: Arc<ConnectConfig>,
174) -> McpConnectAttempt {
175    let outcome = match async {
176        let factory = ctx
177            .oauth_handler_factory
178            .as_ref()
179            .ok_or_else(|| McpError::ConnectionFailed(format!("No OAuth handler factory available for '{name}'")))?;
180        let oauth = config
181            .resolved_oauth()
182            .ok_or_else(|| McpError::ConnectionFailed(format!("OAuth is not available for '{name}'")))?;
183        let handler = factory(OAuthHandlerContext {
184            server_name: name.clone(),
185            callback_port: Some(oauth.callback_port),
186            tx: ctx.event_sender.clone(),
187        })?;
188
189        let auth_client = perform_oauth_flow(
190            &name,
191            &config.transport.uri,
192            handler.as_ref(),
193            aether_auth::OAuthFlowOptions { client_registration: oauth.client_registration, challenge },
194            ctx.oauth_credential_store.clone(),
195        )
196        .await
197        .map_err(|e| McpError::ConnectionFailed(format!("OAuth failed for '{name}': {e}")))?;
198
199        let generation = ctx.next_connection_generation.fetch_add(1, Ordering::Relaxed);
200        let mcp_client = McpClient::new(ctx.client_info.clone(), name.clone(), ctx.event_sender.clone())
201            .with_tool_refresh(ctx.tool_refresh_sender.clone(), generation);
202        McpServerConnection::reconnect_with_auth(&name, config.transport.clone(), auth_client, mcp_client, generation)
203            .await
204    }
205    .await
206    {
207        Ok(conn) => McpConnectOutcome::Connected { conn, reauth_config: Some(config) },
208        Err(error) => McpConnectOutcome::Failed { error },
209    };
210
211    McpConnectAttempt { name, outcome }
212}
213
214impl McpConnectOutcome {
215    fn with_reauth(self, reauth_config: Option<McpHttpConfig>) -> Self {
216        match self {
217            Self::Connected { conn, .. } => Self::Connected { conn, reauth_config },
218            other => other,
219        }
220    }
221}
222
223async fn connect_stdio(
224    server_name: &str,
225    command: String,
226    args: Vec<String>,
227    env: HashMap<String, String>,
228    mcp_client: McpClient,
229    cwd: PathBuf,
230    generation: u64,
231) -> McpConnectOutcome {
232    let mut cmd = Command::new(&command);
233    cmd.args(&args).envs(&env).current_dir(&cwd);
234
235    let (proc, stderr) = match TokioChildProcess::builder(cmd).stderr(Stdio::piped()).spawn() {
236        Ok(parts) => parts,
237        Err(e) => return McpConnectOutcome::Failed { error: McpError::SpawnFailed { command, reason: e.to_string() } },
238    };
239    let stderr_task = stderr.map(|stderr| spawn_stderr_logger(server_name.to_string(), stderr));
240
241    match serve_client_with_lifecycle(mcp_client, proc, client_lifecycle_mode()).await {
242        Ok(client) => McpConnectOutcome::Connected {
243            conn: McpServerConnection::from_parts(client, stderr_task, generation),
244            reauth_config: None,
245        },
246        Err(e) => {
247            if let Some(task) = stderr_task {
248                task.abort();
249            }
250            McpConnectOutcome::Failed { error: McpError::from(e) }
251        }
252    }
253}
254
255fn spawn_stderr_logger(server_name: String, stderr: ChildStderr) -> JoinHandle<()> {
256    tokio::spawn(async move {
257        let mut lines = BufReader::new(stderr).lines();
258        loop {
259            match lines.next_line().await {
260                Ok(Some(line)) => tracing::info!(server = %server_name, stderr = %line, "MCP server stderr"),
261                Ok(None) => break,
262                Err(error) => {
263                    tracing::warn!(server = %server_name, %error, "failed to read MCP server stderr");
264                    break;
265                }
266            }
267        }
268    })
269}
270
271async fn connect_in_memory(
272    name: &str,
273    server: Box<dyn DynService<RoleServer>>,
274    mcp_client: McpClient,
275    generation: u64,
276) -> McpConnectOutcome {
277    match serve_in_memory(server, mcp_client, name).await {
278        Ok((client, handle)) => McpConnectOutcome::Connected {
279            conn: McpServerConnection::from_parts(client, Some(handle), generation),
280            reauth_config: None,
281        },
282        Err(error) => McpConnectOutcome::Failed { error },
283    }
284}
285
286async fn connect_http(
287    name: &str,
288    config: McpHttpConfig,
289    mcp_client: McpClient,
290    oauth_handler_factory: Option<&OAuthHandlerFactory>,
291    oauth_credential_store: Option<&Arc<dyn OAuthCredentialStorage>>,
292    generation: u64,
293) -> McpConnectOutcome {
294    let conn_err = |e| McpError::ConnectionFailed(format!("HTTP MCP server {name}: {e}"));
295    let oauth = config.resolved_oauth();
296    let restored = if let (Some(store), Some(oauth)) = (oauth_credential_store, oauth.as_ref()) {
297        match create_auth_manager_from_store(
298            name,
299            &config.transport.uri,
300            oauth.client_registration.pre_registered_client_id(),
301            &oauth.redirect_uri(),
302            Arc::clone(store),
303        )
304        .await
305        {
306            Ok(manager) => manager,
307            Err(e) => {
308                tracing::warn!(
309                    server = %name,
310                    error = %e,
311                    "Failed to initialize auth manager from stored credentials, proceeding without auth"
312                );
313                None
314            }
315        }
316    } else {
317        None
318    };
319    let result = if let Some(manager) = restored {
320        tracing::debug!("Using OAuth for server '{name}'");
321        let auth_client = AuthClient::new(reqwest::Client::default(), manager);
322        let transport = StreamableHttpClientTransport::with_client(auth_client, config.transport.clone());
323        serve_client_with_lifecycle(mcp_client, transport, client_lifecycle_mode()).await
324    } else {
325        let transport = StreamableHttpClientTransport::from_config(config.transport.clone());
326        serve_client_with_lifecycle(mcp_client, transport, client_lifecycle_mode()).await
327    };
328
329    match result {
330        Ok(client) => McpConnectOutcome::Connected {
331            conn: McpServerConnection::from_parts(client, None, generation),
332            reauth_config: None,
333        },
334        Err(error) => {
335            let challenge = error.auth_challenge().map(str::to_string);
336            let authorization_required = error.is_authorization_required();
337            let error = conn_err(error);
338            tracing::warn!("Failed to connect to MCP server '{name}': {error}");
339            if oauth_handler_factory.is_some() && oauth.is_some() && (authorization_required || challenge.is_some()) {
340                McpConnectOutcome::NeedsOAuth { config, challenge, error }
341            } else {
342                McpConnectOutcome::Failed { error }
343            }
344        }
345    }
346}
347
348fn reauth_config_for(
349    transport: &RuntimeMcpTransport,
350    oauth_handler_factory: Option<&OAuthHandlerFactory>,
351) -> Option<McpHttpConfig> {
352    match transport {
353        RuntimeMcpTransport::Http(config)
354            if oauth_handler_factory.is_some() && config.transport.auth_header.is_none() =>
355        {
356            Some(config.clone())
357        }
358        _ => None,
359    }
360}
361
362async fn serve_in_memory(
363    server: Box<dyn DynService<RoleServer>>,
364    mcp_client: McpClient,
365    label: &str,
366) -> Result<(RunningService<RoleClient, McpClient>, JoinHandle<()>)> {
367    let (client_transport, server_transport) = create_in_memory_transport();
368
369    let server_handle = tokio::spawn(async move {
370        match server.serve(server_transport).await {
371            Ok(_service) => {
372                std::future::pending::<()>().await;
373            }
374            Err(e) => {
375                eprintln!("MCP server error: {e}");
376            }
377        }
378    });
379
380    let client = serve_client_with_lifecycle(mcp_client, client_transport, client_lifecycle_mode())
381        .await
382        .map_err(|e| McpError::ConnectionFailed(format!("Failed to connect to in-memory server '{label}': {e}")))?;
383
384    Ok((client, server_handle))
385}
mcp_utils/client/connection.rs

mcp_utils/client/
connection.rs
