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            // 前置校验（build 早期 bail）保证不可达；防御漂移后 panic
196            (BackendConfig::Stdio { .. }, BackendIsolation::PerSession) => {
197                bail!("stdio 后端不支持 PerSession 隔离（应在前置校验中被拒绝）");
198            }
199        };
200
201        info!(
202            "[StreamServerBuilder] Server created - mcp_id: {}, stateful: {}, isolation: {:?}",
203            mcp_id, self.server_config.stateful_mode, isolation
204        );
205
206        Ok((router, ct, management_handler))
207    }
208
209    async fn create_shared_server(
210        &self,
211        proxy_handler: ProxyHandler,
212    ) -> Result<(axum::Router, CancellationToken)> {
213        let handler = Arc::new(proxy_handler);
214        let ct = CancellationToken::new();
215
216        if self.server_config.stateful_mode {
217            let session_manager = ProxyAwareSessionManager::new(handler.clone())
218                .with_max_sessions(self.server_config.max_sessions);
219            let handler_for_service = handler.clone();
220            let mut server_config = StreamableHttpServerConfig::default();
221            server_config.legacy_session_mode = true;
222            let service = StreamableHttpService::new(
223                move || Ok((*handler_for_service).clone()),
224                session_manager.into(),
225                server_config,
226            );
227            Ok((axum::Router::new().fallback_service(service), ct))
228        } else {
229            use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
230            let handler_for_service = handler.clone();
231            let server_config = StreamableHttpServerConfig::default();
232            let service = StreamableHttpService::new(
233                move || Ok((*handler_for_service).clone()),
234                LocalSessionManager::default().into(),
235                server_config,
236            );
237            Ok((axum::Router::new().fallback_service(service), ct))
238        }
239    }
240
241    async fn create_per_session_server(
242        &self,
243        connector: Arc<dyn BackendConnector>,
244        mcp_id: String,
245        tool_filter: ToolFilter,
246    ) -> Result<(axum::Router, CancellationToken)> {
247        if !self.server_config.stateful_mode {
248            warn!("Per-session isolation is intended for stateful_mode; enabling stateful_mode");
249        }
250        let ct = CancellationToken::new();
251        // Session manager still needs a template handler for version APIs;
252        // use a disconnected per-session template (version stays 0 until connects).
253        let template = Arc::new(ProxyHandler::new_per_session(
254            connector.clone(),
255            mcp_id.clone(),
256            tool_filter.clone(),
257        ));
258        let session_manager = ProxyAwareSessionManager::new(template)
259            .with_max_sessions(self.server_config.max_sessions);
260
261        let mut server_config = StreamableHttpServerConfig::default();
262        server_config.legacy_session_mode = true;
263        let service = StreamableHttpService::new(
264            move || {
265                Ok(ProxyHandler::new_per_session(
266                    connector.clone(),
267                    mcp_id.clone(),
268                    tool_filter.clone(),
269                ))
270            },
271            session_manager.into(),
272            server_config,
273        );
274        Ok((axum::Router::new().fallback_service(service), ct))
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    #[test]
283    fn test_builder_creation() {
284        let builder = StreamServerBuilder::new(BackendConfig::Stdio {
285            command: "echo".into(),
286            args: Some(vec!["hello".into()]),
287            env: None,
288        })
289        .mcp_id("test")
290        .stateful(true);
291
292        assert!(builder.server_config.mcp_id.is_some());
293        assert_eq!(builder.server_config.mcp_id.as_deref(), Some("test"));
294        assert!(builder.server_config.stateful_mode);
295    }
296
297    #[test]
298    fn test_url_backend_config() {
299        let mut headers = HashMap::new();
300        headers.insert("Authorization".into(), "Bearer token123".into());
301        headers.insert("X-Custom".into(), "value".into());
302
303        let builder = StreamServerBuilder::new(BackendConfig::Url {
304            url: "http://localhost:8080/mcp".into(),
305            headers: Some(headers),
306        });
307
308        match &builder.backend_config {
309            BackendConfig::Url { url, headers } => {
310                assert_eq!(url, "http://localhost:8080/mcp");
311                assert!(headers.is_some());
312            }
313            _ => panic!("Expected URL backend"),
314        }
315    }
316
317    #[test]
318    fn default_isolation_by_backend() {
319        assert_eq!(BackendIsolation::for_url(), BackendIsolation::PerSession);
320        assert_eq!(BackendIsolation::for_stdio(), BackendIsolation::Shared);
321    }
322}
mcp_streamable_proxy/server_builder.rs

mcp_streamable_proxy/
server_builder.rs
