winx_code_agent/

http_server.rs

//! Optional remote MCP transport over Streamable HTTP.
//!
//! The endpoint exposes shell and filesystem capabilities, so remote mode is
//! deliberately fail-closed: loopback bind by default, strong file-backed or
//! explicit bearer credentials, bounded request concurrency/rate, and isolated
//! daemon namespaces for independently authenticated principals.

#![allow(clippy::doc_markdown)]

use std::collections::{HashMap, HashSet, VecDeque};
use std::net::{IpAddr, SocketAddr};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use axum::{
    extract::{ConnectInfo, Request, State},
    http::{header, HeaderValue, StatusCode},
    middleware::{self, Next},
    response::{IntoResponse, Response},
    routing::get,
    Json, Router,
};
use rand::RngExt as _;
use rmcp::transport::streamable_http_server::{
    session::local::LocalSessionManager, session::SessionManager as _, StreamableHttpServerConfig,
    StreamableHttpService,
};
use serde::Serialize;
use tokio::sync::{Mutex, Semaphore};
use tower_http::timeout::TimeoutLayer;

use crate::config::{HttpPrincipal, HttpSessionAffinity};
use crate::runtime::{configured_shell_runtime, ShellRuntime};
use crate::server::{SessionIsolation, WinxService};

mod observability;
use observability::{observe_request_entry, PermitBody};

pub type BoxError = Box<dyn std::error::Error + Send + Sync>;

/// The one effective request-body limit. It drives both the declared-length
/// precheck and the SDK's streaming read limit, and at the default concurrency
/// bounds buffered request bodies to 32 x 16 MiB = 512 MiB (bodies are only
/// read after admission). Over HTTP this, not EditFiles' larger aggregate
/// content cap, bounds a tool call; the SDK default it replaces was 4 MiB.
pub const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;
/// Whole-request deadline. Client-profile wait budgets stay below it with margin.
pub(crate) const REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
const RATE_WINDOW: Duration = Duration::from_secs(60);
const MAX_RATE_TRACKED_IPS: usize = 4_096;
const INVALID_AUTH_DELAY: Duration = Duration::from_millis(100);
pub const DEFAULT_MAX_CONCURRENCY: usize = 32;
/// Fixed-window budget of failed authentication attempts per source IP.
/// Authenticated requests never consume it, so junk traffic sharing a tunnel's
/// source IP cannot rate-limit a valid principal.
pub const DEFAULT_REQUESTS_PER_MINUTE: u32 = 120;
/// Fixed-window request budget per authenticated principal. Behind one reverse
/// proxy every client shares a single source IP, so the IP window alone cannot
/// keep one busy principal from starving the others.
pub const DEFAULT_PRINCIPAL_REQUESTS_PER_MINUTE: u32 = 120;

#[derive(Debug)]
struct HealthState {
    ready: AtomicBool,
}

#[derive(Serialize)]
struct HealthResponse {
    status: &'static str,
    ready: bool,
}

/// Opaque per-request correlation shared by the HTTP middleware and MCP handler.
/// It contains no client input and is safe to include in metadata-only usage logs.
#[derive(Clone, Debug)]
pub(crate) struct RequestCorrelation(String);

impl RequestCorrelation {
    fn new() -> Self {
        Self(format!("r_{:032x}", rand::rng().random::<u128>()))
    }

    pub(crate) fn as_str(&self) -> &str {
        &self.0
    }
}

#[derive(Clone, Debug)]
pub struct HttpServerOptions {
    pub bind: String,
    pub principals: Vec<HttpPrincipal>,
    pub extra_hosts: Vec<String>,
    pub allow_query_token: bool,
    pub allow_non_loopback: bool,
    pub session_affinity: HttpSessionAffinity,
    pub max_concurrency: usize,
    /// Failed-authentication attempts allowed per source IP and minute.
    pub requests_per_minute: u32,
    pub principal_requests_per_minute: u32,
}

impl HttpServerOptions {
    pub fn single(
        bind: impl Into<String>,
        token: impl Into<String>,
        extra_hosts: Vec<String>,
        allow_query_token: bool,
    ) -> Result<Self, BoxError> {
        let token = token.into();
        let principal = HttpPrincipal::new("default", token.trim(), false)?;
        Ok(Self {
            bind: bind.into(),
            principals: vec![principal],
            extra_hosts,
            allow_query_token,
            allow_non_loopback: false,
            session_affinity: HttpSessionAffinity::Workspace,
            max_concurrency: DEFAULT_MAX_CONCURRENCY,
            requests_per_minute: DEFAULT_REQUESTS_PER_MINUTE,
            principal_requests_per_minute: DEFAULT_PRINCIPAL_REQUESTS_PER_MINUTE,
        })
    }
}

pub async fn start_http_server(
    bind: &str,
    token: String,
    extra_hosts: Vec<String>,
    allow_query_token: bool,
) -> Result<(), BoxError> {
    let options = HttpServerOptions::single(bind, token, extra_hosts, allow_query_token)?;
    start_http_server_with_options(options).await
}

pub async fn start_http_server_with_options(options: HttpServerOptions) -> Result<(), BoxError> {
    let shell_runtime = configured_shell_runtime().await?;
    start_http_server_with_runtime(options, shell_runtime).await
}

pub async fn start_http_server_with_runtime(
    options: HttpServerOptions,
    shell_runtime: Arc<dyn ShellRuntime>,
) -> Result<(), BoxError> {
    let bind = validate_options(&options)?;
    let principal_names =
        options.principals.iter().map(|principal| principal.name().to_string()).collect::<Vec<_>>();
    let allow_query_token = options.allow_query_token;
    let health_state = Arc::new(HealthState { ready: AtomicBool::new(true) });
    let app = build_app(options, shell_runtime, Arc::clone(&health_state));

    let listener = tokio::net::TcpListener::bind(bind).await?;
    tracing::warn!(
        bind = %bind,
        principals = ?principal_names,
        route = "/mcp",
        "Winx remote MCP is network-reachable; keep it behind authenticated HTTPS"
    );
    if allow_query_token {
        tracing::warn!(
            "query-token authentication is enabled; URLs may expose credentials in proxy logs/history"
        );
    }
    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>())
        .with_graceful_shutdown(shutdown_signal(health_state))
        .await?;
    Ok(())
}

/// Assemble the complete HTTP application for already-validated options.
fn build_app(
    options: HttpServerOptions,
    shell_runtime: Arc<dyn ShellRuntime>,
    health_state: Arc<HealthState>,
) -> Router {
    let shared_concurrency = Arc::new(Semaphore::new(options.max_concurrency));
    let auth_failure_limiter = Arc::new(RateLimiter::new(options.requests_per_minute));
    let principal_rate_limiter = Arc::new(RateLimiter::new(options.principal_requests_per_minute));

    let mut config = StreamableHttpServerConfig::default();
    config.legacy_session_mode = true;
    config.allowed_hosts.extend(options.extra_hosts);
    config.max_request_body_bytes = MAX_BODY_BYTES;

    let shared = WinxService::with_runtime(SessionIsolation::Strict, shell_runtime);
    let session_manager = Arc::new(LocalSessionManager::default());
    let mcp_service =
        StreamableHttpService::new(move || Ok(shared.clone()), session_manager.clone(), config);
    let security = Arc::new(RouteSecurity {
        principals: Arc::new(options.principals),
        allow_query: options.allow_query_token,
        session_affinity: options.session_affinity,
        concurrency: shared_concurrency,
        auth_failure_limiter,
        principal_rate_limiter,
        session_owners: Arc::new(SessionOwners::default()),
        session_manager,
    });
    // The timeout sits INSIDE the security middleware: when it fires, the
    // middleware still receives the 408 response and records the usage event.
    // With the previous order the timed-out future was dropped before logging,
    // so 408s were invisible in the usage log.
    let protected_mcp = Router::new()
        .nest_service("/mcp", mcp_service)
        .layer(TimeoutLayer::with_status_code(StatusCode::REQUEST_TIMEOUT, REQUEST_TIMEOUT))
        .layer(middleware::from_fn_with_state(security, enforce_route_security))
        .layer(middleware::from_fn(reject_declared_oversized_body));
    // Keep OAuth discovery probes and every unrelated path outside bearer auth.
    // Winx does not advertise OAuth metadata, so those routes should be ordinary
    // 404s instead of 401s that encourage clients to start a nonexistent flow.
    let health_routes = Router::new()
        .route("/healthz", get(healthz))
        .route("/readyz", get(readyz))
        .with_state(health_state);
    Router::new().merge(health_routes).merge(protected_mcp)
}

/// Answer 413 for a declared oversized body before authentication or dispatch.
///
/// Streamed bodies without `Content-Length` are bounded by the SDK while it
/// reads them (`max_request_body_bytes`, also a 413). tower-http's
/// `RequestBodyLimitLayer` is deliberately not used: the SDK turns the body
/// error it injects into a 500, so chunked overflows would not answer 413.
async fn reject_declared_oversized_body(request: Request, next: Next) -> Response {
    let declared = request
        .headers()
        .get(header::CONTENT_LENGTH)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.trim().parse::<u64>().ok());
    let limit = u64::try_from(MAX_BODY_BYTES).unwrap_or(u64::MAX);
    if declared.is_some_and(|length| length > limit) {
        return (
            StatusCode::PAYLOAD_TOO_LARGE,
            format!("Payload Too Large: request body exceeds {MAX_BODY_BYTES} bytes"),
        )
            .into_response();
    }
    next.run(request).await
}

async fn healthz() -> impl IntoResponse {
    (StatusCode::OK, Json(HealthResponse { status: "ok", ready: true }))
}

async fn readyz(State(state): State<Arc<HealthState>>) -> impl IntoResponse {
    let ready = state.ready.load(Ordering::Acquire);
    let status = if ready { StatusCode::OK } else { StatusCode::SERVICE_UNAVAILABLE };
    (status, Json(HealthResponse { status: if ready { "ready" } else { "draining" }, ready }))
}

async fn shutdown_signal(state: Arc<HealthState>) {
    #[cfg(unix)]
    {
        let terminate = async {
            match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
                Ok(mut signal) => {
                    signal.recv().await;
                }
                Err(error) => {
                    tracing::warn!(%error, "failed to install SIGTERM handler");
                    std::future::pending::<()>().await;
                }
            }
        };
        tokio::select! {
            result = tokio::signal::ctrl_c() => {
                if let Err(error) = result {
                    tracing::warn!(%error, "failed while waiting for Ctrl-C");
                }
            }
            () = terminate => {}
        }
    }
    #[cfg(not(unix))]
    if let Err(error) = tokio::signal::ctrl_c().await {
        tracing::warn!(%error, "failed while waiting for Ctrl-C");
    }

    state.ready.store(false, Ordering::Release);
    tracing::info!("HTTP shutdown requested; draining in-flight requests");
}

fn validate_options(options: &HttpServerOptions) -> Result<SocketAddr, BoxError> {
    let bind: SocketAddr = options
        .bind
        .parse()
        .map_err(|error| format!("invalid --bind socket address {:?}: {error}", options.bind))?;
    if !bind.ip().is_loopback() && !options.allow_non_loopback {
        return Err(format!(
            "refusing non-loopback HTTP bind {bind}; pass --allow-non-loopback only behind a trusted firewall/tunnel"
        )
        .into());
    }
    if options.max_concurrency == 0 {
        return Err("max_concurrency must be at least 1".into());
    }
    if options.requests_per_minute == 0 {
        return Err("requests_per_minute must be at least 1".into());
    }
    if options.principal_requests_per_minute == 0 {
        return Err("principal_requests_per_minute must be at least 1".into());
    }
    if options.principals.is_empty() {
        return Err("at least one HTTP principal is required".into());
    }
    let mut names = HashSet::new();
    for principal in &options.principals {
        if !names.insert(principal.name()) {
            return Err(format!("duplicate HTTP principal: {}", principal.name()).into());
        }
    }
    Ok(bind)
}

#[derive(Clone)]
struct RouteSecurity {
    principals: Arc<Vec<HttpPrincipal>>,
    allow_query: bool,
    session_affinity: HttpSessionAffinity,
    concurrency: Arc<Semaphore>,
    auth_failure_limiter: Arc<RateLimiter<IpAddr>>,
    principal_rate_limiter: Arc<RateLimiter<String>>,
    session_owners: Arc<SessionOwners>,
    session_manager: Arc<LocalSessionManager>,
}

/// Upper bound on remembered transport sessions. Past it the oldest entry is
/// evicted and its SDK session closed at the same time, so an id never
/// outlives the record that says who may use it.
const MAX_TRACKED_SESSIONS: usize = 65_536;

/// Which principal created each legacy MCP transport session (`Mcp-Session-Id`).
///
/// The SDK's session manager is shared by every principal and treats a known
/// session id as sufficient to resume its stream or delete it. Winx records
/// the creator when the `initialize` response mints the id and refuses any
/// later request that presents the id under another principal's token.
#[derive(Default)]
struct SessionOwners {
    inner: std::sync::Mutex<SessionOwnerTable>,
}

#[derive(Default)]
struct SessionOwnerTable {
    owners: HashMap<String, String>,
    order: VecDeque<String>,
}

impl SessionOwners {
    /// Only the recorded creator may present a minted `session_id`. An id this
    /// process never minted is not an SDK session at all: modern stateless
    /// clients may send one as a conversation identity, and the SDK answers
    /// 404 itself if such an id reaches a legacy session path.
    fn owner_matches(&self, session_id: &str, principal: &str) -> bool {
        let table = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
        table.owners.get(session_id).is_none_or(|owner| owner == principal)
    }

    /// Remember the creator of a freshly minted id. Returns the ids evicted
    /// to stay within [`MAX_TRACKED_SESSIONS`]; the caller closes those SDK
    /// sessions so a forgotten id can never be presented by anyone.
    fn record(&self, session_id: &str, principal: &str) -> Vec<String> {
        let mut table = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
        if table.owners.contains_key(session_id) {
            return Vec::new();
        }
        table.owners.insert(session_id.to_string(), principal.to_string());
        table.order.push_back(session_id.to_string());
        let mut evicted = Vec::new();
        while table.order.len() > MAX_TRACKED_SESSIONS {
            if let Some(oldest) = table.order.pop_front() {
                table.owners.remove(&oldest);
                evicted.push(oldest);
            }
        }
        evicted
    }

    fn forget(&self, session_id: &str) {
        let mut table = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
        if table.owners.remove(session_id).is_some() {
            table.order.retain(|id| id != session_id);
        }
    }
}

fn session_id_header(headers: &header::HeaderMap) -> Option<String> {
    headers
        .get("mcp-session-id")
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToString::to_string)
}

async fn enforce_route_security(
    State(security): State<Arc<RouteSecurity>>,
    ConnectInfo(peer): ConnectInfo<SocketAddr>,
    request: Request,
    next: Next,
) -> Response {
    let path = request.uri().path();
    if path != "/mcp" && !path.starts_with("/mcp/") {
        return next.run(request).await;
    }
    let (started, method, correlation) = observe_request_entry(peer, request.method());
    // Authenticate before touching shared capacity. Behind a tunnel every caller
    // shares one source IP, so junk traffic must not be able to spend the
    // window or the concurrency permits a valid principal needs. The check is a
    // constant-time digest comparison of headers only; the body is never read.
    let Some(principal) =
        authenticate_request(&request, &security.principals, security.allow_query)
    else {
        return reject_failed_auth(&security, peer, &correlation, request.uri().path()).await;
    };

    // A guessable token gets no answer that distinguishes it while its source
    // IP is over the failed-auth budget; strong tokens are not locked out.
    if principal.has_weak_token() && security.auth_failure_limiter.exhausted(&peer.ip()).await {
        tokio::time::sleep(INVALID_AUTH_DELAY).await;
        log_rejection(peer, None, &correlation, "ip_rate_limit", StatusCode::TOO_MANY_REQUESTS);
        return retry_response(StatusCode::TOO_MANY_REQUESTS, "request rate limit exceeded\n");
    }

    let principal_name = principal.name().to_string();
    if !security.principal_rate_limiter.allow(principal_name.clone()).await {
        log_rejection(
            peer,
            Some(&principal_name),
            &correlation,
            "principal_rate_limit",
            StatusCode::TOO_MANY_REQUESTS,
        );
        return retry_response(StatusCode::TOO_MANY_REQUESTS, "principal rate limit exceeded\n");
    }
    // A transport session id belongs to the principal that created it. Any
    // other authenticated principal sees it as unknown, exactly like a stale
    // id: the SDK would otherwise let it resume the stream or close it.
    let presented_session = session_id_header(request.headers());
    if let Some(session_id) = presented_session.as_deref() {
        if !security.session_owners.owner_matches(session_id, &principal_name) {
            tracing::warn!(
                peer = %peer,
                principal = %principal_name,
                "rejected request presenting another principal's MCP session id"
            );
            log_rejection(
                peer,
                Some(&principal_name),
                &correlation,
                "foreign_session",
                StatusCode::NOT_FOUND,
            );
            return (StatusCode::NOT_FOUND, "session not found\n").into_response();
        }
    }
    let Ok(permit) = security.concurrency.clone().try_acquire_owned() else {
        log_rejection(
            peer,
            Some(&principal_name),
            &correlation,
            "concurrency",
            StatusCode::SERVICE_UNAVAILABLE,
        );
        return retry_response(StatusCode::SERVICE_UNAVAILABLE, "too many concurrent requests\n");
    };
    let mut request = request;
    request.extensions_mut().insert(principal);
    request.extensions_mut().insert(security.session_affinity);
    request.extensions_mut().insert(correlation.clone());
    let response = next.run(request).await;
    if response.status().is_success() {
        if let Some(minted) = session_id_header(response.headers()) {
            for evicted in security.session_owners.record(&minted, &principal_name) {
                let manager = security.session_manager.clone();
                tokio::spawn(async move {
                    let id: rmcp::transport::common::server_side_http::SessionId =
                        Arc::from(evicted.as_str());
                    let _ = manager.close_session(&id).await;
                });
            }
        }
        if method == axum::http::Method::DELETE {
            if let Some(session_id) = presented_session.as_deref() {
                security.session_owners.forget(session_id);
            }
        }
    }
    // Preserve the legacy event and duration as the response-ready observation.
    // Body termination is a separate metadata-only event from PermitBody.
    let response_ready = Instant::now();
    tracing::info!(
        target: "winx::usage",
        event = "http_request",
        phase = "response_ready",
        principal = %principal_name,
        request_id = %correlation.as_str(),
        peer = %peer,
        method = %method,
        status = response.status().as_u16(),
        duration_ms = u64::try_from(response_ready.saturating_duration_since(started).as_millis())
            .unwrap_or(u64::MAX),
        "http request"
    );
    // The permit travels with the body, including while a streaming response is
    // open. An already-empty body can now log completion after response-ready.
    response.map(|body| {
        axum::body::Body::new(PermitBody::new(body, permit).observe(
            correlation,
            started,
            response_ready,
        ))
    })
}

/// Failed authentication: the only traffic that spends the per-IP window.
///
/// Every failure is delayed, including those already over budget, so an
/// over-budget guess is never answered faster. No concurrency permit is held
/// while waiting. A strong valid token still succeeds while its IP is over
/// budget (its 32-byte minimum makes guessing infeasible); a weak one does not.
async fn reject_failed_auth(
    security: &RouteSecurity,
    peer: SocketAddr,
    correlation: &RequestCorrelation,
    path: &str,
) -> Response {
    let within_budget = security.auth_failure_limiter.allow(peer.ip()).await;
    tokio::time::sleep(INVALID_AUTH_DELAY).await;
    if !within_budget {
        log_rejection(peer, None, correlation, "ip_rate_limit", StatusCode::TOO_MANY_REQUESTS);
        return retry_response(StatusCode::TOO_MANY_REQUESTS, "request rate limit exceeded\n");
    }
    tracing::warn!(peer = %peer, path = %path, "rejected request with missing or invalid HTTP token");
    log_rejection(peer, None, correlation, "auth", StatusCode::UNAUTHORIZED);
    (StatusCode::UNAUTHORIZED, "missing or invalid token\n").into_response()
}

/// One redacted `winx::usage` event per admission rejection, so overload and
/// auth failures are visible in reports instead of vanishing before the
/// per-request event.
fn log_rejection(
    peer: SocketAddr,
    principal: Option<&str>,
    correlation: &RequestCorrelation,
    reason: &str,
    status: StatusCode,
) {
    tracing::info!(
        target: "winx::usage",
        event = "http_rejected",
        request_id = %correlation.as_str(),
        principal = principal.unwrap_or(""),
        peer = %peer,
        reason,
        status = status.as_u16(),
        "http request rejected"
    );
}
fn retry_response(status: StatusCode, body: &'static str) -> Response {
    let mut response = (status, body).into_response();
    response.headers_mut().insert(header::RETRY_AFTER, HeaderValue::from_static("1"));
    response
}

#[derive(Debug)]
struct RateWindow {
    started_at: Instant,
    last_seen: Instant,
    count: u32,
}

#[derive(Debug)]
/// Fixed-window limiter keyed by source IP (failed-auth budget) or by
/// authenticated principal (fairness behind a shared proxy).
struct RateLimiter<K> {
    requests_per_window: u32,
    windows: Mutex<HashMap<K, RateWindow>>,
}

impl<K: std::hash::Hash + Eq + Clone> RateLimiter<K> {
    fn new(requests_per_window: u32) -> Self {
        Self { requests_per_window, windows: Mutex::new(HashMap::new()) }
    }

    async fn allow(&self, peer: K) -> bool {
        let now = Instant::now();
        let mut windows = self.windows.lock().await;
        windows.retain(|_, window| now.duration_since(window.last_seen) < RATE_WINDOW * 2);
        if !windows.contains_key(&peer) && windows.len() >= MAX_RATE_TRACKED_IPS {
            if let Some(oldest) =
                windows.iter().min_by_key(|(_, window)| window.last_seen).map(|(ip, _)| ip.clone())
            {
                windows.remove(&oldest);
            }
        }
        let window =
            windows.entry(peer).or_insert(RateWindow { started_at: now, last_seen: now, count: 0 });
        if now.duration_since(window.started_at) >= RATE_WINDOW {
            window.started_at = now;
            window.count = 0;
        }
        window.last_seen = now;
        if window.count >= self.requests_per_window {
            return false;
        }
        window.count = window.count.saturating_add(1);
        true
    }

    /// Whether `peer` has spent its budget in the current window, without
    /// consuming any of it.
    async fn exhausted(&self, peer: &K) -> bool {
        let windows = self.windows.lock().await;
        windows.get(peer).is_some_and(|window| {
            Instant::now().duration_since(window.started_at) < RATE_WINDOW
                && window.count >= self.requests_per_window
        })
    }
}

fn authenticate_request(
    request: &Request,
    principals: &[HttpPrincipal],
    allow_query: bool,
) -> Option<HttpPrincipal> {
    let header = bearer_token(request);
    let query = allow_query.then(|| query_token(request)).flatten();
    [header, query].into_iter().flatten().find_map(|presented| {
        principals.iter().find(|principal| principal.matches_token(presented)).cloned()
    })
}

fn bearer_token(request: &Request) -> Option<&str> {
    let value = request.headers().get(header::AUTHORIZATION)?.to_str().ok()?;
    let (scheme, token) = value.split_once(' ')?;
    scheme.eq_ignore_ascii_case("bearer").then_some(token.trim())
}

fn query_token(request: &Request) -> Option<&str> {
    request.uri().query()?.split('&').find_map(|pair| pair.strip_prefix("token="))
}

#[cfg(test)]
mod admission_tests {
    #![allow(clippy::expect_used)]
    use std::sync::Arc;

    use tokio::sync::Semaphore;

    use super::{PermitBody, RateLimiter};

    #[tokio::test]
    async fn principal_windows_are_independent_of_each_other() {
        let limiter: RateLimiter<String> = RateLimiter::new(2);
        assert!(limiter.allow("left".to_string()).await);
        assert!(limiter.allow("left".to_string()).await);
        assert!(!limiter.allow("left".to_string()).await, "left exhausted its window");
        assert!(limiter.allow("right".to_string()).await, "right is unaffected");
    }

    #[tokio::test]
    async fn permit_is_released_only_when_the_body_finishes() {
        let semaphore = Arc::new(Semaphore::new(1));
        let permit = semaphore.clone().try_acquire_owned().expect("permit");
        let body =
            axum::body::Body::new(PermitBody::new(axum::body::Body::from("payload"), permit));
        assert_eq!(semaphore.available_permits(), 0, "held while the body is pending");
        let bytes = axum::body::to_bytes(body, usize::MAX).await.expect("read body");
        assert_eq!(&bytes[..], b"payload");
        assert_eq!(semaphore.available_permits(), 1, "released with the consumed body");

        let permit = semaphore.clone().try_acquire_owned().expect("permit again");
        let body =
            axum::body::Body::new(PermitBody::new(axum::body::Body::from("abandoned"), permit));
        drop(body);
        assert_eq!(semaphore.available_permits(), 1, "released when a client disconnects");
    }
}

#[cfg(test)]
mod session_owner_tests {
    use super::SessionOwners;

    #[test]
    fn minted_session_ids_belong_to_their_first_creator() {
        let owners = SessionOwners::default();
        assert!(owners.owner_matches("conv-a", "left"), "unminted ids are not SDK sessions");
        assert!(owners.record("s1", "left").is_empty(), "nothing evicted");
        assert!(owners.owner_matches("s1", "left"));
        assert!(!owners.owner_matches("s1", "right"), "another principal");
        assert!(owners.record("s1", "right").is_empty(), "nothing evicted");
        assert!(owners.owner_matches("s1", "left"), "first creator wins");
        owners.forget("s1");
        assert!(owners.owner_matches("s1", "right"), "closed ids are plain unknown ids again");
    }

    #[test]
    fn eviction_hands_back_the_ids_whose_sdk_sessions_must_close() {
        let owners = SessionOwners::default();
        assert!(owners.record("oldest", "left").is_empty(), "nothing evicted");
        let mut evicted = Vec::new();
        for index in 0..super::MAX_TRACKED_SESSIONS {
            evicted.extend(owners.record(&format!("s{index}"), "left"));
        }
        assert_eq!(evicted, vec!["oldest".to_string()]);
        assert!(owners.owner_matches(&format!("s{}", super::MAX_TRACKED_SESSIONS - 1), "left"));
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::expect_used)]

    use std::net::{IpAddr, Ipv4Addr};
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::Arc;

    use axum::body::Body;
    use axum::extract::{Request, State};
    use axum::http::StatusCode;
    use axum::response::IntoResponse as _;

    use super::{
        authenticate_request, readyz, validate_options, HealthState, HttpServerOptions, RateLimiter,
    };
    use crate::config::HttpPrincipal;

    fn request(uri: &str, authorization: Option<&str>) -> Request {
        let mut builder = Request::builder().uri(uri);
        if let Some(authorization) = authorization {
            builder = builder.header(axum::http::header::AUTHORIZATION, authorization);
        }
        builder.body(Body::empty()).expect("test request")
    }

    fn principal(name: &str, token_byte: char) -> HttpPrincipal {
        let token = token_byte.to_string().repeat(32);
        HttpPrincipal::new(name, &token, false).expect("valid principal")
    }

    fn strong_options(bind: &str) -> HttpServerOptions {
        HttpServerOptions {
            bind: bind.to_string(),
            principals: vec![principal("default", 'x')],
            extra_hosts: Vec::new(),
            allow_query_token: false,
            allow_non_loopback: false,
            session_affinity: crate::config::HttpSessionAffinity::Workspace,
            max_concurrency: super::DEFAULT_MAX_CONCURRENCY,
            requests_per_minute: super::DEFAULT_REQUESTS_PER_MINUTE,
            principal_requests_per_minute: super::DEFAULT_PRINCIPAL_REQUESTS_PER_MINUTE,
        }
    }

    #[test]
    fn bearer_and_opt_in_query_tokens_select_the_matching_principal() {
        let principals = vec![principal("chatgpt", 'a'), principal("automation", 'b')];
        let header = request("/mcp", Some(&format!("Bearer {}", "a".repeat(32))));
        let matched = authenticate_request(&header, &principals, false).expect("header principal");
        assert_eq!(matched.name(), "chatgpt");

        let query_uri = format!("/mcp?token={}", "b".repeat(32));
        assert!(authenticate_request(&request(&query_uri, None), &principals, false).is_none());
        let matched = authenticate_request(&request(&query_uri, None), &principals, true)
            .expect("query principal");
        assert_eq!(matched.name(), "automation");
    }

    #[test]
    fn non_loopback_binding_requires_explicit_override() {
        assert!(validate_options(&strong_options("127.0.0.1:8000")).is_ok());
        assert!(validate_options(&strong_options("0.0.0.0:8000")).is_err());

        let mut public = strong_options("0.0.0.0:8000");
        public.allow_non_loopback = true;
        assert!(validate_options(&public).is_ok());
    }

    #[tokio::test]
    async fn rate_limiter_rejects_excess_requests() {
        let limiter = RateLimiter::new(2);
        let peer = IpAddr::V4(Ipv4Addr::LOCALHOST);
        assert!(limiter.allow(peer).await);
        assert!(limiter.allow(peer).await);
        assert!(!limiter.allow(peer).await);
    }

    #[test]
    fn duplicate_principal_names_are_rejected() {
        let mut options = strong_options("127.0.0.1:8000");
        options.principals = vec![principal("same", 'a'), principal("same", 'b')];
        assert!(validate_options(&options).is_err());
    }

    #[tokio::test]
    async fn readiness_turns_unavailable_while_draining() {
        let state = Arc::new(HealthState { ready: AtomicBool::new(true) });
        let ready = readyz(State(Arc::clone(&state))).await.into_response();
        assert_eq!(ready.status(), StatusCode::OK);

        state.ready.store(false, Ordering::Release);
        let draining = readyz(State(state)).await.into_response();
        assert_eq!(draining.status(), StatusCode::SERVICE_UNAVAILABLE);
    }
}

#[cfg(test)]
mod observability_tests;

#[cfg(test)]
mod security_tests;