mcp_proxy/
builder.rs

1//! Programmatic proxy builder for library users.
2//!
3//! Constructs a [`ProxyConfig`] via a fluent API, avoiding the need for
4//! TOML files. The resulting config is passed to [`Proxy::from_config()`]
5//! as usual.
6//!
7//! # Example
8//!
9//! ```rust,no_run
10//! use mcp_proxy::builder::ProxyBuilder;
11//!
12//! # async fn example() -> anyhow::Result<()> {
13//! let proxy = ProxyBuilder::new("my-proxy")
14//!     .version("1.0.0")
15//!     .listen("0.0.0.0", 9090)
16//!     .stdio_backend("files", "npx", &["-y", "@modelcontextprotocol/server-filesystem"])
17//!     .http_backend("api", "http://api:8080")
18//!     .build()
19//!     .await?;
20//!
21//! // Embed in an existing axum app
22//! let (router, _session_handle) = proxy.into_router();
23//! # Ok(())
24//! # }
25//! ```
26
27use std::collections::HashMap;
28use std::time::Duration;
29
30use anyhow::Result;
31
32use crate::Proxy;
33use crate::config::*;
34
35/// Fluent builder for constructing an MCP proxy without TOML config files.
36///
37/// Call [`build()`](Self::build) to connect backends and produce a
38/// ready-to-serve [`Proxy`].
39pub struct ProxyBuilder {
40    config: ProxyConfig,
41}
42
43impl ProxyBuilder {
44    /// Create a new proxy builder with the given name.
45    ///
46    /// Defaults: version "0.1.0", separator "/", listen 127.0.0.1:8080.
47    pub fn new(name: impl Into<String>) -> Self {
48        Self {
49            config: ProxyConfig {
50                proxy: ProxySettings {
51                    name: name.into(),
52                    version: "0.1.0".to_string(),
53                    separator: "/".to_string(),
54                    listen: ListenConfig {
55                        host: "127.0.0.1".to_string(),
56                        port: 8080,
57                    },
58                    instructions: None,
59                    shutdown_timeout_seconds: 30,
60                    hot_reload: false,
61                    import_backends: None,
62                    rate_limit: None,
63                    tool_discovery: false,
64                    tool_exposure: crate::config::ToolExposure::default(),
65                },
66                backends: Vec::new(),
67                auth: None,
68                performance: PerformanceConfig::default(),
69                security: SecurityConfig::default(),
70                cache: CacheBackendConfig::default(),
71                observability: ObservabilityConfig::default(),
72                composite_tools: Vec::new(),
73                source_path: None,
74            },
75        }
76    }
77
78    /// Set the proxy version (default: "0.1.0").
79    pub fn version(mut self, version: impl Into<String>) -> Self {
80        self.config.proxy.version = version.into();
81        self
82    }
83
84    /// Set the namespace separator (default: "/").
85    pub fn separator(mut self, separator: impl Into<String>) -> Self {
86        self.config.proxy.separator = separator.into();
87        self
88    }
89
90    /// Set the listen address and port (default: 127.0.0.1:8080).
91    pub fn listen(mut self, host: impl Into<String>, port: u16) -> Self {
92        self.config.proxy.listen = ListenConfig {
93            host: host.into(),
94            port,
95        };
96        self
97    }
98
99    /// Set instructions text sent to MCP clients.
100    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
101        self.config.proxy.instructions = Some(instructions.into());
102        self
103    }
104
105    /// Set the graceful shutdown timeout (default: 30s).
106    pub fn shutdown_timeout(mut self, timeout: Duration) -> Self {
107        self.config.proxy.shutdown_timeout_seconds = timeout.as_secs();
108        self
109    }
110
111    /// Enable hot reload for watching config file changes.
112    pub fn hot_reload(mut self, enabled: bool) -> Self {
113        self.config.proxy.hot_reload = enabled;
114        self
115    }
116
117    /// Set a global rate limit across all backends.
118    pub fn global_rate_limit(mut self, requests: usize, period: Duration) -> Self {
119        self.config.proxy.rate_limit = Some(GlobalRateLimitConfig {
120            requests,
121            period_seconds: period.as_secs(),
122        });
123        self
124    }
125
126    /// Add a stdio backend (subprocess).
127    pub fn stdio_backend(
128        mut self,
129        name: impl Into<String>,
130        command: impl Into<String>,
131        args: &[&str],
132    ) -> Self {
133        self.config.backends.push(BackendConfig {
134            name: name.into(),
135            transport: TransportType::Stdio,
136            command: Some(command.into()),
137            args: args.iter().map(|s| s.to_string()).collect(),
138            url: None,
139            ..default_backend()
140        });
141        self
142    }
143
144    /// Add a stdio backend with environment variables.
145    pub fn stdio_backend_with_env(
146        mut self,
147        name: impl Into<String>,
148        command: impl Into<String>,
149        args: &[&str],
150        env: HashMap<String, String>,
151    ) -> Self {
152        self.config.backends.push(BackendConfig {
153            name: name.into(),
154            transport: TransportType::Stdio,
155            command: Some(command.into()),
156            args: args.iter().map(|s| s.to_string()).collect(),
157            url: None,
158            env,
159            ..default_backend()
160        });
161        self
162    }
163
164    /// Add an HTTP backend.
165    pub fn http_backend(mut self, name: impl Into<String>, url: impl Into<String>) -> Self {
166        self.config.backends.push(BackendConfig {
167            name: name.into(),
168            transport: TransportType::Http,
169            command: None,
170            url: Some(url.into()),
171            ..default_backend()
172        });
173        self
174    }
175
176    /// Add an HTTP backend with a bearer token.
177    pub fn http_backend_with_token(
178        mut self,
179        name: impl Into<String>,
180        url: impl Into<String>,
181        token: impl Into<String>,
182    ) -> Self {
183        self.config.backends.push(BackendConfig {
184            name: name.into(),
185            transport: TransportType::Http,
186            command: None,
187            url: Some(url.into()),
188            bearer_token: Some(token.into()),
189            ..default_backend()
190        });
191        self
192    }
193
194    /// Configure the last added backend with a per-backend modifier.
195    ///
196    /// # Panics
197    ///
198    /// Panics if no backends have been added.
199    pub fn configure_backend(mut self, f: impl FnOnce(&mut BackendConfig)) -> Self {
200        let backend = self
201            .config
202            .backends
203            .last_mut()
204            .expect("configure_backend called with no backends");
205        f(backend);
206        self
207    }
208
209    /// Enable bearer token authentication.
210    /// Enable bearer token authentication.
211    ///
212    /// All tokens in this list have unrestricted access to all tools.
213    /// For per-token tool scoping, use [`scoped_bearer_auth`](Self::scoped_bearer_auth).
214    pub fn bearer_auth(mut self, tokens: Vec<String>) -> Self {
215        self.config.auth = Some(AuthConfig::Bearer {
216            tokens,
217            scoped_tokens: vec![],
218        });
219        self
220    }
221
222    /// Enable bearer token authentication with per-token tool scoping.
223    ///
224    /// Each [`BearerTokenConfig`] can specify
225    /// `allow_tools` or `deny_tools` to restrict which tools that token can access.
226    pub fn scoped_bearer_auth(mut self, scoped_tokens: Vec<BearerTokenConfig>) -> Self {
227        self.config.auth = Some(AuthConfig::Bearer {
228            tokens: vec![],
229            scoped_tokens,
230        });
231        self
232    }
233
234    /// Enable request coalescing.
235    pub fn coalesce_requests(mut self, enabled: bool) -> Self {
236        self.config.performance.coalesce_requests = enabled;
237        self
238    }
239
240    /// Set the maximum argument size for validation.
241    pub fn max_argument_size(mut self, max_bytes: usize) -> Self {
242        self.config.security.max_argument_size = Some(max_bytes);
243        self
244    }
245
246    /// Enable audit logging.
247    pub fn audit_logging(mut self, enabled: bool) -> Self {
248        self.config.observability.audit = enabled;
249        self
250    }
251
252    /// Enable access logging.
253    pub fn access_logging(mut self, enabled: bool) -> Self {
254        self.config.observability.access_log.enabled = enabled;
255        self
256    }
257
258    /// Set the log level (default: "info").
259    pub fn log_level(mut self, level: impl Into<String>) -> Self {
260        self.config.observability.log_level = level.into();
261        self
262    }
263
264    /// Enable structured JSON logging.
265    pub fn json_logs(mut self, enabled: bool) -> Self {
266        self.config.observability.json_logs = enabled;
267        self
268    }
269
270    /// Enable Prometheus metrics.
271    pub fn metrics(mut self, enabled: bool) -> Self {
272        self.config.observability.metrics.enabled = enabled;
273        self
274    }
275
276    /// Set the timeout for the last added backend.
277    ///
278    /// # Panics
279    ///
280    /// Panics if no backends have been added.
281    ///
282    /// # Example
283    ///
284    /// ```rust
285    /// use mcp_proxy::builder::ProxyBuilder;
286    ///
287    /// let config = ProxyBuilder::new("my-proxy")
288    ///     .http_backend("api", "http://api:8080")
289    ///     .timeout(30)
290    ///     .into_config();
291    ///
292    /// assert_eq!(config.backends[0].timeout.as_ref().unwrap().seconds, 30);
293    /// ```
294    pub fn timeout(mut self, seconds: u64) -> Self {
295        let backend = self
296            .config
297            .backends
298            .last_mut()
299            .expect("timeout called with no backends");
300        backend.timeout = Some(TimeoutConfig { seconds });
301        self
302    }
303
304    /// Set the rate limit for the last added backend.
305    ///
306    /// # Panics
307    ///
308    /// Panics if no backends have been added.
309    ///
310    /// # Example
311    ///
312    /// ```rust
313    /// use mcp_proxy::builder::ProxyBuilder;
314    ///
315    /// let config = ProxyBuilder::new("my-proxy")
316    ///     .http_backend("api", "http://api:8080")
317    ///     .rate_limit(100, 1)
318    ///     .into_config();
319    ///
320    /// let rl = config.backends[0].rate_limit.as_ref().unwrap();
321    /// assert_eq!(rl.requests, 100);
322    /// assert_eq!(rl.period_seconds, 1);
323    /// ```
324    pub fn rate_limit(mut self, requests: usize, period_seconds: u64) -> Self {
325        let backend = self
326            .config
327            .backends
328            .last_mut()
329            .expect("rate_limit called with no backends");
330        backend.rate_limit = Some(RateLimitConfig {
331            requests,
332            period_seconds,
333        });
334        self
335    }
336
337    /// Set the circuit breaker for the last added backend.
338    ///
339    /// Uses sensible defaults for other fields: minimum 5 calls,
340    /// 30-second wait duration, and 3 half-open calls.
341    ///
342    /// # Panics
343    ///
344    /// Panics if no backends have been added.
345    ///
346    /// # Example
347    ///
348    /// ```rust
349    /// use mcp_proxy::builder::ProxyBuilder;
350    ///
351    /// let config = ProxyBuilder::new("my-proxy")
352    ///     .http_backend("api", "http://api:8080")
353    ///     .circuit_breaker(0.5)
354    ///     .into_config();
355    ///
356    /// let cb = config.backends[0].circuit_breaker.as_ref().unwrap();
357    /// assert!((cb.failure_rate_threshold - 0.5).abs() < f64::EPSILON);
358    /// ```
359    pub fn circuit_breaker(mut self, failure_rate: f64) -> Self {
360        let backend = self
361            .config
362            .backends
363            .last_mut()
364            .expect("circuit_breaker called with no backends");
365        backend.circuit_breaker = Some(CircuitBreakerConfig {
366            failure_rate_threshold: failure_rate,
367            minimum_calls: 5,
368            wait_duration_seconds: 30,
369            permitted_calls_in_half_open: 3,
370        });
371        self
372    }
373
374    /// Set the tool allowlist for the last added backend.
375    ///
376    /// Only the listed tools will be exposed through the proxy.
377    ///
378    /// # Panics
379    ///
380    /// Panics if no backends have been added.
381    ///
382    /// # Example
383    ///
384    /// ```rust
385    /// use mcp_proxy::builder::ProxyBuilder;
386    ///
387    /// let config = ProxyBuilder::new("my-proxy")
388    ///     .http_backend("api", "http://api:8080")
389    ///     .expose_tools(&["read_file", "list_dir"])
390    ///     .into_config();
391    ///
392    /// assert_eq!(config.backends[0].expose_tools, vec!["read_file", "list_dir"]);
393    /// ```
394    pub fn expose_tools(mut self, tools: &[&str]) -> Self {
395        let backend = self
396            .config
397            .backends
398            .last_mut()
399            .expect("expose_tools called with no backends");
400        backend.expose_tools = tools.iter().map(|s| s.to_string()).collect();
401        self
402    }
403
404    /// Set the tool denylist for the last added backend.
405    ///
406    /// The listed tools will be hidden from clients.
407    ///
408    /// # Panics
409    ///
410    /// Panics if no backends have been added.
411    ///
412    /// # Example
413    ///
414    /// ```rust
415    /// use mcp_proxy::builder::ProxyBuilder;
416    ///
417    /// let config = ProxyBuilder::new("my-proxy")
418    ///     .http_backend("api", "http://api:8080")
419    ///     .hide_tools(&["dangerous_op"])
420    ///     .into_config();
421    ///
422    /// assert_eq!(config.backends[0].hide_tools, vec!["dangerous_op"]);
423    /// ```
424    pub fn hide_tools(mut self, tools: &[&str]) -> Self {
425        let backend = self
426            .config
427            .backends
428            .last_mut()
429            .expect("hide_tools called with no backends");
430        backend.hide_tools = tools.iter().map(|s| s.to_string()).collect();
431        self
432    }
433
434    /// Set the retry policy for the last added backend.
435    ///
436    /// Uses sensible defaults: 100ms initial backoff, 5000ms max backoff,
437    /// no budget limit.
438    ///
439    /// # Panics
440    ///
441    /// Panics if no backends have been added.
442    ///
443    /// # Example
444    ///
445    /// ```rust
446    /// use mcp_proxy::builder::ProxyBuilder;
447    ///
448    /// let config = ProxyBuilder::new("my-proxy")
449    ///     .http_backend("api", "http://api:8080")
450    ///     .retry(3)
451    ///     .into_config();
452    ///
453    /// let retry = config.backends[0].retry.as_ref().unwrap();
454    /// assert_eq!(retry.max_retries, 3);
455    /// ```
456    pub fn retry(mut self, max_retries: u32) -> Self {
457        let backend = self
458            .config
459            .backends
460            .last_mut()
461            .expect("retry called with no backends");
462        backend.retry = Some(RetryConfig {
463            max_retries,
464            initial_backoff_ms: 100,
465            max_backoff_ms: 5000,
466            budget_percent: None,
467            min_retries_per_sec: 10,
468        });
469        self
470    }
471
472    /// Extract the built [`ProxyConfig`] without connecting to backends.
473    ///
474    /// Useful for inspection, serialization, or passing to
475    /// [`Proxy::from_config()`] manually.
476    pub fn into_config(self) -> ProxyConfig {
477        self.config
478    }
479
480    /// Build the proxy: validate config, connect to all backends, and
481    /// construct the middleware stack.
482    pub async fn build(self) -> Result<Proxy> {
483        Proxy::from_config(self.config).await
484    }
485}
486
487/// Default backend config with all optional fields set to `None`/empty.
488fn default_backend() -> BackendConfig {
489    BackendConfig {
490        name: String::new(),
491        transport: TransportType::Stdio,
492        command: None,
493        args: Vec::new(),
494        url: None,
495        env: HashMap::new(),
496        bearer_token: None,
497        forward_auth: false,
498        timeout: None,
499        circuit_breaker: None,
500        rate_limit: None,
501        concurrency: None,
502        retry: None,
503        outlier_detection: None,
504        hedging: None,
505        cache: None,
506        default_args: serde_json::Map::new(),
507        inject_args: Vec::new(),
508        param_overrides: Vec::new(),
509        expose_tools: Vec::new(),
510        hide_tools: Vec::new(),
511        expose_resources: Vec::new(),
512        hide_resources: Vec::new(),
513        expose_prompts: Vec::new(),
514        hide_prompts: Vec::new(),
515        hide_destructive: false,
516        read_only_only: false,
517        failover_for: None,
518        priority: 0,
519        canary_of: None,
520        weight: 100,
521        aliases: Vec::new(),
522        mirror_of: None,
523        mirror_percent: 100,
524    }
525}
526
527#[cfg(test)]
528mod tests {
529    use super::*;
530
531    #[test]
532    fn test_builder_minimal() {
533        let config = ProxyBuilder::new("test-proxy").into_config();
534        assert_eq!(config.proxy.name, "test-proxy");
535        assert_eq!(config.proxy.version, "0.1.0");
536        assert_eq!(config.proxy.separator, "/");
537        assert_eq!(config.proxy.listen.host, "127.0.0.1");
538        assert_eq!(config.proxy.listen.port, 8080);
539        assert!(config.backends.is_empty());
540    }
541
542    #[test]
543    fn test_builder_with_backends() {
544        let config = ProxyBuilder::new("test")
545            .stdio_backend("files", "npx", &["-y", "@mcp/server-files"])
546            .http_backend("api", "http://localhost:8080")
547            .into_config();
548
549        assert_eq!(config.backends.len(), 2);
550        assert_eq!(config.backends[0].name, "files");
551        assert!(matches!(config.backends[0].transport, TransportType::Stdio));
552        assert_eq!(config.backends[0].command.as_deref(), Some("npx"));
553        assert_eq!(config.backends[1].name, "api");
554        assert!(matches!(config.backends[1].transport, TransportType::Http));
555        assert_eq!(
556            config.backends[1].url.as_deref(),
557            Some("http://localhost:8080")
558        );
559    }
560
561    #[test]
562    fn test_builder_configure_backend() {
563        let config = ProxyBuilder::new("test")
564            .http_backend("api", "http://localhost:8080")
565            .configure_backend(|b| {
566                b.timeout = Some(TimeoutConfig { seconds: 30 });
567                b.rate_limit = Some(RateLimitConfig {
568                    requests: 100,
569                    period_seconds: 1,
570                });
571                b.hide_tools = vec!["dangerous_op".to_string()];
572            })
573            .into_config();
574
575        assert!(config.backends[0].timeout.is_some());
576        assert!(config.backends[0].rate_limit.is_some());
577        assert_eq!(config.backends[0].hide_tools, vec!["dangerous_op"]);
578    }
579
580    #[test]
581    fn test_builder_auth_and_observability() {
582        let config = ProxyBuilder::new("test")
583            .bearer_auth(vec!["token1".into(), "token2".into()])
584            .audit_logging(true)
585            .access_logging(true)
586            .metrics(true)
587            .json_logs(true)
588            .log_level("debug")
589            .into_config();
590
591        assert!(config.auth.is_some());
592        assert!(config.observability.audit);
593        assert!(config.observability.access_log.enabled);
594        assert!(config.observability.metrics.enabled);
595        assert!(config.observability.json_logs);
596        assert_eq!(config.observability.log_level, "debug");
597    }
598
599    #[test]
600    fn test_builder_global_rate_limit() {
601        let config = ProxyBuilder::new("test")
602            .global_rate_limit(500, Duration::from_secs(1))
603            .into_config();
604
605        let rl = config.proxy.rate_limit.unwrap();
606        assert_eq!(rl.requests, 500);
607        assert_eq!(rl.period_seconds, 1);
608    }
609
610    #[test]
611    fn test_builder_all_settings() {
612        let config = ProxyBuilder::new("enterprise")
613            .version("2.0.0")
614            .separator("::")
615            .listen("0.0.0.0", 9090)
616            .instructions("Enterprise MCP gateway")
617            .shutdown_timeout(Duration::from_secs(60))
618            .coalesce_requests(true)
619            .max_argument_size(1_048_576)
620            .into_config();
621
622        assert_eq!(config.proxy.name, "enterprise");
623        assert_eq!(config.proxy.version, "2.0.0");
624        assert_eq!(config.proxy.separator, "::");
625        assert_eq!(config.proxy.listen.host, "0.0.0.0");
626        assert_eq!(config.proxy.listen.port, 9090);
627        assert_eq!(
628            config.proxy.instructions.as_deref(),
629            Some("Enterprise MCP gateway")
630        );
631        assert_eq!(config.proxy.shutdown_timeout_seconds, 60);
632        assert!(config.performance.coalesce_requests);
633        assert_eq!(config.security.max_argument_size, Some(1_048_576));
634    }
635
636    #[test]
637    fn test_builder_http_backend_with_token() {
638        let config = ProxyBuilder::new("test")
639            .http_backend_with_token("api", "http://api:8080", "secret")
640            .into_config();
641
642        assert_eq!(config.backends[0].bearer_token.as_deref(), Some("secret"));
643    }
644
645    #[test]
646    fn test_builder_ergonomic_backend_methods() {
647        let config = ProxyBuilder::new("test")
648            .http_backend("api", "http://api:8080")
649            .timeout(30)
650            .rate_limit(100, 1)
651            .circuit_breaker(0.7)
652            .expose_tools(&["read_file", "list_dir"])
653            .retry(5)
654            .stdio_backend("files", "npx", &["-y", "@mcp/server-files"])
655            .hide_tools(&["dangerous_op"])
656            .timeout(60)
657            .into_config();
658
659        // First backend: api
660        let api = &config.backends[0];
661        assert_eq!(api.timeout.as_ref().unwrap().seconds, 30);
662        let rl = api.rate_limit.as_ref().unwrap();
663        assert_eq!(rl.requests, 100);
664        assert_eq!(rl.period_seconds, 1);
665        let cb = api.circuit_breaker.as_ref().unwrap();
666        assert!((cb.failure_rate_threshold - 0.7).abs() < f64::EPSILON);
667        assert_eq!(cb.minimum_calls, 5);
668        assert_eq!(cb.wait_duration_seconds, 30);
669        assert_eq!(cb.permitted_calls_in_half_open, 3);
670        assert_eq!(api.expose_tools, vec!["read_file", "list_dir"]);
671        let retry = api.retry.as_ref().unwrap();
672        assert_eq!(retry.max_retries, 5);
673        assert_eq!(retry.initial_backoff_ms, 100);
674        assert_eq!(retry.max_backoff_ms, 5000);
675        assert!(retry.budget_percent.is_none());
676
677        // Second backend: files
678        let files = &config.backends[1];
679        assert_eq!(files.hide_tools, vec!["dangerous_op"]);
680        assert_eq!(files.timeout.as_ref().unwrap().seconds, 60);
681        assert!(files.circuit_breaker.is_none());
682        assert!(files.rate_limit.is_none());
683    }
684
685    #[test]
686    fn test_builder_stdio_backend_with_env() {
687        let mut env = HashMap::new();
688        env.insert("GITHUB_TOKEN".to_string(), "ghp_xxx".to_string());
689
690        let config = ProxyBuilder::new("test")
691            .stdio_backend_with_env("github", "npx", &["-y", "@mcp/github"], env)
692            .into_config();
693
694        assert_eq!(
695            config.backends[0].env.get("GITHUB_TOKEN").unwrap(),
696            "ghp_xxx"
697        );
698    }
699}
mcp_proxy/builder.rs

mcp_proxy/
builder.rs
