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}