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

mcp_streamable_proxy/
server_builder.rs
