mcp_streamable_proxy/
server_builder.rs

1//! Streamable HTTP Server Builder
2//!
3//! This module provides a high-level Builder API for creating Streamable HTTP MCP servers.
4//! It encapsulates all rmcp-specific types and provides a simple interface for mcp-proxy.
5
6use std::collections::HashMap;
7use std::sync::Arc;
8use std::time::Duration;
9
10use anyhow::{bail, Result};
11use tokio_util::sync::CancellationToken;
12use tracing::{info, warn};
13
14use rmcp::transport::streamable_http_server::{
15    StreamableHttpServerConfig, StreamableHttpService,
16};
17
18use crate::backend_client::UpstreamPeerRegistry;
19use crate::backend_connector::{
20    BackendConnector, BackendIsolation, StdioBackendConnector, UrlBackendConnector,
21};
22use crate::{ProxyAwareSessionManager, ProxyHandler, ToolFilter};
23
24/// Backend configuration for the MCP server
25///
26/// Defines how the proxy connects to the upstream MCP service.
27#[derive(Debug, Clone)]
28pub enum BackendConfig {
29    /// Connect to a local command via stdio
30    Stdio {
31        /// Command to execute (e.g., "npx", "python", etc.)
32        command: String,
33        /// Arguments for the command
34        args: Option<Vec<String>>,
35        /// Environment variables from MCP JSON config
36        env: Option<HashMap<String, String>>,
37    },
38    /// Connect to a remote URL
39    Url {
40        /// URL of the MCP service
41        url: String,
42        /// Custom HTTP headers (including Authorization)
43        headers: Option<HashMap<String, String>>,
44    },
45}
46
47/// Configuration for the Streamable HTTP server
48#[derive(Debug, Clone)]
49pub struct StreamServerConfig {
50    /// Enable stateful mode with session management
51    pub stateful_mode: bool,
52    /// MCP service identifier for logging
53    pub mcp_id: Option<String>,
54    /// Tool filter configuration
55    pub tool_filter: Option<ToolFilter>,
56    /// Backend isolation (`None` = default by backend type)
57    pub backend_isolation: Option<BackendIsolation>,
58    /// Max concurrent upstream sessions (`None` = unlimited)
59    pub max_sessions: Option<usize>,
60    /// Per-session backend connect timeout (URL lazy connect)
61    pub connect_timeout: Duration,
62}
63
64impl Default for StreamServerConfig {
65    fn default() -> Self {
66        Self {
67            stateful_mode: false,
68            mcp_id: None,
69            tool_filter: None,
70            backend_isolation: None,
71            max_sessions: None,
72            connect_timeout: Duration::from_secs(30),
73        }
74    }
75}
76
77/// Builder for creating Streamable HTTP MCP servers
78pub struct StreamServerBuilder {
79    backend_config: BackendConfig,
80    server_config: StreamServerConfig,
81}
82
83impl StreamServerBuilder {
84    /// Create a new builder with the given backend configuration
85    pub fn new(backend: BackendConfig) -> Self {
86        Self {
87            backend_config: backend,
88            server_config: StreamServerConfig::default(),
89        }
90    }
91
92    /// Set whether to enable stateful mode
93    pub fn stateful(mut self, enabled: bool) -> Self {
94        self.server_config.stateful_mode = enabled;
95        self
96    }
97
98    /// Set the MCP service identifier
99    pub fn mcp_id(mut self, id: impl Into<String>) -> Self {
100        self.server_config.mcp_id = Some(id.into());
101        self
102    }
103
104    /// Set the tool filter configuration
105    pub fn tool_filter(mut self, filter: ToolFilter) -> Self {
106        self.server_config.tool_filter = Some(filter);
107        self
108    }
109
110    /// Override backend isolation (URL defaults to per-session, stdio to shared).
111    pub fn backend_isolation(mut self, isolation: BackendIsolation) -> Self {
112        self.server_config.backend_isolation = Some(isolation);
113        self
114    }
115
116    /// Limit concurrent upstream sessions.
117    pub fn max_sessions(mut self, max: usize) -> Self {
118        self.server_config.max_sessions = Some(max);
119        self
120    }
121
122    /// Timeout for per-session lazy backend connect.
123    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
124        self.server_config.connect_timeout = timeout;
125        self
126    }
127
128    /// Build the server and return an axum Router, CancellationToken, and a management ProxyHandler.
129    ///
130    /// For **Shared** isolation the returned handler is the live shared backend.
131    /// For **PerSession** isolation it is a disconnected stub (real backends live
132    /// inside each session's handler); do not rely on `swap_backend` / availability APIs.
133    pub async fn build(self) -> Result<(axum::Router, CancellationToken, ProxyHandler)> {
134        let mcp_id = self
135            .server_config
136            .mcp_id
137            .clone()
138            .unwrap_or_else(|| "stream-proxy".into());
139        let tool_filter = self
140            .server_config
141            .tool_filter
142            .clone()
143            .unwrap_or_default();
144
145        let isolation = self.server_config.backend_isolation.unwrap_or(
146            match &self.backend_config {
147                BackendConfig::Url { .. } => BackendIsolation::for_url(),
148                BackendConfig::Stdio { .. } => BackendIsolation::for_stdio(),
149            },
150        );
151
152        if let (BackendConfig::Stdio { .. }, BackendIsolation::PerSession) =
153            (&self.backend_config, isolation)
154        {
155            bail!("Stdio backend does not support BackendIsolation::PerSession; use Shared");
156        }
157
158        let (router, ct, management_handler) = match (&self.backend_config, isolation) {
159            (BackendConfig::Stdio { command, args, env }, BackendIsolation::Shared) => {
160                let connector = StdioBackendConnector::new(
161                    command.clone(),
162                    args.clone(),
163                    env.clone(),
164                    mcp_id.clone(),
165                );
166                let registry = Arc::new(UpstreamPeerRegistry::new());
167                let client = connector.connect_shared(registry).await?;
168                let handler = ProxyHandler::with_tool_filter(client, mcp_id.clone(), tool_filter);
169                let (router, ct) = self
170                    .create_shared_server(handler.clone())
171                    .await?;
172                (router, ct, handler)
173            }
174            (BackendConfig::Url { url, headers }, BackendIsolation::Shared) => {
175                let connector = Arc::new(
176                    UrlBackendConnector::new(url.clone(), headers.clone())
177                        .with_connect_timeout(self.server_config.connect_timeout),
178                );
179                let registry = Arc::new(UpstreamPeerRegistry::new());
180                let client = connector.connect_shared(registry).await?;
181                let handler = ProxyHandler::with_tool_filter(client, mcp_id.clone(), tool_filter);
182                let (router, ct) = self
183                    .create_shared_server(handler.clone())
184                    .await?;
185                (router, ct, handler)
186            }
187            (BackendConfig::Url { url, headers }, BackendIsolation::PerSession) => {
188                let connector: Arc<dyn BackendConnector> = Arc::new(
189                    UrlBackendConnector::new(url.clone(), headers.clone())
190                        .with_connect_timeout(self.server_config.connect_timeout),
191                );
192                let (router, ct) = self
193                    .create_per_session_server(connector.clone(), mcp_id.clone(), tool_filter.clone())
194                    .await?;
195                // Management stub (disconnected); real backends live per session.
196                let management =
197                    ProxyHandler::new_per_session(connector, mcp_id.clone(), tool_filter);
198                (router, ct, management)
199            }
200            (BackendConfig::Stdio { .. }, BackendIsolation::PerSession) => unreachable!(),
201        };
202
203        info!(
204            "[StreamServerBuilder] Server created - mcp_id: {}, stateful: {}, isolation: {:?}",
205            mcp_id, self.server_config.stateful_mode, isolation
206        );
207
208        Ok((router, ct, management_handler))
209    }
210
211    async fn create_shared_server(
212        &self,
213        proxy_handler: ProxyHandler,
214    ) -> Result<(axum::Router, CancellationToken)> {
215        let handler = Arc::new(proxy_handler);
216        let ct = CancellationToken::new();
217
218        if self.server_config.stateful_mode {
219            let session_manager = ProxyAwareSessionManager::new(handler.clone())
220                .with_max_sessions(self.server_config.max_sessions);
221            let handler_for_service = handler.clone();
222            let mut server_config = StreamableHttpServerConfig::default();
223            server_config.stateful_mode = true;
224            let service = StreamableHttpService::new(
225                move || Ok((*handler_for_service).clone()),
226                session_manager.into(),
227                server_config,
228            );
229            Ok((axum::Router::new().fallback_service(service), ct))
230        } else {
231            use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
232            let handler_for_service = handler.clone();
233            let server_config = StreamableHttpServerConfig::default();
234            let service = StreamableHttpService::new(
235                move || Ok((*handler_for_service).clone()),
236                LocalSessionManager::default().into(),
237                server_config,
238            );
239            Ok((axum::Router::new().fallback_service(service), ct))
240        }
241    }
242
243    async fn create_per_session_server(
244        &self,
245        connector: Arc<dyn BackendConnector>,
246        mcp_id: String,
247        tool_filter: ToolFilter,
248    ) -> Result<(axum::Router, CancellationToken)> {
249        if !self.server_config.stateful_mode {
250            warn!(
251                "Per-session isolation is intended for stateful_mode; enabling stateful_mode"
252            );
253        }
254        let ct = CancellationToken::new();
255        // Session manager still needs a template handler for version APIs;
256        // use a disconnected per-session template (version stays 0 until connects).
257        let template = Arc::new(ProxyHandler::new_per_session(
258            connector.clone(),
259            mcp_id.clone(),
260            tool_filter.clone(),
261        ));
262        let session_manager = ProxyAwareSessionManager::new(template)
263            .with_max_sessions(self.server_config.max_sessions);
264
265        let mut server_config = StreamableHttpServerConfig::default();
266        server_config.stateful_mode = true;
267        let service = StreamableHttpService::new(
268            move || {
269                Ok(ProxyHandler::new_per_session(
270                    connector.clone(),
271                    mcp_id.clone(),
272                    tool_filter.clone(),
273                ))
274            },
275            session_manager.into(),
276            server_config,
277        );
278        Ok((axum::Router::new().fallback_service(service), ct))
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn test_builder_creation() {
288        let builder = StreamServerBuilder::new(BackendConfig::Stdio {
289            command: "echo".into(),
290            args: Some(vec!["hello".into()]),
291            env: None,
292        })
293        .mcp_id("test")
294        .stateful(true);
295
296        assert!(builder.server_config.mcp_id.is_some());
297        assert_eq!(builder.server_config.mcp_id.as_deref(), Some("test"));
298        assert!(builder.server_config.stateful_mode);
299    }
300
301    #[test]
302    fn test_url_backend_config() {
303        let mut headers = HashMap::new();
304        headers.insert("Authorization".into(), "Bearer token123".into());
305        headers.insert("X-Custom".into(), "value".into());
306
307        let builder = StreamServerBuilder::new(BackendConfig::Url {
308            url: "http://localhost:8080/mcp".into(),
309            headers: Some(headers),
310        });
311
312        match &builder.backend_config {
313            BackendConfig::Url { url, headers } => {
314                assert_eq!(url, "http://localhost:8080/mcp");
315                assert!(headers.is_some());
316            }
317            _ => panic!("Expected URL backend"),
318        }
319    }
320
321    #[test]
322    fn default_isolation_by_backend() {
323        assert_eq!(BackendIsolation::for_url(), BackendIsolation::PerSession);
324        assert_eq!(BackendIsolation::for_stdio(), BackendIsolation::Shared);
325    }
326}
mcp_streamable_proxy/server_builder.rs

mcp_streamable_proxy/
server_builder.rs
