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            headers: std::collections::HashMap::new(),
190            ..default_backend()
191        });
192        self
193    }
194
195    /// Configure the last added backend with a per-backend modifier.
196    ///
197    /// # Panics
198    ///
199    /// Panics if no backends have been added.
200    pub fn configure_backend(mut self, f: impl FnOnce(&mut BackendConfig)) -> Self {
201        let backend = self
202            .config
203            .backends
204            .last_mut()
205            .expect("configure_backend called with no backends");
206        f(backend);
207        self
208    }
209
210    /// Enable bearer token authentication.
211    /// Enable bearer token authentication.
212    ///
213    /// All tokens in this list have unrestricted access to all tools.
214    /// For per-token tool scoping, use [`scoped_bearer_auth`](Self::scoped_bearer_auth).
215    pub fn bearer_auth(mut self, tokens: Vec<String>) -> Self {
216        self.config.auth = Some(AuthConfig::Bearer {
217            tokens,
218            scoped_tokens: vec![],
219        });
220        self
221    }
222
223    /// Enable bearer token authentication with per-token tool scoping.
224    ///
225    /// Each [`BearerTokenConfig`] can specify
226    /// `allow_tools` or `deny_tools` to restrict which tools that token can access.
227    pub fn scoped_bearer_auth(mut self, scoped_tokens: Vec<BearerTokenConfig>) -> Self {
228        self.config.auth = Some(AuthConfig::Bearer {
229            tokens: vec![],
230            scoped_tokens,
231        });
232        self
233    }
234
235    /// Enable request coalescing.
236    pub fn coalesce_requests(mut self, enabled: bool) -> Self {
237        self.config.performance.coalesce_requests = enabled;
238        self
239    }
240
241    /// Set the maximum argument size for validation.
242    pub fn max_argument_size(mut self, max_bytes: usize) -> Self {
243        self.config.security.max_argument_size = Some(max_bytes);
244        self
245    }
246
247    /// Enable audit logging.
248    pub fn audit_logging(mut self, enabled: bool) -> Self {
249        self.config.observability.audit = enabled;
250        self
251    }
252
253    /// Enable access logging.
254    pub fn access_logging(mut self, enabled: bool) -> Self {
255        self.config.observability.access_log.enabled = enabled;
256        self
257    }
258
259    /// Set the log level (default: "info").
260    pub fn log_level(mut self, level: impl Into<String>) -> Self {
261        self.config.observability.log_level = level.into();
262        self
263    }
264
265    /// Enable structured JSON logging.
266    pub fn json_logs(mut self, enabled: bool) -> Self {
267        self.config.observability.json_logs = enabled;
268        self
269    }
270
271    /// Enable Prometheus metrics.
272    pub fn metrics(mut self, enabled: bool) -> Self {
273        self.config.observability.metrics.enabled = enabled;
274        self
275    }
276
277    /// Set the timeout for the last added backend.
278    ///
279    /// # Panics
280    ///
281    /// Panics if no backends have been added.
282    ///
283    /// # Example
284    ///
285    /// ```rust
286    /// use mcp_proxy::builder::ProxyBuilder;
287    ///
288    /// let config = ProxyBuilder::new("my-proxy")
289    ///     .http_backend("api", "http://api:8080")
290    ///     .timeout(30)
291    ///     .into_config();
292    ///
293    /// assert_eq!(config.backends[0].timeout.as_ref().unwrap().seconds, 30);
294    /// ```
295    pub fn timeout(mut self, seconds: u64) -> Self {
296        let backend = self
297            .config
298            .backends
299            .last_mut()
300            .expect("timeout called with no backends");
301        backend.timeout = Some(TimeoutConfig { seconds });
302        self
303    }
304
305    /// Set the rate limit for the last added backend.
306    ///
307    /// # Panics
308    ///
309    /// Panics if no backends have been added.
310    ///
311    /// # Example
312    ///
313    /// ```rust
314    /// use mcp_proxy::builder::ProxyBuilder;
315    ///
316    /// let config = ProxyBuilder::new("my-proxy")
317    ///     .http_backend("api", "http://api:8080")
318    ///     .rate_limit(100, 1)
319    ///     .into_config();
320    ///
321    /// let rl = config.backends[0].rate_limit.as_ref().unwrap();
322    /// assert_eq!(rl.requests, 100);
323    /// assert_eq!(rl.period_seconds, 1);
324    /// ```
325    pub fn rate_limit(mut self, requests: usize, period_seconds: u64) -> Self {
326        let backend = self
327            .config
328            .backends
329            .last_mut()
330            .expect("rate_limit called with no backends");
331        backend.rate_limit = Some(RateLimitConfig {
332            requests,
333            period_seconds,
334        });
335        self
336    }
337
338    /// Set the circuit breaker for the last added backend.
339    ///
340    /// Uses sensible defaults for other fields: minimum 5 calls,
341    /// 30-second wait duration, and 3 half-open calls.
342    ///
343    /// # Panics
344    ///
345    /// Panics if no backends have been added.
346    ///
347    /// # Example
348    ///
349    /// ```rust
350    /// use mcp_proxy::builder::ProxyBuilder;
351    ///
352    /// let config = ProxyBuilder::new("my-proxy")
353    ///     .http_backend("api", "http://api:8080")
354    ///     .circuit_breaker(0.5)
355    ///     .into_config();
356    ///
357    /// let cb = config.backends[0].circuit_breaker.as_ref().unwrap();
358    /// assert!((cb.failure_rate_threshold - 0.5).abs() < f64::EPSILON);
359    /// ```
360    pub fn circuit_breaker(mut self, failure_rate: f64) -> Self {
361        let backend = self
362            .config
363            .backends
364            .last_mut()
365            .expect("circuit_breaker called with no backends");
366        backend.circuit_breaker = Some(CircuitBreakerConfig {
367            failure_rate_threshold: failure_rate,
368            minimum_calls: 5,
369            wait_duration_seconds: 30,
370            permitted_calls_in_half_open: 3,
371        });
372        self
373    }
374
375    /// Set the tool allowlist for the last added backend.
376    ///
377    /// Only the listed tools will be exposed through the proxy.
378    ///
379    /// # Panics
380    ///
381    /// Panics if no backends have been added.
382    ///
383    /// # Example
384    ///
385    /// ```rust
386    /// use mcp_proxy::builder::ProxyBuilder;
387    ///
388    /// let config = ProxyBuilder::new("my-proxy")
389    ///     .http_backend("api", "http://api:8080")
390    ///     .expose_tools(&["read_file", "list_dir"])
391    ///     .into_config();
392    ///
393    /// assert_eq!(config.backends[0].expose_tools, vec!["read_file", "list_dir"]);
394    /// ```
395    pub fn expose_tools(mut self, tools: &[&str]) -> Self {
396        let backend = self
397            .config
398            .backends
399            .last_mut()
400            .expect("expose_tools called with no backends");
401        backend.expose_tools = tools.iter().map(|s| s.to_string()).collect();
402        self
403    }
404
405    /// Set the tool denylist for the last added backend.
406    ///
407    /// The listed tools will be hidden from clients.
408    ///
409    /// # Panics
410    ///
411    /// Panics if no backends have been added.
412    ///
413    /// # Example
414    ///
415    /// ```rust
416    /// use mcp_proxy::builder::ProxyBuilder;
417    ///
418    /// let config = ProxyBuilder::new("my-proxy")
419    ///     .http_backend("api", "http://api:8080")
420    ///     .hide_tools(&["dangerous_op"])
421    ///     .into_config();
422    ///
423    /// assert_eq!(config.backends[0].hide_tools, vec!["dangerous_op"]);
424    /// ```
425    pub fn hide_tools(mut self, tools: &[&str]) -> Self {
426        let backend = self
427            .config
428            .backends
429            .last_mut()
430            .expect("hide_tools called with no backends");
431        backend.hide_tools = tools.iter().map(|s| s.to_string()).collect();
432        self
433    }
434
435    /// Set the retry policy for the last added backend.
436    ///
437    /// Uses sensible defaults: 100ms initial backoff, 5000ms max backoff,
438    /// no budget limit.
439    ///
440    /// # Panics
441    ///
442    /// Panics if no backends have been added.
443    ///
444    /// # Example
445    ///
446    /// ```rust
447    /// use mcp_proxy::builder::ProxyBuilder;
448    ///
449    /// let config = ProxyBuilder::new("my-proxy")
450    ///     .http_backend("api", "http://api:8080")
451    ///     .retry(3)
452    ///     .into_config();
453    ///
454    /// let retry = config.backends[0].retry.as_ref().unwrap();
455    /// assert_eq!(retry.max_retries, 3);
456    /// ```
457    pub fn retry(mut self, max_retries: u32) -> Self {
458        let backend = self
459            .config
460            .backends
461            .last_mut()
462            .expect("retry called with no backends");
463        backend.retry = Some(RetryConfig {
464            max_retries,
465            initial_backoff_ms: 100,
466            max_backoff_ms: 5000,
467            budget_percent: None,
468            min_retries_per_sec: 10,
469        });
470        self
471    }
472
473    /// Extract the built [`ProxyConfig`] without connecting to backends.
474    ///
475    /// Useful for inspection, serialization, or passing to
476    /// [`Proxy::from_config()`] manually.
477    pub fn into_config(self) -> ProxyConfig {
478        self.config
479    }
480
481    /// Build the proxy: validate config, connect to all backends, and
482    /// construct the middleware stack.
483    pub async fn build(self) -> Result<Proxy> {
484        Proxy::from_config(self.config).await
485    }
486}
487
488/// Default backend config with all optional fields set to `None`/empty.
489fn default_backend() -> BackendConfig {
490    BackendConfig {
491        name: String::new(),
492        transport: TransportType::Stdio,
493        command: None,
494        args: Vec::new(),
495        url: None,
496        env: HashMap::new(),
497        bearer_token: None,
498        headers: std::collections::HashMap::new(),
499        forward_auth: false,
500        timeout: None,
501        circuit_breaker: None,
502        rate_limit: None,
503        concurrency: None,
504        retry: None,
505        outlier_detection: None,
506        hedging: None,
507        cache: None,
508        default_args: serde_json::Map::new(),
509        inject_args: Vec::new(),
510        param_overrides: Vec::new(),
511        expose_tools: Vec::new(),
512        hide_tools: Vec::new(),
513        expose_resources: Vec::new(),
514        hide_resources: Vec::new(),
515        expose_prompts: Vec::new(),
516        hide_prompts: Vec::new(),
517        hide_destructive: false,
518        read_only_only: false,
519        failover_for: None,
520        priority: 0,
521        canary_of: None,
522        weight: 100,
523        aliases: Vec::new(),
524        mirror_of: None,
525        mirror_percent: 100,
526    }
527}
528
529#[cfg(test)]
530mod tests {
531    use super::*;
532
533    #[test]
534    fn test_builder_minimal() {
535        let config = ProxyBuilder::new("test-proxy").into_config();
536        assert_eq!(config.proxy.name, "test-proxy");
537        assert_eq!(config.proxy.version, "0.1.0");
538        assert_eq!(config.proxy.separator, "/");
539        assert_eq!(config.proxy.listen.host, "127.0.0.1");
540        assert_eq!(config.proxy.listen.port, 8080);
541        assert!(config.backends.is_empty());
542    }
543
544    #[test]
545    fn test_builder_with_backends() {
546        let config = ProxyBuilder::new("test")
547            .stdio_backend("files", "npx", &["-y", "@mcp/server-files"])
548            .http_backend("api", "http://localhost:8080")
549            .into_config();
550
551        assert_eq!(config.backends.len(), 2);
552        assert_eq!(config.backends[0].name, "files");
553        assert!(matches!(config.backends[0].transport, TransportType::Stdio));
554        assert_eq!(config.backends[0].command.as_deref(), Some("npx"));
555        assert_eq!(config.backends[1].name, "api");
556        assert!(matches!(config.backends[1].transport, TransportType::Http));
557        assert_eq!(
558            config.backends[1].url.as_deref(),
559            Some("http://localhost:8080")
560        );
561    }
562
563    #[test]
564    fn test_builder_configure_backend() {
565        let config = ProxyBuilder::new("test")
566            .http_backend("api", "http://localhost:8080")
567            .configure_backend(|b| {
568                b.timeout = Some(TimeoutConfig { seconds: 30 });
569                b.rate_limit = Some(RateLimitConfig {
570                    requests: 100,
571                    period_seconds: 1,
572                });
573                b.hide_tools = vec!["dangerous_op".to_string()];
574            })
575            .into_config();
576
577        assert!(config.backends[0].timeout.is_some());
578        assert!(config.backends[0].rate_limit.is_some());
579        assert_eq!(config.backends[0].hide_tools, vec!["dangerous_op"]);
580    }
581
582    #[test]
583    fn test_builder_auth_and_observability() {
584        let config = ProxyBuilder::new("test")
585            .bearer_auth(vec!["token1".into(), "token2".into()])
586            .audit_logging(true)
587            .access_logging(true)
588            .metrics(true)
589            .json_logs(true)
590            .log_level("debug")
591            .into_config();
592
593        assert!(config.auth.is_some());
594        assert!(config.observability.audit);
595        assert!(config.observability.access_log.enabled);
596        assert!(config.observability.metrics.enabled);
597        assert!(config.observability.json_logs);
598        assert_eq!(config.observability.log_level, "debug");
599    }
600
601    #[test]
602    fn test_builder_global_rate_limit() {
603        let config = ProxyBuilder::new("test")
604            .global_rate_limit(500, Duration::from_secs(1))
605            .into_config();
606
607        let rl = config.proxy.rate_limit.unwrap();
608        assert_eq!(rl.requests, 500);
609        assert_eq!(rl.period_seconds, 1);
610    }
611
612    #[test]
613    fn test_builder_all_settings() {
614        let config = ProxyBuilder::new("enterprise")
615            .version("2.0.0")
616            .separator("::")
617            .listen("0.0.0.0", 9090)
618            .instructions("Enterprise MCP gateway")
619            .shutdown_timeout(Duration::from_secs(60))
620            .coalesce_requests(true)
621            .max_argument_size(1_048_576)
622            .into_config();
623
624        assert_eq!(config.proxy.name, "enterprise");
625        assert_eq!(config.proxy.version, "2.0.0");
626        assert_eq!(config.proxy.separator, "::");
627        assert_eq!(config.proxy.listen.host, "0.0.0.0");
628        assert_eq!(config.proxy.listen.port, 9090);
629        assert_eq!(
630            config.proxy.instructions.as_deref(),
631            Some("Enterprise MCP gateway")
632        );
633        assert_eq!(config.proxy.shutdown_timeout_seconds, 60);
634        assert!(config.performance.coalesce_requests);
635        assert_eq!(config.security.max_argument_size, Some(1_048_576));
636    }
637
638    #[test]
639    fn test_builder_http_backend_with_token() {
640        let config = ProxyBuilder::new("test")
641            .http_backend_with_token("api", "http://api:8080", "secret")
642            .into_config();
643
644        assert_eq!(config.backends[0].bearer_token.as_deref(), Some("secret"));
645    }
646
647    #[test]
648    fn test_builder_ergonomic_backend_methods() {
649        let config = ProxyBuilder::new("test")
650            .http_backend("api", "http://api:8080")
651            .timeout(30)
652            .rate_limit(100, 1)
653            .circuit_breaker(0.7)
654            .expose_tools(&["read_file", "list_dir"])
655            .retry(5)
656            .stdio_backend("files", "npx", &["-y", "@mcp/server-files"])
657            .hide_tools(&["dangerous_op"])
658            .timeout(60)
659            .into_config();
660
661        // First backend: api
662        let api = &config.backends[0];
663        assert_eq!(api.timeout.as_ref().unwrap().seconds, 30);
664        let rl = api.rate_limit.as_ref().unwrap();
665        assert_eq!(rl.requests, 100);
666        assert_eq!(rl.period_seconds, 1);
667        let cb = api.circuit_breaker.as_ref().unwrap();
668        assert!((cb.failure_rate_threshold - 0.7).abs() < f64::EPSILON);
669        assert_eq!(cb.minimum_calls, 5);
670        assert_eq!(cb.wait_duration_seconds, 30);
671        assert_eq!(cb.permitted_calls_in_half_open, 3);
672        assert_eq!(api.expose_tools, vec!["read_file", "list_dir"]);
673        let retry = api.retry.as_ref().unwrap();
674        assert_eq!(retry.max_retries, 5);
675        assert_eq!(retry.initial_backoff_ms, 100);
676        assert_eq!(retry.max_backoff_ms, 5000);
677        assert!(retry.budget_percent.is_none());
678
679        // Second backend: files
680        let files = &config.backends[1];
681        assert_eq!(files.hide_tools, vec!["dangerous_op"]);
682        assert_eq!(files.timeout.as_ref().unwrap().seconds, 60);
683        assert!(files.circuit_breaker.is_none());
684        assert!(files.rate_limit.is_none());
685    }
686
687    #[test]
688    fn test_builder_stdio_backend_with_env() {
689        let mut env = HashMap::new();
690        env.insert("GITHUB_TOKEN".to_string(), "ghp_xxx".to_string());
691
692        let config = ProxyBuilder::new("test")
693            .stdio_backend_with_env("github", "npx", &["-y", "@mcp/github"], env)
694            .into_config();
695
696        assert_eq!(
697            config.backends[0].env.get("GITHUB_TOKEN").unwrap(),
698            "ghp_xxx"
699        );
700    }
701}
mcp_proxy/builder.rs

mcp_proxy/
builder.rs
