mcp_streamable_proxy/
backend_connector.rs

1//! Backend connection factories for shared (stdio) and per-session (URL) isolation.
2
3use std::collections::HashMap;
4use std::process::Stdio;
5use std::sync::Arc;
6use std::time::Duration;
7
8use anyhow::{Context, Result};
9use async_trait::async_trait;
10use process_wrap::tokio::{CommandWrap, KillOnDrop};
11use rmcp::{
12    ServiceExt,
13    transport::{
14        TokioChildProcess,
15        streamable_http_client::{
16            StreamableHttpClientTransport, StreamableHttpClientTransportConfig,
17        },
18    },
19};
20use tracing::info;
21
22#[cfg(unix)]
23use process_wrap::tokio::ProcessGroup;
24
25#[cfg(windows)]
26use process_wrap::tokio::{CreationFlags, JobObject};
27
28use crate::backend_client::{BackendNotificationBridge, UpstreamPeerRegistry, UpstreamPeerSlot};
29use crate::proxy_handler::BackendRunningService;
30
31/// How upstream sessions share (or isolate) backend MCP connections.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
33pub enum BackendIsolation {
34    /// One backend connection shared by all upstream sessions (stdio default).
35    #[default]
36    Shared,
37    /// Each upstream session opens its own backend connection (URL default).
38    PerSession,
39}
40
41impl BackendIsolation {
42    /// Default isolation for a URL backend.
43    pub fn for_url() -> Self {
44        Self::PerSession
45    }
46
47    /// Default isolation for a stdio backend.
48    pub fn for_stdio() -> Self {
49        Self::Shared
50    }
51}
52
53/// Creates backend [`BackendRunningService`] instances for the proxy.
54#[async_trait]
55pub trait BackendConnector: Send + Sync {
56    /// Connect using a shared fan-out registry (stdio / shared isolation).
57    ///
58    /// For stdio, call at most once at process startup.
59    async fn connect_shared(
60        &self,
61        registry: Arc<UpstreamPeerRegistry>,
62    ) -> Result<BackendRunningService>;
63
64    /// Connect for one session with a 1:1 notify slot (URL / per-session).
65    async fn connect_session(&self, slot: UpstreamPeerSlot) -> Result<BackendRunningService>;
66
67    /// Timeout used when awaiting lazy per-session connect.
68    fn connect_timeout(&self) -> Duration {
69        Duration::from_secs(30)
70    }
71}
72
73/// Connect to a remote Streamable HTTP MCP URL (new connection each call).
74#[derive(Debug, Clone)]
75pub struct UrlBackendConnector {
76    pub url: String,
77    pub headers: Option<HashMap<String, String>>,
78    pub connect_timeout: Duration,
79}
80
81impl UrlBackendConnector {
82    /// Create a URL connector.
83    pub fn new(url: impl Into<String>, headers: Option<HashMap<String, String>>) -> Self {
84        Self {
85            url: url.into(),
86            headers,
87            connect_timeout: Duration::from_secs(30),
88        }
89    }
90
91    /// Override connect timeout used by lazy session connect.
92    pub fn with_connect_timeout(mut self, timeout: Duration) -> Self {
93        self.connect_timeout = timeout;
94        self
95    }
96
97    async fn serve_bridge(
98        &self,
99        bridge: BackendNotificationBridge,
100    ) -> Result<BackendRunningService> {
101        let mut req_headers = reqwest::header::HeaderMap::new();
102        let mut auth_header: Option<String> = None;
103
104        if let Some(config_headers) = &self.headers {
105            for (key, value) in config_headers {
106                if key.eq_ignore_ascii_case("Authorization") {
107                    auth_header =
108                        Some(value.strip_prefix("Bearer ").unwrap_or(value).to_string());
109                    continue;
110                }
111                req_headers.insert(
112                    reqwest::header::HeaderName::try_from(key)
113                        .with_context(|| format!("Invalid header name '{key}'"))?,
114                    value
115                        .parse()
116                        .with_context(|| format!("Invalid header value for '{key}'"))?,
117                );
118            }
119        }
120
121        let http_client = reqwest::Client::builder()
122            .default_headers(req_headers)
123            .connect_timeout(self.connect_timeout)
124            .build()
125            .context("Failed to create HTTP client")?;
126
127        let mut config = StreamableHttpClientTransportConfig::with_uri(self.url.clone());
128        config.auth_header = auth_header;
129
130        let transport = StreamableHttpClientTransport::with_client(http_client, config);
131        let client = bridge
132            .serve(transport)
133            .await
134            .context("Failed to connect URL backend")?;
135        info!(url = %self.url, "URL backend connected");
136        Ok(client)
137    }
138}
139
140#[async_trait]
141impl BackendConnector for UrlBackendConnector {
142    async fn connect_shared(
143        &self,
144        registry: Arc<UpstreamPeerRegistry>,
145    ) -> Result<BackendRunningService> {
146        self.serve_bridge(BackendNotificationBridge::with_default_info(registry))
147            .await
148    }
149
150    async fn connect_session(&self, slot: UpstreamPeerSlot) -> Result<BackendRunningService> {
151        self.serve_bridge(BackendNotificationBridge::for_session(slot))
152            .await
153    }
154
155    fn connect_timeout(&self) -> Duration {
156        self.connect_timeout
157    }
158}
159
160/// Stdio child-process backend (shared isolation only).
161pub struct StdioBackendConnector {
162    command: String,
163    args: Option<Vec<String>>,
164    env: Option<HashMap<String, String>>,
165    mcp_id: String,
166}
167
168impl StdioBackendConnector {
169    /// Create a stdio connector.
170    pub fn new(
171        command: impl Into<String>,
172        args: Option<Vec<String>>,
173        env: Option<HashMap<String, String>>,
174        mcp_id: impl Into<String>,
175    ) -> Self {
176        Self {
177            command: command.into(),
178            args,
179            env,
180            mcp_id: mcp_id.into(),
181        }
182    }
183
184    async fn spawn(
185        &self,
186        registry: Arc<UpstreamPeerRegistry>,
187    ) -> Result<BackendRunningService> {
188        let args = self.args.clone();
189        let mut wrapped_cmd = CommandWrap::with_new(&self.command, |cmd| {
190            if let Some(cmd_args) = &args {
191                cmd.args(cmd_args);
192            }
193            if let Some(env_vars) = &self.env {
194                for (k, v) in env_vars {
195                    cmd.env(k, v);
196                }
197            }
198        });
199
200        #[cfg(unix)]
201        wrapped_cmd.wrap(ProcessGroup::leader());
202
203        #[cfg(windows)]
204        {
205            use windows::Win32::System::Threading::{CREATE_NEW_PROCESS_GROUP, CREATE_NO_WINDOW};
206            wrapped_cmd.wrap(CreationFlags(CREATE_NO_WINDOW | CREATE_NEW_PROCESS_GROUP));
207            wrapped_cmd.wrap(JobObject);
208        }
209
210        wrapped_cmd.wrap(KillOnDrop);
211
212        info!(
213            command = %self.command,
214            args = ?args.as_ref().unwrap_or(&vec![]),
215            "Starting stdio backend child process"
216        );
217
218        mcp_common::diagnostic::log_stdio_spawn_context(
219            "StdioBackendConnector",
220            &self.mcp_id,
221            &self.env,
222        );
223
224        let (tokio_process, child_stderr) = TokioChildProcess::builder(wrapped_cmd)
225            .stderr(Stdio::piped())
226            .spawn()
227            .map_err(|e| {
228                anyhow::anyhow!(
229                    "{}",
230                    mcp_common::diagnostic::format_spawn_error(
231                        &self.mcp_id,
232                        &self.command,
233                        &self.args,
234                        e
235                    )
236                )
237            })?;
238
239        if let Some(stderr_pipe) = child_stderr {
240            mcp_common::spawn_stderr_reader(stderr_pipe, self.mcp_id.clone());
241        }
242
243        let bridge = BackendNotificationBridge::with_default_info(registry);
244        let client = bridge
245            .serve(tokio_process)
246            .await
247            .context("Failed to connect stdio backend")?;
248        info!(mcp_id = %self.mcp_id, "Stdio backend connected");
249        Ok(client)
250    }
251}
252
253#[async_trait]
254impl BackendConnector for StdioBackendConnector {
255    async fn connect_shared(
256        &self,
257        registry: Arc<UpstreamPeerRegistry>,
258    ) -> Result<BackendRunningService> {
259        self.spawn(registry).await
260    }
261
262    async fn connect_session(&self, _slot: UpstreamPeerSlot) -> Result<BackendRunningService> {
263        anyhow::bail!(
264            "Stdio backend does not support per-session isolation in this release; \
265             use BackendIsolation::Shared"
266        )
267    }
268}
mcp_streamable_proxy/backend_connector.rs

mcp_streamable_proxy/
backend_connector.rs
