Skip to main content

mcp_proxy/
proxy.rs

1//! Core proxy construction and serving.
2
3use std::collections::HashMap;
4use std::convert::Infallible;
5use std::time::Duration;
6
7use anyhow::{Context, Result};
8use axum::Router;
9use tokio::process::Command;
10use tower::timeout::TimeoutLayer;
11use tower::util::BoxCloneService;
12use tower::{Layer, ServiceExt};
13use tower_mcp::SessionHandle;
14use tower_mcp::auth::{AuthLayer, StaticBearerValidator};
15use tower_mcp::client::StdioClientTransport;
16use tower_mcp::proxy::McpProxy;
17use tower_mcp::{RouterRequest, RouterResponse};
18use tower_resilience::retry::RetryLayer;
19
20use crate::admin::BackendMeta;
21use crate::alias;
22use crate::cache;
23use crate::coalesce;
24use crate::config::{AuthConfig, ProxyConfig, TransportType};
25use crate::filter::CapabilityFilterService;
26#[cfg(feature = "oauth")]
27use crate::rbac::{RbacConfig, RbacService};
28use crate::validation::{ValidationConfig, ValidationService};
29
30/// A fully constructed MCP proxy ready to serve or embed.
31pub struct Proxy {
32    router: Router,
33    session_handle: SessionHandle,
34    inner: McpProxy,
35    config: ProxyConfig,
36    #[cfg(feature = "discovery")]
37    discovery_index: Option<crate::discovery::SharedDiscoveryIndex>,
38}
39
40impl Proxy {
41    /// Build a proxy from a [`ProxyConfig`].
42    ///
43    /// Connects to all backends, builds the middleware stack, and prepares
44    /// the axum router. Call [`serve()`](Self::serve) to run standalone or
45    /// [`into_router()`](Self::into_router) to embed in an existing app.
46    pub async fn from_config(config: ProxyConfig) -> Result<Self> {
47        let (mcp_proxy, cb_handles) = build_mcp_proxy(&config).await?;
48        let proxy_for_admin = mcp_proxy.clone();
49        // Only the discovery index build mutates this clone.
50        #[cfg_attr(not(feature = "discovery"), allow(unused_mut))]
51        let mut proxy_for_caller = mcp_proxy.clone();
52        let proxy_for_management = mcp_proxy.clone();
53
54        // Install Prometheus metrics recorder (must happen before middleware)
55        #[cfg(feature = "metrics")]
56        let metrics_handle = if config.observability.metrics.enabled {
57            tracing::info!("Prometheus metrics enabled at /admin/metrics");
58            let builder = metrics_exporter_prometheus::PrometheusBuilder::new();
59            let handle = builder
60                .install_recorder()
61                .context("installing Prometheus metrics recorder")?;
62            Some(handle)
63        } else {
64            None
65        };
66        #[cfg(not(feature = "metrics"))]
67        let metrics_handle = None;
68
69        let (service, cache_handle) = build_middleware_stack(&config, mcp_proxy)?;
70
71        let (router, session_handle) =
72            tower_mcp::transport::http::HttpTransport::from_service(service)
73                .into_router_with_handle();
74
75        // Inbound authentication (axum-level middleware)
76        let router = apply_auth(&config, router).await?;
77
78        // Collect backend metadata for the health checker
79        let backend_meta: std::collections::HashMap<String, BackendMeta> = config
80            .backends
81            .iter()
82            .map(|b| {
83                (
84                    b.name.clone(),
85                    BackendMeta {
86                        transport: format!("{:?}", b.transport).to_lowercase(),
87                    },
88                )
89            })
90            .collect();
91
92        // Admin API
93        let admin_state = crate::admin::spawn_health_checker(
94            proxy_for_admin,
95            config.proxy.name.clone(),
96            config.proxy.version.clone(),
97            config.backends.len(),
98            backend_meta,
99        );
100        let router = router.nest(
101            "/admin",
102            crate::admin::admin_router(
103                admin_state.clone(),
104                metrics_handle,
105                session_handle.clone(),
106                cache_handle,
107                proxy_for_management,
108                &config,
109                config.source_path.clone(),
110                cb_handles,
111            ),
112        );
113        tracing::info!("Admin API enabled at /admin/backends");
114
115        // Build discovery index if enabled (search mode implies discovery)
116        #[cfg(feature = "discovery")]
117        let discovery_enabled = config.proxy.tool_discovery
118            || config.proxy.tool_exposure == crate::config::ToolExposure::Search;
119        #[cfg(feature = "discovery")]
120        let (discovery_index, discovery_tools) = if discovery_enabled {
121            let index =
122                crate::discovery::build_index(&mut proxy_for_caller, &config.proxy.separator).await;
123            let tools = crate::discovery::build_discovery_tools(index.clone());
124            (Some(index), Some(tools))
125        } else {
126            (None, None)
127        };
128        #[cfg(not(feature = "discovery"))]
129        let discovery_tools: Option<Vec<tower_mcp::Tool>> = None;
130
131        // MCP admin tools (proxy/ namespace)
132        if let Err(e) = crate::admin_tools::register_admin_tools(
133            &proxy_for_caller,
134            admin_state,
135            session_handle.clone(),
136            &config,
137            discovery_tools,
138        )
139        .await
140        {
141            tracing::warn!("Failed to register admin tools: {e}");
142        } else {
143            tracing::info!("MCP admin tools registered under proxy/ namespace");
144        }
145
146        Ok(Self {
147            router,
148            session_handle,
149            inner: proxy_for_caller,
150            config,
151            #[cfg(feature = "discovery")]
152            discovery_index,
153        })
154    }
155
156    /// Get a reference to the session handle for monitoring active sessions.
157    pub fn session_handle(&self) -> &SessionHandle {
158        &self.session_handle
159    }
160
161    /// Get a reference to the underlying [`McpProxy`] for dynamic operations.
162    ///
163    /// Use this to add backends dynamically via [`McpProxy::add_backend()`].
164    pub fn mcp_proxy(&self) -> &McpProxy {
165        &self.inner
166    }
167
168    /// Enable hot reload by watching the given config file path.
169    ///
170    /// New backends added to the config file will be connected dynamically
171    /// without restarting the proxy.
172    pub fn enable_hot_reload(&self, config_path: std::path::PathBuf) {
173        tracing::info!("Hot reload enabled, watching config file for changes");
174        crate::reload::spawn_config_watcher(
175            config_path,
176            self.inner.clone(),
177            #[cfg(feature = "discovery")]
178            self.discovery_index
179                .as_ref()
180                .map(|idx| (idx.clone(), self.config.proxy.separator.clone())),
181        );
182    }
183
184    /// Consume the proxy and return the axum Router and SessionHandle.
185    ///
186    /// Use this to embed the proxy in an existing axum application:
187    ///
188    /// ```rust,ignore
189    /// let (proxy_router, session_handle) = proxy.into_router();
190    ///
191    /// let app = Router::new()
192    ///     .nest("/mcp", proxy_router)
193    ///     .route("/health", get(|| async { "ok" }));
194    /// ```
195    pub fn into_router(self) -> (Router, SessionHandle) {
196        (self.router, self.session_handle)
197    }
198
199    /// Serve the proxy on the configured listen address.
200    ///
201    /// Blocks until a shutdown signal (SIGTERM/SIGINT) is received,
202    /// then drains connections for the configured timeout period.
203    pub async fn serve(self) -> Result<()> {
204        let addr = format!(
205            "{}:{}",
206            self.config.proxy.listen.host, self.config.proxy.listen.port
207        );
208
209        tracing::info!(listen = %addr, "Proxy ready");
210
211        let listener = tokio::net::TcpListener::bind(&addr)
212            .await
213            .with_context(|| format!("binding to {}", addr))?;
214
215        let shutdown_timeout = Duration::from_secs(self.config.proxy.shutdown_timeout_seconds);
216        axum::serve(listener, self.router)
217            .with_graceful_shutdown(shutdown_signal(shutdown_timeout))
218            .await
219            .context("server error")?;
220
221        tracing::info!("Proxy shut down");
222        Ok(())
223    }
224}
225
226/// Circuit breaker handle type alias.
227pub type CbHandle = tower_resilience::circuitbreaker::CircuitBreakerHandle;
228
229/// Build the circuit breaker layer for one backend from its config.
230///
231/// The count-based breaker only evaluates once its sliding window is full,
232/// and the window defaults to 100 calls. Aligning the window with
233/// `minimum_calls` makes the failure rate cover the most recent
234/// `minimum_calls` calls and evaluation begin exactly at the documented
235/// threshold (#220).
236fn build_breaker_layer(
237    cb: &crate::config::CircuitBreakerConfig,
238    backend_name: &str,
239) -> (
240    tower_resilience::circuitbreaker::CircuitBreakerLayer,
241    CbHandle,
242) {
243    tower_resilience::circuitbreaker::CircuitBreakerLayer::builder()
244        .failure_rate_threshold(cb.failure_rate_threshold)
245        .minimum_number_of_calls(cb.minimum_calls)
246        .sliding_window_size(cb.minimum_calls)
247        .wait_duration_in_open(Duration::from_secs(cb.wait_duration_seconds))
248        .permitted_calls_in_half_open(cb.permitted_calls_in_half_open)
249        .name(format!("{backend_name}-cb"))
250        .build_with_handle()
251}
252
253type InfallibleService = BoxCloneService<RouterRequest, RouterResponse, Infallible>;
254type MwService = BoxCloneService<RouterRequest, RouterResponse, tower::BoxError>;
255type MwStage = Box<dyn Fn(MwService) -> MwService + Send + Sync>;
256
257/// Every configured per-backend middleware composed into ONE tower-mcp
258/// backend layer. tower-mcp's `backend_layer` replaces the previously applied
259/// layer rather than stacking (joshrotenberg/tower-mcp#1173), so handing it
260/// the middlewares one call at a time silently keeps only the last.
261///
262/// Composition happens in three zones matching the middlewares' type
263/// contracts: retry and hedging operate on the raw `Error = Infallible`
264/// backend service (both classify failures from the response, and hedging
265/// requires a `Clone` error type); concurrency, rate limit, timeout, and
266/// circuit breaker operate in a widened `BoxError` plane where they can
267/// produce service-level errors; a `CatchError` fold then converts those
268/// errors into JSON-RPC error responses, and outlier detection sits on top,
269/// observing the response plane (its `Service` impl requires
270/// `Error = Infallible`).
271#[derive(Default)]
272struct BackendMiddlewareLayer {
273    /// Innermost: retries against the raw backend service.
274    retry: Option<RetryLayer<RouterRequest, RouterResponse, Infallible>>,
275    /// Outside retry, still in the Infallible zone.
276    hedge: Option<tower_resilience::hedge::HedgeLayer>,
277    /// BoxError zone, applied outward in push order: concurrency, rate
278    /// limit, timeout, circuit breaker.
279    stages: Vec<MwStage>,
280    /// Outermost, after the error fold: observes response-plane errors.
281    outlier: Option<crate::outlier::OutlierDetectionLayer>,
282}
283
284impl BackendMiddlewareLayer {
285    fn is_empty(&self) -> bool {
286        self.retry.is_none()
287            && self.hedge.is_none()
288            && self.stages.is_empty()
289            && self.outlier.is_none()
290    }
291
292    fn stage<L>(&mut self, layer: L)
293    where
294        L: Layer<MwService> + Send + Sync + 'static,
295        L::Service:
296            tower::Service<RouterRequest, Response = RouterResponse> + Clone + Send + 'static,
297        <L::Service as tower::Service<RouterRequest>>::Error: std::fmt::Display,
298        <L::Service as tower::Service<RouterRequest>>::Future: Send + 'static,
299    {
300        self.stages.push(Box::new(move |inner| {
301            BoxCloneService::new(
302                layer
303                    .layer(inner)
304                    .map_err(|e| -> tower::BoxError { e.to_string().into() }),
305            )
306        }));
307    }
308}
309
310impl<S> Layer<S> for BackendMiddlewareLayer
311where
312    S: tower::Service<RouterRequest, Response = RouterResponse, Error = Infallible>
313        + Clone
314        + Send
315        + 'static,
316    S::Future: Send + 'static,
317{
318    type Service = InfallibleService;
319
320    fn layer(&self, base: S) -> InfallibleService {
321        fn widen(e: Infallible) -> tower::BoxError {
322            match e {}
323        }
324
325        // Zone A: retry classifies failures from responses over the
326        // Infallible base (hedging also needs a Clone inner error, which
327        // Infallible satisfies).
328        let after_retry: InfallibleService = match &self.retry {
329            Some(retry) => BoxCloneService::new(retry.layer(base)),
330            None => BoxCloneService::new(base),
331        };
332
333        // Zone B: error-producing middlewares in the BoxError plane. Hedging
334        // wraps the inner error type (`HedgeError<Infallible>`), so it forms
335        // the floor of this zone.
336        let mut svc: MwService = match &self.hedge {
337            Some(hedge) => BoxCloneService::new(
338                hedge
339                    .layer(after_retry)
340                    .map_err(|e| -> tower::BoxError { e.to_string().into() }),
341            ),
342            None => BoxCloneService::new(after_retry.map_err(widen)),
343        };
344        for stage in &self.stages {
345            svc = stage(svc);
346        }
347
348        // Fold middleware errors into JSON-RPC error responses.
349        let folded: InfallibleService =
350            BoxCloneService::new(tower_mcp::transport::CatchError::new(svc));
351
352        // Zone C: outlier detection observes the response plane.
353        match &self.outlier {
354            Some(outlier) => BoxCloneService::new(outlier.layer(folded)),
355            None => folded,
356        }
357    }
358}
359
360/// Build the McpProxy with all backends and per-backend middleware.
361/// Returns the proxy and a map of backend name -> circuit breaker handle.
362async fn build_mcp_proxy(config: &ProxyConfig) -> Result<(McpProxy, HashMap<String, CbHandle>)> {
363    let mut builder = McpProxy::builder(&config.proxy.name, &config.proxy.version)
364        .separator(&config.proxy.separator);
365    let mut cb_handles: HashMap<String, CbHandle> = HashMap::new();
366
367    if let Some(instructions) = &config.proxy.instructions {
368        builder = builder.instructions(instructions);
369    }
370
371    // Create shared outlier detector if any backend has outlier_detection configured.
372    // Use the max of all max_ejection_percent values.
373    let outlier_detector = {
374        let max_pct = config
375            .backends
376            .iter()
377            .filter_map(|b| b.outlier_detection.as_ref())
378            .map(|od| od.max_ejection_percent)
379            .max();
380        max_pct.map(crate::outlier::OutlierDetector::new)
381    };
382
383    for backend in &config.backends {
384        tracing::info!(name = %backend.name, transport = ?backend.transport, "Adding backend");
385
386        match backend.transport {
387            TransportType::Stdio => {
388                let command = backend.command.as_deref().unwrap();
389                let args: Vec<&str> = backend.args.iter().map(|s| s.as_str()).collect();
390
391                let mut cmd = Command::new(command);
392                cmd.args(&args);
393
394                for (key, value) in &backend.env {
395                    cmd.env(key, value);
396                }
397
398                let transport = StdioClientTransport::spawn_command(&mut cmd)
399                    .await
400                    .with_context(|| format!("spawning backend '{}'", backend.name))?;
401
402                builder = builder.backend(&backend.name, transport).await;
403            }
404            TransportType::Http => {
405                let url = backend.url.as_deref().unwrap();
406                let mut transport = tower_mcp::client::HttpClientTransport::new(url);
407                if let Some(token) = &backend.bearer_token {
408                    transport = transport.bearer_token(token);
409                }
410
411                builder = builder.backend(&backend.name, transport).await;
412            }
413            #[cfg(feature = "websocket")]
414            TransportType::Websocket => {
415                let url = backend.url.as_deref().unwrap();
416                tracing::info!(url = %url, "Connecting to WebSocket backend");
417                let transport = if let Some(token) = &backend.bearer_token {
418                    crate::ws_transport::WebSocketClientTransport::connect_with_bearer_token(
419                        url, token,
420                    )
421                    .await
422                    .with_context(|| {
423                        format!("connecting to WebSocket backend '{}'", backend.name)
424                    })?
425                } else {
426                    crate::ws_transport::WebSocketClientTransport::connect(url)
427                        .await
428                        .with_context(|| {
429                            format!("connecting to WebSocket backend '{}'", backend.name)
430                        })?
431                };
432
433                builder = builder.backend(&backend.name, transport).await;
434            }
435            #[cfg(not(feature = "websocket"))]
436            TransportType::Websocket => {
437                anyhow::bail!(
438                    "WebSocket transport requires the 'websocket' feature. \
439                     Rebuild with: cargo install mcp-proxy --features websocket"
440                );
441            }
442        }
443
444        // Per-backend middleware stack (applied in order: inner -> outer).
445        // Collected into one BackendMiddlewareLayer and handed to tower-mcp
446        // as a single backend_layer call; see the type's doc comment.
447        let mut mw = BackendMiddlewareLayer::default();
448
449        // Retry (innermost -- retries happen before other middleware)
450        if let Some(retry_cfg) = &backend.retry {
451            tracing::info!(
452                backend = %backend.name,
453                max_retries = retry_cfg.max_retries,
454                initial_backoff_ms = retry_cfg.initial_backoff_ms,
455                max_backoff_ms = retry_cfg.max_backoff_ms,
456                "Applying retry policy"
457            );
458            mw.retry = Some(crate::retry::build_retry_layer(retry_cfg, &backend.name));
459        }
460
461        // Hedging (after retry, before concurrency -- hedges are separate requests)
462        if let Some(hedge_cfg) = &backend.hedging {
463            let delay = Duration::from_millis(hedge_cfg.delay_ms);
464            let max_attempts = hedge_cfg.max_hedges + 1; // +1 for the primary request
465            tracing::info!(
466                backend = %backend.name,
467                delay_ms = hedge_cfg.delay_ms,
468                max_hedges = hedge_cfg.max_hedges,
469                "Applying request hedging"
470            );
471            let layer = if delay.is_zero() {
472                tower_resilience::hedge::HedgeLayer::builder()
473                    .no_delay()
474                    .max_hedged_attempts(max_attempts)
475                    .name(format!("{}-hedge", backend.name))
476                    .build()
477            } else {
478                tower_resilience::hedge::HedgeLayer::builder()
479                    .delay(delay)
480                    .max_hedged_attempts(max_attempts)
481                    .name(format!("{}-hedge", backend.name))
482                    .build()
483            };
484            mw.hedge = Some(layer);
485        }
486
487        // Concurrency limit
488        if let Some(cc) = &backend.concurrency {
489            tracing::info!(
490                backend = %backend.name,
491                max = cc.max_concurrent,
492                "Applying concurrency limit"
493            );
494            mw.stage(tower::limit::ConcurrencyLimitLayer::new(cc.max_concurrent));
495        }
496
497        // Rate limit
498        if let Some(rl) = &backend.rate_limit {
499            tracing::info!(
500                backend = %backend.name,
501                requests = rl.requests,
502                period_seconds = rl.period_seconds,
503                "Applying rate limit"
504            );
505            let layer = tower_resilience::ratelimiter::RateLimiterLayer::builder()
506                .limit_for_period(rl.requests)
507                .refresh_period(Duration::from_secs(rl.period_seconds))
508                .name(format!("{}-ratelimit", backend.name))
509                .build();
510            mw.stage(layer);
511        }
512
513        // Timeout
514        if let Some(timeout) = &backend.timeout {
515            tracing::info!(
516                backend = %backend.name,
517                seconds = timeout.seconds,
518                "Applying timeout"
519            );
520            mw.stage(TimeoutLayer::new(Duration::from_secs(timeout.seconds)));
521        }
522
523        // Circuit breaker
524        if let Some(cb) = &backend.circuit_breaker {
525            tracing::info!(
526                backend = %backend.name,
527                failure_rate = cb.failure_rate_threshold,
528                wait_seconds = cb.wait_duration_seconds,
529                "Applying circuit breaker"
530            );
531            let (layer, handle) = build_breaker_layer(cb, &backend.name);
532            cb_handles.insert(backend.name.clone(), handle);
533            mw.stage(layer);
534        }
535
536        // Outlier detection (outermost -- observes errors after all other middleware)
537        if let Some(od) = &backend.outlier_detection
538            && let Some(ref detector) = outlier_detector
539        {
540            tracing::info!(
541                backend = %backend.name,
542                consecutive_errors = od.consecutive_errors,
543                base_ejection_seconds = od.base_ejection_seconds,
544                max_ejection_percent = od.max_ejection_percent,
545                "Applying outlier detection"
546            );
547            let layer = crate::outlier::OutlierDetectionLayer::new(
548                backend.name.clone(),
549                od.clone(),
550                detector.clone(),
551            );
552            mw.outlier = Some(layer);
553        }
554
555        if !mw.is_empty() {
556            builder = builder.backend_layer(mw);
557        }
558    }
559
560    let result = builder.build().await?;
561
562    if !result.skipped.is_empty() {
563        for s in &result.skipped {
564            tracing::warn!("Skipped backend: {s}");
565        }
566    }
567
568    Ok((result.proxy, cb_handles))
569}
570
571/// Build a scope-enforcement layer from configured OAuth `required_scopes`.
572///
573/// Returns `None` when no scopes are required (the layer would be a no-op).
574/// Otherwise returns a [`ScopeEnforcementLayer`](tower_mcp::oauth::ScopeEnforcementLayer)
575/// whose default policy requires *all* of `required_scopes` to be present in the
576/// token (AND semantics) for every request.
577#[cfg(feature = "oauth")]
578fn oauth_scope_layer(
579    required_scopes: &[String],
580) -> Option<tower_mcp::oauth::ScopeEnforcementLayer> {
581    if required_scopes.is_empty() {
582        return None;
583    }
584    let policy = tower_mcp::oauth::ScopePolicy::new().default_scopes(
585        tower_mcp::oauth::ScopeRequirement::all(required_scopes.iter().cloned()),
586    );
587    Some(tower_mcp::oauth::ScopeEnforcementLayer::new(policy))
588}
589
590/// Build the MCP-level middleware stack around the proxy.
591fn build_middleware_stack(
592    config: &ProxyConfig,
593    proxy: McpProxy,
594) -> Result<(
595    BoxCloneService<RouterRequest, RouterResponse, Infallible>,
596    Option<cache::CacheHandle>,
597)> {
598    let mut service: BoxCloneService<RouterRequest, RouterResponse, Infallible> =
599        BoxCloneService::new(proxy);
600    let mut cache_handle: Option<cache::CacheHandle> = None;
601
602    // Argument injection (innermost -- merges default/per-tool args into CallTool requests)
603    let injection_rules: Vec<_> = config
604        .backends
605        .iter()
606        .filter(|b| !b.default_args.is_empty() || !b.inject_args.is_empty())
607        .map(|b| {
608            let namespace = format!("{}{}", b.name, config.proxy.separator);
609            tracing::info!(
610                backend = %b.name,
611                default_args = b.default_args.len(),
612                tool_rules = b.inject_args.len(),
613                "Applying argument injection"
614            );
615            crate::inject::InjectionRules::new(
616                namespace,
617                b.default_args.clone(),
618                b.inject_args.clone(),
619            )
620        })
621        .collect();
622
623    if !injection_rules.is_empty() {
624        service = BoxCloneService::new(crate::inject::InjectArgsService::new(
625            service,
626            injection_rules,
627        ));
628    }
629
630    // Parameter overrides (after inject, before filter -- hides/renames tool params)
631    let param_overrides: Vec<_> = config
632        .backends
633        .iter()
634        .filter(|b| !b.param_overrides.is_empty())
635        .flat_map(|b| {
636            let namespace = format!("{}{}", b.name, config.proxy.separator);
637            tracing::info!(
638                backend = %b.name,
639                overrides = b.param_overrides.len(),
640                "Applying parameter overrides"
641            );
642            b.param_overrides
643                .iter()
644                .map(move |c| crate::param_override::ToolOverride::new(&namespace, c))
645        })
646        .collect();
647
648    if !param_overrides.is_empty() {
649        service = BoxCloneService::new(crate::param_override::ParamOverrideService::new(
650            service,
651            param_overrides,
652        ));
653    }
654
655    // Canary routing (rewrites requests from primary to canary namespace based on weight)
656    let canary_mappings: std::collections::HashMap<String, (String, u32, u32)> = config
657        .backends
658        .iter()
659        .filter_map(|b| {
660            b.canary_of.as_ref().map(|primary_name| {
661                // Find the primary backend's weight
662                let primary_weight = config
663                    .backends
664                    .iter()
665                    .find(|p| p.name == *primary_name)
666                    .map(|p| p.weight)
667                    .unwrap_or(100);
668                (
669                    primary_name.clone(),
670                    (b.name.clone(), primary_weight, b.weight),
671                )
672            })
673        })
674        .collect();
675
676    if !canary_mappings.is_empty() {
677        for (primary, (canary, pw, cw)) in &canary_mappings {
678            tracing::info!(
679                primary = %primary,
680                canary = %canary,
681                primary_weight = pw,
682                canary_weight = cw,
683                "Enabling canary routing"
684            );
685        }
686        service = BoxCloneService::new(crate::canary::CanaryService::new(
687            service,
688            canary_mappings,
689            &config.proxy.separator,
690        ));
691    }
692
693    // Failover routing (deterministic fallback on primary error)
694    // Collect failover backends grouped by primary, sorted by priority (ascending).
695    let mut failover_groups: std::collections::HashMap<String, Vec<(u32, String)>> =
696        std::collections::HashMap::new();
697    for b in &config.backends {
698        if let Some(ref primary) = b.failover_for {
699            failover_groups
700                .entry(primary.clone())
701                .or_default()
702                .push((b.priority, b.name.clone()));
703        }
704    }
705    // Sort each group by priority (lower = preferred)
706    let failover_mappings: std::collections::HashMap<String, Vec<String>> = failover_groups
707        .into_iter()
708        .map(|(primary, mut backends)| {
709            backends.sort_by_key(|(priority, _)| *priority);
710            let names: Vec<String> = backends.into_iter().map(|(_, name)| name).collect();
711            (primary, names)
712        })
713        .collect();
714
715    if !failover_mappings.is_empty() {
716        for (primary, failovers) in &failover_mappings {
717            tracing::info!(
718                primary = %primary,
719                failovers = ?failovers,
720                "Enabling failover routing"
721            );
722        }
723        service = BoxCloneService::new(crate::failover::FailoverService::new(
724            service,
725            failover_mappings,
726            &config.proxy.separator,
727        ));
728    }
729
730    // Traffic mirroring (sends cloned requests through the proxy)
731    let mirror_mappings: std::collections::HashMap<String, (String, u32)> = config
732        .backends
733        .iter()
734        .filter_map(|b| {
735            b.mirror_of
736                .as_ref()
737                .map(|source| (source.clone(), (b.name.clone(), b.mirror_percent)))
738        })
739        .collect();
740
741    if !mirror_mappings.is_empty() {
742        for (source, (mirror, pct)) in &mirror_mappings {
743            tracing::info!(
744                source = %source,
745                mirror = %mirror,
746                percent = pct,
747                "Enabling traffic mirroring"
748            );
749        }
750        service = BoxCloneService::new(crate::mirror::MirrorService::new(
751            service,
752            mirror_mappings,
753            &config.proxy.separator,
754        ));
755    }
756
757    // Response caching
758    let cache_configs: Vec<_> = config
759        .backends
760        .iter()
761        .filter_map(|b| {
762            b.cache
763                .as_ref()
764                .map(|c| (format!("{}{}", b.name, config.proxy.separator), c))
765        })
766        .collect();
767
768    if !cache_configs.is_empty() {
769        for (ns, cfg) in &cache_configs {
770            tracing::info!(
771                backend = %ns.trim_end_matches(&config.proxy.separator),
772                resource_ttl = cfg.resource_ttl_seconds,
773                tool_ttl = cfg.tool_ttl_seconds,
774                max_entries = cfg.max_entries,
775                "Applying response cache"
776            );
777        }
778        let (cache_svc, handle) = cache::CacheService::new(service, cache_configs, &config.cache);
779        service = BoxCloneService::new(cache_svc);
780        cache_handle = Some(handle);
781    }
782
783    // Request coalescing
784    if config.performance.coalesce_requests {
785        tracing::info!("Request coalescing enabled");
786        service = BoxCloneService::new(coalesce::CoalesceService::new(service));
787    }
788
789    // Request validation
790    if config.security.max_argument_size.is_some() {
791        let validation = ValidationConfig {
792            max_argument_size: config.security.max_argument_size,
793        };
794        if let Some(max) = validation.max_argument_size {
795            tracing::info!(max_argument_size = max, "Applying request validation");
796        }
797        service = BoxCloneService::new(ValidationService::new(service, validation));
798    }
799
800    // Static capability filtering
801    let filters: Vec<_> = config
802        .backends
803        .iter()
804        .filter_map(|b| b.build_filter(&config.proxy.separator).transpose())
805        .collect::<anyhow::Result<Vec<_>>>()?;
806
807    if !filters.is_empty() {
808        for f in &filters {
809            tracing::info!(
810                backend = %f.namespace.trim_end_matches(&config.proxy.separator),
811                tool_filter = ?f.tool_filter,
812                resource_filter = ?f.resource_filter,
813                prompt_filter = ?f.prompt_filter,
814                "Applying capability filter"
815            );
816        }
817        service = BoxCloneService::new(CapabilityFilterService::new(service, filters));
818    }
819
820    // Search-mode filtering: hide all tools except proxy/ namespace
821    if config.proxy.tool_exposure == crate::config::ToolExposure::Search {
822        let prefix = format!("proxy{}", config.proxy.separator);
823        tracing::info!(
824            prefix = %prefix,
825            "Search mode: ListTools will only show proxy/ namespace tools"
826        );
827        service =
828            BoxCloneService::new(crate::filter::SearchModeFilterService::new(service, prefix));
829    }
830
831    // Tool aliasing
832    let alias_mappings: Vec<_> = config
833        .backends
834        .iter()
835        .flat_map(|b| {
836            let ns = format!("{}{}", b.name, config.proxy.separator);
837            b.aliases
838                .iter()
839                .map(move |a| (ns.clone(), a.from.clone(), a.to.clone()))
840        })
841        .collect();
842
843    if let Some(alias_map) = alias::AliasMap::new(alias_mappings) {
844        let count = alias_map.forward.len();
845        tracing::info!(aliases = count, "Applying tool aliases");
846        service = BoxCloneService::new(alias::AliasService::new(service, alias_map));
847    }
848
849    // Composite tools (fan-out to multiple backend tools)
850    if !config.composite_tools.is_empty() {
851        let count = config.composite_tools.len();
852        tracing::info!(composite_tools = count, "Applying composite tool fan-out");
853        service = BoxCloneService::new(crate::composite::CompositeService::new(
854            service,
855            config.composite_tools.clone(),
856        ));
857    }
858
859    // Bearer token scoping (per-token allow/deny lists)
860    #[cfg(feature = "oauth")]
861    if matches!(
862        &config.auth,
863        Some(AuthConfig::Bearer {
864            scoped_tokens,
865            ..
866        }) if !scoped_tokens.is_empty()
867    ) {
868        tracing::info!("Enabling bearer token scoping middleware");
869        service = BoxCloneService::new(crate::bearer_scope::BearerScopingService::new(service));
870    }
871
872    // RBAC (JWT auth only)
873    #[cfg(feature = "oauth")]
874    {
875        let rbac_config = match &config.auth {
876            Some(
877                AuthConfig::Jwt {
878                    roles,
879                    role_mapping: Some(mapping),
880                    ..
881                }
882                | AuthConfig::OAuth {
883                    roles,
884                    role_mapping: Some(mapping),
885                    ..
886                },
887            ) if !roles.is_empty() => {
888                tracing::info!(
889                    roles = roles.len(),
890                    claim = %mapping.claim,
891                    "Enabling RBAC"
892                );
893                Some(RbacConfig::new(roles, mapping))
894            }
895            _ => None,
896        };
897
898        if let Some(rbac) = rbac_config {
899            service = BoxCloneService::new(RbacService::new(service, rbac));
900        }
901
902        // OAuth `required_scopes` enforcement: a coarse global gate that rejects
903        // any token missing one of the configured scopes (AND semantics). Runs
904        // outside RBAC so a token lacking the required scopes is denied for every
905        // operation, including `tools/list`. Reads TokenClaims injected by the
906        // OAuth auth layer; requests without claims pass through (already rejected
907        // upstream by the HTTP auth layer when auth is enabled).
908        let required_scopes: &[String] = match &config.auth {
909            Some(AuthConfig::OAuth {
910                required_scopes, ..
911            }) => required_scopes,
912            _ => &[],
913        };
914        if let Some(layer) = oauth_scope_layer(required_scopes) {
915            tracing::info!(
916                scopes = ?required_scopes,
917                "Enabling OAuth required_scopes enforcement"
918            );
919            service = BoxCloneService::new(tower::Layer::layer(&layer, service));
920        }
921
922        // Token passthrough (inject ClientToken for forward_auth backends)
923        let forward_namespaces: std::collections::HashSet<String> = config
924            .backends
925            .iter()
926            .filter(|b| b.forward_auth)
927            .map(|b| format!("{}{}", b.name, config.proxy.separator))
928            .collect();
929
930        if !forward_namespaces.is_empty() {
931            tracing::info!(
932                backends = ?forward_namespaces,
933                "Enabling token passthrough for forward_auth backends"
934            );
935            service = BoxCloneService::new(crate::token::TokenPassthroughService::new(
936                service,
937                forward_namespaces,
938            ));
939        }
940    }
941
942    // Metrics
943    #[cfg(feature = "metrics")]
944    if config.observability.metrics.enabled {
945        service = BoxCloneService::new(crate::metrics::MetricsService::new(service));
946    }
947
948    // Structured access logging
949    if config.observability.access_log.enabled {
950        tracing::info!("Access logging enabled (target: mcp::access)");
951        service = BoxCloneService::new(crate::access_log::AccessLogService::new(
952            service,
953            &config.proxy.separator,
954        ));
955    }
956
957    // Audit logging
958    if config.observability.audit {
959        tracing::info!("Audit logging enabled (target: mcp::audit)");
960        let audited = tower::Layer::layer(&tower_mcp::AuditLayer::new(), service);
961        service = BoxCloneService::new(tower_mcp::CatchError::new(audited));
962    }
963
964    // Global rate limit (outermost -- protects entire proxy)
965    if let Some(ref rl) = config.proxy.rate_limit {
966        tracing::info!(
967            requests = rl.requests,
968            period_seconds = rl.period_seconds,
969            "Applying global rate limit"
970        );
971        let layer = tower_resilience::ratelimiter::RateLimiterLayer::builder()
972            .limit_for_period(rl.requests)
973            .refresh_period(Duration::from_secs(rl.period_seconds))
974            .name("global-ratelimit")
975            .build();
976        let limited = tower::Layer::layer(&layer, service);
977        service = BoxCloneService::new(tower_mcp::CatchError::new(limited));
978    }
979
980    Ok((service, cache_handle))
981}
982
983/// Apply inbound authentication middleware to the router.
984async fn apply_auth(config: &ProxyConfig, router: Router) -> Result<Router> {
985    let router = if let Some(auth) = &config.auth {
986        match auth {
987            AuthConfig::Bearer {
988                tokens,
989                scoped_tokens,
990            } => {
991                let total = tokens.len() + scoped_tokens.len();
992                if scoped_tokens.is_empty() {
993                    // Simple bearer auth: use StaticBearerValidator
994                    tracing::info!(token_count = total, "Enabling bearer token auth");
995                    let validator = StaticBearerValidator::new(tokens.iter().cloned());
996                    let layer = AuthLayer::new(validator);
997                    router.layer(layer)
998                } else {
999                    // Scoped bearer auth: use custom layer that injects TokenClaims
1000                    #[cfg(feature = "oauth")]
1001                    {
1002                        tracing::info!(
1003                            token_count = total,
1004                            scoped = scoped_tokens.len(),
1005                            "Enabling bearer token auth with per-token scoping"
1006                        );
1007                        let layer =
1008                            crate::bearer_scope::ScopedBearerAuthLayer::new(tokens, scoped_tokens);
1009                        router.layer(layer)
1010                    }
1011                    #[cfg(not(feature = "oauth"))]
1012                    {
1013                        anyhow::bail!(
1014                            "Per-token tool scoping requires the 'oauth' feature. \
1015                             Rebuild with: cargo install mcp-proxy --features oauth"
1016                        );
1017                    }
1018                }
1019            }
1020            #[cfg(feature = "oauth")]
1021            AuthConfig::Jwt {
1022                issuer,
1023                audience,
1024                jwks_uri,
1025                ..
1026            } => {
1027                tracing::info!(
1028                    issuer = %issuer,
1029                    audience = %audience,
1030                    jwks_uri = %jwks_uri,
1031                    "Enabling JWT auth (JWKS)"
1032                );
1033                let validator = tower_mcp::oauth::JwksValidator::builder(jwks_uri)
1034                    .expected_audience(audience)
1035                    .expected_issuer(issuer)
1036                    .build()
1037                    .await
1038                    .context("building JWKS validator")?;
1039
1040                let addr = format!(
1041                    "http://{}:{}",
1042                    config.proxy.listen.host, config.proxy.listen.port
1043                );
1044                let metadata = tower_mcp::oauth::ProtectedResourceMetadata::new(&addr)
1045                    .authorization_server(issuer);
1046
1047                let layer = tower_mcp::oauth::OAuthLayer::new(validator, metadata);
1048                router.layer(layer)
1049            }
1050            #[cfg(not(feature = "oauth"))]
1051            AuthConfig::Jwt { .. } => {
1052                anyhow::bail!(
1053                    "JWT auth requires the 'oauth' feature. Rebuild with: cargo install mcp-proxy --features oauth"
1054                );
1055            }
1056            #[cfg(feature = "oauth")]
1057            AuthConfig::OAuth {
1058                issuer,
1059                audience,
1060                token_validation,
1061                jwks_uri,
1062                introspection_endpoint,
1063                client_id,
1064                client_secret,
1065                ..
1066            } => {
1067                use crate::config::TokenValidationStrategy;
1068
1069                tracing::info!(
1070                    issuer = %issuer,
1071                    audience = %audience,
1072                    strategy = ?token_validation,
1073                    "Enabling OAuth 2.1 auth"
1074                );
1075
1076                // Auto-discover endpoints from issuer if not overridden
1077                let discovered = crate::introspection::discover_auth_server(issuer)
1078                    .await
1079                    .context("discovering OAuth authorization server")?;
1080
1081                let effective_jwks_uri = jwks_uri
1082                    .as_deref()
1083                    .or(discovered.jwks_uri.as_deref())
1084                    .ok_or_else(|| {
1085                        anyhow::anyhow!(
1086                            "JWKS URI not found via discovery and not configured manually"
1087                        )
1088                    })?;
1089
1090                let effective_introspection = introspection_endpoint
1091                    .as_deref()
1092                    .or(discovered.introspection_endpoint.as_deref());
1093
1094                let addr = format!(
1095                    "http://{}:{}",
1096                    config.proxy.listen.host, config.proxy.listen.port
1097                );
1098                let metadata = tower_mcp::oauth::ProtectedResourceMetadata::new(&addr)
1099                    .authorization_server(issuer);
1100
1101                match token_validation {
1102                    TokenValidationStrategy::Jwt => {
1103                        let validator =
1104                            tower_mcp::oauth::JwksValidator::builder(effective_jwks_uri)
1105                                .expected_audience(audience)
1106                                .expected_issuer(issuer)
1107                                .build()
1108                                .await
1109                                .context("building JWKS validator")?;
1110                        let layer = tower_mcp::oauth::OAuthLayer::new(validator, metadata);
1111                        router.layer(layer)
1112                    }
1113                    TokenValidationStrategy::Introspection => {
1114                        let endpoint = effective_introspection.ok_or_else(|| {
1115                            anyhow::anyhow!(
1116                                "introspection endpoint not found via discovery and not configured"
1117                            )
1118                        })?;
1119                        let validator = crate::introspection::IntrospectionValidator::new(
1120                            endpoint,
1121                            client_id.as_deref().unwrap(),
1122                            client_secret.as_deref().unwrap(),
1123                        )
1124                        .expected_audience(audience);
1125                        let layer = tower_mcp::oauth::OAuthLayer::new(validator, metadata);
1126                        router.layer(layer)
1127                    }
1128                    TokenValidationStrategy::Both => {
1129                        let endpoint = effective_introspection.ok_or_else(|| {
1130                            anyhow::anyhow!(
1131                                "introspection endpoint not found via discovery and not configured"
1132                            )
1133                        })?;
1134                        let jwt_validator =
1135                            tower_mcp::oauth::JwksValidator::builder(effective_jwks_uri)
1136                                .expected_audience(audience)
1137                                .expected_issuer(issuer)
1138                                .build()
1139                                .await
1140                                .context("building JWKS validator")?;
1141                        let introspection_validator =
1142                            crate::introspection::IntrospectionValidator::new(
1143                                endpoint,
1144                                client_id.as_deref().unwrap(),
1145                                client_secret.as_deref().unwrap(),
1146                            )
1147                            .expected_audience(audience);
1148                        let fallback = crate::introspection::FallbackValidator::new(
1149                            jwt_validator,
1150                            introspection_validator,
1151                        );
1152                        let layer = tower_mcp::oauth::OAuthLayer::new(fallback, metadata);
1153                        router.layer(layer)
1154                    }
1155                }
1156            }
1157            #[cfg(not(feature = "oauth"))]
1158            AuthConfig::OAuth { .. } => {
1159                anyhow::bail!(
1160                    "OAuth auth requires the 'oauth' feature. Rebuild with: cargo install mcp-proxy --features oauth"
1161                );
1162            }
1163        }
1164    } else {
1165        router
1166    };
1167    Ok(router)
1168}
1169
1170/// Wait for SIGTERM or SIGINT, then log and return.
1171pub async fn shutdown_signal(timeout: Duration) {
1172    let ctrl_c = tokio::signal::ctrl_c();
1173    #[cfg(unix)]
1174    {
1175        let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
1176            .expect("SIGTERM handler");
1177        tokio::select! {
1178            _ = ctrl_c => {},
1179            _ = sigterm.recv() => {},
1180        }
1181    }
1182    #[cfg(not(unix))]
1183    {
1184        ctrl_c.await.ok();
1185    }
1186    tracing::info!(
1187        timeout_seconds = timeout.as_secs(),
1188        "Shutdown signal received, draining connections"
1189    );
1190}
1191
1192#[cfg(all(test, feature = "oauth"))]
1193mod scope_enforcement_tests {
1194    use std::collections::HashMap;
1195
1196    use tower::{Layer, Service};
1197    use tower_mcp::oauth::token::TokenClaims;
1198    use tower_mcp::protocol::{CallToolParams, McpRequest, RequestId};
1199    use tower_mcp::router::Extensions;
1200
1201    use super::oauth_scope_layer;
1202    use crate::test_util::MockService;
1203
1204    /// Build a `tools/call` request, optionally carrying a token with `scope`.
1205    fn call_with_scope(scope: Option<&str>) -> tower_mcp::RouterRequest {
1206        let mut extensions = Extensions::new();
1207        if let Some(scope) = scope {
1208            extensions.insert(TokenClaims {
1209                sub: Some("user".into()),
1210                iss: None,
1211                aud: None,
1212                exp: None,
1213                scope: Some(scope.to_string()),
1214                client_id: None,
1215                extra: HashMap::new(),
1216            });
1217        }
1218        tower_mcp::RouterRequest {
1219            id: RequestId::Number(1),
1220            inner: McpRequest::CallTool(CallToolParams {
1221                name: "fs/read".into(),
1222                arguments: serde_json::json!({}),
1223                input_responses: None,
1224                request_state: None,
1225                meta: None,
1226                task: None,
1227            }),
1228            extensions,
1229        }
1230    }
1231
1232    #[test]
1233    fn no_required_scopes_yields_no_layer() {
1234        assert!(oauth_scope_layer(&[]).is_none());
1235    }
1236
1237    #[tokio::test]
1238    async fn token_missing_required_scope_is_rejected() {
1239        let required = vec!["mcp:access".to_string()];
1240        let layer = oauth_scope_layer(&required).expect("layer for non-empty scopes");
1241        let mut svc = layer.layer(MockService::with_tools(&["fs/read"]));
1242
1243        // Token carries a different scope -> missing the required one.
1244        let resp = svc
1245            .call(call_with_scope(Some("other:scope")))
1246            .await
1247            .unwrap();
1248        let err = resp.inner.unwrap_err();
1249        assert!(
1250            err.message.to_lowercase().contains("scope"),
1251            "expected insufficient-scope error, got: {}",
1252            err.message
1253        );
1254    }
1255
1256    #[tokio::test]
1257    async fn token_with_all_required_scopes_is_allowed() {
1258        let required = vec!["mcp:access".to_string(), "mcp:read".to_string()];
1259        let layer = oauth_scope_layer(&required).expect("layer for non-empty scopes");
1260        let mut svc = layer.layer(MockService::with_tools(&["fs/read"]));
1261
1262        let resp = svc
1263            .call(call_with_scope(Some("mcp:access mcp:read mcp:extra")))
1264            .await
1265            .unwrap();
1266        assert!(
1267            resp.inner.is_ok(),
1268            "token carrying all required scopes should be allowed"
1269        );
1270    }
1271}
1272
1273#[cfg(test)]
1274mod middleware_stack_tests {
1275    //! Regression tests for #218: the per-backend middleware chain must
1276    //! compose as one stack. Under tower-mcp's last-wins `backend_layer`
1277    //! semantics (joshrotenberg/tower-mcp#1173), only the final middleware
1278    //! survived and the first test here fails at "call 1 should time out".
1279
1280    use std::pin::Pin;
1281    use std::sync::Arc;
1282    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
1283    use std::task::{Context, Poll};
1284
1285    use tower::Service;
1286    use tower_mcp::CallToolResult;
1287    use tower_mcp::protocol::{CallToolParams, McpRequest, McpResponse, RequestId};
1288    use tower_mcp::router::Extensions;
1289    use tower_mcp_types::JsonRpcError;
1290
1291    use super::*;
1292
1293    /// Infallible backend with a hit counter, switchable slowness, and a
1294    /// countdown of response-plane failures to emit before succeeding.
1295    #[derive(Clone)]
1296    struct FlakyBackend {
1297        hits: Arc<AtomicUsize>,
1298        slow: Arc<AtomicBool>,
1299        fail_responses: Arc<AtomicUsize>,
1300    }
1301
1302    impl FlakyBackend {
1303        fn new() -> (Self, Arc<AtomicUsize>, Arc<AtomicBool>, Arc<AtomicUsize>) {
1304            let hits = Arc::new(AtomicUsize::new(0));
1305            let slow = Arc::new(AtomicBool::new(false));
1306            let fail_responses = Arc::new(AtomicUsize::new(0));
1307            (
1308                Self {
1309                    hits: hits.clone(),
1310                    slow: slow.clone(),
1311                    fail_responses: fail_responses.clone(),
1312                },
1313                hits,
1314                slow,
1315                fail_responses,
1316            )
1317        }
1318    }
1319
1320    impl tower::Service<RouterRequest> for FlakyBackend {
1321        type Response = RouterResponse;
1322        type Error = Infallible;
1323        type Future = Pin<Box<dyn Future<Output = Result<RouterResponse, Infallible>> + Send>>;
1324
1325        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
1326            Poll::Ready(Ok(()))
1327        }
1328
1329        fn call(&mut self, req: RouterRequest) -> Self::Future {
1330            let hits = self.hits.clone();
1331            let slow = self.slow.clone();
1332            let fail_responses = self.fail_responses.clone();
1333            Box::pin(async move {
1334                hits.fetch_add(1, Ordering::SeqCst);
1335                if slow.load(Ordering::SeqCst) {
1336                    tokio::time::sleep(Duration::from_millis(100)).await;
1337                }
1338                let inner = if fail_responses
1339                    .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
1340                    .is_ok()
1341                {
1342                    Err(JsonRpcError::internal_error("transient backend failure"))
1343                } else {
1344                    Ok(McpResponse::CallTool(CallToolResult::text("pong")))
1345                };
1346                Ok(RouterResponse { id: req.id, inner })
1347            })
1348        }
1349    }
1350
1351    fn req() -> RouterRequest {
1352        RouterRequest {
1353            id: RequestId::Number(1),
1354            inner: McpRequest::CallTool(CallToolParams {
1355                name: "ping".to_string(),
1356                arguments: serde_json::json!({}),
1357                input_responses: None,
1358                request_state: None,
1359                meta: None,
1360                task: None,
1361            }),
1362            extensions: Extensions::new(),
1363        }
1364    }
1365
1366    async fn drive(svc: &mut InfallibleService) -> RouterResponse {
1367        svc.ready()
1368            .await
1369            .expect("infallible")
1370            .call(req())
1371            .await
1372            .expect("infallible")
1373    }
1374
1375    fn test_breaker(
1376        minimum_calls: usize,
1377        wait_in_open: Duration,
1378        permitted_in_half_open: usize,
1379    ) -> tower_resilience::circuitbreaker::CircuitBreakerLayer {
1380        let (layer, _handle) = tower_resilience::circuitbreaker::CircuitBreakerLayer::builder()
1381            .failure_rate_threshold(0.5)
1382            .minimum_number_of_calls(minimum_calls)
1383            // The count-based window must be full before the breaker
1384            // evaluates; align it with minimum_calls (see #220 for the
1385            // production-side mapping fix).
1386            .sliding_window_size(minimum_calls)
1387            .wait_duration_in_open(wait_in_open)
1388            .permitted_calls_in_half_open(permitted_in_half_open)
1389            .name("test-cb")
1390            .build_with_handle();
1391        layer
1392    }
1393
1394    /// Timeout and circuit breaker both apply; timeouts count as failures,
1395    /// the breaker opens, rejects without reaching the backend, and closes
1396    /// again through half-open probes once the fault clears.
1397    #[tokio::test]
1398    async fn timeout_and_breaker_compose_and_breaker_trips() {
1399        let (backend, hits, slow, _fail) = FlakyBackend::new();
1400        slow.store(true, Ordering::SeqCst);
1401
1402        let mut mw = BackendMiddlewareLayer::default();
1403        mw.stage(TimeoutLayer::new(Duration::from_millis(10)));
1404        mw.stage(test_breaker(4, Duration::from_millis(200), 2));
1405
1406        let mut svc = mw.layer(backend);
1407
1408        for i in 1..=4 {
1409            let resp = drive(&mut svc).await;
1410            assert!(resp.inner.is_err(), "call {i} should time out");
1411        }
1412        assert_eq!(
1413            hits.load(Ordering::SeqCst),
1414            4,
1415            "all four calls reach the backend"
1416        );
1417
1418        let resp = drive(&mut svc).await;
1419        assert!(resp.inner.is_err(), "open breaker rejects");
1420        assert_eq!(
1421            hits.load(Ordering::SeqCst),
1422            4,
1423            "rejected call must not reach the backend"
1424        );
1425
1426        tokio::time::sleep(Duration::from_millis(300)).await;
1427        slow.store(false, Ordering::SeqCst);
1428        for i in 1..=3 {
1429            let resp = drive(&mut svc).await;
1430            assert!(resp.inner.is_ok(), "recovered call {i} succeeds");
1431        }
1432        assert_eq!(hits.load(Ordering::SeqCst), 7);
1433    }
1434
1435    /// Retry (Infallible zone) and timeout (BoxError zone) coexist: the
1436    /// retry policy absorbs transient response-plane failures while the
1437    /// timeout stays armed above it.
1438    #[tokio::test]
1439    async fn retry_and_timeout_compose() {
1440        let (backend, hits, _slow, fail) = FlakyBackend::new();
1441        fail.store(2, Ordering::SeqCst);
1442
1443        let cfg = crate::config::RetryConfig {
1444            max_retries: 3,
1445            initial_backoff_ms: 1,
1446            max_backoff_ms: 5,
1447            budget_percent: None,
1448            min_retries_per_sec: 10,
1449        };
1450        let mut mw = BackendMiddlewareLayer {
1451            retry: Some(crate::retry::build_retry_layer(&cfg, "test")),
1452            ..Default::default()
1453        };
1454        mw.stage(TimeoutLayer::new(Duration::from_millis(500)));
1455
1456        let mut svc = mw.layer(backend);
1457
1458        let resp = drive(&mut svc).await;
1459        assert!(resp.inner.is_ok(), "retries absorb the transient failures");
1460        assert_eq!(
1461            hits.load(Ordering::SeqCst),
1462            3,
1463            "two failures plus the successful attempt"
1464        );
1465    }
1466
1467    /// The empty layer is the identity.
1468    #[tokio::test]
1469    async fn empty_middleware_layer_is_identity() {
1470        let (backend, hits, _slow, _fail) = FlakyBackend::new();
1471        let mw = BackendMiddlewareLayer::default();
1472        assert!(mw.is_empty());
1473
1474        let mut svc = mw.layer(backend);
1475        let resp = drive(&mut svc).await;
1476        assert!(resp.inner.is_ok());
1477        assert_eq!(hits.load(Ordering::SeqCst), 1);
1478    }
1479
1480    /// Regression for #220: a breaker built through the production config
1481    /// mapping opens after exactly `minimum_calls` failing calls, not after
1482    /// the tower-resilience default 100-call window fills.
1483    #[tokio::test]
1484    async fn config_mapped_breaker_opens_at_minimum_calls() {
1485        let (backend, hits, slow, _fail) = FlakyBackend::new();
1486        slow.store(true, Ordering::SeqCst);
1487
1488        let cfg = crate::config::CircuitBreakerConfig {
1489            failure_rate_threshold: 0.5,
1490            minimum_calls: 4,
1491            wait_duration_seconds: 60,
1492            permitted_calls_in_half_open: 1,
1493        };
1494        let (breaker, _handle) = build_breaker_layer(&cfg, "test");
1495
1496        let mut mw = BackendMiddlewareLayer::default();
1497        mw.stage(TimeoutLayer::new(Duration::from_millis(10)));
1498        mw.stage(breaker);
1499
1500        let mut svc = mw.layer(backend);
1501
1502        for i in 1..=4 {
1503            let resp = drive(&mut svc).await;
1504            assert!(resp.inner.is_err(), "call {i} should time out");
1505        }
1506        assert_eq!(hits.load(Ordering::SeqCst), 4);
1507
1508        let resp = drive(&mut svc).await;
1509        assert!(resp.inner.is_err(), "open breaker rejects");
1510        assert_eq!(
1511            hits.load(Ordering::SeqCst),
1512            4,
1513            "rejected call must not reach the backend"
1514        );
1515    }
1516}