tower-mcp 0.11.0

//! Streamable HTTP transport for MCP
//!
//! Implements the Streamable HTTP transport from MCP specification 2025-11-25,
//! with version-gated support for the 2026-07-28 stateless protocol (SEP-2575 /
//! SEP-2567) when the `stateless` feature is compiled in.
//!
//! ## Features
//!
//! - Single endpoint for POST (requests) and GET (SSE notifications)
//! - Session management via `MCP-Session-Id` header
//! - SSE streaming for server notifications and progress updates
//! - SSE event IDs and stream resumption via `Last-Event-ID` header (SEP-1699)
//! - Configurable session TTL and cleanup
//! - **Sampling support**: Server-to-client LLM requests via SSE + POST
//! - **Stateless mode** (`stateless` feature): version-gated dispatch for
//!   2026-07-28+ clients with per-request `_meta` and no session handshake
//!
//! ## Stateless mode (2026-07-28 protocol)
//!
//! When the `stateless` feature is compiled in, the transport handles two
//! distinct stateless paths:
//!
//! ### Automatic version-gated path (2026-07-28+)
//!
//! Any request that arrives with `MCP-Protocol-Version: 2026-07-28` (or
//! later) and no `mcp-session-id` header is dispatched statelessly, regardless
//! of whether [`HttpTransport::stateless()`] was called. The client is fully
//! self-identifying: every request carries its protocol version, client info,
//! and client capabilities in the `_meta` object; no initialize handshake is
//! needed. Handlers access this data via
//! [`RequestContext::per_request_meta()`](crate::context::RequestContext::per_request_meta).
//!
//! This path runs before the legacy SEP-1442 opt-in path, so 2026-07-28
//! clients are always handled correctly even on transports that never call
//! `HttpTransport::stateless()`.
//!
//! ### Legacy SEP-1442 opt-in path
//!
//! Calling [`HttpTransport::stateless()`] with a [`crate::stateless::StatelessConfig`]
//! activates the older SEP-1442-style opt-in stateless behavior for clients that
//! do not carry `MCP-Protocol-Version: 2026-07-28`. See
//! [`crate::stateless::StatelessConfig`] for details on what this path controls.
//!
//! Stateful clients (those sending an `mcp-session-id`) continue to work
//! normally on the same transport alongside both stateless paths.
//!
//! ## `messages/listen` SSE stream
//!
//! Clients using the 2026-07-28 protocol open a server-to-client notification
//! stream by POSTing a `messages/listen` JSON-RPC request. The server responds
//! with `Content-Type: text/event-stream` and streams zero or more
//! `notifications/*` events until the client disconnects.
//!
//! This replaces the `GET /` SSE endpoint used by the 2025-11-25 protocol. The
//! `GET /` endpoint is still supported for 2025-11-25 sessions; `messages/listen`
//! is only available for 2026-07-28+ clients.
//!
//! ```text
//! Client (2026-07-28)                          Server
//!   |                                            |
//!   |-- POST / {method: "messages/listen",       |
//!   |           MCP-Protocol-Version: 2026-07-28} -->|
//!   |<-- 200 Content-Type: text/event-stream ----|
//!   |<-- event: message (notification) ----------|
//!   |<-- event: message (notification) ----------|
//!   |   (client disconnects)                     |
//! ```
//!
//! ## SEP-2243 HTTP headers
//!
//! SEP-2243 defines HTTP headers that let load balancers, proxies, and
//! observability tools inspect MCP traffic without parsing the JSON-RPC body:
//!
//! | Header | Required when | Description |
//! |--------|---------------|-------------|
//! | `Mcp-Method` | All POST requests (strict mode) | Mirrors the JSON-RPC `method` field |
//! | `Mcp-Name` | `tools/call`, `prompts/get`, `resources/read` (strict mode) | Mirrors `params.name` or `params.uri` |
//! | `MCP-Protocol-Version` | All requests (strict mode) | The protocol version in use |
//!
//! Validation is **lenient** for 2025-11-25 clients: headers present in the
//! request are validated for consistency with the body, but missing headers are
//! not an error. Validation is **strict** for 2026-07-28+ clients: `Mcp-Method`
//! must be present on every POST and `Mcp-Name` must be present for the three
//! named methods. Violations return `-32001` (HeaderMismatch).
//!
//! The public constants [`MCP_METHOD_HEADER`], [`MCP_NAME_HEADER`], and
//! [`MCP_PARAM_HEADER_PREFIX`] hold the canonical lowercase header names.
//!
//! ## Sampling (Server-to-Client Requests)
//!
//! When using `HttpTransport::new(router).with_sampling()`, tool handlers can request
//! LLM completions from the client. The flow is:
//!
//! 1. Tool handler calls `ctx.sample(params)`
//! 2. Server sends the sampling request on the SSE stream
//! 3. Client receives the request and processes it
//! 4. Client sends the response as a POST to the MCP endpoint
//! 5. Server routes the response back to the waiting handler
//!
//! This follows the MCP spec which states servers MAY send JSON-RPC requests
//! on the SSE stream.
//!
//! ## Session Reconnection
//!
//! When a session is not found (e.g., after server restart or session expiration),
//! the server returns a JSON-RPC error with code `-32005` (SessionNotFound).
//! Clients should handle this by re-initializing the connection:
//!
//! ```text
//! Client                          Server
//!   |                               |
//!   |-- tools/list (old session) -->|
//!   |<-- error: SessionNotFound ----|
//!   |                               |
//!   |-- initialize --------------->|
//!   |<-- result + new session id ---|
//!   |                               |
//!   |-- tools/list (new session) -->|
//!   |<-- result -------------------|
//! ```
//!
//! ## SSE Stream Resumption (SEP-1699)
//!
//! Each SSE event includes a unique, monotonically increasing event ID. If a
//! client disconnects and reconnects, it can include the `Last-Event-ID` header
//! with the ID of the last event it received. The server will replay any buffered
//! events with IDs greater than the provided ID before continuing with live events.
//!
//! ```text
//! Client                              Server
//!   |-- GET / (Accept: text/event-stream) -->|
//!   |<-- id:0, data:{progress...} -----------|
//!   |<-- id:1, data:{progress...} -----------|
//!   |<-- id:2, data:{progress...} -----------|
//!   |                                        |
//!   |  ** Client disconnects **              |
//!   |                                        |
//!   |                   (server buffers id:3, id:4, id:5)
//!   |                                        |
//!   |-- GET / (Last-Event-ID: 2) ----------->|
//!   |<-- id:3, data:{...} (replayed) --------|
//!   |<-- id:4, data:{...} (replayed) --------|
//!   |<-- id:5, data:{...} (replayed) --------|
//!   |<-- id:6, data:{...} (live) ------------|
//! ```
//!
//! The server buffers up to 1000 events per session by default.
//!
//! ## Error Codes
//!
//! | Code    | Name                      | Description                                        |
//! |---------|---------------------------|----------------------------------------------------|
//! | -32001  | HeaderMismatch            | Required HTTP header missing or inconsistent with body (SEP-2243, strict mode) |
//! | -32004  | UnsupportedProtocolVersion| Server does not support the requested protocol version (SEP-2575) |
//! | -32005  | SessionNotFound           | Session expired or server restarted                |
//! | -32006  | SessionRequired           | MCP-Session-Id header missing                      |
//!
//! ## Session Handling
//!
//! By default, sessions are optional: requests without an `mcp-session-id`
//! header are allowed and receive a transient, pre-initialized session. This
//! ensures compatibility with clients (Codex CLI, Cursor, etc.) that don't
//! carry the session ID forward after initialization.
//!
//! Clients that do send session IDs continue to work normally.
//!
//! To require strict session management (reject requests without a session ID),
//! use [`HttpTransport::require_sessions()`]:
//!
//! ```rust,ignore
//! let transport = HttpTransport::new(router).require_sessions();
//! ```
//!
//! ## CORS Support
//!
//! Browser-based MCP clients require CORS headers. Since [`HttpTransport::into_router()`]
//! returns a standard [`axum::Router`], you can add CORS support using
//! `tower_http::cors::CorsLayer`:
//!
//! ```rust,ignore
//! use tower_mcp::McpRouter;
//! use tower_mcp::transport::http::HttpTransport;
//! use tower_http::cors::{CorsLayer, Any};
//! use http::Method;
//!
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! let router = McpRouter::new().server_info("my-server", "1.0.0");
//! let transport = HttpTransport::new(router);
//!
//! // Wrap the axum router with CORS middleware
//! let app = transport.into_router().layer(
//!     CorsLayer::new()
//!         .allow_origin(Any)
//!         .allow_methods([Method::GET, Method::POST, Method::DELETE])
//!         .allow_headers(Any)
//!         .expose_headers(Any),
//! );
//!
//! let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await?;
//! axum::serve(listener, app).await?;
//! # Ok(())
//! # }
//! ```
//!
//! For production, replace `Any` origins with your specific allowed origins.
//!
//! **Note:** [`HttpTransport::layer()`] applies middleware at the *MCP request* level
//! (inside the JSON-RPC service). CORS must be applied at the *HTTP* level using
//! `into_router().layer(...)` as shown above.
//!
//! # Example
//!
//! ```rust,no_run
//! use tower_mcp::{BoxError, McpRouter, ToolBuilder, CallToolResult};
//! use tower_mcp::transport::http::HttpTransport;
//! use schemars::JsonSchema;
//! use serde::Deserialize;
//!
//! #[derive(Debug, Deserialize, JsonSchema)]
//! struct Input { value: String }
//!
//! #[tokio::main]
//! async fn main() -> Result<(), BoxError> {
//!     let tool = ToolBuilder::new("echo")
//!         .handler(|i: Input| async move { Ok(CallToolResult::text(i.value)) })
//!         .build();
//!
//!     let router = McpRouter::new()
//!         .server_info("my-server", "1.0.0")
//!         .tool(tool);
//!
//!     let transport = HttpTransport::new(router);
//!
//!     // Run on localhost:3000
//!     transport.serve("127.0.0.1:3000").await?;
//!     Ok(())
//! }
//! ```

use std::collections::HashMap;
use std::convert::Infallible;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use axum::{
    Router,
    extract::State,
    http::{HeaderMap, HeaderValue, StatusCode, header},
    response::{IntoResponse, Response, Sse, sse::Event},
    routing::{delete, get, post},
};
use tokio::sync::{Mutex, RwLock, broadcast, oneshot};
use tokio_stream::StreamExt;
use tokio_stream::wrappers::BroadcastStream;

use crate::context::{
    ChannelClientRequester, ClientRequesterHandle, NotificationReceiver, OutgoingRequestReceiver,
    notification_channel, outgoing_request_channel,
};
use crate::error::{Error, JsonRpcError, Result};
use crate::jsonrpc::JsonRpcService;
use crate::protocol::{
    ClientCapabilities, Implementation, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse,
    LATEST_PROTOCOL_VERSION, McpNotification, RequestId, SUPPORTED_PROTOCOL_VERSIONS,
    UPCOMING_PROTOCOL_VERSION,
};
use crate::router::{McpRouter, RouterRequest, RouterResponse};
use crate::transport::service::{
    CatchError, InjectAnnotations, McpBoxService, ServiceFactory, identity_factory,
};
use tower::util::BoxCloneService;

/// SEP-2575 per-request `_meta` extraction. Pulls `StatelessRequestMeta` from
/// the parsed request params and inserts it into the per-request `Extensions`
/// so handlers can read it via `ctx.per_request_meta()`. No-op if the request
/// has no `_meta`, params aren't an object, or the meta can't deserialize.
#[cfg(feature = "stateless")]
fn stash_per_request_meta(req: &JsonRpcRequest, ext: &mut crate::router::Extensions) {
    if let Some(params) = req.params.as_ref()
        && let Some(meta) = crate::stateless::StatelessRequestMeta::from_params(params)
    {
        ext.insert(meta);
    }
}

/// Header name for MCP session ID
pub const MCP_SESSION_ID_HEADER: &str = "mcp-session-id";

/// Header name for MCP protocol version
pub const MCP_PROTOCOL_VERSION_HEADER: &str = "mcp-protocol-version";

/// SEP-2243: header that mirrors the JSON-RPC `method` field for HTTP
/// intermediaries (load balancers, observability) so they can route or
/// classify MCP traffic without parsing the body. Required on all POST
/// requests when the negotiated protocol version implements SEP-2243.
pub const MCP_METHOD_HEADER: &str = "mcp-method";

/// SEP-2243: header that mirrors `params.name` (for `tools/call` and
/// `prompts/get`) or `params.uri` (for `resources/read`). Required for
/// those three methods when the negotiated protocol version implements
/// SEP-2243.
pub const MCP_NAME_HEADER: &str = "mcp-name";

/// SEP-2243: prefix for custom headers derived from tool parameters
/// marked with the `x-mcp-header` JSON Schema extension. The full header
/// name is `Mcp-Param-{Name}`.
pub const MCP_PARAM_HEADER_PREFIX: &str = "mcp-param-";

/// First MCP protocol version that mandates SEP-2243 HTTP header
/// validation. Prior versions are treated leniently: present headers are
/// still validated for body consistency, but missing headers are not an
/// error.
pub(super) const SEP_2243_MIN_PROTOCOL_VERSION: &str = "2026-07-28";

/// SSE event type for JSON-RPC messages
const SSE_MESSAGE_EVENT: &str = "message";

/// Header name for Last-Event-ID (for SSE stream resumption per SEP-1699)
const LAST_EVENT_ID_HEADER: &str = "last-event-id";

/// Pending request waiting for a response from the client
struct PendingRequest {
    response_tx: oneshot::Sender<Result<serde_json::Value>>,
}

/// Session state for HTTP transport
/// How a session produces its MCP service for request processing.
enum SessionServiceSource {
    /// Session was created from an McpRouter with a factory for middleware wrapping.
    Router {
        router: McpRouter,
        factory: ServiceFactory,
    },
    /// Session was created from a pre-built boxed service (e.g., McpProxy).
    /// Wrapped in Mutex because BoxCloneService is Send but not Sync,
    /// and Session must be Sync for Arc<Session> to be Send.
    Boxed(std::sync::Mutex<McpBoxService>),
}

struct Session {
    /// Session ID
    id: String,
    /// Source for creating the MCP service
    service_source: SessionServiceSource,
    /// Broadcast channel for SSE notifications and outgoing requests
    notifications_tx: broadcast::Sender<String>,
    /// When this session was created
    created_at: Instant,
    /// Last time this session was accessed
    last_accessed: RwLock<Instant>,
    /// Pending outgoing requests waiting for responses
    pending_requests: Mutex<HashMap<RequestId, PendingRequest>>,
    /// Receiver for outgoing requests (used by SSE stream)
    request_rx: Mutex<Option<OutgoingRequestReceiver>>,
    /// Negotiated protocol version (set after initialize)
    protocol_version: RwLock<String>,
    /// Client implementation info advertised in the `initialize` request.
    ///
    /// Populated by `handle_post` after a successful initialize response,
    /// and restored from a [`SessionRecord`](crate::session_store::SessionRecord)
    /// when a session is rebuilt from the persistent store. `None` until the
    /// first initialize completes.
    client_info: RwLock<Option<Implementation>>,
    /// Client capabilities advertised in the `initialize` request.
    ///
    /// Populated by `handle_post` after a successful initialize response,
    /// and restored from a [`SessionRecord`](crate::session_store::SessionRecord)
    /// when a session is rebuilt from the persistent store. `None` until the
    /// first initialize completes.
    client_capabilities: RwLock<Option<ClientCapabilities>>,
    /// Counter for SSE event IDs (for stream resumption per SEP-1699)
    event_counter: AtomicU64,
    /// Pluggable store for SSE events (enables cross-instance replay)
    event_store: Arc<dyn crate::event_store::EventStore>,
}

impl Session {
    fn new(
        router: McpRouter,
        sampling_enabled: bool,
        service_factory: ServiceFactory,
        event_store: Arc<dyn crate::event_store::EventStore>,
    ) -> Self {
        let (notifications_tx, _) = broadcast::channel(100);

        // Set up notification forwarding: mpsc -> broadcast
        // The router sends notifications (progress, log, resource updates) to
        // an mpsc channel. We bridge these to the session's broadcast channel
        // so they reach connected SSE clients.
        let (notif_sender, mut notif_receiver) = notification_channel(256);
        let router = router.with_notification_sender(notif_sender);

        let broadcast_tx = notifications_tx.clone();
        tokio::spawn(async move {
            while let Some(notification) = notif_receiver.recv().await {
                if let Some(json) = crate::transport::stdio::serialize_notification(&notification) {
                    // Best effort: if no subscribers, the message is dropped
                    let _ = broadcast_tx.send(json);
                }
            }
        });

        // Set up client requester if sampling is enabled
        let (router, request_rx) = if sampling_enabled {
            let (request_tx, request_rx) = outgoing_request_channel(32);
            let client_requester: ClientRequesterHandle =
                Arc::new(ChannelClientRequester::new(request_tx));
            let router = router.with_client_requester(client_requester);
            (router, Some(request_rx))
        } else {
            (router, None)
        };

        let now = Instant::now();
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            service_source: SessionServiceSource::Router {
                router,
                factory: service_factory,
            },
            notifications_tx,
            created_at: now,
            last_accessed: RwLock::new(now),
            pending_requests: Mutex::new(HashMap::new()),
            request_rx: Mutex::new(request_rx),
            protocol_version: RwLock::new(LATEST_PROTOCOL_VERSION.to_string()),
            client_info: RwLock::new(None),
            client_capabilities: RwLock::new(None),
            event_counter: AtomicU64::new(0),
            event_store,
        }
    }

    /// Create a session from a pre-built boxed service.
    ///
    /// This is used when the transport is created via [`HttpTransport::from_service()`].
    /// Notification bridging and sampling setup are skipped — the caller is
    /// responsible for configuring these on the service before passing it in.
    fn from_service(
        service: McpBoxService,
        event_store: Arc<dyn crate::event_store::EventStore>,
    ) -> Self {
        let (notifications_tx, _) = broadcast::channel(100);

        let now = Instant::now();
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            service_source: SessionServiceSource::Boxed(std::sync::Mutex::new(service)),
            notifications_tx,
            created_at: now,
            last_accessed: RwLock::new(now),
            pending_requests: Mutex::new(HashMap::new()),
            request_rx: Mutex::new(None),
            protocol_version: RwLock::new(LATEST_PROTOCOL_VERSION.to_string()),
            client_info: RwLock::new(None),
            client_capabilities: RwLock::new(None),
            event_counter: AtomicU64::new(0),
            event_store,
        }
    }

    /// Rebuild a session from a [`SessionRecord`] so a request for an
    /// unknown session ID can be served transparently.
    ///
    /// The router is pre-marked initialized and the protocol version is
    /// restored from the record. Runtime state (broadcast channels,
    /// pending-request table) is freshly allocated — in-flight state from
    /// before the rebuild is not recovered. The `event_counter` is left at
    /// zero; the [`SessionRegistry`] seeds it from the event store so
    /// future event IDs don't collide with buffered ones.
    fn restored(
        record: &crate::session_store::SessionRecord,
        router: McpRouter,
        sampling_enabled: bool,
        service_factory: ServiceFactory,
        event_store: Arc<dyn crate::event_store::EventStore>,
    ) -> Self {
        // Skip the Initializing intermediate state — this session was
        // already initialized on the original instance.
        router.session().mark_initialized();

        let (notifications_tx, _) = broadcast::channel(100);
        let (notif_sender, mut notif_receiver) = notification_channel(256);
        let router = router.with_notification_sender(notif_sender);

        let broadcast_tx = notifications_tx.clone();
        tokio::spawn(async move {
            while let Some(notification) = notif_receiver.recv().await {
                if let Some(json) = crate::transport::stdio::serialize_notification(&notification) {
                    let _ = broadcast_tx.send(json);
                }
            }
        });

        let (router, request_rx) = if sampling_enabled {
            let (request_tx, request_rx) = outgoing_request_channel(32);
            let client_requester: ClientRequesterHandle =
                Arc::new(ChannelClientRequester::new(request_tx));
            let router = router.with_client_requester(client_requester);
            (router, Some(request_rx))
        } else {
            (router, None)
        };

        let now = Instant::now();
        Self {
            id: record.id.clone(),
            service_source: SessionServiceSource::Router {
                router,
                factory: service_factory,
            },
            notifications_tx,
            created_at: now,
            last_accessed: RwLock::new(now),
            pending_requests: Mutex::new(HashMap::new()),
            request_rx: Mutex::new(request_rx),
            protocol_version: RwLock::new(record.protocol_version.clone()),
            client_info: RwLock::new(record.client_info.clone()),
            client_capabilities: RwLock::new(record.client_capabilities.clone()),
            event_counter: AtomicU64::new(0),
            event_store,
        }
    }

    /// Rebuild a session from a [`SessionRecord`] for transports built
    /// with [`HttpTransport::from_service`]. The service's internal state
    /// (if any) is not restored — the caller is responsible for anything
    /// beyond the metadata in the record.
    fn from_service_restored(
        service: McpBoxService,
        record: &crate::session_store::SessionRecord,
        event_store: Arc<dyn crate::event_store::EventStore>,
    ) -> Self {
        let (notifications_tx, _) = broadcast::channel(100);
        let now = Instant::now();
        Self {
            id: record.id.clone(),
            service_source: SessionServiceSource::Boxed(std::sync::Mutex::new(service)),
            notifications_tx,
            created_at: now,
            last_accessed: RwLock::new(now),
            pending_requests: Mutex::new(HashMap::new()),
            request_rx: Mutex::new(None),
            protocol_version: RwLock::new(record.protocol_version.clone()),
            client_info: RwLock::new(record.client_info.clone()),
            client_capabilities: RwLock::new(record.client_capabilities.clone()),
            event_counter: AtomicU64::new(0),
            event_store,
        }
    }

    /// Create a middleware-wrapped service from this session's service source.
    fn make_service(&self) -> McpBoxService {
        match &self.service_source {
            SessionServiceSource::Router { router, factory } => (factory)(router.clone()),
            SessionServiceSource::Boxed(mutex) => mutex.lock().unwrap().clone(),
        }
    }

    /// Handle a client notification (fire-and-forget, no response).
    ///
    /// For router-based sessions, delegates to the router's notification handler.
    /// For service-based sessions, notifications are logged but not processed
    /// (the service should handle its own notification needs).
    fn handle_notification(&self, notification: McpNotification) {
        match &self.service_source {
            SessionServiceSource::Router { router, .. } => {
                router.handle_notification(notification);
            }
            SessionServiceSource::Boxed(_) => {
                tracing::debug!(
                    notification = ?notification,
                    "Notification received on service-based session (not forwarded)"
                );
            }
        }
    }

    /// Get the next SSE event ID for this session.
    ///
    /// Event IDs are monotonically increasing per session, enabling
    /// stream resumption via the Last-Event-ID header (SEP-1699).
    fn next_event_id(&self) -> u64 {
        self.event_counter.fetch_add(1, Ordering::SeqCst)
    }

    /// Buffer an event for potential replay (SEP-1699).
    ///
    /// Delegates to the configured [`EventStore`](crate::event_store::EventStore).
    /// Store errors are logged but non-fatal — the transport continues
    /// serving the client even if the external event buffer is unavailable,
    /// since the event has already been sent on the live SSE stream.
    async fn buffer_event(&self, id: u64, data: String) {
        let record = crate::event_store::EventRecord::new(id, data);
        if let Err(e) = self.event_store.append(&self.id, record).await {
            tracing::warn!(session_id = %self.id, event_id = id, error = %e, "Failed to append event to event store");
        }
    }

    /// Get buffered events after the given event ID.
    ///
    /// Returns events with IDs greater than `after_id`, in order. Used for
    /// stream resumption when a client reconnects with the `Last-Event-ID`
    /// header. Store errors produce an empty replay list and are logged.
    async fn get_events_after(&self, after_id: u64) -> Vec<crate::event_store::EventRecord> {
        match self.event_store.replay_after(&self.id, after_id).await {
            Ok(events) => events,
            Err(e) => {
                tracing::warn!(session_id = %self.id, error = %e, "Failed to replay events from event store");
                Vec::new()
            }
        }
    }

    /// Update the last accessed time
    async fn touch(&self) {
        *self.last_accessed.write().await = Instant::now();
    }

    /// Check if the session has expired
    async fn is_expired(&self, ttl: Duration) -> bool {
        self.last_accessed.read().await.elapsed() > ttl
    }

    /// Store a pending request
    async fn add_pending_request(
        &self,
        id: RequestId,
        response_tx: oneshot::Sender<Result<serde_json::Value>>,
    ) {
        let mut pending = self.pending_requests.lock().await;
        pending.insert(id, PendingRequest { response_tx });
    }

    /// Complete a pending request with a response
    async fn complete_pending_request(
        &self,
        id: &RequestId,
        result: Result<serde_json::Value>,
    ) -> bool {
        let pending = {
            let mut pending_requests = self.pending_requests.lock().await;
            pending_requests.remove(id)
        };

        match pending {
            Some(pending) => {
                // Send result to waiter (ignore if they've dropped the receiver)
                let _ = pending.response_tx.send(result);
                true
            }
            None => false,
        }
    }
}

/// Default session TTL (30 minutes)
pub const DEFAULT_SESSION_TTL: Duration = Duration::from_secs(30 * 60);

/// Default cleanup interval (1 minute)
const DEFAULT_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);

/// Configuration for session management
#[derive(Debug, Clone)]
pub struct SessionConfig {
    /// Time-to-live for inactive sessions
    pub ttl: Duration,
    /// Maximum number of sessions (None = unlimited)
    pub max_sessions: Option<usize>,
    /// How often to run the cleanup task
    pub cleanup_interval: Duration,
}

impl Default for SessionConfig {
    fn default() -> Self {
        Self {
            ttl: DEFAULT_SESSION_TTL,
            max_sessions: None,
            cleanup_interval: DEFAULT_CLEANUP_INTERVAL,
        }
    }
}

impl SessionConfig {
    /// Create a new session config with the given TTL
    pub fn with_ttl(ttl: Duration) -> Self {
        Self {
            ttl,
            ..Default::default()
        }
    }

    /// Set the maximum number of sessions
    pub fn max_sessions(mut self, max: usize) -> Self {
        self.max_sessions = Some(max);
        self
    }

    /// Set the cleanup interval
    pub fn cleanup_interval(mut self, interval: Duration) -> Self {
        self.cleanup_interval = interval;
        self
    }
}

/// Registry coordinating live session runtime state with a pluggable
/// persistent [`SessionStore`](crate::session_store::SessionStore).
///
/// - Runtime state (broadcast channels, pending requests, live services) is
///   kept in the in-process `sessions` map and cannot be serialized.
/// - Persistent metadata (IDs, timestamps, protocol version) is mirrored into
///   the caller-supplied [`SessionStore`]. The default
///   [`MemorySessionStore`](crate::session_store::MemorySessionStore) keeps
///   metadata in-process (same behavior as before this trait existed).
struct SessionRegistry {
    sessions: RwLock<HashMap<String, Arc<Session>>>,
    config: SessionConfig,
    sampling_enabled: bool,
    persistent: Arc<dyn crate::session_store::SessionStore>,
    events: Arc<dyn crate::event_store::EventStore>,
    /// Source for rebuilding services when restoring a session.
    service_source: ServiceSource,
    /// If `true`, a request for an unknown session ID whose record is not
    /// in the persistent store spins up a new session with synthetic
    /// client info instead of returning 404 (see anubis-mcp #125 for the
    /// precedent).
    auto_reinit: bool,
}

impl SessionRegistry {
    fn new(
        config: SessionConfig,
        sampling_enabled: bool,
        persistent: Arc<dyn crate::session_store::SessionStore>,
        events: Arc<dyn crate::event_store::EventStore>,
        service_source: ServiceSource,
        auto_reinit: bool,
    ) -> Self {
        Self {
            sessions: RwLock::new(HashMap::new()),
            config,
            sampling_enabled,
            persistent,
            events,
            service_source,
            auto_reinit,
        }
    }

    /// Build a SessionRecord reflecting the given live Session.
    async fn record_for(&self, session: &Session) -> crate::session_store::SessionRecord {
        let protocol_version = session.protocol_version.read().await.clone();
        let last_accessed = session.last_accessed.read().await;
        let mut record = crate::session_store::SessionRecord::new(
            session.id.clone(),
            protocol_version,
            self.config.ttl,
        );
        // Populate the client identity / capabilities advertised at
        // initialize time so persisted records faithfully describe the
        // session. These remain `None` until a successful initialize.
        record.client_info = session.client_info.read().await.clone();
        record.client_capabilities = session.client_capabilities.read().await.clone();
        // Convert from monotonic Instant to SystemTime approximation.
        let now = std::time::SystemTime::now();
        let created_ago = session.created_at.elapsed();
        let last_accessed_ago = last_accessed.elapsed();
        record.created_at = now.checked_sub(created_ago).unwrap_or(now);
        record.last_accessed = now.checked_sub(last_accessed_ago).unwrap_or(now);
        record.expires_at = record.last_accessed + self.config.ttl;
        record
    }

    /// Persist metadata for a newly created session, logging on failure.
    ///
    /// Persistence errors are intentionally non-fatal: the live runtime
    /// session is already registered locally, so the transport can continue
    /// serving requests even if the external store is briefly unavailable.
    async fn persist_new(&self, session: &Session) {
        let record = self.record_for(session).await;
        if let Err(e) = self.persistent.create(&mut record.clone()).await {
            tracing::warn!(session_id = %session.id, error = %e, "Failed to persist session record");
        }
    }

    /// Persist an update to an existing session's record (upsert).
    ///
    /// Called after the session's state changes in a way that should be
    /// reflected in the persistent store -- notably after a successful
    /// `initialize` so the stored record carries the client's advertised
    /// `client_info` and `capabilities` (rather than the defaults captured
    /// at create time). Failures are logged but non-fatal.
    async fn save_record(&self, session: &Session) {
        let record = self.record_for(session).await;
        if let Err(e) = self.persistent.save(&record).await {
            tracing::warn!(session_id = %session.id, error = %e, "Failed to save session record");
        }
    }

    async fn create(
        &self,
        router: McpRouter,
        service_factory: ServiceFactory,
    ) -> Option<Arc<Session>> {
        let session = {
            let mut sessions = self.sessions.write().await;

            // Check max sessions limit
            if let Some(max) = self.config.max_sessions
                && sessions.len() >= max
            {
                tracing::warn!(
                    max_sessions = max,
                    current = sessions.len(),
                    "Session limit reached, rejecting new session"
                );
                return None;
            }

            let session = Arc::new(Session::new(
                router,
                self.sampling_enabled,
                service_factory,
                self.events.clone(),
            ));
            sessions.insert(session.id.clone(), session.clone());
            tracing::debug!(session_id = %session.id, sampling = self.sampling_enabled, "Created new session");
            session
        };
        self.persist_new(&session).await;
        Some(session)
    }

    async fn create_from_service(&self, service: McpBoxService) -> Option<Arc<Session>> {
        let session = {
            let mut sessions = self.sessions.write().await;

            if let Some(max) = self.config.max_sessions
                && sessions.len() >= max
            {
                tracing::warn!(
                    max_sessions = max,
                    current = sessions.len(),
                    "Session limit reached, rejecting new session"
                );
                return None;
            }

            let session = Arc::new(Session::from_service(service, self.events.clone()));
            sessions.insert(session.id.clone(), session.clone());
            tracing::debug!(session_id = %session.id, "Created new session from service");
            session
        };
        self.persist_new(&session).await;
        Some(session)
    }

    /// Create a new session with its router already marked as initialized.
    ///
    /// Used by the optional-sessions feature to serve requests from clients
    /// that skip the initialize handshake.
    async fn create_initialized(
        &self,
        router: McpRouter,
        service_factory: ServiceFactory,
    ) -> Option<Arc<Session>> {
        // Pre-initialize the router's session state so it won't reject requests
        router.session().mark_initialized();

        let session = {
            let mut sessions = self.sessions.write().await;

            if let Some(max) = self.config.max_sessions
                && sessions.len() >= max
            {
                return None;
            }

            let session = Arc::new(Session::new(
                router,
                self.sampling_enabled,
                service_factory,
                self.events.clone(),
            ));
            sessions.insert(session.id.clone(), session.clone());
            tracing::debug!(session_id = %session.id, "Created pre-initialized session (optional_sessions)");
            session
        };
        self.persist_new(&session).await;
        Some(session)
    }

    /// Create a pre-initialized session from a boxed service.
    async fn create_initialized_from_service(
        &self,
        service: McpBoxService,
    ) -> Option<Arc<Session>> {
        let session = {
            let mut sessions = self.sessions.write().await;

            if let Some(max) = self.config.max_sessions
                && sessions.len() >= max
            {
                return None;
            }

            let session = Arc::new(Session::from_service(service, self.events.clone()));
            sessions.insert(session.id.clone(), session.clone());
            tracing::debug!(session_id = %session.id, "Created pre-initialized session from service (optional_sessions)");
            session
        };
        self.persist_new(&session).await;
        Some(session)
    }

    async fn get(&self, id: &str) -> Option<Arc<Session>> {
        // Fast path: the session is live in this process.
        {
            let sessions = self.sessions.read().await;
            if let Some(s) = sessions.get(id).cloned() {
                s.touch().await;
                return Some(s);
            }
        }

        // Slow path #1: the session is unknown locally but the persistent
        // store has a record — rebuild it.
        match self.persistent.load(id).await {
            Ok(Some(record)) => {
                tracing::info!(session_id = %id, "Restoring session from persistent store");
                if let Some(session) = self.restore_from_record(record).await {
                    return Some(session);
                }
            }
            Ok(None) => {}
            Err(e) => {
                tracing::warn!(session_id = %id, error = %e, "Failed to load session record");
            }
        }

        // Slow path #2 (opt-in): auto-reinitialize with synthetic client
        // info so the client can continue without a re-handshake. Useful
        // for single-instance restarts where no external store is
        // configured; loses original client identity.
        if self.auto_reinit {
            tracing::info!(session_id = %id, "Auto-reinitializing unknown session");
            return self.auto_reinitialize(id).await;
        }

        None
    }

    /// Restore a live [`Session`] from a persisted [`SessionRecord`].
    ///
    /// The caller must ensure the record's ID is not already live locally;
    /// on success the session is inserted into the local registry, the
    /// event counter is seeded so new event IDs don't collide with
    /// buffered ones, and the record's `last_accessed` is refreshed and
    /// saved back to the store.
    async fn restore_from_record(
        &self,
        record: crate::session_store::SessionRecord,
    ) -> Option<Arc<Session>> {
        let session = {
            let mut sessions = self.sessions.write().await;

            if let Some(max) = self.config.max_sessions
                && sessions.len() >= max
            {
                tracing::warn!(
                    max_sessions = max,
                    "Session limit reached, cannot restore session"
                );
                return None;
            }

            // Guard against a concurrent create that beat us here.
            if let Some(existing) = sessions.get(&record.id).cloned() {
                existing.touch().await;
                return Some(existing);
            }

            let session: Arc<Session> = match &self.service_source {
                ServiceSource::Router { router, factory } => Arc::new(Session::restored(
                    &record,
                    router.with_fresh_session(),
                    self.sampling_enabled,
                    factory.clone(),
                    self.events.clone(),
                )),
                ServiceSource::Service(svc) => {
                    let service = svc.lock().unwrap().clone();
                    Arc::new(Session::from_service_restored(
                        service,
                        &record,
                        self.events.clone(),
                    ))
                }
            };

            sessions.insert(record.id.clone(), session.clone());
            tracing::debug!(session_id = %session.id, "Restored session into local registry");
            session
        };

        // Seed the event counter past the highest buffered event ID so new
        // SSE events don't collide with ones the client may still replay.
        if let Ok(events) = self.events.replay_after(&record.id, 0).await
            && let Some(max_id) = events.iter().map(|e| e.id).max()
        {
            session
                .event_counter
                .store(max_id + 1, std::sync::atomic::Ordering::SeqCst);
        }

        // Refresh last_accessed in the store so the record doesn't expire
        // immediately after restore.
        let mut refreshed = record;
        refreshed.touch(self.config.ttl);
        if let Err(e) = self.persistent.save(&refreshed).await {
            tracing::warn!(session_id = %refreshed.id, error = %e, "Failed to refresh restored session record");
        }

        Some(session)
    }

    /// Create a new session with the requested ID and synthetic client
    /// info, skipping the initialize handshake. Used when `auto_reinit`
    /// is enabled and no stored record exists.
    ///
    /// Loses the original client's identity and capabilities — the server
    /// sees a session from client `"auto-recovered"`.
    async fn auto_reinitialize(&self, id: &str) -> Option<Arc<Session>> {
        let mut record = crate::session_store::SessionRecord::new(
            id.to_string(),
            LATEST_PROTOCOL_VERSION.to_string(),
            self.config.ttl,
        );
        record.client_info = Some(crate::protocol::Implementation {
            name: "auto-recovered".into(),
            version: "unknown".into(),
            title: None,
            description: None,
            icons: None,
            website_url: None,
            meta: None,
        });
        record.client_capabilities = Some(crate::protocol::ClientCapabilities::default());

        // Persist first so a concurrent request sees the record. Ignore
        // persistence errors; the in-memory session will still work.
        if let Err(e) = self.persistent.create(&mut record).await {
            tracing::warn!(session_id = %id, error = %e, "Failed to persist auto-reinitialized session");
        }

        self.restore_from_record(record).await
    }

    async fn remove(&self, id: &str) -> bool {
        let removed = {
            let mut sessions = self.sessions.write().await;
            sessions.remove(id).is_some()
        };
        if removed {
            tracing::debug!(session_id = %id, "Removed session");
            if let Err(e) = self.persistent.delete(id).await {
                tracing::warn!(session_id = %id, error = %e, "Failed to delete session record");
            }
            if let Err(e) = self.events.purge_session(id).await {
                tracing::warn!(session_id = %id, error = %e, "Failed to purge session events");
            }
        }
        removed
    }

    /// Send a pre-serialized JSON notification to every live session's SSE
    /// broadcast channel.
    ///
    /// Used by the external-notification fan-out task. Failures to send
    /// (no SSE subscribers attached to a session yet) are silent — the
    /// broadcast channel drops the message naturally.
    async fn broadcast_to_all(&self, json: &str) {
        let sessions = self.sessions.read().await;
        for session in sessions.values() {
            let _ = session.notifications_tx.send(json.to_string());
        }
    }

    /// Remove expired sessions, returns count of removed sessions
    async fn cleanup_expired(&self) -> usize {
        let expired = {
            let mut sessions = self.sessions.write().await;
            let ttl = self.config.ttl;

            let mut expired = Vec::new();
            for (id, session) in sessions.iter() {
                if session.is_expired(ttl).await {
                    expired.push(id.clone());
                }
            }

            for id in &expired {
                sessions.remove(id);
                tracing::debug!(session_id = %id, "Expired session removed");
            }

            if !expired.is_empty() {
                tracing::info!(
                    expired_count = expired.len(),
                    remaining = sessions.len(),
                    "Session cleanup completed"
                );
            }
            expired
        };

        for id in &expired {
            if let Err(e) = self.persistent.delete(id).await {
                tracing::warn!(session_id = %id, error = %e, "Failed to delete expired session record");
            }
            if let Err(e) = self.events.purge_session(id).await {
                tracing::warn!(session_id = %id, error = %e, "Failed to purge expired session events");
            }
        }

        expired.len()
    }
}

/// Metadata about an active session.
///
/// Returned by [`SessionHandle::list_sessions()`].
#[derive(Debug, Clone)]
pub struct SessionInfo {
    /// The session ID.
    pub id: String,
    /// How long ago this session was created.
    pub created_at: Duration,
    /// How long ago this session was last accessed.
    pub last_activity: Duration,
}

/// A handle for querying and managing HTTP transport sessions.
///
/// Obtained from [`HttpTransport::into_router_with_handle()`] or
/// [`HttpTransport::into_router_at_with_handle()`]. The handle is cheap to
/// clone and can be shared across threads.
///
/// # Example
///
/// ```rust,ignore
/// use tower_mcp::transport::http::HttpTransport;
///
/// let transport = HttpTransport::new(router);
/// let (router, handle) = transport.into_router_with_handle();
///
/// // Later, in an admin endpoint:
/// let count = handle.session_count().await;
/// for info in handle.list_sessions().await {
///     println!("{}: created {:?} ago", info.id, info.created_at);
/// }
/// handle.terminate_session("session-id").await;
/// ```
#[derive(Clone)]
pub struct SessionHandle {
    store: Arc<SessionRegistry>,
}

impl SessionHandle {
    /// Returns the number of currently active sessions.
    pub async fn session_count(&self) -> usize {
        self.store.sessions.read().await.len()
    }

    /// Returns metadata for all active sessions.
    pub async fn list_sessions(&self) -> Vec<SessionInfo> {
        let sessions = self.store.sessions.read().await;
        let mut infos = Vec::with_capacity(sessions.len());
        for session in sessions.values() {
            let last_accessed = session.last_accessed.read().await;
            infos.push(SessionInfo {
                id: session.id.clone(),
                created_at: session.created_at.elapsed(),
                last_activity: last_accessed.elapsed(),
            });
        }
        infos
    }

    /// Terminates a session by ID, returning `true` if the session existed.
    pub async fn terminate_session(&self, id: &str) -> bool {
        self.store.remove(id).await
    }
}

/// The source of the MCP service for session creation.
#[derive(Clone)]
enum ServiceSource {
    /// Created from an McpRouter with a factory for middleware wrapping.
    Router {
        router: McpRouter,
        factory: ServiceFactory,
    },
    /// Created from a pre-built boxed service (e.g., McpProxy).
    /// Wrapped in Arc<Mutex<_>> because BoxCloneService is Send but not Sync.
    Service(Arc<std::sync::Mutex<McpBoxService>>),
}

/// Shared state for the HTTP transport
struct AppState {
    /// Source for creating new session services
    service_source: ServiceSource,
    /// Session store
    sessions: Arc<SessionRegistry>,
    /// Whether to validate Origin header
    validate_origin: bool,
    /// Allowed origins (if validation is enabled)
    allowed_origins: Vec<String>,
    /// Whether to validate Host header (defense against direct DNS rebinding)
    validate_host: bool,
    /// Allowed hosts (host:port). Localhost variants are always allowed.
    allowed_hosts: Vec<String>,
    /// Whether sampling is enabled
    sampling_enabled: bool,
    /// Whether sessions are optional (for clients that don't track session IDs)
    optional_sessions: bool,
    /// SEP-1442 stateless mode configuration
    #[cfg(feature = "stateless")]
    stateless_config: Option<crate::stateless::StatelessConfig>,
}

/// Configuration for OAuth 2.1 Protected Resource Metadata.
///
/// When set on [`HttpTransport`], a `GET /.well-known/oauth-protected-resource`
/// endpoint is added that returns the metadata JSON, enabling OAuth client
/// discovery per RFC 9728.
#[cfg(feature = "oauth")]
#[derive(Clone)]
pub(crate) struct OAuthConfig {
    /// Protected Resource Metadata to serve at the well-known endpoint.
    pub(crate) metadata: crate::oauth::ProtectedResourceMetadata,
}

/// HTTP transport for MCP servers
///
/// Implements the Streamable HTTP transport from the MCP specification.
///
/// # Construction
///
/// There are two ways to create an `HttpTransport`:
///
/// - [`HttpTransport::new(router)`](HttpTransport::new) — wraps an [`McpRouter`], with full
///   support for per-session notification bridging, sampling, and `.layer()` middleware.
///
/// - [`HttpTransport::from_service(service)`](HttpTransport::from_service) — wraps any
///   `Service<RouterRequest>` (e.g., [`McpProxy`](crate::proxy::McpProxy)). The service is
///   cloned for each session. Notification bridging and sampling are not set up automatically;
///   the caller should configure these on the service before passing it in.
///   `.layer()` is not supported in this mode.
pub struct HttpTransport {
    service_source: ServiceSource,
    validate_origin: bool,
    allowed_origins: Vec<String>,
    validate_host: bool,
    allowed_hosts: Vec<String>,
    session_config: SessionConfig,
    sampling_enabled: bool,
    optional_sessions: bool,
    session_store: Arc<dyn crate::session_store::SessionStore>,
    event_store: Arc<dyn crate::event_store::EventStore>,
    auto_reinit_sessions: bool,
    /// Caller-owned receiver for notifications pushed from outside any
    /// request handler. Drained by a background task and fanned out to
    /// every live session's SSE stream.
    external_notifications: Option<NotificationReceiver>,
    #[cfg(feature = "stateless")]
    stateless_config: Option<crate::stateless::StatelessConfig>,
    #[cfg(feature = "oauth")]
    oauth_config: Option<OAuthConfig>,
}

impl HttpTransport {
    /// Create a new HTTP transport wrapping an MCP router.
    ///
    /// Supports per-session notification bridging, sampling, and `.layer()` middleware.
    pub fn new(router: McpRouter) -> Self {
        Self {
            service_source: ServiceSource::Router {
                router,
                factory: identity_factory(),
            },
            validate_origin: true,
            allowed_origins: vec![],
            validate_host: true,
            allowed_hosts: vec![],
            session_config: SessionConfig::default(),
            sampling_enabled: false,
            optional_sessions: true,
            session_store: Arc::new(crate::session_store::MemorySessionStore::new()),
            event_store: Arc::new(crate::event_store::MemoryEventStore::new()),
            auto_reinit_sessions: false,
            external_notifications: None,
            #[cfg(feature = "stateless")]
            stateless_config: None,
            #[cfg(feature = "oauth")]
            oauth_config: None,
        }
    }

    /// Create an HTTP transport from a pre-built service.
    ///
    /// This accepts any `Service<RouterRequest>` implementation, such as
    /// [`McpProxy`](crate::proxy::McpProxy). The service is cloned for each
    /// HTTP session.
    ///
    /// Notification bridging and sampling are **not** set up automatically.
    /// The caller should configure these on the service before passing it in.
    ///
    /// `.layer()` is not supported when using `from_service()` — wrap the
    /// service with middleware before passing it in.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use tower_mcp::transport::http::HttpTransport;
    /// use tower_mcp::proxy::McpProxy;
    ///
    /// let proxy: McpProxy = /* ... */;
    /// let transport = HttpTransport::from_service(proxy);
    /// transport.serve("127.0.0.1:3000").await?;
    /// ```
    pub fn from_service<S>(service: S) -> Self
    where
        S: tower::Service<
                RouterRequest,
                Response = RouterResponse,
                Error = std::convert::Infallible,
            > + Clone
            + Send
            + 'static,
        S::Future: Send,
    {
        Self {
            service_source: ServiceSource::Service(Arc::new(std::sync::Mutex::new(
                BoxCloneService::new(service),
            ))),
            validate_origin: true,
            allowed_origins: vec![],
            validate_host: true,
            allowed_hosts: vec![],
            session_config: SessionConfig::default(),
            sampling_enabled: false,
            optional_sessions: true,
            session_store: Arc::new(crate::session_store::MemorySessionStore::new()),
            event_store: Arc::new(crate::event_store::MemoryEventStore::new()),
            auto_reinit_sessions: false,
            external_notifications: None,
            #[cfg(feature = "stateless")]
            stateless_config: None,
            #[cfg(feature = "oauth")]
            oauth_config: None,
        }
    }

    /// Create an HTTP transport that drains a caller-owned notification
    /// channel and fans the items out to every live session's SSE stream.
    ///
    /// This mirrors [`GenericStdioTransport::with_notifications`](crate::transport::stdio::GenericStdioTransport::with_notifications)
    /// and is the supported way to push server-originated notifications
    /// (e.g. `notifications/resources/updated`) from outside any request
    /// handler — background tasks, lifecycle hooks, anything async that
    /// needs to notify subscribed clients.
    ///
    /// Per-session notification channels (in-handler `ctx.send_log()`,
    /// progress updates) are unaffected. The external channel runs in
    /// parallel and broadcasts to every active session; MCP clients are
    /// expected to ignore notifications they didn't subscribe to.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use tower_mcp::{BoxError, McpRouter};
    /// use tower_mcp::context::{ServerNotification, notification_channel};
    /// use tower_mcp::transport::http::HttpTransport;
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), BoxError> {
    ///     let (notif_tx, notif_rx) = notification_channel(256);
    ///
    ///     let router = McpRouter::new().server_info("my-server", "1.0.0");
    ///
    ///     // Hold onto notif_tx in your application state so background tasks
    ///     // can push notifications. tx is `Clone`.
    ///     let pusher = notif_tx.clone();
    ///     tokio::spawn(async move {
    ///         let _ = pusher.send(ServerNotification::ResourceUpdated {
    ///             uri: "claude://chats/123".to_string(),
    ///         }).await;
    ///     });
    ///
    ///     let transport = HttpTransport::with_notifications(router, notif_rx);
    ///     transport.serve("127.0.0.1:3000").await?;
    ///     Ok(())
    /// }
    /// ```
    pub fn with_notifications(router: McpRouter, notification_rx: NotificationReceiver) -> Self {
        Self {
            external_notifications: Some(notification_rx),
            ..Self::new(router)
        }
    }

    /// Attach a caller-owned notification receiver after construction.
    ///
    /// Useful when wrapping a pre-built service via
    /// [`from_service`](Self::from_service), where setting a sender on the
    /// router isn't part of the flow. See [`with_notifications`](Self::with_notifications)
    /// for the typical router-based path.
    pub fn external_notifications(mut self, notification_rx: NotificationReceiver) -> Self {
        self.external_notifications = Some(notification_rx);
        self
    }

    /// Enable sampling support for this transport.
    ///
    /// When sampling is enabled, tool handlers can use `ctx.sample()` to
    /// request LLM completions from connected clients. The server sends
    /// sampling requests on the SSE stream, and clients respond via POST.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use tower_mcp::{BoxError, McpRouter, ToolBuilder, CallToolResult, CreateMessageParams, SamplingMessage};
    /// use tower_mcp::extract::{Context, RawArgs};
    /// use tower_mcp::transport::http::HttpTransport;
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), BoxError> {
    ///     let tool = ToolBuilder::new("ai-tool")
    ///         .extractor_handler((), |ctx: Context, RawArgs(_): RawArgs| async move {
    ///             // Request LLM completion from client
    ///             let params = CreateMessageParams::new(
    ///                 vec![SamplingMessage::user("Summarize this...")],
    ///                 500,
    ///             );
    ///             let result = ctx.sample(params).await?;
    ///             Ok(CallToolResult::text(format!("{:?}", result.content)))
    ///         })
    ///         .build();
    ///
    ///     let router = McpRouter::new()
    ///         .server_info("my-server", "1.0.0")
    ///         .tool(tool);
    ///
    ///     let transport = HttpTransport::new(router).with_sampling();
    ///     transport.serve("127.0.0.1:3000").await?;
    ///     Ok(())
    /// }
    /// ```
    pub fn with_sampling(mut self) -> Self {
        self.sampling_enabled = true;
        self
    }

    /// Require strict session management.
    ///
    /// When enabled, requests without an `mcp-session-id` header are rejected
    /// with a `SessionRequired` error (-32006). Clients must complete the
    /// `initialize` handshake and include the session ID on all subsequent
    /// requests, as specified by the MCP 2025-11-25 spec.
    ///
    /// By default, sessions are optional for compatibility with clients
    /// (Codex CLI, Cursor, etc.) that don't carry the session ID forward
    /// after initialization.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use tower_mcp::McpRouter;
    /// use tower_mcp::transport::http::HttpTransport;
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
    ///     let router = McpRouter::new().server_info("my-server", "1.0.0");
    ///     let transport = HttpTransport::new(router).require_sessions();
    ///     transport.serve("127.0.0.1:3000").await?;
    ///     Ok(())
    /// }
    /// ```
    pub fn require_sessions(mut self) -> Self {
        self.optional_sessions = false;
        self
    }

    /// Enable the legacy SEP-1442 stateless opt-in path.
    ///
    /// This activates the SEP-1442-style stateless behavior for clients that
    /// do NOT send `MCP-Protocol-Version: 2026-07-28`. Specifically, when a
    /// [`crate::stateless::StatelessConfig`] is set:
    ///
    /// - Requests without a session ID can be served without an initialize
    ///   handshake (if [`crate::stateless::StatelessConfig::optional_sessions`]
    ///   is `true`).
    /// - The `server/discover` RPC is enabled (if
    ///   [`crate::stateless::StatelessConfig::enable_discover`] is `true`).
    /// - Protocol version may be required in every request body (if
    ///   [`crate::stateless::StatelessConfig::require_protocol_version`] is `true`).
    ///
    /// **Note:** this method does NOT control the automatic version-gated
    /// stateless path for 2026-07-28+ clients. When the `stateless` feature
    /// is compiled in, any request with `MCP-Protocol-Version: 2026-07-28`
    /// and no `mcp-session-id` is dispatched statelessly regardless of
    /// whether this method is called. See the [`crate::stateless`] module
    /// documentation for the full two-path explanation.
    ///
    /// Stateful clients (those that send `mcp-session-id`) continue to work
    /// normally on the same transport.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use tower_mcp::McpRouter;
    /// use tower_mcp::transport::http::HttpTransport;
    /// use tower_mcp::stateless::StatelessConfig;
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
    ///     let router = McpRouter::new().server_info("my-server", "1.0.0");
    ///     // Enables the SEP-1442 opt-in path. 2026-07-28 clients are
    ///     // handled statelessly regardless of this call.
    ///     let transport = HttpTransport::new(router)
    ///         .stateless(StatelessConfig::new());
    ///     transport.serve("127.0.0.1:3000").await?;
    ///     Ok(())
    /// }
    /// ```
    #[cfg(feature = "stateless")]
    pub fn stateless(mut self, config: crate::stateless::StatelessConfig) -> Self {
        self.stateless_config = Some(config);
        self
    }

    /// Disable Origin header validation (not recommended for production)
    pub fn disable_origin_validation(mut self) -> Self {
        self.validate_origin = false;
        self
    }

    /// Set allowed origins for CORS/security validation
    pub fn allowed_origins(mut self, origins: Vec<String>) -> Self {
        self.allowed_origins = origins;
        self
    }

    /// Disable Host header validation (not recommended when binding to a
    /// non-loopback interface).
    ///
    /// Host validation is the defense-in-depth pair to Origin validation: it
    /// rejects requests whose `Host` header doesn't match the server's
    /// expected hostname, blocking direct DNS-rebinding attacks where a
    /// malicious site resolves its own domain to `127.0.0.1`.
    pub fn disable_host_validation(mut self) -> Self {
        self.validate_host = false;
        self
    }

    /// Set allowed hosts for the `Host` header allowlist.
    ///
    /// Each entry should be a `host:port` pair (e.g. `"api.example.com"`,
    /// `"api.example.com:8443"`). Localhost variants (`localhost`,
    /// `127.0.0.1`, `::1`, with any port) are always accepted regardless
    /// of this list.
    ///
    /// When the `Host` header is missing, the validator falls back to the
    /// HTTP/2 `:authority` pseudo-header from `request.uri().authority()`,
    /// since middleware like `axum::Router::nest` can strip the synthesized
    /// `Host` header before it reaches our handler.
    pub fn allowed_hosts(mut self, hosts: Vec<String>) -> Self {
        self.allowed_hosts = hosts;
        self
    }

    /// Configure session management (TTL, max sessions, cleanup interval)
    pub fn session_config(mut self, config: SessionConfig) -> Self {
        self.session_config = config;
        self
    }

    /// Set session TTL (convenience method)
    pub fn session_ttl(mut self, ttl: Duration) -> Self {
        self.session_config.ttl = ttl;
        self
    }

    /// Set maximum number of concurrent sessions (convenience method)
    pub fn max_sessions(mut self, max: usize) -> Self {
        self.session_config.max_sessions = Some(max);
        self
    }

    /// Configure a pluggable [`SessionStore`](crate::session_store::SessionStore)
    /// for persisting session metadata.
    ///
    /// The default is an in-process
    /// [`MemorySessionStore`](crate::session_store::MemorySessionStore) —
    /// supply an external store (Redis, Postgres, etc.) to share session
    /// metadata across server instances behind a load balancer.
    ///
    /// Runtime state (broadcast channels, pending requests, service
    /// instances) is always kept per-instance; only persistent metadata is
    /// mirrored to the store.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use std::sync::Arc;
    /// use tower_mcp::{HttpTransport, McpRouter};
    /// use tower_mcp::session_store::{MemorySessionStore, SessionStore};
    ///
    /// let router = McpRouter::new();
    /// let store: Arc<dyn SessionStore> = Arc::new(MemorySessionStore::new());
    /// let transport = HttpTransport::new(router).session_store(store);
    /// ```
    pub fn session_store(mut self, store: Arc<dyn crate::session_store::SessionStore>) -> Self {
        self.session_store = store;
        self
    }

    /// Configure a pluggable [`EventStore`](crate::event_store::EventStore)
    /// for SSE event buffering and stream resumption.
    ///
    /// The default is an in-process
    /// [`MemoryEventStore`](crate::event_store::MemoryEventStore) with a
    /// 1000-event ring buffer per session — supply an external store (Redis,
    /// etc.) so clients can resume SSE streams after reconnecting to a
    /// different server instance behind a load balancer (SEP-1699).
    ///
    /// Typically paired with a matching
    /// [`session_store`](Self::session_store) so both session metadata and
    /// buffered events survive across instances.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use std::sync::Arc;
    /// use tower_mcp::{HttpTransport, McpRouter};
    /// use tower_mcp::event_store::{EventStore, MemoryEventStore};
    ///
    /// let router = McpRouter::new();
    /// let store: Arc<dyn EventStore> = Arc::new(MemoryEventStore::new());
    /// let transport = HttpTransport::new(router).event_store(store);
    /// ```
    pub fn event_store(mut self, store: Arc<dyn crate::event_store::EventStore>) -> Self {
        self.event_store = store;
        self
    }

    /// Enable auto-reinitialization for unknown session IDs.
    ///
    /// When a request arrives with an `mcp-session-id` that is not live
    /// locally and has no record in the configured
    /// [`session_store`](Self::session_store), the transport normally
    /// returns a session-not-found error. With this flag enabled, the
    /// transport instead spins up a new session claiming that ID and
    /// completes the initialize handshake internally with synthetic
    /// client info (`name = "auto-recovered"`, empty capabilities).
    ///
    /// This lets tolerant clients continue after a server restart without
    /// repeating the handshake, at the cost of losing the original
    /// client's identity and negotiated capabilities. Prefer pairing this
    /// with a real [`session_store`](Self::session_store) — the store
    /// path runs first and preserves full identity when a record exists.
    ///
    /// Disabled by default. This is the pattern established by
    /// [anubis-mcp #125](https://github.com/zoedsoupe/anubis-mcp/pull/125).
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use tower_mcp::{HttpTransport, McpRouter};
    ///
    /// let router = McpRouter::new();
    /// let transport = HttpTransport::new(router).auto_reinitialize_sessions(true);
    /// ```
    pub fn auto_reinitialize_sessions(mut self, enabled: bool) -> Self {
        self.auto_reinit_sessions = enabled;
        self
    }

    /// Configure OAuth 2.1 Protected Resource Metadata for this transport.
    ///
    /// When set, adds a `GET /.well-known/oauth-protected-resource` endpoint
    /// that returns the metadata JSON, enabling OAuth client discovery per
    /// RFC 9728.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use tower_mcp::oauth::ProtectedResourceMetadata;
    /// use tower_mcp::transport::http::HttpTransport;
    /// use tower_mcp::McpRouter;
    ///
    /// let metadata = ProtectedResourceMetadata::new("https://mcp.example.com")
    ///     .authorization_server("https://auth.example.com")
    ///     .scope("mcp:read");
    ///
    /// let router = McpRouter::new().server_info("my-server", "1.0.0");
    /// let transport = HttpTransport::new(router).oauth(metadata);
    /// ```
    #[cfg(feature = "oauth")]
    pub fn oauth(mut self, metadata: crate::oauth::ProtectedResourceMetadata) -> Self {
        self.oauth_config = Some(OAuthConfig { metadata });
        self
    }

    /// Apply a tower middleware layer to MCP request processing.
    ///
    /// # Panics
    ///
    /// Panics if this transport was created via [`from_service()`](Self::from_service).
    /// When using `from_service()`, wrap the service with middleware before passing it in.
    pub fn layer<L>(mut self, layer: L) -> Self
    where
        L: tower::Layer<McpRouter> + Send + Sync + 'static,
        L::Service:
            tower::Service<RouterRequest, Response = RouterResponse> + Clone + Send + 'static,
        <L::Service as tower::Service<RouterRequest>>::Error: std::fmt::Display + Send,
        <L::Service as tower::Service<RouterRequest>>::Future: Send,
    {
        match &mut self.service_source {
            ServiceSource::Router { factory, .. } => {
                *factory = Arc::new(move |router: McpRouter| {
                    let annotations = router.tool_annotations_map();
                    let wrapped = layer.layer(router);
                    tower::util::BoxCloneService::new(InjectAnnotations::new(
                        CatchError::new(wrapped),
                        annotations,
                    ))
                });
            }
            ServiceSource::Service(_) => {
                panic!(
                    "layer() cannot be used with from_service() — \
                     wrap the service with middleware before passing it in"
                );
            }
        }
        self
    }

    fn build_state(&self) -> Arc<AppState> {
        let sessions = Arc::new(SessionRegistry::new(
            self.session_config.clone(),
            self.sampling_enabled,
            self.session_store.clone(),
            self.event_store.clone(),
            self.service_source.clone(),
            self.auto_reinit_sessions,
        ));

        // Spawn cleanup task
        let cleanup_sessions = sessions.clone();
        let cleanup_interval = self.session_config.cleanup_interval;
        tokio::spawn(async move {
            loop {
                tokio::time::sleep(cleanup_interval).await;
                cleanup_sessions.cleanup_expired().await;
            }
        });

        Arc::new(AppState {
            service_source: self.service_source.clone(),
            sessions,
            validate_origin: self.validate_origin,
            allowed_origins: self.allowed_origins.clone(),
            validate_host: self.validate_host,
            allowed_hosts: self.allowed_hosts.clone(),
            sampling_enabled: self.sampling_enabled,
            optional_sessions: self.optional_sessions,
            #[cfg(feature = "stateless")]
            stateless_config: self.stateless_config.clone(),
        })
    }

    /// Build the axum router for this transport.
    pub fn into_router(self) -> Router {
        let (router, _handle) = self.into_router_with_handle();
        router
    }

    /// Build the axum router and return a [`SessionHandle`] for querying
    /// session metrics (e.g., active session count).
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let transport = HttpTransport::new(router);
    /// let (router, handle) = transport.into_router_with_handle();
    ///
    /// // Use handle in an admin endpoint
    /// let count = handle.session_count().await;
    /// ```
    pub fn into_router_with_handle(mut self) -> (Router, SessionHandle) {
        let external_rx = self.external_notifications.take();
        let state = self.build_state();
        let handle = SessionHandle {
            store: state.sessions.clone(),
        };

        spawn_external_notification_fanout(external_rx, state.sessions.clone());

        let router = Router::new()
            .route("/", post(handle_post))
            .route("/", get(handle_get))
            .route("/", delete(handle_delete))
            .route("/health", get(handle_health))
            .with_state(state);

        #[cfg(feature = "oauth")]
        let router = self.add_oauth_route(router, "");

        (router, handle)
    }

    /// Build an axum router mounted at a specific path.
    pub fn into_router_at(self, path: &str) -> Router {
        let (router, _handle) = self.into_router_at_with_handle(path);
        router
    }

    /// Build an axum router mounted at a specific path and return a
    /// [`SessionHandle`] for querying session metrics.
    pub fn into_router_at_with_handle(mut self, path: &str) -> (Router, SessionHandle) {
        let external_rx = self.external_notifications.take();
        let state = self.build_state();
        let handle = SessionHandle {
            store: state.sessions.clone(),
        };

        spawn_external_notification_fanout(external_rx, state.sessions.clone());

        let mcp_router = Router::new()
            .route("/", post(handle_post))
            .route("/", get(handle_get))
            .route("/", delete(handle_delete))
            .route("/health", get(handle_health))
            .with_state(state);

        let router = Router::new().nest(path, mcp_router);

        #[cfg(feature = "oauth")]
        let router = self.add_oauth_route(router, path);

        (router, handle)
    }

    /// Serve the transport on the given address
    ///
    /// This is a convenience method that creates a TCP listener and serves the transport.
    pub async fn serve(self, addr: &str) -> Result<()> {
        let listener = tokio::net::TcpListener::bind(addr)
            .await
            .map_err(|e| Error::Transport(format!("Failed to bind to {}: {}", addr, e)))?;

        tracing::info!("MCP HTTP transport listening on {}", addr);

        let router = self.into_router();
        axum::serve(listener, router)
            .await
            .map_err(|e| Error::Transport(format!("Server error: {}", e)))?;

        Ok(())
    }

    /// Add the OAuth Protected Resource Metadata well-known route if configured.
    #[cfg(feature = "oauth")]
    fn add_oauth_route(&self, router: Router, base_path: &str) -> Router {
        if let Some(ref config) = self.oauth_config {
            let metadata = config.metadata.clone();
            let well_known_path = if base_path.is_empty() {
                crate::oauth::ProtectedResourceMetadata::well_known_path().to_string()
            } else {
                format!(
                    "{}{}",
                    base_path.trim_end_matches('/'),
                    crate::oauth::ProtectedResourceMetadata::well_known_path()
                )
            };
            router.route(
                &well_known_path,
                get(move || {
                    let m = metadata.clone();
                    async move { axum::Json(m) }
                }),
            )
        } else {
            router
        }
    }
}

/// Check if an origin is a localhost origin (safe from DNS rebinding).
/// Drain a caller-supplied notification channel and fan items out to every
/// live session's SSE broadcast.
///
/// No-op when `rx` is `None`. When present, spawns a long-running task that
/// runs for the lifetime of the transport (until the channel closes).
fn spawn_external_notification_fanout(
    rx: Option<NotificationReceiver>,
    sessions: Arc<SessionRegistry>,
) {
    let Some(mut rx) = rx else {
        return;
    };
    tokio::spawn(async move {
        while let Some(notification) = rx.recv().await {
            if let Some(json) = crate::transport::stdio::serialize_notification(&notification) {
                sessions.broadcast_to_all(&json).await;
            }
        }
        tracing::debug!("External notification channel closed; fan-out task exiting");
    });
}

fn is_localhost_origin(origin: &str) -> bool {
    // Parse the origin to extract the host
    if let Some(rest) = origin
        .strip_prefix("http://")
        .or_else(|| origin.strip_prefix("https://"))
    {
        is_localhost_host(rest)
    } else {
        false
    }
}

/// Check if a `host:port` (or `[ipv6]:port`) value refers to localhost.
///
/// Used by both Origin validation (after stripping the `http(s)://` scheme)
/// and Host validation (where there's no scheme to begin with).
fn is_localhost_host(host: &str) -> bool {
    let host_only = if host.starts_with('[') {
        // Bracketed IPv6: [::1]:3000 -> ::1
        host.split(']')
            .next()
            .unwrap_or(host)
            .trim_start_matches('[')
    } else {
        // Strip port if present
        host.split(':').next().unwrap_or(host)
    };
    matches!(host_only, "localhost" | "127.0.0.1" | "::1")
}

/// Resolve the effective host for validation.
///
/// Prefers the `Host` header, falling back to the HTTP/2 `:authority`
/// pseudo-header (`request.uri().authority()`) when the header is missing.
/// This matters behind middleware like `axum::Router::nest`, which can
/// strip Hyper's synthesized `Host` before our handler sees it.
fn effective_host<'a>(headers: &'a HeaderMap, uri: &'a axum::http::Uri) -> Option<&'a str> {
    if let Some(value) = headers.get(header::HOST)
        && let Ok(s) = value.to_str()
    {
        return Some(s);
    }
    uri.authority().map(|a| a.as_str())
}

/// Validate the `Host` header (defense-in-depth alongside Origin).
///
/// Returns Some(Response) if validation fails, None if it passes.
fn validate_host(headers: &HeaderMap, uri: &axum::http::Uri, state: &AppState) -> Option<Response> {
    if !state.validate_host {
        return None;
    }

    let Some(host) = effective_host(headers, uri) else {
        if state.allowed_hosts.is_empty() {
            // No Host header and no allowlist: fall back to permissive
            // behavior matching pre-validation defaults so we don't break
            // existing deployments. (Origin already protects browsers.)
            return None;
        }
        tracing::warn!("Rejecting request: missing Host header and no :authority fallback");
        return Some((StatusCode::BAD_REQUEST, "Missing Host header").into_response());
    };

    if is_localhost_host(host) {
        return None;
    }

    if state.allowed_hosts.is_empty() {
        // Non-localhost host with no explicit allowlist: keep accepting it.
        // Operators who want strict Host validation must opt in via
        // `.allowed_hosts(...)`. This preserves the historical behavior of
        // not enforcing Host on non-loopback deployments by default.
        return None;
    }

    if state.allowed_hosts.iter().any(|h| h == host) {
        return None;
    }

    tracing::warn!(host = %host, "Rejecting request: Host not in allowlist");
    Some((StatusCode::BAD_REQUEST, "Host not allowed").into_response())
}

/// Validate Origin header for security.
///
/// When origin validation is enabled:
/// - Requests without an Origin header are allowed (same-origin)
/// - Localhost origins are always allowed (DNS rebinding protection)
/// - If `allowed_origins` is non-empty, non-localhost origins must match
/// - If `allowed_origins` is empty, non-localhost origins are rejected
///
/// Returns Some(Response) if validation fails, None if it passes.
fn validate_origin(headers: &HeaderMap, state: &AppState) -> Option<Response> {
    if !state.validate_origin {
        return None;
    }

    if let Some(origin) = headers.get(header::ORIGIN) {
        let origin_str = origin.to_str().unwrap_or("");

        // Always allow localhost origins (DNS rebinding protection allows these)
        if is_localhost_origin(origin_str) {
            return None;
        }

        // Non-localhost origin: check against allowed list
        if state.allowed_origins.is_empty() {
            tracing::warn!(
                origin = %origin_str,
                "Rejecting request: cross-origin not allowed (no allowlist configured)"
            );
            return Some(
                (StatusCode::FORBIDDEN, "Cross-origin requests not allowed").into_response(),
            );
        }

        if !state
            .allowed_origins
            .iter()
            .any(|o| o == origin_str || o == "*")
        {
            tracing::warn!(origin = %origin_str, "Rejecting request: Origin not in allowlist");
            return Some((StatusCode::FORBIDDEN, "Origin not allowed").into_response());
        }
    }

    None
}

/// Extract and validate session ID from headers
fn get_session_id(headers: &HeaderMap) -> Option<String> {
    headers
        .get(MCP_SESSION_ID_HEADER)
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string())
}

/// Extract protocol version from headers
fn get_protocol_version(headers: &HeaderMap) -> Option<String> {
    headers
        .get(MCP_PROTOCOL_VERSION_HEADER)
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string())
}

/// Extract Last-Event-ID from headers for SSE stream resumption (SEP-1699)
fn get_last_event_id(headers: &HeaderMap) -> Option<u64> {
    headers
        .get(LAST_EVENT_ID_HEADER)
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.parse::<u64>().ok())
}

/// Check if the request is an initialize request
fn is_initialize_request(body: &serde_json::Value) -> bool {
    body.get("method")
        .and_then(|m| m.as_str())
        .map(|m| m == "initialize")
        .unwrap_or(false)
}

/// Check if this is a response to one of our outgoing requests
fn is_response(parsed: &serde_json::Value) -> bool {
    parsed.get("method").is_none()
        && (parsed.get("result").is_some() || parsed.get("error").is_some())
}

/// Extract request ID from a JSON value
fn extract_request_id(parsed: &serde_json::Value) -> Option<RequestId> {
    parsed.get("id").and_then(|id| {
        if let Some(n) = id.as_i64() {
            Some(RequestId::Number(n))
        } else {
            id.as_str().map(|s| RequestId::String(s.to_string()))
        }
    })
}

/// Handle POST requests (JSON-RPC messages from client)
async fn handle_post(
    State(state): State<Arc<AppState>>,
    request: axum::extract::Request,
) -> Response {
    let (parts, body_bytes) = request.into_parts();
    let headers = parts.headers;
    let uri = parts.uri.clone();

    // Validate Host (DNS rebinding defense, complement to Origin)
    if let Some(resp) = validate_host(&headers, &uri, &state) {
        return resp;
    }

    // Validate Origin
    if let Some(resp) = validate_origin(&headers, &state) {
        return resp;
    }

    let body = match axum::body::to_bytes(body_bytes, usize::MAX).await {
        Ok(bytes) => match String::from_utf8(bytes.to_vec()) {
            Ok(s) => s,
            Err(e) => {
                return json_rpc_error_response(
                    None,
                    JsonRpcError::parse_error(format!("Invalid UTF-8: {}", e)),
                );
            }
        },
        Err(e) => {
            return json_rpc_error_response(
                None,
                JsonRpcError::parse_error(format!("Failed to read body: {}", e)),
            );
        }
    };

    // Bridge TokenClaims from HTTP extensions to MCP extensions (if present)
    #[cfg(feature = "oauth")]
    let http_extensions = parts.extensions;
    #[cfg(not(feature = "oauth"))]
    let _ = parts.extensions;

    // Parse the request body
    let parsed: serde_json::Value = match serde_json::from_str(&body) {
        Ok(v) => v,
        Err(e) => {
            return json_rpc_error_response(
                None,
                JsonRpcError::parse_error(format!("Invalid JSON: {}", e)),
            );
        }
    };

    // Check if this is an initialize request (creates new session)
    let is_init = is_initialize_request(&parsed);

    // SEP-2575 / SEP-2567: version-gated stateless mode for 2026-07-28+ clients.
    //
    // When the requested (or carried) protocol version is >= 2026-07-28 and the
    // request has no mcp-session-id, every request -- including `initialize` --
    // is served without creating or looking up a session. Each request is fully
    // self-contained; client identity and capabilities flow through per-request
    // `_meta` rather than a session handshake.
    //
    // This block runs before the legacy SEP-1442 stateless path so that
    // 2026-07-28 requests are handled here regardless of whether
    // `stateless_config` is set on the transport.
    #[cfg(feature = "stateless")]
    {
        let version_in_play: Option<String> = if is_init {
            // For `initialize`, read the version the client is requesting from
            // the params object.
            parsed
                .get("params")
                .and_then(|p| p.get("protocolVersion"))
                .and_then(|v| v.as_str())
                .map(|s| s.to_string())
        } else {
            // For non-init requests, only the HTTP-level `MCP-Protocol-Version`
            // header gates stateless mode. Body-level `_meta.protocolVersion` is
            // plumbed to handlers via `stash_per_request_meta` in both paths.
            get_protocol_version(&headers)
        };

        if let Some(ref version) = version_in_play
            && is_stateless_protocol_version(version)
            && get_session_id(&headers).is_none()
            // `messages/listen` opens an SSE stream; let it fall through to the
            // dedicated intercept below rather than handling it as a plain RPC call.
            && parsed.get("method").and_then(|m| m.as_str()) != Some("messages/listen")
        {
            // Notifications and responses are fire-and-forget; no dispatch needed.
            if !is_init && (parsed.get("id").is_none() || is_response(&parsed)) {
                return StatusCode::ACCEPTED.into_response();
            }

            // SEP-2243 validation before `parsed` is consumed by deserialization.
            // 2026-07-28 falls into strict mode, so missing Mcp-Method is an error.
            let sep_2243_mode = super::http_headers::mode_for_version(version);
            if let Err(err) = super::http_headers::validate(&headers, &parsed, sep_2243_mode) {
                tracing::warn!(
                    mode = ?sep_2243_mode,
                    version = %version,
                    error = %err.message,
                    "Rejecting stateless request: SEP-2243 header validation failed",
                );
                let id = extract_request_id(&parsed);
                let mut resp = json_rpc_error_response(id, err);
                *resp.status_mut() = StatusCode::BAD_REQUEST;
                return resp;
            }

            let request: JsonRpcRequest = match serde_json::from_value(parsed) {
                Ok(r) => r,
                Err(e) => {
                    return json_rpc_error_response(
                        None,
                        JsonRpcError::parse_error(format!("Invalid request: {}", e)),
                    );
                }
            };

            // Ephemeral pre-initialized service -- no session is stored or created.
            let mut service = match &state.service_source {
                ServiceSource::Router { router, factory } => {
                    let ephemeral = router.with_fresh_session();
                    ephemeral.session().mark_initialized();
                    JsonRpcService::new(factory(ephemeral))
                }
                ServiceSource::Service(mutex) => JsonRpcService::new(mutex.lock().unwrap().clone()),
            };

            let mut ext = crate::router::Extensions::new();
            #[cfg(feature = "oauth")]
            if let Some(claims) = http_extensions.get::<crate::oauth::token::TokenClaims>() {
                ext.insert(claims.clone());
            }
            stash_per_request_meta(&request, &mut ext);
            if !ext.is_empty() {
                service = service.with_extensions(ext);
            }

            let mut response = match service.call_single(request).await {
                Ok(resp) => resp,
                Err(e) => {
                    return json_rpc_error_response(
                        None,
                        JsonRpcError::internal_error(e.to_string()),
                    );
                }
            };

            // For `initialize`: the router's version negotiation falls back to
            // `LATEST_PROTOCOL_VERSION` ("2025-11-25") because 2026-07-28 is not
            // yet in `SUPPORTED_PROTOCOL_VERSIONS` at the types layer. Patch the
            // response to reflect the version the transport is actually serving so
            // the client sees the version it requested.
            if is_init
                && let JsonRpcResponse::Result(ref mut result) = response
                && let Some(pv) = result.result.get_mut("protocolVersion")
            {
                *pv = serde_json::Value::String(version.clone());
            }

            let mut resp = axum::Json(response).into_response();
            resp.headers_mut().insert(
                MCP_PROTOCOL_VERSION_HEADER,
                HeaderValue::from_str(version).unwrap(),
            );
            // Intentionally NO `mcp-session-id` header for 2026-07-28+ clients.
            return resp;
        }
    }

    // SEP-1442: Handle stateless requests (no session needed).
    // Stateless requests have a protocol version but no session ID and are not
    // initialize requests. They are processed with an ephemeral service and
    // return immediately without storing any session state.
    #[cfg(feature = "stateless")]
    if !is_init && state.stateless_config.is_some() && get_session_id(&headers).is_none() {
        let version_from_header = get_protocol_version(&headers);
        let params = parsed.get("params").unwrap_or(&parsed);
        let version_from_meta = crate::stateless::StatelessRequestMeta::from_params(params)
            .and_then(|m| m.protocol_version);

        if let Some(version) = version_from_header.or(version_from_meta) {
            if let Err(err) = crate::stateless::validate_protocol_version(&version) {
                return json_rpc_error_response(None, err);
            }

            // Notifications and responses don't make sense without a session
            if parsed.get("id").is_none() || is_response(&parsed) {
                return StatusCode::ACCEPTED.into_response();
            }

            let request: JsonRpcRequest = match serde_json::from_value(parsed) {
                Ok(r) => r,
                Err(e) => {
                    return json_rpc_error_response(
                        None,
                        JsonRpcError::parse_error(format!("Invalid request: {}", e)),
                    );
                }
            };

            // Ephemeral pre-initialized service -- no session stored
            let mut service = match &state.service_source {
                ServiceSource::Router { router, factory } => {
                    let ephemeral = router.with_fresh_session();
                    ephemeral.session().mark_initialized();
                    JsonRpcService::new(factory(ephemeral))
                }
                ServiceSource::Service(mutex) => JsonRpcService::new(mutex.lock().unwrap().clone()),
            };

            let mut ext = crate::router::Extensions::new();
            #[cfg(feature = "oauth")]
            if let Some(claims) = http_extensions.get::<crate::oauth::token::TokenClaims>() {
                ext.insert(claims.clone());
            }
            #[cfg(feature = "stateless")]
            stash_per_request_meta(&request, &mut ext);
            if !ext.is_empty() {
                service = service.with_extensions(ext);
            }

            let response = match service.call_single(request).await {
                Ok(resp) => resp,
                Err(e) => {
                    return json_rpc_error_response(
                        None,
                        JsonRpcError::internal_error(e.to_string()),
                    );
                }
            };

            let mut resp = axum::Json(response).into_response();
            resp.headers_mut().insert(
                MCP_PROTOCOL_VERSION_HEADER,
                HeaderValue::from_str(&version).unwrap(),
            );
            return resp;
        }
    }

    // Get or create session
    let session = if is_init {
        // Create new session for initialize
        let create_result = match &state.service_source {
            ServiceSource::Router { router, factory } => {
                // Use with_fresh_session() to ensure each session has its own state
                state
                    .sessions
                    .create(router.with_fresh_session(), factory.clone())
                    .await
            }
            ServiceSource::Service(mutex) => {
                let service = mutex.lock().unwrap().clone();
                state.sessions.create_from_service(service).await
            }
        };
        match create_result {
            Some(s) => s,
            None => {
                return (
                    StatusCode::SERVICE_UNAVAILABLE,
                    "Maximum session limit reached",
                )
                    .into_response();
            }
        }
    } else if let Some(session_id) = get_session_id(&headers) {
        // Client sent a session ID -- look it up
        match state.sessions.get(&session_id).await {
            Some(s) => s,
            None => {
                // Return JSON-RPC error with session info so clients know to re-initialize
                return json_rpc_error_response(
                    None,
                    JsonRpcError::session_not_found_with_id(&session_id),
                );
            }
        }
    } else if state.optional_sessions {
        // No session ID, but sessions are optional -- create a transient,
        // pre-initialized session so the router won't reject the request.
        // This supports clients (Codex CLI, Cursor, etc.) that perform
        // initialize + tools/list during setup but don't carry the session
        // ID forward to subsequent requests.
        let create_result = match &state.service_source {
            ServiceSource::Router { router, factory } => {
                state
                    .sessions
                    .create_initialized(router.with_fresh_session(), factory.clone())
                    .await
            }
            ServiceSource::Service(mutex) => {
                let service = mutex.lock().unwrap().clone();
                state
                    .sessions
                    .create_initialized_from_service(service)
                    .await
            }
        };
        match create_result {
            Some(s) => s,
            None => {
                return (
                    StatusCode::SERVICE_UNAVAILABLE,
                    "Maximum session limit reached",
                )
                    .into_response();
            }
        }
    } else {
        // No session ID and sessions are required
        return json_rpc_error_response(None, JsonRpcError::session_required());
    };

    // SEP-2575 / SEP-2567: intercept `messages/listen` before the standard
    // version validation. `messages/listen` is only available when the
    // effective protocol version is >= 2026-07-28; otherwise we return a
    // proper JSON-RPC error rather than silently falling through to the
    // router (which would return `MethodNotFound` anyway, but without the
    // protocol-version context).
    //
    // We check the Mcp-Protocol-Version header first (per-request override),
    // falling back to the session-negotiated version. Intercepting here
    // also prevents the version-validation guard below from rejecting the
    // 2026-07-28 header before we can inspect it.
    {
        let method_str = parsed.get("method").and_then(|m| m.as_str()).unwrap_or("");
        if method_str == "messages/listen" {
            let req_id = extract_request_id(&parsed);
            let effective_version = if let Some(v) = get_protocol_version(&headers) {
                v
            } else {
                session.protocol_version.read().await.clone()
            };
            if version_supports_messages_listen(&effective_version) {
                return handle_messages_listen_sse(session).await;
            } else {
                return json_rpc_error_response(
                    req_id,
                    JsonRpcError::method_not_found("messages/listen"),
                );
            }
        }
    }

    // Validate protocol version (if present and not init request).
    // Per SEP-2575, unsupported versions get a JSON-RPC error with code
    // -32004 and `{ supported, requested }` data, not a plain-text 400.
    if !is_init
        && let Some(version) = get_protocol_version(&headers)
        && !SUPPORTED_PROTOCOL_VERSIONS.contains(&version.as_str())
    {
        let id = extract_request_id(&parsed);
        return json_rpc_error_response(
            id,
            JsonRpcError::unsupported_protocol_version(
                version,
                SUPPORTED_PROTOCOL_VERSIONS.iter().copied(),
            ),
        );
    }

    // SEP-2243: validate the standardized HTTP headers (Mcp-Method,
    // Mcp-Name, Mcp-Param-*) against the body. Mode is "strict" only
    // when the negotiated protocol version is at or beyond the
    // SEP-2243-inclusion version; otherwise present headers are still
    // checked for body consistency but missing headers are allowed.
    //
    // For `initialize` requests the session's protocol version hasn't
    // been negotiated yet, so we fall back to the version the client
    // requested in the body. For all other requests we use the session's
    // negotiated version (which is also reflected back in the response
    // `Mcp-Protocol-Version` header).
    let sep_2243_version = if is_init {
        match parsed
            .get("params")
            .and_then(|p| p.get("protocolVersion"))
            .and_then(|v| v.as_str())
        {
            Some(v) => v.to_string(),
            None => session.protocol_version.read().await.clone(),
        }
    } else {
        session.protocol_version.read().await.clone()
    };
    let sep_2243_mode = super::http_headers::mode_for_version(&sep_2243_version);
    if let Err(err) = super::http_headers::validate(&headers, &parsed, sep_2243_mode) {
        tracing::warn!(
            mode = ?sep_2243_mode,
            version = %sep_2243_version,
            error = %err.message,
            "Rejecting request: SEP-2243 header validation failed",
        );
        let id = extract_request_id(&parsed);
        let mut resp = json_rpc_error_response(id, err);
        // Per SEP-2243 §"Error Code" the HTTP status MUST be 400.
        *resp.status_mut() = StatusCode::BAD_REQUEST;
        return resp;
    }

    // Check if this is a response to one of our outgoing requests (sampling)
    if is_response(&parsed) {
        if let Some(id) = extract_request_id(&parsed) {
            let result = if let Some(error) = parsed.get("error") {
                let code = error.get("code").and_then(|c| c.as_i64()).unwrap_or(-1);
                let message = error
                    .get("message")
                    .and_then(|m| m.as_str())
                    .unwrap_or("Unknown error");
                Err(Error::Internal(format!(
                    "Client error ({}): {}",
                    code, message
                )))
            } else if let Some(result) = parsed.get("result") {
                Ok(result.clone())
            } else {
                Err(Error::Internal(
                    "Response has neither result nor error".to_string(),
                ))
            };

            if session.complete_pending_request(&id, result).await {
                tracing::debug!(request_id = ?id, "Completed pending request");
            } else {
                tracing::warn!(request_id = ?id, "Received response for unknown request");
            }
        }
        return StatusCode::ACCEPTED.into_response();
    }

    // Check if this is a notification (no id field)
    if parsed.get("id").is_none() {
        // Handle notification
        if let Ok(notification) = serde_json::from_value::<JsonRpcNotification>(parsed)
            && let Ok(mcp_notification) = McpNotification::from_jsonrpc(&notification)
        {
            session.handle_notification(mcp_notification);
        }
        return StatusCode::ACCEPTED.into_response();
    }

    // For initialize requests, capture the advertised client info /
    // capabilities from the raw params before `parsed` is consumed by
    // deserialization. These are stashed onto the live `Session` after a
    // successful initialize so the persisted SessionRecord faithfully
    // describes the client (rather than carrying the defaults set at
    // session-create time).
    let init_client_metadata: Option<(Option<Implementation>, Option<ClientCapabilities>)> =
        if is_init {
            let params = parsed.get("params");
            let client_info = params
                .and_then(|p| p.get("clientInfo"))
                .and_then(|v| serde_json::from_value::<Implementation>(v.clone()).ok());
            let client_capabilities = params
                .and_then(|p| p.get("capabilities"))
                .and_then(|v| serde_json::from_value::<ClientCapabilities>(v.clone()).ok());
            Some((client_info, client_capabilities))
        } else {
            None
        };

    // Handle as JSON-RPC request
    let request: JsonRpcRequest = match serde_json::from_value(parsed) {
        Ok(r) => r,
        Err(e) => {
            return json_rpc_error_response(
                None,
                JsonRpcError::parse_error(format!("Invalid request: {}", e)),
            );
        }
    };

    // Process the request through the middleware-wrapped service
    let mut service = JsonRpcService::new(session.make_service());

    // Bridge per-request data from HTTP into MCP Extensions: OAuth claims,
    // SEP-2575 `_meta` (clientInfo, clientCapabilities, etc.). Empty ext is
    // skipped to avoid pointless allocation.
    #[allow(unused_mut)]
    let mut ext = crate::router::Extensions::new();
    #[cfg(feature = "oauth")]
    if let Some(claims) = http_extensions.get::<crate::oauth::token::TokenClaims>() {
        ext.insert(claims.clone());
    }
    #[cfg(feature = "stateless")]
    stash_per_request_meta(&request, &mut ext);
    if !ext.is_empty() {
        service = service.with_extensions(ext);
    }
    let response = match service.call_single(request).await {
        Ok(resp) => resp,
        Err(e) => {
            return json_rpc_error_response(None, JsonRpcError::internal_error(e.to_string()));
        }
    };

    // For successful initialize responses, extract and store the negotiated
    // protocol version, stash the client's advertised identity / capabilities
    // on the live session, and persist the now-complete record to the session
    // store so a restore from a peer instance sees the original client info
    // instead of defaults.
    if is_init && let JsonRpcResponse::Result(ref result) = response {
        if let Some(version) = result
            .result
            .get("protocolVersion")
            .and_then(|v| v.as_str())
        {
            *session.protocol_version.write().await = version.to_string();
        }
        if let Some((client_info, client_capabilities)) = init_client_metadata {
            *session.client_info.write().await = client_info;
            *session.client_capabilities.write().await = client_capabilities;
        }
        state.sessions.save_record(&session).await;
    }

    // Build response with headers
    let mut resp = axum::Json(response).into_response();

    if is_init {
        resp.headers_mut().insert(
            MCP_SESSION_ID_HEADER,
            HeaderValue::from_str(&session.id).unwrap(),
        );
    }

    // Always include the negotiated protocol version header
    let version = session.protocol_version.read().await;
    resp.headers_mut().insert(
        MCP_PROTOCOL_VERSION_HEADER,
        HeaderValue::from_str(&version).unwrap(),
    );

    resp
}

/// Returns `true` when the given protocol version string enables `messages/listen`.
///
/// `messages/listen` is part of the 2026-07-28 spec (SEP-2575 / SEP-2567).
/// Version strings are YYYY-MM-DD dates, so lexicographic comparison is correct.
fn version_supports_messages_listen(version: &str) -> bool {
    version >= UPCOMING_PROTOCOL_VERSION
}

/// Returns `true` when the given protocol version string enables stateless
/// (sessionless) mode for the HTTP transport.
///
/// Stateless mode is introduced in the 2026-07-28 protocol (SEP-2575 /
/// SEP-2567). Version strings are YYYY-MM-DD; lexicographic comparison is
/// correct for date-ordered MCP versions.
#[cfg(feature = "stateless")]
fn is_stateless_protocol_version(version: &str) -> bool {
    version >= UPCOMING_PROTOCOL_VERSION
}

/// Serve a `messages/listen` request as an SSE stream.
///
/// Subscribes to the session's notification broadcast channel and returns a
/// streaming `text/event-stream` response. The stream closes naturally when:
/// - The client disconnects (axum drops the response body).
/// - The broadcast channel closes (server shutdown / session expiry).
///
/// Each notification is assigned a monotonically increasing event ID for
/// potential stream resumption (SEP-1699).
async fn handle_messages_listen_sse(session: Arc<Session>) -> Response {
    let rx = session.notifications_tx.subscribe();
    let session_clone = session.clone();

    let stream = BroadcastStream::new(rx)
        .then(move |result: std::result::Result<String, _>| {
            let session = session_clone.clone();
            async move {
                match result {
                    Ok(msg) => {
                        let event_id = session.next_event_id();
                        // Buffer the event for potential replay (SEP-1699)
                        session.buffer_event(event_id, msg.clone()).await;
                        Some(Ok::<_, Infallible>(
                            Event::default()
                                .id(event_id.to_string())
                                .event(SSE_MESSAGE_EVENT)
                                .data(msg),
                        ))
                    }
                    Err(_) => None,
                }
            }
        })
        .filter_map(|x| x);

    Sse::new(stream)
        .keep_alive(
            axum::response::sse::KeepAlive::new()
                .interval(Duration::from_secs(30))
                .text("ping"),
        )
        .into_response()
}

/// Handle GET requests (SSE stream for server notifications and outgoing requests)
async fn handle_get(
    State(state): State<Arc<AppState>>,
    request: axum::extract::Request,
) -> Response {
    let (parts, _body) = request.into_parts();
    let headers = parts.headers;
    let uri = parts.uri.clone();

    // Validate Host (DNS rebinding defense, complement to Origin)
    if let Some(resp) = validate_host(&headers, &uri, &state) {
        return resp;
    }

    // Validate Origin
    if let Some(resp) = validate_origin(&headers, &state) {
        return resp;
    }

    // Check Accept header
    let accept = headers
        .get(header::ACCEPT)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");

    if !accept.contains("text/event-stream") {
        return (
            StatusCode::NOT_ACCEPTABLE,
            "Accept header must include text/event-stream",
        )
            .into_response();
    }

    // Get session
    let session_id = match get_session_id(&headers) {
        Some(id) => id,
        None => {
            return json_rpc_error_response(None, JsonRpcError::session_required());
        }
    };

    let session = match state.sessions.get(&session_id).await {
        Some(s) => s,
        None => {
            return json_rpc_error_response(
                None,
                JsonRpcError::session_not_found_with_id(&session_id),
            );
        }
    };

    // Check for Last-Event-ID header for stream resumption (SEP-1699)
    let last_event_id = get_last_event_id(&headers);

    // If sampling is enabled, use bidirectional stream
    if state.sampling_enabled {
        return handle_get_bidirectional(session, last_event_id).await;
    }

    // Simple mode: just notifications
    let rx = session.notifications_tx.subscribe();
    let session_clone = session.clone();

    // Replay buffered events if Last-Event-ID was provided (SEP-1699)
    let replay_events: Vec<_> = if let Some(after_id) = last_event_id {
        let events = session.get_events_after(after_id).await;
        tracing::debug!(
            after_id = after_id,
            replay_count = events.len(),
            "Replaying buffered events for stream resumption"
        );
        events
            .into_iter()
            .map(|e| {
                Ok::<_, Infallible>(
                    Event::default()
                        .id(e.id.to_string())
                        .event(SSE_MESSAGE_EVENT)
                        .data(e.data),
                )
            })
            .collect()
    } else {
        Vec::new()
    };

    // Create replay stream from buffered events
    let replay_stream = tokio_stream::iter(replay_events);

    // Create live stream for new events
    // Use `then` for async processing, then `filter_map` to remove errors
    let live_stream = BroadcastStream::new(rx)
        .then(move |result: std::result::Result<String, _>| {
            let session = session_clone.clone();
            async move {
                match result {
                    Ok(msg) => {
                        let event_id = session.next_event_id();
                        // Buffer the event for potential replay (SEP-1699)
                        session.buffer_event(event_id, msg.clone()).await;
                        Some(Ok::<_, Infallible>(
                            Event::default()
                                .id(event_id.to_string())
                                .event(SSE_MESSAGE_EVENT)
                                .data(msg),
                        ))
                    }
                    Err(_) => None,
                }
            }
        })
        .filter_map(|x| x);

    // Chain replay stream with live stream
    let stream = replay_stream.chain(live_stream);

    Sse::new(stream)
        .keep_alive(
            axum::response::sse::KeepAlive::new()
                .interval(Duration::from_secs(30))
                .text("ping"),
        )
        .into_response()
}

/// Handle GET requests with bidirectional support (sampling enabled)
async fn handle_get_bidirectional(session: Arc<Session>, last_event_id: Option<u64>) -> Response {
    // Take ownership of the request receiver for this SSE connection
    let request_rx = {
        let mut rx_guard = session.request_rx.lock().await;
        rx_guard.take()
    };

    // Create a channel for the stream
    let (tx, rx) = tokio::sync::mpsc::channel::<std::result::Result<Event, Infallible>>(100);

    // Replay buffered events if Last-Event-ID was provided (SEP-1699)
    if let Some(after_id) = last_event_id {
        let events = session.get_events_after(after_id).await;
        tracing::debug!(
            after_id = after_id,
            replay_count = events.len(),
            "Replaying buffered events for bidirectional stream resumption"
        );
        for event in events {
            let sse_event = Event::default()
                .id(event.id.to_string())
                .event(SSE_MESSAGE_EVENT)
                .data(event.data);
            if tx.send(Ok(sse_event)).await.is_err() {
                // Client disconnected before replay completed
                return Sse::new(tokio_stream::wrappers::ReceiverStream::new(rx))
                    .keep_alive(
                        axum::response::sse::KeepAlive::new()
                            .interval(Duration::from_secs(30))
                            .text("ping"),
                    )
                    .into_response();
            }
        }
    }

    // Spawn task to multiplex notifications and outgoing requests
    let session_clone = session.clone();
    tokio::spawn(async move {
        let mut notification_rx = session_clone.notifications_tx.subscribe();

        // If we have a request receiver, use select! to handle both
        if let Some(mut req_rx) = request_rx {
            loop {
                tokio::select! {
                    // Handle notifications
                    result = notification_rx.recv() => {
                        match result {
                            Ok(msg) => {
                                let event_id = session_clone.next_event_id();
                                // Buffer the event for potential replay (SEP-1699)
                                session_clone.buffer_event(event_id, msg.clone()).await;
                                let event = Event::default()
                                    .id(event_id.to_string())
                                    .event(SSE_MESSAGE_EVENT)
                                    .data(msg);
                                if tx.send(Ok(event)).await.is_err() {
                                    break; // Client disconnected
                                }
                            }
                            Err(broadcast::error::RecvError::Closed) => break,
                            Err(broadcast::error::RecvError::Lagged(_)) => continue,
                        }
                    }

                    // Handle outgoing requests (sampling)
                    Some(outgoing) = req_rx.recv() => {
                        // Build JSON-RPC request
                        let request = JsonRpcRequest {
                            jsonrpc: "2.0".to_string(),
                            id: outgoing.id.clone(),
                            method: outgoing.method,
                            params: Some(outgoing.params),
                        };

                        match serde_json::to_string(&request) {
                            Ok(request_json) => {
                                tracing::debug!(output = %request_json, "Sending request to client via SSE");

                                // Store pending request
                                session_clone.add_pending_request(
                                    outgoing.id,
                                    outgoing.response_tx,
                                ).await;

                                // Send on SSE stream
                                let event_id = session_clone.next_event_id();
                                // Buffer the event for potential replay (SEP-1699)
                                session_clone.buffer_event(event_id, request_json.clone()).await;
                                let event = Event::default()
                                    .id(event_id.to_string())
                                    .event(SSE_MESSAGE_EVENT)
                                    .data(request_json);
                                if tx.send(Ok(event)).await.is_err() {
                                    break; // Client disconnected
                                }
                            }
                            Err(e) => {
                                tracing::error!(error = %e, "Failed to serialize outgoing request");
                                // Notify the waiter of the error
                                let _ = outgoing.response_tx.send(Err(Error::Internal(
                                    format!("Failed to serialize request: {}", e),
                                )));
                            }
                        }
                    }
                }
            }
        } else {
            // No request receiver, just handle notifications
            loop {
                match notification_rx.recv().await {
                    Ok(msg) => {
                        let event_id = session_clone.next_event_id();
                        // Buffer the event for potential replay (SEP-1699)
                        session_clone.buffer_event(event_id, msg.clone()).await;
                        let event = Event::default()
                            .id(event_id.to_string())
                            .event(SSE_MESSAGE_EVENT)
                            .data(msg);
                        if tx.send(Ok(event)).await.is_err() {
                            break;
                        }
                    }
                    Err(broadcast::error::RecvError::Closed) => break,
                    Err(broadcast::error::RecvError::Lagged(_)) => continue,
                }
            }
        }
    });

    // Convert the receiver into a stream
    let stream = tokio_stream::wrappers::ReceiverStream::new(rx);

    Sse::new(stream)
        .keep_alive(
            axum::response::sse::KeepAlive::new()
                .interval(Duration::from_secs(30))
                .text("ping"),
        )
        .into_response()
}

/// Handle DELETE requests (session termination)
async fn handle_delete(
    State(state): State<Arc<AppState>>,
    request: axum::extract::Request,
) -> Response {
    let (parts, _body) = request.into_parts();
    let headers = parts.headers;
    let uri = parts.uri.clone();

    // Validate Host (DNS rebinding defense, complement to Origin)
    if let Some(resp) = validate_host(&headers, &uri, &state) {
        return resp;
    }

    // Validate Origin
    if let Some(resp) = validate_origin(&headers, &state) {
        return resp;
    }

    let session_id = match get_session_id(&headers) {
        Some(id) => id,
        None => {
            return json_rpc_error_response(None, JsonRpcError::session_required());
        }
    };

    if state.sessions.remove(&session_id).await {
        tracing::info!(session_id = %session_id, "Session terminated");
        StatusCode::OK.into_response()
    } else {
        // For DELETE, it's okay if the session doesn't exist - it's already gone
        // Return OK instead of an error for idempotency
        tracing::debug!(session_id = %session_id, "Session already removed or never existed");
        StatusCode::OK.into_response()
    }
}

/// Handle GET /health requests
///
/// Returns a simple 200 OK response for health checks.
/// Does not require authentication or session state.
async fn handle_health() -> Response {
    StatusCode::OK.into_response()
}

/// Create a JSON-RPC error response
fn json_rpc_error_response(
    id: Option<crate::protocol::RequestId>,
    error: JsonRpcError,
) -> Response {
    let response = JsonRpcResponse::error(id, error);
    axum::Json(response).into_response()
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::Request;
    use tower::ServiceExt;

    fn create_test_router() -> McpRouter {
        McpRouter::new().server_info("test-server", "1.0.0")
    }

    #[tokio::test]
    async fn test_initialize_creates_session() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();

        assert_eq!(response.status(), StatusCode::OK);
        assert!(response.headers().contains_key(MCP_SESSION_ID_HEADER));
        // Verify protocol version header is present on initialize response
        assert_eq!(
            response
                .headers()
                .get(MCP_PROTOCOL_VERSION_HEADER)
                .and_then(|v| v.to_str().ok()),
            Some("2025-11-25")
        );
    }

    #[tokio::test]
    async fn test_protocol_version_header_on_subsequent_requests() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        // Initialize
        let init_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-03-26",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let init_response = app.clone().oneshot(init_request).await.unwrap();
        let session_id = init_response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Verify init response has negotiated version (2025-03-26, not latest)
        assert_eq!(
            init_response
                .headers()
                .get(MCP_PROTOCOL_VERSION_HEADER)
                .and_then(|v| v.to_str().ok()),
            Some("2025-03-26")
        );

        // Send initialized notification
        let initialized_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .header(MCP_PROTOCOL_VERSION_HEADER, "2025-03-26")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": "notifications/initialized"
                })
                .to_string(),
            ))
            .unwrap();

        app.clone().oneshot(initialized_request).await.unwrap();

        // Send tools/list and check for protocol version header
        let list_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .header(MCP_PROTOCOL_VERSION_HEADER, "2025-03-26")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(list_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get(MCP_PROTOCOL_VERSION_HEADER)
                .and_then(|v| v.to_str().ok()),
            Some("2025-03-26")
        );
    }

    #[tokio::test]
    async fn unsupported_protocol_version_returns_spec_shape_error() {
        // SEP-2575: requests carrying an unrecognized MCP-Protocol-Version
        // header (post-initialize) get a JSON-RPC error with code -32004 and
        // data `{ supported: [...], requested: "..." }`.
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        // Initialize first so we're past the init exemption.
        let init_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "t", "version": "0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let init_response = app.clone().oneshot(init_request).await.unwrap();
        let session_id = init_response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Now send a request with a bogus version header.
        let bad = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .header(MCP_PROTOCOL_VERSION_HEADER, "1999-01-01")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 99,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(bad).await.unwrap();
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["error"]["code"].as_i64().unwrap(), -32004);
        assert_eq!(json["error"]["data"]["requested"], "1999-01-01");
        let supported = json["error"]["data"]["supported"]
            .as_array()
            .expect("supported must be an array");
        assert!(supported.contains(&serde_json::json!("2025-11-25")));
        // The request id must be echoed (we have one in the body).
        assert_eq!(json["id"], 99);
        // Field name must be `supported`, NOT `supportedVersions` (SEP-2575 shape).
        assert!(
            json["error"]["data"].get("supportedVersions").is_none(),
            "error data must use 'supported', not 'supportedVersions': {:?}",
            json["error"]["data"]
        );
        // The supported set must exactly match SUPPORTED_PROTOCOL_VERSIONS --
        // no extras, none missing.
        let expected: Vec<serde_json::Value> = SUPPORTED_PROTOCOL_VERSIONS
            .iter()
            .map(|v| serde_json::json!(v))
            .collect();
        assert_eq!(
            supported, &expected,
            "data.supported must exactly match SUPPORTED_PROTOCOL_VERSIONS"
        );
    }

    /// When a request (no session) arrives with an invalid `Mcp-Protocol-Version`
    /// header, the transport must return -32004 with the correct SEP-2575 wire
    /// shape: `{ supported: [...], requested: "..." }`. This verifies the
    /// version-validation path fires without requiring a session.
    #[cfg(feature = "stateless")]
    #[tokio::test]
    async fn stateless_unsupported_protocol_version_returns_spec_shape_error() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        // "1999-01-01" is below UPCOMING_PROTOCOL_VERSION ("2026-07-28"), so it
        // does not enter the version-gated stateless block. It is also not in
        // SUPPORTED_PROTOCOL_VERSIONS, so it triggers the -32004 version check
        // at lines ~2370-2382. No session header is sent, exercising the
        // sessionless request path through version validation.
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_PROTOCOL_VERSION_HEADER, "1999-01-01")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 42,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(
            json["error"]["code"].as_i64().unwrap(),
            -32004,
            "must return UnsupportedProtocolVersion (-32004): {json}"
        );
        assert_eq!(
            json["error"]["data"]["requested"], "1999-01-01",
            "data.requested must echo the version: {json}"
        );
        // Field name must be `supported`, not `supportedVersions`.
        assert!(
            json["error"]["data"].get("supportedVersions").is_none(),
            "error data must use 'supported', not 'supportedVersions': {json}"
        );
        let supported = json["error"]["data"]["supported"]
            .as_array()
            .expect("data.supported must be an array");
        let expected: Vec<serde_json::Value> = SUPPORTED_PROTOCOL_VERSIONS
            .iter()
            .map(|v| serde_json::json!(v))
            .collect();
        assert_eq!(
            supported, &expected,
            "data.supported must exactly match SUPPORTED_PROTOCOL_VERSIONS"
        );
    }

    // =========================================================================
    // SEP-2243: HTTP header standardization (Mcp-Method, Mcp-Name, Mcp-Param-*)
    // =========================================================================

    /// In lenient mode (negotiated protocol version < 2026-07-28) a
    /// request without any SEP-2243 headers must still succeed — older
    /// clients that haven't opted in must keep working.
    #[tokio::test]
    async fn sep_2243_lenient_mode_accepts_missing_headers() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        // Initialize (no SEP-2243 headers) negotiates 2025-11-25 — lenient.
        let init = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "t", "version": "0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let init_response = app.clone().oneshot(init).await.unwrap();
        assert_eq!(init_response.status(), StatusCode::OK);
        let session_id = init_response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // tools/list with no Mcp-Method header — must succeed in lenient mode.
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
    }

    /// In lenient mode, if the client opts in by sending Mcp-Method,
    /// the server still validates against the body and rejects a
    /// mismatch with -32001.
    #[tokio::test]
    async fn sep_2243_lenient_mode_validates_present_headers() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        let init = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header(MCP_METHOD_HEADER, "initialize")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "t", "version": "0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let init_response = app.clone().oneshot(init).await.unwrap();
        assert_eq!(init_response.status(), StatusCode::OK);
        let session_id = init_response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // tools/list with a deliberately-wrong Mcp-Method header.
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .header(MCP_METHOD_HEADER, "ping")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["error"]["code"].as_i64().unwrap(), -32001);
        assert!(
            json["error"]["message"]
                .as_str()
                .unwrap()
                .contains("Mcp-Method")
        );
        assert_eq!(json["id"], 2);
    }

    /// tools/call with matching Mcp-Method + Mcp-Name passes validation
    /// even in strict mode. Driven via initialize with the upcoming
    /// 2026-07-28 protocol version so we exercise the strict branch.
    ///
    /// Gated to `not(stateless)` because with the stateless feature enabled,
    /// initialize requests for 2026-07-28 are handled without a session (chunk 5).
    /// The stateless-mode equivalent is `stateless_v2026_tools_call_without_session_succeeds`.
    #[tokio::test]
    #[cfg(not(feature = "stateless"))]
    async fn sep_2243_strict_mode_tools_call_with_matching_headers() {
        use crate::{CallToolResult, ToolBuilder};

        let router = McpRouter::new().server_info("t", "1.0.0").tool(
            ToolBuilder::new("echo")
                .description("echo")
                .handler(|args: serde_json::Value| async move {
                    Ok(CallToolResult::text(args.to_string()))
                })
                .build(),
        );
        let transport = HttpTransport::new(router).disable_origin_validation();
        let app = transport.into_router();

        // Initialize requesting 2026-07-28 so the session falls into
        // strict mode. The server will negotiate the actual returned
        // version against SUPPORTED_PROTOCOL_VERSIONS, but for SEP-2243
        // gating on init the requested version is what counts.
        let init = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header(MCP_METHOD_HEADER, "initialize")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2026-07-28",
                        "capabilities": {},
                        "clientInfo": { "name": "t", "version": "0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let init_response = app.clone().oneshot(init).await.unwrap();
        assert_eq!(init_response.status(), StatusCode::OK);
        let session_id = init_response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();
        let negotiated_version = init_response
            .headers()
            .get(MCP_PROTOCOL_VERSION_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // For subsequent requests, the session's negotiated protocol
        // version is what the validator gates on. If the server did
        // NOT honor 2026-07-28 (because it isn't in SUPPORTED yet) the
        // session will be lenient — which is fine, we just want to
        // confirm the happy path works. If it IS honored, the strict
        // branch is exercised.
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .header(MCP_PROTOCOL_VERSION_HEADER, &negotiated_version)
            .header(MCP_METHOD_HEADER, "tools/call")
            .header(MCP_NAME_HEADER, "echo")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/call",
                    "params": {
                        "name": "echo",
                        "arguments": {"message": "hi"}
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
    }

    /// tools/call with mismatched Mcp-Name vs body params.name MUST be
    /// rejected with -32001 and HTTP 400 even in lenient mode.
    #[tokio::test]
    async fn sep_2243_tools_call_mcp_name_mismatch_rejected() {
        use crate::{CallToolResult, ToolBuilder};

        let router = McpRouter::new().server_info("t", "1.0.0").tool(
            ToolBuilder::new("echo")
                .description("echo")
                .handler(|args: serde_json::Value| async move {
                    Ok(CallToolResult::text(args.to_string()))
                })
                .build(),
        );
        let transport = HttpTransport::new(router).disable_origin_validation();
        let app = transport.into_router();

        let init = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "t", "version": "0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let init_response = app.clone().oneshot(init).await.unwrap();
        let session_id = init_response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .header(MCP_METHOD_HEADER, "tools/call")
            .header(MCP_NAME_HEADER, "not-echo")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 7,
                    "method": "tools/call",
                    "params": {
                        "name": "echo",
                        "arguments": {"message": "hi"}
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["error"]["code"].as_i64().unwrap(), -32001);
        let msg = json["error"]["message"].as_str().unwrap();
        assert!(msg.contains("Mcp-Name"), "got: {msg}");
        assert_eq!(json["id"], 7);
    }

    /// Mcp-Param-* with a Base64-encoded value that decodes to the
    /// body argument must pass validation.
    #[tokio::test]
    async fn sep_2243_mcp_param_base64_decoded_and_matched() {
        use crate::{CallToolResult, ToolBuilder};

        let router = McpRouter::new().server_info("t", "1.0.0").tool(
            ToolBuilder::new("echo")
                .description("echo")
                .handler(|args: serde_json::Value| async move {
                    Ok(CallToolResult::text(args.to_string()))
                })
                .build(),
        );
        let transport = HttpTransport::new(router).disable_origin_validation();
        let app = transport.into_router();

        let init = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "t", "version": "0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let init_response = app.clone().oneshot(init).await.unwrap();
        let session_id = init_response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Body argument is "Hello"; header is "=?base64?SGVsbG8=?=".
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .header(MCP_METHOD_HEADER, "tools/call")
            .header(MCP_NAME_HEADER, "echo")
            .header("mcp-param-message", "=?base64?SGVsbG8=?=")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 5,
                    "method": "tools/call",
                    "params": {
                        "name": "echo",
                        "arguments": {"message": "Hello"}
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
    }

    /// notifications/initialized still receives ACCEPTED when SEP-2243
    /// headers match (regression for the notification fast path).
    #[tokio::test]
    async fn sep_2243_notification_with_matching_method_header_accepted() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        let init = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "t", "version": "0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let init_response = app.clone().oneshot(init).await.unwrap();
        let session_id = init_response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .header(MCP_METHOD_HEADER, "notifications/initialized")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": "notifications/initialized"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::ACCEPTED);
    }

    #[tokio::test]
    async fn test_request_without_session_fails() {
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .require_sessions();
        let app = transport.into_router();

        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();

        // We now return JSON-RPC errors for session issues
        assert_eq!(response.status(), StatusCode::OK);

        // Verify it's a JSON-RPC error response
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(json.get("error").is_some());
        assert_eq!(json["error"]["code"], -32006); // SessionRequired
    }

    #[tokio::test]
    async fn test_delete_session() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        // First, initialize to get a session
        let init_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.clone().oneshot(init_request).await.unwrap();
        let session_id = response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Delete the session
        let delete_request = Request::builder()
            .method("DELETE")
            .uri("/")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .body(Body::empty())
            .unwrap();

        let response = app.clone().oneshot(delete_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        // Verify session is gone
        let list_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(list_request).await.unwrap();
        // We now return JSON-RPC errors for session issues
        assert_eq!(response.status(), StatusCode::OK);

        // Verify it's a JSON-RPC error response
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(json.get("error").is_some());
        assert_eq!(json["error"]["code"], -32005); // SessionNotFound
    }

    #[tokio::test]
    async fn test_custom_session_store_receives_create_and_delete() {
        use crate::session_store::{MemorySessionStore, SessionStore as PublicSessionStore};

        let store = Arc::new(MemorySessionStore::new());
        let store_dyn: Arc<dyn PublicSessionStore> = store.clone();

        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .session_store(store_dyn);
        let (app, handle) = transport.into_router_with_handle();

        // Initialize to create a session.
        let init_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test-client", "version": "1.0.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.clone().oneshot(init_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let session_id = response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Custom store should have the record.
        assert_eq!(store.len().await, 1);
        let record = store
            .load(&session_id)
            .await
            .unwrap()
            .expect("expected session to be persisted");
        assert_eq!(record.id, session_id);

        // After initialize completes the record must carry the client's
        // advertised identity / capabilities (issue #786). Previously these
        // were left as `None` because the record was created before
        // initialize ran.
        let client_info = record
            .client_info
            .expect("client_info should be populated after initialize");
        assert_eq!(client_info.name, "test-client");
        assert_eq!(client_info.version, "1.0.0");
        assert!(
            record.client_capabilities.is_some(),
            "client_capabilities should be populated after initialize"
        );

        // Terminate session via the handle -- store should be cleared.
        assert!(handle.terminate_session(&session_id).await);
        assert_eq!(store.len().await, 0);
        assert!(store.load(&session_id).await.unwrap().is_none());
    }

    #[tokio::test]
    async fn test_session_store_record_carries_negotiated_protocol_version() {
        // Issue #786: stored record should reflect the negotiated protocol
        // version (taken from the initialize response), not the default the
        // session was created with.
        use crate::session_store::{MemorySessionStore, SessionStore as PublicSessionStore};

        let store = Arc::new(MemorySessionStore::new());
        let store_dyn: Arc<dyn PublicSessionStore> = store.clone();

        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .session_store(store_dyn);
        let app = transport.into_router();

        // Initialize using an older supported protocol version so we can
        // tell the persisted version apart from `LATEST_PROTOCOL_VERSION`.
        let init_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-03-26",
                        "capabilities": {},
                        "clientInfo": { "name": "v-client", "version": "2.0.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(init_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let session_id = response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        let record = store
            .load(&session_id)
            .await
            .unwrap()
            .expect("session should be persisted");
        assert_eq!(record.protocol_version, "2025-03-26");
        let client_info = record.client_info.expect("client_info should be populated");
        assert_eq!(client_info.name, "v-client");
    }

    #[tokio::test]
    async fn test_restored_session_exposes_original_client_info() {
        // Issue #786: a session restored from the persistent store on a
        // peer instance should retain the original client's identity and
        // capabilities, not the synthetic defaults used for auto-reinit.
        use crate::session_store::{MemorySessionStore, SessionStore as PublicSessionStore};

        let store = Arc::new(MemorySessionStore::new());
        let store_dyn: Arc<dyn PublicSessionStore> = store.clone();

        // First "instance": initialize, then drop the transport so the
        // local registry is gone but the persistent record survives.
        let session_id = {
            let transport = HttpTransport::new(create_test_router())
                .disable_origin_validation()
                .session_store(store_dyn.clone());
            let app = transport.into_router();

            let init_request = Request::builder()
                .method("POST")
                .uri("/")
                .header("Content-Type", "application/json")
                .header("Accept", "application/json, text/event-stream")
                .body(Body::from(
                    serde_json::json!({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "initialize",
                        "params": {
                            "protocolVersion": "2025-11-25",
                            "capabilities": { "roots": {} },
                            "clientInfo": {
                                "name": "original-client",
                                "version": "3.1.4"
                            }
                        }
                    })
                    .to_string(),
                ))
                .unwrap();
            let response = app.oneshot(init_request).await.unwrap();
            assert_eq!(response.status(), StatusCode::OK);
            response
                .headers()
                .get(MCP_SESSION_ID_HEADER)
                .unwrap()
                .to_str()
                .unwrap()
                .to_string()
        };

        // Sanity check: the persisted record now carries the client info.
        let stored = store
            .load(&session_id)
            .await
            .unwrap()
            .expect("record should survive transport drop");
        assert_eq!(
            stored.client_info.as_ref().map(|c| c.name.as_str()),
            Some("original-client")
        );

        // Second "instance": brand new transport, same store. A request
        // with the existing session id triggers restore_from_record.
        let transport2 = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .session_store(store_dyn);
        let app2 = transport2.into_router();

        let list_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app2.oneshot(list_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(
            json.get("result").is_some(),
            "expected tools/list result, got {json}"
        );

        // After the restore, the record in the store still carries the
        // original client info (refreshed expiry, not a synthetic
        // "auto-recovered" identity).
        let after_restore = store
            .load(&session_id)
            .await
            .unwrap()
            .expect("record should still be present after restore");
        let client_info = after_restore
            .client_info
            .expect("restored record should retain client_info");
        assert_eq!(client_info.name, "original-client");
        assert_eq!(client_info.version, "3.1.4");
        assert!(
            after_restore.client_capabilities.is_some(),
            "restored record should retain client_capabilities"
        );
    }

    #[tokio::test]
    async fn test_auto_reinitialize_marks_synthetic_client_info() {
        // Companion to the restored-client-info test: the auto-reinit
        // path must continue to flag the client as `"auto-recovered"` so
        // the two paths remain distinguishable on inspection of the
        // persisted record.
        use crate::session_store::{MemorySessionStore, SessionStore as PublicSessionStore};

        let store = Arc::new(MemorySessionStore::new());
        let store_dyn: Arc<dyn PublicSessionStore> = store.clone();

        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .session_store(store_dyn)
            .auto_reinitialize_sessions(true);
        let app = transport.into_router();

        let list_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header(MCP_SESSION_ID_HEADER, "made-up-id")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(list_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        let record = store
            .load("made-up-id")
            .await
            .unwrap()
            .expect("auto-reinitialize should persist a record");
        assert_eq!(
            record.client_info.as_ref().map(|c| c.name.as_str()),
            Some("auto-recovered")
        );
    }

    #[tokio::test]
    async fn test_custom_event_store_buffers_and_purges() {
        use crate::event_store::{EventStore as PublicEventStore, MemoryEventStore};

        let events = Arc::new(MemoryEventStore::new());
        let events_dyn: Arc<dyn PublicEventStore> = events.clone();

        // Build a session directly so we can exercise buffer_event/get_events_after
        // without needing a live SSE subscriber.
        let session = Arc::new(Session::new(
            create_test_router(),
            false,
            identity_factory(),
            events_dyn,
        ));

        session.buffer_event(0, "first".to_string()).await;
        session.buffer_event(1, "second".to_string()).await;

        // Custom store should have both events.
        assert_eq!(events.total_events().await, 2);
        let replayed = events.replay_after(&session.id, 0).await.unwrap();
        assert_eq!(replayed.len(), 1);
        assert_eq!(replayed[0].id, 1);
        assert_eq!(replayed[0].data, "second");

        // Purging should clear the session's log.
        events.purge_session(&session.id).await.unwrap();
        assert_eq!(events.total_events().await, 0);
    }

    #[tokio::test]
    async fn test_restore_from_store_serves_unknown_session_id() {
        use crate::session_store::{MemorySessionStore, SessionRecord, SessionStore};

        // Two transports share a single session store (simulating two
        // server instances behind a load balancer).
        let store = Arc::new(MemorySessionStore::new());
        let store_dyn: Arc<dyn SessionStore> = store.clone();

        // Seed the store with a record as if a peer instance had created it.
        let mut seeded = SessionRecord::new(
            "shared-session".to_string(),
            "2025-11-25".to_string(),
            Duration::from_secs(60),
        );
        store.create(&mut seeded).await.unwrap();
        let seeded_id = seeded.id;

        // This transport has never seen the session locally.
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .session_store(store_dyn);
        let app = transport.into_router();

        let list_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header(MCP_SESSION_ID_HEADER, &seeded_id)
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(list_request).await.unwrap();
        // Without restore this would produce a SessionNotFound JSON-RPC
        // error; with restore the request is served normally.
        assert_eq!(response.status(), StatusCode::OK);

        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(
            json.get("result").is_some(),
            "expected tools/list result, got {json}"
        );
    }

    #[tokio::test]
    async fn test_auto_reinitialize_serves_unknown_session_without_store_record() {
        // No seeded store record — the client just shows up with a
        // session ID the server has never heard of. With auto-reinit
        // enabled the transport spins up a synthetic session.
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .auto_reinitialize_sessions(true);
        let app = transport.into_router();

        let list_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header(MCP_SESSION_ID_HEADER, "client-made-up-id")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(list_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(
            json.get("result").is_some(),
            "expected tools/list result, got {json}"
        );
    }

    #[tokio::test]
    async fn test_unknown_session_without_restore_or_auto_reinit_returns_error() {
        // Default transport: no store seeded, no auto-reinit.
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        let list_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header(MCP_SESSION_ID_HEADER, "never-seen-before")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(list_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(json.get("error").is_some(), "expected error, got {json}");
        assert_eq!(json["error"]["code"], -32005); // SessionNotFound
    }

    #[tokio::test]
    async fn test_session_expiration() {
        // Create transport with very short TTL
        let config = SessionConfig::with_ttl(Duration::from_millis(50))
            .cleanup_interval(Duration::from_millis(10));
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .session_config(config);
        let app = transport.into_router();

        // Initialize to get a session
        let init_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.clone().oneshot(init_request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let session_id = response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Wait for session to expire and cleanup to run
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Session should be expired now
        let list_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(list_request).await.unwrap();
        // We now return JSON-RPC errors for session issues
        assert_eq!(response.status(), StatusCode::OK);

        // Verify it's a JSON-RPC error response
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(json.get("error").is_some());
        assert_eq!(json["error"]["code"], -32005); // SessionNotFound
    }

    #[tokio::test]
    async fn test_layer_with_identity() {
        // Verify that .layer() compiles and produces a working transport
        // using a no-op layer (tower::layer::Identity)
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .layer(tower::layer::util::Identity::new());
        let app = transport.into_router();

        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert!(response.headers().contains_key(MCP_SESSION_ID_HEADER));
    }

    #[tokio::test]
    async fn test_layer_with_timeout() {
        // Verify that .layer() works with TimeoutLayer
        use std::time::Duration;
        use tower::timeout::TimeoutLayer;

        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .layer(TimeoutLayer::new(Duration::from_secs(30)));
        let app = transport.into_router();

        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert!(response.headers().contains_key(MCP_SESSION_ID_HEADER));
    }

    #[tokio::test]
    async fn test_layer_middleware_error_produces_jsonrpc_error() {
        // Use an extremely short timeout to force an error.
        // The CatchError wrapper should convert it to a JSON-RPC error response.
        use std::time::Duration;
        use tower::timeout::TimeoutLayer;

        let slow_tool = crate::tool::ToolBuilder::new("slow")
            .description("A slow tool")
            .handler(|_: serde_json::Value| async move {
                tokio::time::sleep(Duration::from_secs(10)).await;
                Ok(crate::CallToolResult::text("done"))
            })
            .build();

        let router = McpRouter::new()
            .server_info("test-server", "1.0.0")
            .tool(slow_tool);

        // 1ms timeout will definitely expire before the tool completes
        let transport = HttpTransport::new(router)
            .disable_origin_validation()
            .layer(TimeoutLayer::new(Duration::from_millis(1)));
        let app = transport.into_router();

        // Initialize first
        let init_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.clone().oneshot(init_request).await.unwrap();
        let session_id = response
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Call the slow tool -- should timeout and return a JSON-RPC error
        let tool_request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/call",
                    "params": {
                        "name": "slow",
                        "arguments": {}
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(tool_request).await.unwrap();
        // Should still return 200 with a JSON-RPC error body
        assert_eq!(response.status(), StatusCode::OK);

        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(
            json.get("error").is_some(),
            "Expected JSON-RPC error response, got: {}",
            json
        );
    }

    #[tokio::test]
    async fn test_max_sessions_limit() {
        // Create transport with max 1 session
        let config = SessionConfig::default().max_sessions(1);
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .session_config(config);
        let app = transport.into_router();

        // First initialize should succeed
        let init_request1 = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.clone().oneshot(init_request1).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        // Second initialize should fail (max sessions reached)
        let init_request2 = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client-2",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(init_request2).await.unwrap();
        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
    }

    #[tokio::test]
    async fn test_session_event_buffering() {
        // Test that events are buffered and can be retrieved for replay (SEP-1699)
        let session = Session::new(
            create_test_router(),
            false,
            identity_factory(),
            Arc::new(crate::event_store::MemoryEventStore::new()),
        );

        // Buffer some events
        session.buffer_event(0, "event0".to_string()).await;
        session.buffer_event(1, "event1".to_string()).await;
        session.buffer_event(2, "event2".to_string()).await;

        // Get events after event 0
        let events = session.get_events_after(0).await;
        assert_eq!(events.len(), 2);
        assert_eq!(events[0].id, 1);
        assert_eq!(events[0].data, "event1");
        assert_eq!(events[1].id, 2);
        assert_eq!(events[1].data, "event2");

        // Get events after event 1
        let events = session.get_events_after(1).await;
        assert_eq!(events.len(), 1);
        assert_eq!(events[0].id, 2);

        // Get events after event 2 (none)
        let events = session.get_events_after(2).await;
        assert!(events.is_empty());
    }

    #[tokio::test]
    async fn test_session_event_counter_increments() {
        // Test that event IDs increment monotonically (SEP-1699)
        let session = Session::new(
            create_test_router(),
            false,
            identity_factory(),
            Arc::new(crate::event_store::MemoryEventStore::new()),
        );

        assert_eq!(session.next_event_id(), 0);
        assert_eq!(session.next_event_id(), 1);
        assert_eq!(session.next_event_id(), 2);
    }

    #[tokio::test]
    async fn test_session_event_buffer_limit() {
        // Test that buffer respects max size limit
        // Create a session - buffer limit is DEFAULT_MAX_BUFFERED_EVENTS (1000)
        let session = Session::new(
            create_test_router(),
            false,
            identity_factory(),
            Arc::new(crate::event_store::MemoryEventStore::new()),
        );

        // Buffer more events than we can test practically, but verify the mechanism works
        // by checking that old events are evicted when we exceed the limit
        for i in 0..10 {
            session.buffer_event(i, format!("event{}", i)).await;
        }

        // All 10 events should be present
        let events = session.get_events_after(0).await;
        // Events after 0 should be 1-9 (9 events)
        assert_eq!(events.len(), 9);
    }

    #[tokio::test]
    async fn test_session_handle_count() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let (app, handle) = transport.into_router_with_handle();

        // No sessions initially
        assert_eq!(handle.session_count().await, 0);

        // Initialize to create a session
        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), 200);

        // Now we should have 1 session
        assert_eq!(handle.session_count().await, 1);
    }

    #[tokio::test]
    async fn test_session_handle_list_and_terminate() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let (app, handle) = transport.into_router_with_handle();

        // No sessions initially
        assert!(handle.list_sessions().await.is_empty());

        // Initialize to create a session
        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": {
                            "name": "test-client",
                            "version": "1.0.0"
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), 200);

        // list_sessions should return 1 session with valid metadata
        let sessions = handle.list_sessions().await;
        assert_eq!(sessions.len(), 1);
        assert!(!sessions[0].id.is_empty());

        // Terminate the session
        let session_id = sessions[0].id.clone();
        assert!(handle.terminate_session(&session_id).await);
        assert_eq!(handle.session_count().await, 0);

        // Terminating again returns false
        assert!(!handle.terminate_session(&session_id).await);
    }

    #[tokio::test]
    async fn test_request_without_session_id_rejected() {
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .require_sessions();
        let app = transport.into_router();

        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            // No mcp-session-id header
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list",
                    "params": {}
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK); // JSON-RPC errors still 200
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        // Should return session required error
        assert!(json["error"].is_object());
    }

    #[tokio::test]
    async fn test_invalid_session_id_returns_error() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header("mcp-session-id", "nonexistent-session-id")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list",
                    "params": {}
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["error"]["code"].as_i64().unwrap(), -32005); // SessionNotFound
    }

    #[tokio::test]
    async fn test_notification_returns_accepted() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        // First initialize to get a session
        let init_req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.clone().oneshot(init_req).await.unwrap();
        let session_id = resp
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Send a notification (no id field) -- should return 202 Accepted
        let notif = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("mcp-session-id", &session_id)
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": "notifications/initialized"
                })
                .to_string(),
            ))
            .unwrap();

        let response = app.oneshot(notif).await.unwrap();
        assert_eq!(response.status(), StatusCode::ACCEPTED);
    }

    #[tokio::test]
    async fn test_invalid_json_returns_parse_error() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        let request = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .body(Body::from("not valid json{{{"))
            .unwrap();

        let response = app.oneshot(request).await.unwrap();
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        tower_mcp_types::testing::assert_jsonrpc_error_response(&json);
        assert!(
            json["id"].is_null(),
            "id must be null on parse error: {json}"
        );
        assert_eq!(json["error"]["code"].as_i64().unwrap(), -32700);
    }

    #[tokio::test]
    async fn test_session_config_max_sessions() {
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .session_config(SessionConfig::default().max_sessions(1));
        let app = transport.into_router();

        // First initialize succeeds
        let init1 = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test1", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp1 = app.clone().oneshot(init1).await.unwrap();
        assert_eq!(resp1.status(), StatusCode::OK);

        // Second initialize should fail (max 1 session)
        let init2 = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test2", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp2 = app.oneshot(init2).await.unwrap();
        assert_eq!(resp2.status(), StatusCode::SERVICE_UNAVAILABLE);
    }

    #[tokio::test]
    async fn test_delete_terminates_session() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        // Initialize
        let init_req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.clone().oneshot(init_req).await.unwrap();
        let session_id = resp
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // DELETE should terminate the session
        let delete_req = Request::builder()
            .method("DELETE")
            .uri("/")
            .header("mcp-session-id", &session_id)
            .body(Body::empty())
            .unwrap();

        let resp = app.clone().oneshot(delete_req).await.unwrap();
        assert!(resp.status().is_success());

        // Subsequent request with that session ID should fail
        let list_req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header("mcp-session-id", &session_id)
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list",
                    "params": {}
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.oneshot(list_req).await.unwrap();
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["error"]["code"].as_i64().unwrap(), -32005);
    }

    // -----------------------------------------------------------------------
    // Origin validation / DNS rebinding protection
    // -----------------------------------------------------------------------

    #[test]
    fn test_is_localhost_origin_http() {
        assert!(is_localhost_origin("http://localhost"));
        assert!(is_localhost_origin("http://localhost:3000"));
        assert!(is_localhost_origin("http://127.0.0.1"));
        assert!(is_localhost_origin("http://127.0.0.1:8080"));
        assert!(is_localhost_origin("http://[::1]"));
        assert!(is_localhost_origin("http://[::1]:3000"));
    }

    #[test]
    fn test_is_localhost_origin_https() {
        assert!(is_localhost_origin("https://localhost"));
        assert!(is_localhost_origin("https://127.0.0.1:443"));
    }

    #[test]
    fn test_is_not_localhost_origin() {
        assert!(!is_localhost_origin("http://example.com"));
        assert!(!is_localhost_origin("http://evil-localhost.com"));
        assert!(!is_localhost_origin("http://localhost.evil.com"));
        assert!(!is_localhost_origin("ftp://localhost"));
        assert!(!is_localhost_origin("localhost"));
        assert!(!is_localhost_origin(""));
    }

    #[tokio::test]
    async fn test_origin_validation_rejects_cross_origin() {
        let transport = HttpTransport::new(create_test_router());
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Origin", "http://evil.com")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn test_origin_validation_allows_localhost() {
        let transport = HttpTransport::new(create_test_router());
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Origin", "http://localhost:3000")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_origin_validation_allows_configured_origin() {
        let transport = HttpTransport::new(create_test_router())
            .allowed_origins(vec!["https://my-app.example.com".to_string()]);
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Origin", "https://my-app.example.com")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_origin_validation_rejects_unconfigured_origin() {
        let transport = HttpTransport::new(create_test_router())
            .allowed_origins(vec!["https://my-app.example.com".to_string()]);
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Origin", "https://other-app.example.com")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn test_origin_validation_no_header_allowed() {
        // Requests without Origin header should be allowed (same-origin)
        let transport = HttpTransport::new(create_test_router());
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            // No Origin header
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_disabled_origin_validation_allows_any() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Origin", "http://evil.com")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    // =========================================================================
    // Host header validation (DNS rebinding defense complement to Origin)
    // =========================================================================

    fn initialize_body() -> Body {
        Body::from(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": {},
                    "clientInfo": { "name": "test", "version": "1.0" }
                }
            })
            .to_string(),
        )
    }

    #[test]
    fn test_is_localhost_host_variants() {
        assert!(is_localhost_host("localhost"));
        assert!(is_localhost_host("localhost:3000"));
        assert!(is_localhost_host("127.0.0.1"));
        assert!(is_localhost_host("127.0.0.1:8080"));
        assert!(is_localhost_host("[::1]"));
        assert!(is_localhost_host("[::1]:3000"));

        assert!(!is_localhost_host("evil.com"));
        assert!(!is_localhost_host("api.example.com:8443"));
        assert!(!is_localhost_host("10.0.0.1"));
    }

    #[tokio::test]
    async fn test_host_validation_allows_localhost() {
        let transport = HttpTransport::new(create_test_router())
            .allowed_hosts(vec!["api.example.com".to_string()]);
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Host", "127.0.0.1:3000")
            .body(initialize_body())
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_host_validation_allows_configured_host() {
        let transport = HttpTransport::new(create_test_router())
            .allowed_hosts(vec!["api.example.com".to_string()]);
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Host", "api.example.com")
            .body(initialize_body())
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_host_validation_rejects_unconfigured_host() {
        let transport = HttpTransport::new(create_test_router())
            .allowed_hosts(vec!["api.example.com".to_string()]);
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Host", "evil.com")
            .body(initialize_body())
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }

    #[tokio::test]
    async fn test_host_validation_no_allowlist_accepts_any_host() {
        // Existing deployments that haven't opted into Host validation
        // (no `.allowed_hosts(...)`) should keep accepting non-localhost
        // hosts; Origin still protects browsers.
        let transport = HttpTransport::new(create_test_router());
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Host", "any.example.com")
            .body(initialize_body())
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_disabled_host_validation_allows_any_with_allowlist() {
        let transport = HttpTransport::new(create_test_router())
            .disable_host_validation()
            .allowed_hosts(vec!["api.example.com".to_string()]);
        let app = transport.into_router();

        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("Host", "evil.com")
            .body(initialize_body())
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[test]
    fn test_effective_host_prefers_header() {
        let mut headers = HeaderMap::new();
        headers.insert(header::HOST, HeaderValue::from_static("api.example.com"));
        let uri: axum::http::Uri = "http://other.example.com/path".parse().unwrap();
        assert_eq!(effective_host(&headers, &uri), Some("api.example.com"));
    }

    #[test]
    fn test_effective_host_falls_back_to_authority() {
        // When Host header is missing (HTTP/2 + middleware that strips it),
        // we should fall back to the URI authority.
        let headers = HeaderMap::new();
        let uri: axum::http::Uri = "http://api.example.com/path".parse().unwrap();
        assert_eq!(effective_host(&headers, &uri), Some("api.example.com"));
    }

    #[test]
    fn test_effective_host_returns_none_when_both_missing() {
        let headers = HeaderMap::new();
        let uri: axum::http::Uri = "/path".parse().unwrap();
        assert_eq!(effective_host(&headers, &uri), None);
    }

    // =========================================================================
    // External notification fan-out
    // =========================================================================

    /// Initialize a session against `app` and return its session id.
    async fn init_session(app: &Router) -> String {
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        resp.headers()
            .get(MCP_SESSION_ID_HEADER)
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string())
            .expect("initialize must return a session id")
    }

    #[tokio::test]
    async fn test_external_notification_reaches_single_session() {
        let (notif_tx, notif_rx) = notification_channel(8);
        let transport = HttpTransport::with_notifications(create_test_router(), notif_rx);
        let (app, session_handle) = transport.into_router_with_handle();

        let session_id = init_session(&app).await;

        // Subscribe to the session's broadcast channel before firing.
        let mut rx = {
            let sessions = session_handle.store.sessions.read().await;
            let session = sessions
                .get(&session_id)
                .expect("session should be registered");
            session.notifications_tx.subscribe()
        };

        notif_tx
            .send(crate::context::ServerNotification::ResourceUpdated {
                uri: "claude://chats/abc".to_string(),
            })
            .await
            .unwrap();

        let json = tokio::time::timeout(Duration::from_secs(1), rx.recv())
            .await
            .expect("notification should arrive within timeout")
            .expect("broadcast channel closed");
        assert!(json.contains("notifications/resources/updated"));
        assert!(json.contains("claude://chats/abc"));
    }

    #[tokio::test]
    async fn test_external_notification_fans_out_to_all_sessions() {
        let (notif_tx, notif_rx) = notification_channel(8);
        let transport = HttpTransport::with_notifications(create_test_router(), notif_rx);
        let (app, session_handle) = transport.into_router_with_handle();

        let session_a = init_session(&app).await;
        let session_b = init_session(&app).await;
        assert_ne!(session_a, session_b);

        let (mut rx_a, mut rx_b) = {
            let sessions = session_handle.store.sessions.read().await;
            let a = sessions.get(&session_a).unwrap();
            let b = sessions.get(&session_b).unwrap();
            (
                a.notifications_tx.subscribe(),
                b.notifications_tx.subscribe(),
            )
        };

        notif_tx
            .send(crate::context::ServerNotification::ResourcesListChanged)
            .await
            .unwrap();

        let json_a = tokio::time::timeout(Duration::from_secs(1), rx_a.recv())
            .await
            .unwrap()
            .unwrap();
        let json_b = tokio::time::timeout(Duration::from_secs(1), rx_b.recv())
            .await
            .unwrap()
            .unwrap();
        assert!(json_a.contains("notifications/resources/list_changed"));
        assert!(json_b.contains("notifications/resources/list_changed"));
    }

    #[tokio::test]
    async fn test_external_notifications_builder_method() {
        // `external_notifications` should be equivalent to the constructor.
        let (notif_tx, notif_rx) = notification_channel(8);
        let transport = HttpTransport::new(create_test_router()).external_notifications(notif_rx);
        let (app, session_handle) = transport.into_router_with_handle();

        let session_id = init_session(&app).await;
        let mut rx = {
            let sessions = session_handle.store.sessions.read().await;
            sessions
                .get(&session_id)
                .unwrap()
                .notifications_tx
                .subscribe()
        };

        notif_tx
            .send(crate::context::ServerNotification::ToolsListChanged)
            .await
            .unwrap();

        let json = tokio::time::timeout(Duration::from_secs(1), rx.recv())
            .await
            .unwrap()
            .unwrap();
        assert!(json.contains("notifications/tools/list_changed"));
    }

    #[tokio::test]
    async fn test_default_transport_has_no_external_fanout_task() {
        // Smoke test: a transport without external notifications builds and
        // serves normally. (Verifying the fan-out task is *not* spawned is
        // hard to do directly; this just confirms we didn't accidentally
        // gate the happy path on the channel being present.)
        let transport = HttpTransport::new(create_test_router());
        let (app, _handle) = transport.into_router_with_handle();
        let _session_id = init_session(&app).await;
    }

    // =========================================================================
    // Chunk 5: version-gated stateless mode for 2026-07-28+ clients
    // =========================================================================

    /// 2026-07-28 initialize must NOT return mcp-session-id.
    #[tokio::test]
    #[cfg(feature = "stateless")]
    async fn stateless_v2026_initialize_omits_session_id() {
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .disable_host_validation();
        let app = transport.into_router();
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header(MCP_METHOD_HEADER, "initialize")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2026-07-28",
                        "capabilities": {},
                        "clientInfo": { "name": "sc", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert!(
            !response.headers().contains_key(MCP_SESSION_ID_HEADER),
            "2026-07-28 initialize must not return mcp-session-id"
        );
        assert_eq!(
            response
                .headers()
                .get(MCP_PROTOCOL_VERSION_HEADER)
                .and_then(|v| v.to_str().ok()),
            Some("2026-07-28")
        );
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(
            json["result"]["protocolVersion"], "2026-07-28",
            "initialize result must carry 2026-07-28 protocol version"
        );
    }

    /// 2026-07-28 tools/call without session header succeeds.
    #[tokio::test]
    #[cfg(feature = "stateless")]
    async fn stateless_v2026_tools_call_without_session_succeeds() {
        use crate::{CallToolResult, ToolBuilder};
        let router = McpRouter::new().server_info("t", "1.0.0").tool(
            ToolBuilder::new("echo")
                .description("echo")
                .handler(|args: serde_json::Value| async move {
                    Ok(CallToolResult::text(args.to_string()))
                })
                .build(),
        );
        let transport = HttpTransport::new(router).disable_origin_validation();
        let app = transport.into_router();
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_PROTOCOL_VERSION_HEADER, "2026-07-28")
            .header(MCP_METHOD_HEADER, "tools/call")
            .header(MCP_NAME_HEADER, "echo")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/call",
                    "params": {
                        "name": "echo",
                        "arguments": {"message": "hello"},
                        "_meta": {
                            "io.modelcontextprotocol/protocolVersion": "2026-07-28",
                            "io.modelcontextprotocol/clientInfo": {
                                "name": "sc", "version": "1.0"
                            },
                            "io.modelcontextprotocol/clientCapabilities": {}
                        }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert!(
            !response.headers().contains_key(MCP_SESSION_ID_HEADER),
            "stateless tools/call must not set mcp-session-id"
        );
        assert_eq!(
            response
                .headers()
                .get(MCP_PROTOCOL_VERSION_HEADER)
                .and_then(|v| v.to_str().ok()),
            Some("2026-07-28")
        );
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(
            json.get("result").is_some(),
            "expected tools/call result, got: {json}"
        );
    }

    /// 2025-11-25 initialize still returns mcp-session-id (unchanged).
    #[tokio::test]
    async fn stateless_v2025_initialize_still_gets_session_id() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0", "id": 1, "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "old-client", "version": "1.0" }
                    }
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert!(
            response.headers().contains_key(MCP_SESSION_ID_HEADER),
            "2025-11-25 initialize must return mcp-session-id"
        );
    }

    /// With require_sessions(), 2025-11-25 tools/list without session
    /// header fails with SessionRequired (-32006) -- behavior unchanged.
    #[tokio::test]
    async fn stateless_v2025_tools_list_without_session_rejected() {
        let transport = HttpTransport::new(create_test_router())
            .disable_origin_validation()
            .require_sessions();
        let app = transport.into_router();
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_PROTOCOL_VERSION_HEADER, "2025-11-25")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0", "id": 1, "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(json.get("error").is_some(), "expected error, got: {json}");
        assert_eq!(
            json["error"]["code"].as_i64().unwrap(),
            -32006,
            "expected SessionRequired (-32006)"
        );
    }

    /// 2026-07-28 tools/list without session header succeeds (#856).
    #[tokio::test]
    #[cfg(feature = "stateless")]
    async fn stateless_v2026_tools_list_without_session_succeeds() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_PROTOCOL_VERSION_HEADER, "2026-07-28")
            .header(MCP_METHOD_HEADER, "tools/list")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert!(
            !response.headers().contains_key(MCP_SESSION_ID_HEADER),
            "stateless tools/list must not set mcp-session-id"
        );
        assert_eq!(
            response
                .headers()
                .get(MCP_PROTOCOL_VERSION_HEADER)
                .and_then(|v| v.to_str().ok()),
            Some("2026-07-28")
        );
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(
            json["result"]["tools"].is_array(),
            "expected tools array in result, got: {json}"
        );
    }

    /// 2026-07-28 stateless notification (no id) returns 202 and creates no session (#857).
    #[tokio::test]
    #[cfg(feature = "stateless")]
    async fn stateless_v2026_notification_returns_202_no_session() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let (app, handle) = transport.into_router_with_handle();
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_PROTOCOL_VERSION_HEADER, "2026-07-28")
            .header(MCP_METHOD_HEADER, "notifications/initialized")
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": "notifications/initialized"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(
            response.status(),
            StatusCode::ACCEPTED,
            "stateless notification must return 202 ACCEPTED"
        );
        assert!(
            !response.headers().contains_key(MCP_SESSION_ID_HEADER),
            "stateless notification must not set mcp-session-id"
        );
        assert_eq!(
            handle.session_count().await,
            0,
            "stateless notification must not create a session"
        );
    }

    /// 2026-07-28 stateless request missing Mcp-Method returns -32001 + HTTP 400 (#859).
    #[tokio::test]
    #[cfg(feature = "stateless")]
    async fn stateless_v2026_missing_mcp_method_returns_400() {
        let transport = HttpTransport::new(create_test_router()).disable_origin_validation();
        let app = transport.into_router();
        let req = Request::builder()
            .method("POST")
            .uri("/")
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header(MCP_PROTOCOL_VERSION_HEADER, "2026-07-28")
            // Intentionally NO Mcp-Method header
            .body(Body::from(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                })
                .to_string(),
            ))
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(
            response.status(),
            StatusCode::BAD_REQUEST,
            "missing Mcp-Method must return HTTP 400"
        );
        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(json.get("error").is_some(), "expected error, got: {json}");
        assert_eq!(
            json["error"]["code"].as_i64().unwrap(),
            -32001,
            "expected InvalidRequest (-32001)"
        );
        assert!(
            json["error"]["message"]
                .as_str()
                .unwrap_or("")
                .contains("Mcp-Method"),
            "error message must mention Mcp-Method, got: {json}"
        );
    }
}