mcp_utils/client/
connection.rs

1use super::{
2    McpClientEvent, McpError, OAuthHandlerFactory, Result,
3    config::{McpHttpConfig, 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<McpHttpConfig> },
83    NeedsOAuth { config: McpHttpConfig, 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: McpHttpConfig,
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 {
165            server_name: name.clone(),
166            callback_port: config.callback_port(),
167            tx: ctx.event_sender.clone(),
168        })?;
169
170        let auth_client = perform_oauth_flow(
171            &name,
172            &config.transport.uri,
173            handler.as_ref(),
174            config.oauth_client_id(),
175            ctx.oauth_credential_store.clone(),
176        )
177        .await
178        .map_err(|e| McpError::ConnectionFailed(format!("OAuth failed for '{name}': {e}")))?;
179
180        let mcp_client = McpClient::new(ctx.client_info.clone(), name.clone(), ctx.event_sender.clone());
181        McpServerConnection::reconnect_with_auth(&name, config.transport.clone(), auth_client, mcp_client).await
182    }
183    .await
184    {
185        Ok(conn) => McpConnectOutcome::Connected { conn, reauth_config: Some(config) },
186        Err(error) => McpConnectOutcome::Failed { error },
187    };
188
189    McpConnectAttempt { name, proxied, outcome }
190}
191
192impl McpConnectOutcome {
193    fn with_reauth(self, reauth_config: Option<McpHttpConfig>) -> Self {
194        match self {
195            Self::Connected { conn, .. } => Self::Connected { conn, reauth_config },
196            other => other,
197        }
198    }
199}
200
201async fn connect_stdio(
202    server_name: &str,
203    command: String,
204    args: Vec<String>,
205    env: HashMap<String, String>,
206    mcp_client: McpClient,
207    cwd: PathBuf,
208) -> McpConnectOutcome {
209    let mut cmd = Command::new(&command);
210    cmd.args(&args).envs(&env).current_dir(&cwd);
211
212    let (proc, stderr) = match TokioChildProcess::builder(cmd).stderr(Stdio::piped()).spawn() {
213        Ok(parts) => parts,
214        Err(e) => return McpConnectOutcome::Failed { error: McpError::SpawnFailed { command, reason: e.to_string() } },
215    };
216    let stderr_task = stderr.map(|stderr| spawn_stderr_logger(server_name.to_string(), stderr));
217
218    match serve_client(mcp_client, proc).await {
219        Ok(client) => McpConnectOutcome::Connected {
220            conn: McpServerConnection::from_parts(client, stderr_task),
221            reauth_config: None,
222        },
223        Err(e) => {
224            if let Some(task) = stderr_task {
225                task.abort();
226            }
227            McpConnectOutcome::Failed { error: McpError::from(e) }
228        }
229    }
230}
231
232fn spawn_stderr_logger(server_name: String, stderr: ChildStderr) -> JoinHandle<()> {
233    tokio::spawn(async move {
234        let mut lines = BufReader::new(stderr).lines();
235        loop {
236            match lines.next_line().await {
237                Ok(Some(line)) => tracing::info!(server = %server_name, stderr = %line, "MCP server stderr"),
238                Ok(None) => break,
239                Err(error) => {
240                    tracing::warn!(server = %server_name, %error, "failed to read MCP server stderr");
241                    break;
242                }
243            }
244        }
245    })
246}
247
248async fn connect_in_memory(
249    name: &str,
250    server: Box<dyn DynService<RoleServer>>,
251    mcp_client: McpClient,
252) -> McpConnectOutcome {
253    match serve_in_memory(server, mcp_client, name).await {
254        Ok((client, handle)) => McpConnectOutcome::Connected {
255            conn: McpServerConnection::from_parts(client, Some(handle)),
256            reauth_config: None,
257        },
258        Err(error) => McpConnectOutcome::Failed { error },
259    }
260}
261
262async fn connect_http(
263    name: &str,
264    config: McpHttpConfig,
265    mcp_client: McpClient,
266    oauth_handler_factory: Option<&OAuthHandlerFactory>,
267    oauth_credential_store: Option<&Arc<dyn OAuthCredentialStorage>>,
268) -> McpConnectOutcome {
269    let conn_err = |e| McpError::ConnectionFailed(format!("HTTP MCP server {name}: {e}"));
270    let stored_auth_manager = if let Some(store) = oauth_credential_store
271        && config.transport.auth_header.is_none()
272    {
273        match create_auth_manager_from_store(name, &config.transport.uri, config.oauth_client_id(), Arc::clone(store))
274            .await
275        {
276            Ok(manager) => manager,
277            Err(e) => {
278                tracing::warn!(
279                    server = %name,
280                    error = %e,
281                    "Failed to initialize auth manager from stored credentials, proceeding without auth"
282                );
283                None
284            }
285        }
286    } else {
287        None
288    };
289    let result = if let Some(auth_manager) = stored_auth_manager {
290        tracing::debug!("Using OAuth for server '{name}'");
291        let auth_client = AuthClient::new(reqwest::Client::default(), auth_manager);
292        let transport = StreamableHttpClientTransport::with_client(auth_client, config.transport.clone());
293        serve_client(mcp_client, transport).await.map_err(conn_err)
294    } else {
295        let transport = StreamableHttpClientTransport::from_config(config.transport.clone());
296        serve_client(mcp_client, transport).await.map_err(conn_err)
297    };
298
299    match result {
300        Ok(client) => {
301            McpConnectOutcome::Connected { conn: McpServerConnection::from_parts(client, None), reauth_config: None }
302        }
303        Err(error) => {
304            tracing::warn!("Failed to connect to MCP server '{name}': {error}");
305            if oauth_handler_factory.is_some() && config.transport.auth_header.is_none() {
306                McpConnectOutcome::NeedsOAuth { config, error }
307            } else {
308                McpConnectOutcome::Failed { error }
309            }
310        }
311    }
312}
313
314fn reauth_config_for(
315    transport: &McpTransport,
316    oauth_handler_factory: Option<&OAuthHandlerFactory>,
317) -> Option<McpHttpConfig> {
318    match transport {
319        McpTransport::Http(config) if oauth_handler_factory.is_some() && config.transport.auth_header.is_none() => {
320            Some(config.clone())
321        }
322        _ => None,
323    }
324}
325
326async fn serve_in_memory(
327    server: Box<dyn DynService<RoleServer>>,
328    mcp_client: McpClient,
329    label: &str,
330) -> Result<(RunningService<RoleClient, McpClient>, JoinHandle<()>)> {
331    let (client_transport, server_transport) = create_in_memory_transport();
332
333    let server_handle = tokio::spawn(async move {
334        match server.serve(server_transport).await {
335            Ok(_service) => {
336                std::future::pending::<()>().await;
337            }
338            Err(e) => {
339                eprintln!("MCP server error: {e}");
340            }
341        }
342    });
343
344    let client = serve_client(mcp_client, client_transport)
345        .await
346        .map_err(|e| McpError::ConnectionFailed(format!("Failed to connect to in-memory server '{label}': {e}")))?;
347
348    Ok((client, server_handle))
349}
mcp_utils/client/connection.rs

mcp_utils/client/
connection.rs
