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