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