mcp_utils/client/
connection.rs

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

mcp_utils/client/
connection.rs
