turbomcp-server 3.6.0

//! Server Configuration
//!
//! This module provides configuration options for MCP servers including:
//! - Protocol version negotiation
//! - Rate limiting
//! - Connection limits
//! - Capability requirements

use std::collections::HashSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};

use parking_lot::Mutex;
use serde::{Deserialize, Serialize};

// Re-export from core (single source of truth - DRY)
pub use turbomcp_core::SUPPORTED_VERSIONS as SUPPORTED_PROTOCOL_VERSIONS;
pub use turbomcp_types::ProtocolVersion;

/// Default maximum connections for TCP transport.
pub const DEFAULT_MAX_CONNECTIONS: usize = 1000;

/// Default rate limit (requests per second).
pub const DEFAULT_RATE_LIMIT: u32 = 100;

/// Default rate limit window.
pub const DEFAULT_RATE_LIMIT_WINDOW: Duration = Duration::from_secs(1);

/// Default maximum message size (10MB).
pub const DEFAULT_MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024;

/// Default idle lifetime of a Streamable HTTP session (1 hour).
pub const DEFAULT_HTTP_SESSION_IDLE_TIMEOUT: Duration = Duration::from_secs(60 * 60);

/// Default cap on concurrent Streamable HTTP sessions.
pub const DEFAULT_MAX_HTTP_SESSIONS: usize = 10_000;

/// Origin validation configuration for HTTP transports.
#[derive(Debug, Clone)]
pub struct OriginValidationConfig {
    /// Explicitly allowed origins.
    pub allowed_origins: HashSet<String>,
    /// Whether to allow localhost/browser-dev origins.
    pub allow_localhost: bool,
    /// Whether to disable origin checks entirely.
    pub allow_any: bool,
    /// Trusted reverse-proxy IPs / CIDRs whose `X-Forwarded-For`,
    /// `X-Real-IP`, `CF-Connecting-IP`, and `X-Client-IP` headers we will
    /// honour. **Empty means trust nothing** — direct clients can no longer
    /// spoof their source IP via headers. Entries accept CIDR notation
    /// (`10.0.0.0/8`) or bare addresses (`10.0.0.5`).
    pub trusted_proxies: Vec<String>,
    /// Whether to accept requests with no `Origin` header from non-loopback
    /// clients.
    ///
    /// Non-browser clients — CLIs, agent runtimes, other servers — send no
    /// `Origin`, so a server reachable over the network refuses all of them by
    /// default. This admits them while still validating every `Origin` that is
    /// present, which `allow_any` does not. Pair it with authentication.
    pub allow_missing_origin: bool,
    /// Whether the HTTP transport answers CORS preflights and attaches CORS
    /// headers for the origins this policy accepts.
    ///
    /// Off by default. A browser-based client on an allowlisted origin needs it
    /// to read responses at all, and in particular to see `Mcp-Session-Id`.
    pub cors: bool,
}

impl Default for OriginValidationConfig {
    fn default() -> Self {
        Self {
            allowed_origins: HashSet::new(),
            allow_localhost: true,
            allow_any: false,
            trusted_proxies: Vec::new(),
            allow_missing_origin: false,
            cors: false,
        }
    }
}

impl OriginValidationConfig {
    /// Create a new origin validation configuration.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

/// Session lifetime policy for the Streamable HTTP transport.
///
/// A session holds memory for as long as it exists — replay history, log
/// level, pending server requests — and a client that initializes and walks
/// away never says so. §Session Management lets the server end a session at
/// any time; requests naming it then get 404 and a conforming client starts
/// a new one.
#[derive(Debug, Clone, Copy)]
pub struct HttpSessionConfig {
    /// How long a session may sit idle before it is reaped.
    ///
    /// A session is idle while it has no request in flight and no attached
    /// stream, so a client holding a GET stream open is never reaped however
    /// quiet it is.
    pub idle_timeout: Duration,
    /// Maximum concurrent sessions. An `initialize` beyond this is refused
    /// with 503 Service Unavailable, after idle sessions have been reaped.
    pub max_sessions: usize,
}

impl Default for HttpSessionConfig {
    fn default() -> Self {
        Self {
            idle_timeout: DEFAULT_HTTP_SESSION_IDLE_TIMEOUT,
            max_sessions: DEFAULT_MAX_HTTP_SESSIONS,
        }
    }
}

/// Server configuration.
#[derive(Debug, Clone)]
pub struct ServerConfig {
    /// Protocol version configuration.
    pub protocol: ProtocolConfig,
    /// Rate limiting configuration.
    pub rate_limit: Option<RateLimitConfig>,
    /// Connection limits.
    pub connection_limits: ConnectionLimits,
    /// Required client capabilities.
    pub required_capabilities: RequiredCapabilities,
    /// Maximum message size in bytes (default: 10MB).
    pub max_message_size: usize,
    /// HTTP origin validation policy.
    pub origin_validation: OriginValidationConfig,
    /// Streamable HTTP session lifetime policy.
    pub http_sessions: HttpSessionConfig,
    /// MCP authorization for the Streamable HTTP transport. `None` serves
    /// every request unauthenticated.
    #[cfg(feature = "http")]
    pub authorization: Option<crate::HttpAuthorization>,
}

impl Default for ServerConfig {
    fn default() -> Self {
        Self {
            protocol: ProtocolConfig::default(),
            rate_limit: None,
            connection_limits: ConnectionLimits::default(),
            required_capabilities: RequiredCapabilities::default(),
            max_message_size: DEFAULT_MAX_MESSAGE_SIZE,
            origin_validation: OriginValidationConfig::default(),
            http_sessions: HttpSessionConfig::default(),
            #[cfg(feature = "http")]
            authorization: None,
        }
    }
}

impl ServerConfig {
    /// Create a new server configuration with defaults.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a builder for server configuration.
    #[must_use]
    pub fn builder() -> ServerConfigBuilder {
        ServerConfigBuilder::default()
    }
}

/// Builder for server configuration.
#[derive(Debug, Clone, Default)]
pub struct ServerConfigBuilder {
    protocol: Option<ProtocolConfig>,
    rate_limit: Option<RateLimitConfig>,
    connection_limits: Option<ConnectionLimits>,
    required_capabilities: Option<RequiredCapabilities>,
    max_message_size: Option<usize>,
    origin_validation: Option<OriginValidationConfig>,
    http_sessions: Option<HttpSessionConfig>,
    #[cfg(feature = "http")]
    authorization: Option<crate::HttpAuthorization>,
}

impl ServerConfigBuilder {
    /// Set protocol configuration.
    #[must_use]
    pub fn protocol(mut self, config: ProtocolConfig) -> Self {
        self.protocol = Some(config);
        self
    }

    /// Set rate limiting configuration.
    #[must_use]
    pub fn rate_limit(mut self, config: RateLimitConfig) -> Self {
        self.rate_limit = Some(config);
        self
    }

    /// Set connection limits.
    #[must_use]
    pub fn connection_limits(mut self, limits: ConnectionLimits) -> Self {
        self.connection_limits = Some(limits);
        self
    }

    /// Set required client capabilities.
    #[must_use]
    pub fn required_capabilities(mut self, caps: RequiredCapabilities) -> Self {
        self.required_capabilities = Some(caps);
        self
    }

    /// Set maximum message size in bytes.
    ///
    /// Messages exceeding this size will be rejected.
    /// Default: 10MB.
    #[must_use]
    pub fn max_message_size(mut self, size: usize) -> Self {
        self.max_message_size = Some(size);
        self
    }

    /// Set HTTP origin validation configuration.
    #[must_use]
    pub fn origin_validation(mut self, config: OriginValidationConfig) -> Self {
        self.origin_validation = Some(config);
        self
    }

    /// Add a single allowed origin for HTTP transports.
    #[must_use]
    pub fn allow_origin(mut self, origin: impl Into<String>) -> Self {
        self.origin_validation
            .get_or_insert_with(OriginValidationConfig::default)
            .allowed_origins
            .insert(origin.into());
        self
    }

    /// Add multiple allowed origins for HTTP transports.
    #[must_use]
    pub fn allow_origins<I, S>(mut self, origins: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        let config = self
            .origin_validation
            .get_or_insert_with(OriginValidationConfig::default);
        config
            .allowed_origins
            .extend(origins.into_iter().map(Into::into));
        self
    }

    /// Control whether localhost origins are accepted.
    #[must_use]
    pub fn allow_localhost_origins(mut self, allow: bool) -> Self {
        self.origin_validation
            .get_or_insert_with(OriginValidationConfig::default)
            .allow_localhost = allow;
        self
    }

    /// Disable origin checks entirely.
    #[must_use]
    pub fn allow_any_origin(mut self, allow: bool) -> Self {
        self.origin_validation
            .get_or_insert_with(OriginValidationConfig::default)
            .allow_any = allow;
        self
    }

    /// Accept requests with no `Origin` header from non-loopback clients,
    /// while still validating every `Origin` that is present.
    ///
    /// See [`OriginValidationConfig::allow_missing_origin`].
    #[must_use]
    pub fn allow_missing_origin(mut self, allow: bool) -> Self {
        self.origin_validation
            .get_or_insert_with(OriginValidationConfig::default)
            .allow_missing_origin = allow;
        self
    }

    /// Answer CORS preflights and attach CORS headers for accepted origins.
    ///
    /// See [`OriginValidationConfig::cors`].
    #[must_use]
    pub fn cors(mut self, enabled: bool) -> Self {
        self.origin_validation
            .get_or_insert_with(OriginValidationConfig::default)
            .cors = enabled;
        self
    }

    /// Require MCP authorization on the Streamable HTTP transport.
    ///
    /// See [`HttpAuthorization`](crate::HttpAuthorization): the server
    /// publishes its protected-resource metadata, and every MCP request must
    /// carry a bearer token the configured validator accepts.
    #[cfg(feature = "http")]
    #[must_use]
    pub fn authorization(mut self, authorization: crate::HttpAuthorization) -> Self {
        self.authorization = Some(authorization);
        self
    }

    /// Set the Streamable HTTP session policy.
    #[must_use]
    pub fn http_sessions(mut self, config: HttpSessionConfig) -> Self {
        self.http_sessions = Some(config);
        self
    }

    /// Set how long a Streamable HTTP session may sit idle before it is reaped.
    ///
    /// Default: 1 hour.
    #[must_use]
    pub fn http_session_idle_timeout(mut self, timeout: Duration) -> Self {
        self.http_sessions
            .get_or_insert_with(HttpSessionConfig::default)
            .idle_timeout = timeout;
        self
    }

    /// Set the maximum number of concurrent Streamable HTTP sessions.
    ///
    /// Default: 10,000.
    #[must_use]
    pub fn max_http_sessions(mut self, max: usize) -> Self {
        self.http_sessions
            .get_or_insert_with(HttpSessionConfig::default)
            .max_sessions = max;
        self
    }

    /// Build the server configuration with sensible defaults.
    ///
    /// This method always succeeds and uses defaults for any unset fields.
    /// For strict validation, use [`try_build()`](Self::try_build).
    #[must_use]
    pub fn build(self) -> ServerConfig {
        ServerConfig {
            protocol: self.protocol.unwrap_or_default(),
            rate_limit: self.rate_limit,
            connection_limits: self.connection_limits.unwrap_or_default(),
            required_capabilities: self.required_capabilities.unwrap_or_default(),
            max_message_size: self.max_message_size.unwrap_or(DEFAULT_MAX_MESSAGE_SIZE),
            origin_validation: self.origin_validation.unwrap_or_default(),
            http_sessions: self.http_sessions.unwrap_or_default(),
            #[cfg(feature = "http")]
            authorization: self.authorization,
        }
    }

    /// Build the server configuration with validation.
    ///
    /// This method validates the configuration and returns an error if any
    /// constraints are violated. Use this for stricter configuration checking
    /// in enterprise deployments.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - `max_message_size` is less than 1024 bytes (minimum viable message size)
    /// - Rate limit `max_requests` is 0
    /// - Rate limit `window` is zero
    /// - Connection limits have all values set to 0
    /// - `max_http_sessions` is 0
    ///
    /// # Example
    ///
    /// ```rust
    /// use turbomcp_server::ServerConfig;
    ///
    /// // Validated build - catches configuration errors
    /// let config = ServerConfig::builder()
    ///     .max_message_size(1024 * 1024) // 1MB
    ///     .try_build()
    ///     .expect("Invalid configuration");
    /// ```
    pub fn try_build(self) -> Result<ServerConfig, ConfigValidationError> {
        let max_message_size = self.max_message_size.unwrap_or(DEFAULT_MAX_MESSAGE_SIZE);

        // Validate message size
        if max_message_size < 1024 {
            return Err(ConfigValidationError::InvalidMessageSize {
                size: max_message_size,
                min: 1024,
            });
        }

        // Validate rate limit if provided
        if let Some(ref rate_limit) = self.rate_limit {
            if rate_limit.max_requests == 0 {
                return Err(ConfigValidationError::InvalidRateLimit {
                    reason: "max_requests cannot be 0".to_string(),
                });
            }
            if rate_limit.window.is_zero() {
                return Err(ConfigValidationError::InvalidRateLimit {
                    reason: "rate limit window cannot be zero".to_string(),
                });
            }
        }

        // Validate connection limits
        let connection_limits = self.connection_limits.unwrap_or_default();
        if connection_limits.max_tcp_connections == 0
            && connection_limits.max_websocket_connections == 0
            && connection_limits.max_http_concurrent == 0
            && connection_limits.max_unix_connections == 0
        {
            return Err(ConfigValidationError::InvalidConnectionLimits {
                reason: "at least one connection limit must be non-zero".to_string(),
            });
        }

        let http_sessions = self.http_sessions.unwrap_or_default();
        if http_sessions.max_sessions == 0 {
            return Err(ConfigValidationError::InvalidConnectionLimits {
                reason: "max_http_sessions cannot be 0".to_string(),
            });
        }

        Ok(ServerConfig {
            protocol: self.protocol.unwrap_or_default(),
            rate_limit: self.rate_limit,
            connection_limits,
            required_capabilities: self.required_capabilities.unwrap_or_default(),
            max_message_size,
            origin_validation: self.origin_validation.unwrap_or_default(),
            http_sessions,
            #[cfg(feature = "http")]
            authorization: self.authorization,
        })
    }
}

/// Every field carries over. The destructuring is exhaustive on purpose: a
/// field added to `ServerConfig` and not handled here is a compile error.
/// `ServerBuilder::with_config` used to copy fields one by one and silently
/// dropped `http_sessions` and `authorization` — the latter leaving a server
/// configured for authorization serving HTTP unauthenticated.
impl From<ServerConfig> for ServerConfigBuilder {
    fn from(config: ServerConfig) -> Self {
        let ServerConfig {
            protocol,
            rate_limit,
            connection_limits,
            required_capabilities,
            max_message_size,
            origin_validation,
            http_sessions,
            #[cfg(feature = "http")]
            authorization,
        } = config;
        Self {
            protocol: Some(protocol),
            rate_limit,
            connection_limits: Some(connection_limits),
            required_capabilities: Some(required_capabilities),
            max_message_size: Some(max_message_size),
            origin_validation: Some(origin_validation),
            http_sessions: Some(http_sessions),
            #[cfg(feature = "http")]
            authorization,
        }
    }
}

/// Errors that can occur during configuration validation.
#[derive(Debug, Clone, thiserror::Error)]
pub enum ConfigValidationError {
    /// Invalid message size configuration.
    #[error("Invalid max_message_size: {size} bytes is below minimum of {min} bytes")]
    InvalidMessageSize {
        /// The configured size.
        size: usize,
        /// The minimum allowed size.
        min: usize,
    },

    /// Invalid rate limit configuration.
    #[error("Invalid rate limit: {reason}")]
    InvalidRateLimit {
        /// Description of the validation failure.
        reason: String,
    },

    /// Invalid connection limits configuration.
    #[error("Invalid connection limits: {reason}")]
    InvalidConnectionLimits {
        /// Description of the validation failure.
        reason: String,
    },
}

/// Protocol version configuration.
#[derive(Debug, Clone)]
pub struct ProtocolConfig {
    /// Preferred protocol version.
    pub preferred_version: ProtocolVersion,
    /// Supported protocol versions.
    pub supported_versions: Vec<ProtocolVersion>,
    /// Answer an unsupported requested version with the preferred one.
    ///
    /// The lifecycle spec requires this, so it should stay `true`. With
    /// `false` an unsupported request is refused with -32602 instead.
    pub allow_fallback: bool,
}

impl Default for ProtocolConfig {
    /// Multi-version by default in v3.1.
    ///
    /// Pre-3.1 the default was exact-match against `LATEST`, which silently
    /// rejected clients on older stable spec versions even though we ship
    /// version adapters for them. The spec permits the server to choose a
    /// different version than the client requested; defaulting to the full
    /// stable set unblocks older clients while still preferring the newest.
    /// Use [`Self::strict`] to restore the old single-version behavior.
    ///
    /// `allow_fallback` is on because the lifecycle spec makes it mandatory:
    /// "If the server supports the requested protocol version, it MUST respond
    /// with the same version. Otherwise, the server MUST respond with another
    /// protocol version it supports." Refusing the handshake outright is not
    /// one of the permitted outcomes — the client is the party that decides
    /// whether the offered version is acceptable.
    fn default() -> Self {
        Self {
            preferred_version: ProtocolVersion::LATEST.clone(),
            supported_versions: ProtocolVersion::STABLE.to_vec(),
            allow_fallback: true,
        }
    }
}

impl ProtocolConfig {
    /// Create a configuration that speaks only the specified version.
    ///
    /// A client asking for any other version is answered with this one — the
    /// lifecycle spec's required response to an unsupported request — and
    /// decides for itself whether to continue. It used to be refused outright,
    /// which the spec does not permit.
    #[must_use]
    pub fn strict(version: impl Into<ProtocolVersion>) -> Self {
        let v = version.into();
        Self {
            preferred_version: v.clone(),
            supported_versions: vec![v],
            allow_fallback: true,
        }
    }

    /// Create a multi-version configuration that accepts all stable versions.
    ///
    /// The preferred version is the latest stable. Older clients are accepted
    /// and responses are filtered through the appropriate version adapter.
    /// A version outside the stable set is answered with the preferred one
    /// rather than an error, as the lifecycle spec requires.
    #[must_use]
    pub fn multi_version() -> Self {
        Self {
            preferred_version: ProtocolVersion::LATEST.clone(),
            supported_versions: ProtocolVersion::STABLE.to_vec(),
            allow_fallback: true,
        }
    }

    /// Check if a protocol version is supported.
    #[must_use]
    pub fn is_supported(&self, version: &ProtocolVersion) -> bool {
        self.supported_versions.contains(version)
    }

    /// Negotiate protocol version with client.
    ///
    /// Returns the negotiated version or None if no compatible version found.
    #[must_use]
    pub fn negotiate(&self, client_version: Option<&str>) -> Option<ProtocolVersion> {
        match client_version {
            Some(version_str) => {
                let version = ProtocolVersion::from(version_str);
                if self.is_supported(&version) {
                    Some(version)
                } else if self.allow_fallback {
                    Some(self.preferred_version.clone())
                } else {
                    None
                }
            }
            None => Some(self.preferred_version.clone()),
        }
    }
}

/// Rate limiting configuration.
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    /// Maximum requests per window.
    pub max_requests: u32,
    /// Time window for rate limiting.
    pub window: Duration,
    /// Whether to rate limit per client (by user_id or IP).
    pub per_client: bool,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            max_requests: DEFAULT_RATE_LIMIT,
            window: DEFAULT_RATE_LIMIT_WINDOW,
            per_client: true,
        }
    }
}

impl RateLimitConfig {
    /// Create a new rate limit configuration.
    #[must_use]
    pub fn new(max_requests: u32, window: Duration) -> Self {
        Self {
            max_requests,
            window,
            per_client: true,
        }
    }

    /// Set per-client rate limiting.
    #[must_use]
    pub fn per_client(mut self, enabled: bool) -> Self {
        self.per_client = enabled;
        self
    }
}

/// Connection limits.
#[derive(Debug, Clone)]
pub struct ConnectionLimits {
    /// Maximum concurrent TCP connections.
    pub max_tcp_connections: usize,
    /// Maximum concurrent WebSocket connections.
    pub max_websocket_connections: usize,
    /// Maximum concurrent HTTP requests.
    pub max_http_concurrent: usize,
    /// Maximum concurrent Unix socket connections.
    pub max_unix_connections: usize,
}

impl Default for ConnectionLimits {
    fn default() -> Self {
        Self {
            max_tcp_connections: DEFAULT_MAX_CONNECTIONS,
            max_websocket_connections: DEFAULT_MAX_CONNECTIONS,
            max_http_concurrent: DEFAULT_MAX_CONNECTIONS,
            max_unix_connections: DEFAULT_MAX_CONNECTIONS,
        }
    }
}

impl ConnectionLimits {
    /// Create a new connection limits configuration.
    #[must_use]
    pub fn new(max_connections: usize) -> Self {
        Self {
            max_tcp_connections: max_connections,
            max_websocket_connections: max_connections,
            max_http_concurrent: max_connections,
            max_unix_connections: max_connections,
        }
    }
}

/// Required client capabilities.
///
/// Specifies which client capabilities the server requires.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RequiredCapabilities {
    /// Require roots capability.
    #[serde(default)]
    pub roots: bool,
    /// Require sampling capability.
    #[serde(default)]
    pub sampling: bool,
    /// Require draft extensions.
    #[serde(default)]
    pub extensions: HashSet<String>,
    /// Require experimental capabilities.
    #[serde(default)]
    pub experimental: HashSet<String>,
}

impl RequiredCapabilities {
    /// Create empty required capabilities (no requirements).
    #[must_use]
    pub fn none() -> Self {
        Self::default()
    }

    /// Require roots capability.
    #[must_use]
    pub fn with_roots(mut self) -> Self {
        self.roots = true;
        self
    }

    /// Require sampling capability.
    #[must_use]
    pub fn with_sampling(mut self) -> Self {
        self.sampling = true;
        self
    }

    /// Require a draft extension.
    #[must_use]
    pub fn with_extension(mut self, name: impl Into<String>) -> Self {
        self.extensions.insert(name.into());
        self
    }

    /// Require an experimental capability.
    #[must_use]
    pub fn with_experimental(mut self, name: impl Into<String>) -> Self {
        self.experimental.insert(name.into());
        self
    }

    /// Check if all required capabilities are present in client capabilities.
    #[must_use]
    pub fn validate(&self, client_caps: &ClientCapabilities) -> CapabilityValidation {
        let mut missing = Vec::new();

        if self.roots && !client_caps.roots {
            missing.push("roots".to_string());
        }

        if self.sampling && !client_caps.sampling {
            missing.push("sampling".to_string());
        }

        for extension in &self.extensions {
            if !client_caps.extensions.contains(extension) {
                missing.push(format!("extensions/{}", extension));
            }
        }

        for exp in &self.experimental {
            if !client_caps.experimental.contains(exp) {
                missing.push(format!("experimental/{}", exp));
            }
        }

        if missing.is_empty() {
            CapabilityValidation::Valid
        } else {
            CapabilityValidation::Missing(missing)
        }
    }
}

/// Client capabilities received during initialization.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ClientCapabilities {
    /// Client supports roots.
    #[serde(default)]
    pub roots: bool,
    /// Client supports sampling.
    #[serde(default)]
    pub sampling: bool,
    /// Client draft extensions.
    #[serde(default)]
    pub extensions: HashSet<String>,
    /// Client experimental capabilities.
    #[serde(default)]
    pub experimental: HashSet<String>,
}

impl ClientCapabilities {
    /// Parse client capabilities from initialize request params.
    #[must_use]
    pub fn from_params(params: &serde_json::Value) -> Self {
        let caps = params.get("capabilities").cloned().unwrap_or_default();

        Self {
            roots: caps.get("roots").map(|v| !v.is_null()).unwrap_or(false),
            sampling: caps.get("sampling").map(|v| !v.is_null()).unwrap_or(false),
            extensions: caps
                .get("extensions")
                .and_then(|v| v.as_object())
                .map(|obj| obj.keys().cloned().collect())
                .unwrap_or_default(),
            experimental: caps
                .get("experimental")
                .and_then(|v| v.as_object())
                .map(|obj| obj.keys().cloned().collect())
                .unwrap_or_default(),
        }
    }
}

/// Result of capability validation.
#[derive(Debug, Clone)]
pub enum CapabilityValidation {
    /// All required capabilities are present.
    Valid,
    /// Some required capabilities are missing.
    Missing(Vec<String>),
}

impl CapabilityValidation {
    /// Check if validation passed.
    #[must_use]
    pub fn is_valid(&self) -> bool {
        matches!(self, Self::Valid)
    }

    /// Get missing capabilities if any.
    #[must_use]
    pub fn missing(&self) -> Option<&[String]> {
        match self {
            Self::Valid => None,
            Self::Missing(caps) => Some(caps),
        }
    }
}

/// Rate limiter using token bucket algorithm.
#[derive(Debug)]
pub struct RateLimiter {
    config: RateLimitConfig,
    /// Global bucket for non-per-client limiting.
    global_bucket: Mutex<TokenBucket>,
    /// Per-client buckets (keyed by client ID).
    client_buckets: Mutex<std::collections::HashMap<String, TokenBucket>>,
    /// Last cleanup timestamp for automatic cleanup.
    last_cleanup: Mutex<Instant>,
}

impl RateLimiter {
    /// Create a new rate limiter.
    #[must_use]
    pub fn new(config: RateLimitConfig) -> Self {
        Self {
            global_bucket: Mutex::new(TokenBucket::new(config.max_requests, config.window)),
            client_buckets: Mutex::new(std::collections::HashMap::new()),
            last_cleanup: Mutex::new(Instant::now()),
            config,
        }
    }

    /// Check if a request is allowed.
    ///
    /// Returns `true` if allowed, `false` if rate limited.
    pub fn check(&self, client_id: Option<&str>) -> bool {
        // Periodic cleanup of stale client buckets (avoid unbounded growth)
        let needs_cleanup = {
            let last = self.last_cleanup.lock();
            last.elapsed() > Duration::from_secs(60)
        };
        if needs_cleanup {
            self.cleanup(Duration::from_secs(300));
            *self.last_cleanup.lock() = Instant::now();
        }

        if self.config.per_client {
            if let Some(id) = client_id {
                let mut buckets = self.client_buckets.lock();
                let bucket = buckets.entry(id.to_string()).or_insert_with(|| {
                    TokenBucket::new(self.config.max_requests, self.config.window)
                });
                bucket.try_acquire()
            } else {
                // No client ID, use global bucket
                self.global_bucket.lock().try_acquire()
            }
        } else {
            self.global_bucket.lock().try_acquire()
        }
    }

    /// Clean up old client buckets to prevent memory growth.
    pub fn cleanup(&self, max_age: Duration) {
        let mut buckets = self.client_buckets.lock();
        let now = Instant::now();
        buckets.retain(|_, bucket| now.duration_since(bucket.last_access) < max_age);
    }

    /// Get the current number of tracked client buckets.
    #[must_use]
    pub fn client_bucket_count(&self) -> usize {
        self.client_buckets.lock().len()
    }
}

/// Token bucket for rate limiting.
#[derive(Debug)]
struct TokenBucket {
    tokens: f64,
    max_tokens: f64,
    refill_rate: f64, // tokens per second
    last_refill: Instant,
    last_access: Instant,
}

impl TokenBucket {
    fn new(max_requests: u32, window: Duration) -> Self {
        let max_tokens = max_requests as f64;
        let refill_rate = max_tokens / window.as_secs_f64();
        Self {
            tokens: max_tokens,
            max_tokens,
            refill_rate,
            last_refill: Instant::now(),
            last_access: Instant::now(),
        }
    }

    fn try_acquire(&mut self) -> bool {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill);

        // Only refill if meaningful time has passed (reduces syscalls on burst traffic)
        if elapsed >= Duration::from_millis(10) {
            self.tokens =
                (self.tokens + elapsed.as_secs_f64() * self.refill_rate).min(self.max_tokens);
            self.last_refill = now;
        }

        self.last_access = now;

        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            true
        } else {
            false
        }
    }
}

/// Connection counter for tracking active connections.
///
/// This is designed to be wrapped in `Arc` and shared across async tasks.
/// Use `try_acquire_arc` to get a guard that can be moved into spawned tasks.
#[derive(Debug)]
pub struct ConnectionCounter {
    current: AtomicUsize,
    max: usize,
}

impl ConnectionCounter {
    /// Create a new connection counter.
    #[must_use]
    pub fn new(max: usize) -> Self {
        Self {
            current: AtomicUsize::new(0),
            max,
        }
    }

    /// Try to acquire a connection slot (for use when counter is in Arc).
    ///
    /// Returns a guard that releases the slot when dropped, or None if at capacity.
    /// The guard is `Send + 'static` and can be moved into spawned async tasks.
    ///
    /// The CAS loop is unbounded — under genuine contention the standard pattern
    /// is to keep retrying until either the slot is acquired or capacity is hit;
    /// progress is guaranteed because each iteration either advances `current`
    /// (acquire success) or witnesses another thread's advance (their CAS won).
    pub fn try_acquire_arc(self: &Arc<Self>) -> Option<ConnectionGuard> {
        loop {
            let current = self.current.load(Ordering::Relaxed);
            if current >= self.max {
                return None;
            }
            if self
                .current
                .compare_exchange(current, current + 1, Ordering::SeqCst, Ordering::Relaxed)
                .is_ok()
            {
                return Some(ConnectionGuard {
                    counter: Arc::clone(self),
                });
            }
            // Hint to the CPU that we're spinning (avoids pipeline stalls).
            std::hint::spin_loop();
        }
    }

    /// Get current connection count.
    #[must_use]
    pub fn current(&self) -> usize {
        self.current.load(Ordering::Relaxed)
    }

    /// Get maximum connections.
    #[must_use]
    pub fn max(&self) -> usize {
        self.max
    }

    fn release(&self) {
        self.current.fetch_sub(1, Ordering::SeqCst);
    }
}

/// Guard that releases a connection slot when dropped.
///
/// This guard is `Send + 'static` and can be safely moved into spawned async tasks.
#[derive(Debug)]
pub struct ConnectionGuard {
    counter: Arc<ConnectionCounter>,
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.counter.release();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_protocol_negotiation_exact_match() {
        let config = ProtocolConfig::default();
        assert_eq!(
            config.negotiate(Some("2025-11-25")),
            Some(ProtocolVersion::V2025_11_25)
        );
    }

    #[test]
    fn test_protocol_negotiation_default_accepts_stable_versions() {
        // v3.1: default is multi-version (all stable versions). Older spec
        // versions are accepted; older clients route through version adapters.
        // Use ProtocolConfig::strict(LATEST) to restore exact-match behavior.
        let config = ProtocolConfig::default();
        assert_eq!(
            config.negotiate(Some("2025-06-18")),
            Some(ProtocolVersion::V2025_06_18)
        );
    }

    /// `strict` speaks one version, and offers it to a client asking for
    /// another: refusing the handshake is not an answer the lifecycle spec
    /// permits.
    #[test]
    fn test_protocol_negotiation_strict_offers_its_version() {
        let config = ProtocolConfig::strict(ProtocolVersion::LATEST.clone());
        assert_eq!(
            config.negotiate(Some("2025-06-18")),
            Some(ProtocolVersion::LATEST)
        );
    }

    #[test]
    fn test_protocol_negotiation_multi_version_accepts_older() {
        let config = ProtocolConfig::multi_version();
        assert_eq!(
            config.negotiate(Some("2025-06-18")),
            Some(ProtocolVersion::V2025_06_18)
        );
        assert_eq!(
            config.negotiate(Some("2025-11-25")),
            Some(ProtocolVersion::V2025_11_25)
        );
    }

    #[test]
    fn test_protocol_negotiation_none_returns_preferred() {
        let config = ProtocolConfig::default();
        assert_eq!(config.negotiate(None), Some(ProtocolVersion::V2025_11_25));
    }

    /// The lifecycle spec gives a server two choices for a version it does not
    /// support, and refusing the handshake is not one of them: it MUST answer
    /// with a version it does support and let the client decide.
    #[test]
    fn test_protocol_negotiation_unknown_version_falls_back() {
        let config = ProtocolConfig::default();
        assert_eq!(
            config.negotiate(Some("unknown-version")),
            Some(ProtocolVersion::V2025_11_25)
        );
    }

    /// The real-world case: clients still pinned to retired spec revisions.
    /// Before 3.5.0 these could not connect to a TurboMCP server at all.
    #[test]
    fn test_protocol_negotiation_retired_versions_fall_back() {
        let config = ProtocolConfig::default();
        for retired in ["2024-11-05", "2025-03-26"] {
            assert_eq!(
                config.negotiate(Some(retired)),
                Some(ProtocolVersion::V2025_11_25),
                "{retired} should be offered the latest supported version"
            );
        }
    }

    /// `strict` narrows what is spoken, not whether the handshake is
    /// answered; only `allow_fallback: false` refuses.
    #[test]
    fn test_protocol_negotiation_strict() {
        let config = ProtocolConfig::strict("2025-11-25");
        assert!(!config.is_supported(&ProtocolVersion::V2025_06_18));

        let refusing = ProtocolConfig {
            allow_fallback: false,
            ..config
        };
        assert_eq!(refusing.negotiate(Some("2025-06-18")), None);
    }

    #[test]
    fn test_capability_validation() {
        let required = RequiredCapabilities::none().with_roots();
        let client = ClientCapabilities {
            roots: true,
            ..Default::default()
        };
        assert!(required.validate(&client).is_valid());

        let client_missing = ClientCapabilities::default();
        assert!(!required.validate(&client_missing).is_valid());
    }

    #[test]
    fn test_extension_capability_validation() {
        let required = RequiredCapabilities::none().with_extension("trace");
        let client = ClientCapabilities {
            extensions: ["trace".to_string()].into_iter().collect(),
            ..Default::default()
        };
        assert!(required.validate(&client).is_valid());

        let missing = ClientCapabilities::default();
        let validation = required.validate(&missing);
        assert!(!validation.is_valid());
        assert_eq!(
            validation.missing(),
            Some(&["extensions/trace".to_string()][..])
        );
    }

    #[test]
    fn test_client_capabilities_parse_extensions() {
        let params = serde_json::json!({
            "capabilities": {
                "extensions": {
                    "trace": {"version": "1"},
                    "handoff": {}
                }
            }
        });

        let caps = ClientCapabilities::from_params(&params);
        assert!(caps.extensions.contains("trace"));
        assert!(caps.extensions.contains("handoff"));
    }

    #[test]
    fn test_rate_limiter() {
        let config = RateLimitConfig::new(2, Duration::from_secs(1));
        let limiter = RateLimiter::new(config);

        assert!(limiter.check(None));
        assert!(limiter.check(None));
        assert!(!limiter.check(None)); // Should be rate limited
    }

    #[test]
    fn test_connection_counter() {
        let counter = Arc::new(ConnectionCounter::new(2));

        let guard1 = counter.try_acquire_arc();
        assert!(guard1.is_some());
        assert_eq!(counter.current(), 1);

        let guard2 = counter.try_acquire_arc();
        assert!(guard2.is_some());
        assert_eq!(counter.current(), 2);

        let guard3 = counter.try_acquire_arc();
        assert!(guard3.is_none()); // At capacity

        drop(guard1);
        assert_eq!(counter.current(), 1);

        let guard4 = counter.try_acquire_arc();
        assert!(guard4.is_some());
    }

    // =========================================================================
    // Builder validation tests
    // =========================================================================

    #[test]
    fn test_builder_default_succeeds() {
        // Default configuration should always succeed
        let config = ServerConfig::builder().build();
        assert_eq!(config.max_message_size, DEFAULT_MAX_MESSAGE_SIZE);
        assert!(config.origin_validation.allow_localhost);
        assert!(config.origin_validation.allowed_origins.is_empty());
    }

    #[test]
    fn test_builder_origin_validation_overrides() {
        let config = ServerConfig::builder()
            .allow_origin("https://app.example.com")
            .allow_localhost_origins(false)
            .build();

        assert!(!config.origin_validation.allow_localhost);
        assert!(
            config
                .origin_validation
                .allowed_origins
                .contains("https://app.example.com")
        );
    }

    #[test]
    fn test_builder_try_build_valid() {
        let result = ServerConfig::builder()
            .max_message_size(1024 * 1024)
            .try_build();
        assert!(result.is_ok());
    }

    #[test]
    fn test_builder_try_build_invalid_message_size() {
        let result = ServerConfig::builder()
            .max_message_size(100) // Below minimum
            .try_build();
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            ConfigValidationError::InvalidMessageSize { .. }
        ));
    }

    #[test]
    fn test_builder_try_build_invalid_rate_limit() {
        let result = ServerConfig::builder()
            .rate_limit(RateLimitConfig {
                max_requests: 0, // Invalid
                window: Duration::from_secs(1),
                per_client: true,
            })
            .try_build();
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            ConfigValidationError::InvalidRateLimit { .. }
        ));
    }

    #[test]
    fn test_builder_try_build_zero_window() {
        let result = ServerConfig::builder()
            .rate_limit(RateLimitConfig {
                max_requests: 100,
                window: Duration::ZERO, // Invalid
                per_client: true,
            })
            .try_build();
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            ConfigValidationError::InvalidRateLimit { .. }
        ));
    }

    #[test]
    fn test_builder_try_build_invalid_connection_limits() {
        let result = ServerConfig::builder()
            .connection_limits(ConnectionLimits {
                max_tcp_connections: 0,
                max_websocket_connections: 0,
                max_http_concurrent: 0,
                max_unix_connections: 0,
            })
            .try_build();
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            ConfigValidationError::InvalidConnectionLimits { .. }
        ));
    }
}

#[cfg(test)]
mod proptest_tests {
    use super::*;
    use proptest::prelude::*;

    proptest! {
        #[test]
        fn config_builder_never_panics(
            max_msg_size in 0usize..10_000_000,
        ) {
            // Builder should never panic, just return errors for invalid inputs
            let _ = ServerConfig::builder()
                .max_message_size(max_msg_size)
                .try_build();
        }

        #[test]
        fn connection_counter_bounded(max in 1usize..10000) {
            let counter = Arc::new(ConnectionCounter::new(max));
            let mut guards = Vec::new();
            // Should never acquire more than max
            for _ in 0..max + 10 {
                if let Some(guard) = counter.try_acquire_arc() {
                    guards.push(guard);
                }
            }
            assert_eq!(guards.len(), max);
            assert_eq!(counter.current(), max);
        }
    }
}