Skip to main content

mcp_utils/client/
transport.rs

1use super::error::{ClientError, Result};
2use super::sigv4::SigV4HttpClient;
3use super::{
4    ClientOptions,
5    handler::Handler,
6    oauth::{self, OAuthSettings},
7};
8use crate::server::McpServer;
9use aether_auth::SigV4Signer;
10use reqwest::header::HeaderMap;
11use rmcp::{
12    ClientLifecycleMode, RoleClient,
13    model::ProtocolVersion,
14    serve_client_with_lifecycle,
15    service::{ClientInitializeError, RunningService},
16    transport::{
17        IntoTransport, StreamableHttpClientTransport, TokioChildProcess,
18        streamable_http_client::StreamableHttpClientTransportConfig,
19    },
20};
21use std::collections::HashMap;
22use std::path::PathBuf;
23use std::process::Stdio;
24use std::sync::Arc;
25use tokio::{
26    io::{AsyncBufReadExt, BufReader},
27    net::UnixStream,
28    process::{ChildStderr, Command},
29};
30
31#[derive(Debug, Clone)]
32pub enum Transport {
33    Stdio { command: String, args: Vec<String>, env: HashMap<String, String>, cwd: Option<PathBuf> },
34    Http { url: String, headers: HeaderMap, auth: HttpAuth },
35    InProcess(McpServer),
36    Unix(PathBuf),
37}
38
39#[derive(Debug, Clone)]
40pub enum HttpAuth {
41    None,
42    OAuth(OAuthSettings),
43    SigV4(SigV4Signer),
44}
45
46pub(super) type ClientService = RunningService<RoleClient, Arc<Handler>>;
47
48impl Transport {
49    pub(crate) fn supports_oauth(&self, options: &ClientOptions) -> bool {
50        options.oauth_prompts().is_some() && matches!(self, Self::Http { auth: HttpAuth::OAuth(_), .. })
51    }
52
53    pub(super) async fn connect(self, handler: Handler, options: &ClientOptions) -> Result<ClientService> {
54        let supports_oauth = self.supports_oauth(options);
55        let handler = Arc::new(handler);
56        let connected = match self {
57            Self::Stdio { command, args, env, cwd } => connect_stdio(handler, command, args, env, cwd).await,
58            Self::Http { url, headers, auth } => connect_http(handler, url, headers, auth, options).await,
59            Self::InProcess(server) => serve(handler, server.serve_in_memory()).await,
60            Self::Unix(path) => connect_unix(handler, path).await,
61        };
62        match connected {
63            Err(ClientError::Connect { server, source }) if supports_oauth && requires_authorization(&source) => {
64                Err(ClientError::AuthRequired { server, challenge: source.auth_challenge().map(str::to_string) })
65            }
66            connected => connected,
67        }
68    }
69
70    pub(super) async fn authorize(
71        self,
72        handler: Handler,
73        options: &ClientOptions,
74        challenge: Option<String>,
75    ) -> Result<ClientService> {
76        let server = handler.server().to_string();
77        let Self::Http { url, headers, auth: HttpAuth::OAuth(oauth) } = self else {
78            return Err(ClientError::OAuthUnavailable { server });
79        };
80        let auth_client = oauth::authorize(&server, &url, oauth, options, challenge).await?;
81        let transport = StreamableHttpClientTransport::with_client(auth_client, http_config(url, headers));
82        serve(Arc::new(handler), transport).await
83    }
84}
85
86fn client_lifecycle_mode() -> ClientLifecycleMode {
87    ClientLifecycleMode::Auto {
88        preferred_versions: vec![
89            ProtocolVersion::V_2026_07_28,
90            ProtocolVersion::V_2025_11_25,
91            ProtocolVersion::V_2025_06_18,
92            ProtocolVersion::V_2025_03_26,
93            ProtocolVersion::V_2024_11_05,
94        ],
95        legacy_version: Some(ProtocolVersion::V_2025_11_25),
96    }
97}
98
99async fn connect_stdio(
100    handler: Arc<Handler>,
101    command: String,
102    args: Vec<String>,
103    env: HashMap<String, String>,
104    cwd: Option<PathBuf>,
105) -> Result<ClientService> {
106    let mut cmd = Command::new(&command);
107    cmd.args(&args).envs(&env);
108    if let Some(cwd) = cwd {
109        cmd.current_dir(cwd);
110    }
111    let (process, stderr) = TokioChildProcess::builder(cmd)
112        .stderr(Stdio::piped())
113        .spawn()
114        .map_err(|source| ClientError::Spawn { command, source })?;
115
116    if let Some(stderr) = stderr {
117        spawn_stderr_logger(handler.server().to_string(), stderr);
118    }
119
120    serve(handler, process).await
121}
122
123async fn connect_http(
124    handler: Arc<Handler>,
125    url: String,
126    headers: HeaderMap,
127    auth: HttpAuth,
128    options: &ClientOptions,
129) -> Result<ClientService> {
130    let config = http_config(url, headers);
131    match auth {
132        HttpAuth::None => serve(handler, StreamableHttpClientTransport::from_config(config)).await,
133        HttpAuth::OAuth(oauth) => match oauth::restore(handler.server(), &config.uri, &oauth, options).await {
134            Some(auth_client) => serve(handler, StreamableHttpClientTransport::with_client(auth_client, config)).await,
135            None => serve(handler, StreamableHttpClientTransport::from_config(config)).await,
136        },
137        HttpAuth::SigV4(signer) => {
138            serve(handler, StreamableHttpClientTransport::with_client(SigV4HttpClient::new(signer), config)).await
139        }
140    }
141}
142
143async fn connect_unix(handler: Arc<Handler>, path: PathBuf) -> Result<ClientService> {
144    let stream = UnixStream::connect(&path).await.map_err(|source| ClientError::SocketConnect { path, source })?;
145    serve(handler, stream).await
146}
147
148async fn serve<T, U, V>(handler: Arc<Handler>, transport: T) -> Result<ClientService>
149where
150    T: IntoTransport<RoleClient, U, V>,
151    U: std::error::Error + Send + Sync + 'static,
152{
153    let server = handler.server().to_string();
154    serve_client_with_lifecycle(handler, transport, client_lifecycle_mode())
155        .await
156        .map_err(|source| ClientError::connect(server, source))
157}
158
159fn requires_authorization(error: &ClientInitializeError) -> bool {
160    error.is_authorization_required() || error.auth_challenge().is_some()
161}
162
163fn http_config(url: String, headers: HeaderMap) -> StreamableHttpClientTransportConfig {
164    let headers = headers.into_iter().filter_map(|(name, value)| Some((name?, value))).collect();
165    StreamableHttpClientTransportConfig::with_uri(url).custom_headers(headers)
166}
167
168fn spawn_stderr_logger(server: String, stderr: ChildStderr) {
169    tokio::spawn(async move {
170        let mut lines = BufReader::new(stderr).lines();
171        loop {
172            match lines.next_line().await {
173                Ok(Some(line)) => tracing::info!(server = %server, stderr = %line, "MCP server stderr"),
174                Ok(None) => break,
175                Err(error) => {
176                    tracing::warn!(server = %server, %error, "failed to read MCP server stderr");
177                    break;
178                }
179            }
180        }
181    });
182}