1use std::collections::BTreeMap;
2use std::process::Stdio;
3#[cfg(feature = "oauth")]
4use std::sync::Arc;
5use std::sync::Once;
6
7use rmcp::model::{CallToolRequestParams, GetPromptRequestParams, ReadResourceRequestParams, Tool};
8use rmcp::service::RunningService;
9use rmcp::transport::{
10 streamable_http_client::StreamableHttpClientTransportConfig, StreamableHttpClientTransport,
11 TokioChildProcess,
12};
13use rmcp::{RoleClient, ServiceExt};
14use serde_json::{Map, Value};
15use tokio::io::AsyncReadExt;
16use tokio::process::Command;
17
18use crate::config::UpstreamConfig;
19#[cfg(feature = "oauth")]
20use crate::oauth::UpstreamOAuthProvider;
21use crate::process::guard::SpawnGuard;
22use crate::upstream::http_body_cap::BodyCappedHttpClient;
23use crate::upstream::http_client::{decide_http_transport, HttpTransportDecision};
24use crate::upstream::transport::websocket::{
25 connect as connect_websocket_transport, WebSocketTransportConfig,
26};
27use crate::upstream::{
28 CapScope, PromptDescriptor, ResourceDescriptor, ResponseCaps, ToolDescriptor, TransportKind,
29 UpstreamError, UpstreamSnapshot,
30};
31
32#[derive(Clone)]
33pub(super) struct LiveConnectContext<'a> {
34 response_caps: &'a ResponseCaps,
35 #[cfg(feature = "oauth")]
36 oauth: Option<LiveOauthContext<'a>>,
37}
38
39#[cfg(feature = "oauth")]
40#[derive(Clone)]
41pub(super) struct LiveOauthContext<'a> {
42 pub subject: &'a str,
43 pub provider: Arc<dyn UpstreamOAuthProvider>,
44}
45
46impl<'a> LiveConnectContext<'a> {
47 pub(super) fn shared(response_caps: &'a ResponseCaps) -> Self {
48 Self {
49 response_caps,
50 #[cfg(feature = "oauth")]
51 oauth: None,
52 }
53 }
54
55 #[cfg(feature = "oauth")]
56 pub(super) fn oauth(
57 response_caps: &'a ResponseCaps,
58 subject: &'a str,
59 provider: Arc<dyn UpstreamOAuthProvider>,
60 ) -> Self {
61 Self {
62 response_caps,
63 oauth: Some(LiveOauthContext { subject, provider }),
64 }
65 }
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub(super) enum LiveKind {
70 Http(TransportKind),
71 Stdio,
72 WebSocket,
73}
74
75pub(super) struct LiveUpstream {
76 _service: RunningService<RoleClient, ()>,
77 peer: rmcp::service::Peer<RoleClient>,
78}
79
80impl LiveUpstream {
81 pub(super) fn peer(&self) -> rmcp::service::Peer<RoleClient> {
82 self.peer.clone()
83 }
84}
85
86pub(super) async fn connect_live(
87 config: &UpstreamConfig,
88 guard: &SpawnGuard,
89 context: LiveConnectContext<'_>,
90) -> Result<(LiveUpstream, UpstreamSnapshot), UpstreamError> {
91 let (service, peer, kind) = if let Some(url) = config.url.as_deref() {
92 match decide_http_transport(url) {
93 HttpTransportDecision::WebSocket => connect_websocket(config, url).await?,
94 HttpTransportDecision::Json | HttpTransportDecision::Sse => {
95 connect_http(config, url, context).await?
96 }
97 }
98 } else if let Some(command) = config.command.as_deref() {
99 connect_stdio(config, command, guard).await?
100 } else {
101 return Err(UpstreamError::Unsupported {
102 upstream: config.name.clone(),
103 capability: "transport",
104 });
105 };
106
107 let tools = peer
108 .list_all_tools()
109 .await
110 .map_err(|error| UpstreamError::connect(config, error))?;
111 let resources = if config.proxy_resources {
112 list_resources_or_empty(config, &peer).await?
113 } else {
114 Vec::new()
115 };
116 let prompts = if config.proxy_prompts {
117 list_prompts_or_empty(config, &peer).await?
118 } else {
119 Vec::new()
120 };
121
122 let mut snapshot = UpstreamSnapshot::empty(
123 config.name.clone(),
124 match kind {
125 LiveKind::Http(transport) => transport,
126 LiveKind::Stdio => TransportKind::Stdio,
127 LiveKind::WebSocket => TransportKind::WebSocket,
128 },
129 );
130 snapshot.tools = tools.into_iter().map(tool_descriptor).collect();
131 snapshot.resources = resources.into_iter().map(resource_descriptor).collect();
132 snapshot.prompts = prompts.into_iter().map(prompt_descriptor).collect();
133 Ok((
134 LiveUpstream {
135 _service: service,
136 peer,
137 },
138 snapshot,
139 ))
140}
141
142pub(super) async fn call_live_tool(
143 upstream: &str,
144 peer: rmcp::service::Peer<RoleClient>,
145 tool: String,
146 params: Value,
147) -> Result<Value, UpstreamError> {
148 let Value::Object(args) = params else {
149 return Err(UpstreamError::ParamsMustBeObject);
150 };
151 let result = peer
152 .call_tool(CallToolRequestParams::new(tool).with_arguments(args))
153 .await
154 .map_err(|error| UpstreamError::LiveCall {
155 upstream: upstream.to_owned(),
156 operation: "tools/call",
157 message: error.to_string(),
158 })?;
159 if let Some(value) = result.structured_content.clone() {
160 return Ok(value);
161 }
162 Ok(serde_json::to_value(result).unwrap_or(Value::Null))
163}
164
165pub(super) async fn read_live_resource(
166 upstream: &str,
167 peer: rmcp::service::Peer<RoleClient>,
168 uri: String,
169) -> Result<Value, UpstreamError> {
170 let result = peer
171 .read_resource(ReadResourceRequestParams::new(uri))
172 .await
173 .map_err(|error| UpstreamError::LiveCall {
174 upstream: upstream.to_owned(),
175 operation: "resources/read",
176 message: error.to_string(),
177 })?;
178 serde_json::to_value(result).map_err(|error| UpstreamError::LiveCall {
179 upstream: upstream.to_owned(),
180 operation: "resources/read",
181 message: error.to_string(),
182 })
183}
184
185pub(super) async fn get_live_prompt(
186 upstream: &str,
187 peer: rmcp::service::Peer<RoleClient>,
188 name: String,
189 arguments: Option<Map<String, Value>>,
190) -> Result<Value, UpstreamError> {
191 let mut params = GetPromptRequestParams::new(name);
192 params.arguments = arguments;
193 let result = peer
194 .get_prompt(params)
195 .await
196 .map_err(|error| UpstreamError::LiveCall {
197 upstream: upstream.to_owned(),
198 operation: "prompts/get",
199 message: error.to_string(),
200 })?;
201 serde_json::to_value(result).map_err(|error| UpstreamError::LiveCall {
202 upstream: upstream.to_owned(),
203 operation: "prompts/get",
204 message: error.to_string(),
205 })
206}
207
208async fn connect_http(
209 config: &UpstreamConfig,
210 url: &str,
211 context: LiveConnectContext<'_>,
212) -> Result<
213 (
214 RunningService<RoleClient, ()>,
215 rmcp::service::Peer<RoleClient>,
216 LiveKind,
217 ),
218 UpstreamError,
219> {
220 ensure_rustls_crypto_provider();
221 let transport_kind = match decide_http_transport(url) {
222 HttpTransportDecision::Json => TransportKind::HttpJson,
223 HttpTransportDecision::Sse => TransportKind::HttpSse,
224 HttpTransportDecision::WebSocket => TransportKind::WebSocket,
225 };
226 let mut transport_config = StreamableHttpClientTransportConfig::with_uri(url.to_owned());
227 #[cfg(feature = "oauth")]
228 if config.oauth.is_some() {
229 let oauth = context.oauth.ok_or_else(|| UpstreamError::LiveConnect {
230 upstream: config.name.clone(),
231 message: "oauth upstream requires subject-scoped connection context".to_owned(),
232 })?;
233 let client = BodyCappedHttpClient::default_with_caps(
234 context.response_caps.limit_for(CapScope::HttpJson),
235 context.response_caps.limit_for(CapScope::HttpSseEvent),
236 );
237 let auth_client = oauth
238 .provider
239 .authenticated_http_client(config, oauth.subject, client)
240 .await
241 .map_err(|error| UpstreamError::LiveConnect {
242 upstream: config.name.clone(),
243 message: error.to_string(),
244 })?;
245 let transport = StreamableHttpClientTransport::with_client(auth_client, transport_config);
246 let service = ().serve(transport).await.map_err(|error| {
247 UpstreamError::connect(config, format!("oauth http connect failed: {error}"))
248 })?;
249 let peer = service.peer().clone();
250 return Ok((service, peer, LiveKind::Http(transport_kind)));
251 }
252 #[cfg(not(feature = "oauth"))]
253 if config.oauth.is_some() {
254 return Err(UpstreamError::LiveConnect {
255 upstream: config.name.clone(),
256 message: "oauth upstream support is not compiled into soma-mcp-client".to_owned(),
257 });
258 }
259 if let Some(token) = bearer_token_from_env(config) {
260 transport_config = transport_config.auth_header(token);
261 }
262 let client = BodyCappedHttpClient::default_with_caps(
263 context.response_caps.limit_for(CapScope::HttpJson),
264 context.response_caps.limit_for(CapScope::HttpSseEvent),
265 );
266 let transport = StreamableHttpClientTransport::with_client(client, transport_config);
267 let service = ()
268 .serve(transport)
269 .await
270 .map_err(|error| UpstreamError::connect(config, format!("http connect failed: {error}")))?;
271 let peer = service.peer().clone();
272 Ok((service, peer, LiveKind::Http(transport_kind)))
273}
274
275async fn connect_websocket(
276 config: &UpstreamConfig,
277 url: &str,
278) -> Result<
279 (
280 RunningService<RoleClient, ()>,
281 rmcp::service::Peer<RoleClient>,
282 LiveKind,
283 ),
284 UpstreamError,
285> {
286 ensure_rustls_crypto_provider();
287 let transport_config = WebSocketTransportConfig::new(url.to_owned())
288 .with_authorization(websocket_authorization(config));
289 let service =
290 ().serve(connect_websocket_transport(transport_config))
291 .await
292 .map_err(|error| {
293 UpstreamError::connect(config, format!("websocket connect failed: {error}"))
294 })?;
295 let peer = service.peer().clone();
296 Ok((service, peer, LiveKind::WebSocket))
297}
298
299async fn connect_stdio(
300 config: &UpstreamConfig,
301 command: &str,
302 guard: &SpawnGuard,
303) -> Result<
304 (
305 RunningService<RoleClient, ()>,
306 rmcp::service::Peer<RoleClient>,
307 LiveKind,
308 ),
309 UpstreamError,
310> {
311 let spec = crate::upstream::pool::connect_stdio::plan_stdio_connection(config, guard).map_err(
312 |error| UpstreamError::LiveConnect {
313 upstream: config.name.clone(),
314 message: error.to_string(),
315 },
316 )?;
317 let mut cmd = Command::new(command);
318 cmd.args(&spec.args)
319 .env_clear()
320 .envs(stdio_env())
321 .envs(spec.env.iter())
322 .stderr(Stdio::piped());
323 if let Some(env_name) = config.bearer_token_env.as_deref() {
324 if let Ok(token) = std::env::var(env_name) {
325 cmd.env(env_name, token);
326 }
327 }
328
329 #[cfg(unix)]
330 let command = {
331 use process_wrap::tokio::{CommandWrap, ProcessGroup};
332 let mut wrapped = CommandWrap::from(cmd);
333 wrapped.wrap(ProcessGroup::leader());
334 wrapped
335 };
336 #[cfg(not(unix))]
337 let command = cmd;
338
339 let (transport, stderr) = TokioChildProcess::builder(command)
340 .spawn()
341 .map_err(|error| UpstreamError::LiveConnect {
342 upstream: config.name.clone(),
343 message: format!("stdio spawn failed: {error}"),
344 })?;
345 drain_stderr(config.name.clone(), stderr);
346 let service = ().serve(transport).await.map_err(|error| {
347 UpstreamError::connect(config, format!("stdio initialize failed: {error}"))
348 })?;
349 let peer = service.peer().clone();
350 Ok((service, peer, LiveKind::Stdio))
351}
352
353async fn list_resources_or_empty(
354 config: &UpstreamConfig,
355 peer: &rmcp::service::Peer<RoleClient>,
356) -> Result<Vec<rmcp::model::Resource>, UpstreamError> {
357 match peer.list_all_resources().await {
358 Ok(resources) => Ok(resources),
359 Err(error) if capability_is_absent(&error.to_string()) => Ok(Vec::new()),
360 Err(error) => Err(UpstreamError::LiveConnect {
361 upstream: config.name.clone(),
362 message: format!("resources/list failed: {error}"),
363 }),
364 }
365}
366
367async fn list_prompts_or_empty(
368 config: &UpstreamConfig,
369 peer: &rmcp::service::Peer<RoleClient>,
370) -> Result<Vec<rmcp::model::Prompt>, UpstreamError> {
371 match peer.list_all_prompts().await {
372 Ok(prompts) => Ok(prompts),
373 Err(error) if capability_is_absent(&error.to_string()) => Ok(Vec::new()),
374 Err(error) => Err(UpstreamError::LiveConnect {
375 upstream: config.name.clone(),
376 message: format!("prompts/list failed: {error}"),
377 }),
378 }
379}
380
381fn tool_descriptor(tool: Tool) -> ToolDescriptor {
382 ToolDescriptor {
383 name: tool.name.to_string(),
384 description: tool.description.map(|value| value.to_string()),
385 input_schema: Some(Value::Object((*tool.input_schema).clone())),
386 output_schema: tool
387 .output_schema
388 .map(|schema| Value::Object((*schema).clone())),
389 destructive: tool
390 .annotations
391 .as_ref()
392 .and_then(|annotations| annotations.destructive_hint)
393 .unwrap_or(true),
394 }
395}
396
397fn resource_descriptor(resource: rmcp::model::Resource) -> ResourceDescriptor {
398 ResourceDescriptor {
399 uri: resource.uri,
400 name: Some(resource.name),
401 }
402}
403
404fn prompt_descriptor(prompt: rmcp::model::Prompt) -> PromptDescriptor {
405 PromptDescriptor {
406 name: prompt.name,
407 description: prompt.description,
408 }
409}
410
411fn normalize_bearer_value(token: &str) -> String {
412 token
413 .trim()
414 .strip_prefix("Bearer ")
415 .unwrap_or_else(|| token.trim())
416 .to_owned()
417}
418
419fn websocket_authorization(config: &UpstreamConfig) -> Option<String> {
420 bearer_token_from_env(config).map(|token| format!("Bearer {token}"))
421}
422
423fn bearer_token_from_env(config: &UpstreamConfig) -> Option<String> {
424 let env_name = config.bearer_token_env.as_deref()?;
425 let token = std::env::var(env_name).ok()?;
426 let token = normalize_bearer_value(&token);
427 (!token.is_empty()).then_some(token)
428}
429
430fn capability_is_absent(error: &str) -> bool {
431 error.contains("-32601")
432 || error.contains("Method not found")
433 || error.contains("method not found")
434}
435
436fn ensure_rustls_crypto_provider() {
437 static INSTALL: Once = Once::new();
438 INSTALL.call_once(|| {
439 let _ = rustls::crypto::ring::default_provider().install_default();
440 });
441}
442
443fn stdio_env() -> BTreeMap<String, String> {
444 const ALLOWLIST: &[&str] = &[
445 "PATH",
446 "HOME",
447 "USER",
448 "LOGNAME",
449 "TERM",
450 "TZ",
451 "TMPDIR",
452 "TMP",
453 "TEMP",
454 "LANG",
455 "LC_ALL",
456 "XDG_CACHE_HOME",
457 "XDG_CONFIG_HOME",
458 "XDG_DATA_HOME",
459 "SSL_CERT_FILE",
460 "SSL_CERT_DIR",
461 "NODE_EXTRA_CA_CERTS",
462 "REQUESTS_CA_BUNDLE",
463 "CURL_CA_BUNDLE",
464 ];
465 ALLOWLIST
466 .iter()
467 .filter_map(|key| {
468 std::env::var(key)
469 .ok()
470 .map(|value| ((*key).to_owned(), value))
471 })
472 .collect()
473}
474
475fn drain_stderr(upstream: String, stderr: Option<tokio::process::ChildStderr>) {
476 let Some(mut stderr) = stderr else {
477 return;
478 };
479 tokio::spawn(async move {
480 let mut bytes = Vec::new();
481 if stderr.read_to_end(&mut bytes).await.is_ok() && !bytes.is_empty() {
482 tracing::debug!(
483 upstream,
484 stderr = %String::from_utf8_lossy(&bytes),
485 "upstream stdio stderr"
486 );
487 }
488 });
489}
490
491#[cfg(test)]
492#[path = "live_tests.rs"]
493mod tests;