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 = Some(value.strip_prefix("Bearer ").unwrap_or(value).to_string());
108                    continue;
109                }
110                req_headers.insert(
111                    reqwest::header::HeaderName::try_from(key)
112                        .with_context(|| format!("Invalid header name '{key}'"))?,
113                    value
114                        .parse()
115                        .with_context(|| format!("Invalid header value for '{key}'"))?,
116                );
117            }
118        }
119
120        let http_client = reqwest::Client::builder()
121            .default_headers(req_headers)
122            .connect_timeout(self.connect_timeout)
123            .build()
124            .context("Failed to create HTTP client")?;
125
126        let mut config = StreamableHttpClientTransportConfig::with_uri(self.url.clone());
127        config.auth_header = auth_header;
128
129        let transport = StreamableHttpClientTransport::with_client(http_client, config);
130        let client = bridge
131            .serve(transport)
132            .await
133            .context("Failed to connect URL backend")?;
134        info!(url = %self.url, "URL backend connected");
135        Ok(client)
136    }
137}
138
139#[async_trait]
140impl BackendConnector for UrlBackendConnector {
141    async fn connect_shared(
142        &self,
143        registry: Arc<UpstreamPeerRegistry>,
144    ) -> Result<BackendRunningService> {
145        self.serve_bridge(BackendNotificationBridge::with_default_info(registry))
146            .await
147    }
148
149    async fn connect_session(&self, slot: UpstreamPeerSlot) -> Result<BackendRunningService> {
150        self.serve_bridge(BackendNotificationBridge::for_session(slot))
151            .await
152    }
153
154    fn connect_timeout(&self) -> Duration {
155        self.connect_timeout
156    }
157}
158
159/// Stdio child-process backend (shared isolation only).
160pub struct StdioBackendConnector {
161    command: String,
162    args: Option<Vec<String>>,
163    env: Option<HashMap<String, String>>,
164    mcp_id: String,
165}
166
167impl StdioBackendConnector {
168    /// Create a stdio connector.
169    pub fn new(
170        command: impl Into<String>,
171        args: Option<Vec<String>>,
172        env: Option<HashMap<String, String>>,
173        mcp_id: impl Into<String>,
174    ) -> Self {
175        Self {
176            command: command.into(),
177            args,
178            env,
179            mcp_id: mcp_id.into(),
180        }
181    }
182
183    async fn spawn(&self, registry: Arc<UpstreamPeerRegistry>) -> Result<BackendRunningService> {
184        let args = self.args.clone();
185        let mut wrapped_cmd = CommandWrap::with_new(&self.command, |cmd| {
186            if let Some(cmd_args) = &args {
187                cmd.args(cmd_args);
188            }
189            if let Some(env_vars) = &self.env {
190                for (k, v) in env_vars {
191                    cmd.env(k, v);
192                }
193            }
194        });
195
196        #[cfg(unix)]
197        wrapped_cmd.wrap(ProcessGroup::leader());
198
199        #[cfg(windows)]
200        {
201            use windows::Win32::System::Threading::{CREATE_NEW_PROCESS_GROUP, CREATE_NO_WINDOW};
202            wrapped_cmd.wrap(CreationFlags(CREATE_NO_WINDOW | CREATE_NEW_PROCESS_GROUP));
203            wrapped_cmd.wrap(JobObject);
204        }
205
206        wrapped_cmd.wrap(KillOnDrop);
207
208        info!(
209            command = %self.command,
210            args = ?args.as_ref().unwrap_or(&vec![]),
211            "Starting stdio backend child process"
212        );
213
214        mcp_common::diagnostic::log_stdio_spawn_context(
215            "StdioBackendConnector",
216            &self.mcp_id,
217            &self.env,
218        );
219
220        let (tokio_process, child_stderr) = TokioChildProcess::builder(wrapped_cmd)
221            .stderr(Stdio::piped())
222            .spawn()
223            .map_err(|e| {
224                anyhow::anyhow!(
225                    "{}",
226                    mcp_common::diagnostic::format_spawn_error(
227                        &self.mcp_id,
228                        &self.command,
229                        &self.args,
230                        e
231                    )
232                )
233            })?;
234
235        if let Some(stderr_pipe) = child_stderr {
236            mcp_common::spawn_stderr_reader(stderr_pipe, self.mcp_id.clone());
237        }
238
239        let bridge = BackendNotificationBridge::with_default_info(registry);
240        let client = bridge
241            .serve(tokio_process)
242            .await
243            .context("Failed to connect stdio backend")?;
244        info!(mcp_id = %self.mcp_id, "Stdio backend connected");
245        Ok(client)
246    }
247}
248
249#[async_trait]
250impl BackendConnector for StdioBackendConnector {
251    async fn connect_shared(
252        &self,
253        registry: Arc<UpstreamPeerRegistry>,
254    ) -> Result<BackendRunningService> {
255        self.spawn(registry).await
256    }
257
258    async fn connect_session(&self, _slot: UpstreamPeerSlot) -> Result<BackendRunningService> {
259        anyhow::bail!(
260            "Stdio backend does not support per-session isolation in this release; \
261             use BackendIsolation::Shared"
262        )
263    }
264}
mcp_streamable_proxy/backend_connector.rs

mcp_streamable_proxy/
backend_connector.rs
