rmcp_soddygo/transport/streamable_http_server/

tower.rs

use std::{
    borrow::Cow,
    collections::HashMap,
    convert::Infallible,
    fmt::Display,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::Duration,
};

use bytes::Bytes;
use futures::{Stream, StreamExt, future::BoxFuture};
use http::{HeaderMap, Method, Request, Response, header::ALLOW};
use http_body::Body;
use http_body_util::{BodyExt, Full, combinators::BoxBody};
use pin_project_lite::pin_project;
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;

use super::session::{
    RestoreOutcome, SessionId, SessionManager, SessionRestoreMarker, SessionState, SessionStore,
};
use crate::{
    RoleServer,
    model::{
        ClientCapabilities, ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorCode,
        ErrorData, GetExtensions, GetMeta, Implementation, InitializeRequest,
        InitializeRequestParams, InitializedNotification, JsonObject, JsonRpcError,
        ProtocolVersion, RequestId, ServerJsonRpcMessage,
    },
    serve_server,
    service::serve_directly_with_ct,
    transport::{
        OneshotTransport, TransportAdapterIdentity,
        common::{
            http_header::{
                EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_MCP_PROTOCOL_VERSION,
                HEADER_SESSION_ID, JSON_MIME_TYPE,
            },
            mcp_headers,
            server_side_http::{
                BoxResponse, ServerSseMessage, accepted_response, expect_json,
                internal_error_response, sse_stream_response, unexpected_message_response,
            },
        },
    },
};

/// Default maximum POST request body size (4 MiB).
pub(crate) const DEFAULT_MAX_REQUEST_BODY_BYTES: usize = 4 * 1024 * 1024;

#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct StreamableHttpServerConfig {
    /// The ping message duration for SSE connections.
    pub sse_keep_alive: Option<Duration>,
    /// The retry interval for SSE priming events.
    pub sse_retry: Option<Duration>,
    /// If true, the server will create a session for each request and keep it alive.
    /// When enabled, SSE priming events are sent to enable client reconnection.
    pub stateful_mode: bool,
    /// When true and `stateful_mode` is false, the server prefers
    /// `Content-Type: application/json` for simple request-response tools.
    /// If the handler emits a notification or request before the final response,
    /// the server falls back to `text/event-stream` so no message is lost.
    pub json_response: bool,
    /// Cancellation token for the Streamable HTTP server.
    ///
    /// When this token is cancelled, all active sessions are terminated and
    /// the server stops accepting new requests.
    pub cancellation_token: CancellationToken,
    /// Allowed hostnames or `host:port` authorities for inbound `Host` validation.
    ///
    /// By default, Streamable HTTP servers only accept loopback hosts to
    /// prevent DNS rebinding attacks against locally running servers. Public
    /// deployments should override this list with their own hostnames.
    /// examples:
    ///     allowed_hosts = ["localhost", "127.0.0.1", "0.0.0.0"]
    /// or with ports:
    ///     allowed_hosts = ["example.com", "example.com:8080"]
    pub allowed_hosts: Vec<String>,
    /// Allowed browser origins for inbound `Origin` validation.
    ///
    /// Defaults to an empty list, which disables Origin validation. When
    /// non-empty, requests carrying an `Origin` header must match per RFC 6454
    /// `(scheme, host, port)`; missing-`Origin` requests still pass. Entries
    /// must include a scheme; `"null"` matches the browser's `Origin: null`.
    /// examples:
    ///     allowed_origins = ["https://app.example.com", "http://localhost:8080"]
    pub allowed_origins: Vec<String>,
    /// Optional external session store for cross-instance recovery.
    ///
    /// When set, [`SessionState`] (the client's `initialize` parameters) is
    /// persisted after a successful handshake and deleted when the session
    /// closes. On any subsequent request that arrives at an instance with no
    /// in-memory session, the store is consulted: if an entry is found the
    /// session is transparently restored so the client does not need to
    /// re-initialize.
    ///
    /// # Example
    /// ```rust,ignore
    /// use std::sync::Arc;
    /// use rmcp::transport::streamable_http_server::{
    ///     StreamableHttpServerConfig, session::SessionStore,
    /// };
    ///
    /// let config = StreamableHttpServerConfig {
    ///     session_store: Some(Arc::new(MyRedisStore::new())),
    ///     ..Default::default()
    /// };
    /// ```
    pub session_store: Option<Arc<dyn SessionStore>>,
    /// Maximum POST request body size in bytes.
    ///
    /// Enforced while streaming the body, independent of `Content-Length`,
    /// chunked transfer encoding, or HTTP version. Oversized payloads receive
    /// a `413 Payload Too Large` response.
    pub max_request_body_bytes: usize,
}

impl std::fmt::Debug for dyn SessionStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("<SessionStore>")
    }
}

impl Default for StreamableHttpServerConfig {
    fn default() -> Self {
        Self {
            sse_keep_alive: Some(Duration::from_secs(15)),
            sse_retry: Some(Duration::from_secs(3)),
            stateful_mode: true,
            json_response: false,
            cancellation_token: CancellationToken::new(),
            allowed_hosts: vec!["localhost".into(), "127.0.0.1".into(), "::1".into()],
            allowed_origins: vec![],
            session_store: None,
            max_request_body_bytes: DEFAULT_MAX_REQUEST_BODY_BYTES,
        }
    }
}

impl StreamableHttpServerConfig {
    pub fn with_allowed_hosts(
        mut self,
        allowed_hosts: impl IntoIterator<Item = impl Into<String>>,
    ) -> Self {
        self.allowed_hosts = allowed_hosts.into_iter().map(Into::into).collect();
        self
    }
    /// Disable allowed hosts. This will allow requests with any `Host` header, which is NOT recommended for public deployments.
    pub fn disable_allowed_hosts(mut self) -> Self {
        self.allowed_hosts.clear();
        self
    }
    pub fn with_allowed_origins(
        mut self,
        allowed_origins: impl IntoIterator<Item = impl Into<String>>,
    ) -> Self {
        self.allowed_origins = allowed_origins.into_iter().map(Into::into).collect();
        self
    }
    /// Disable Origin validation, reverting to the default ignore-Origin behavior.
    pub fn disable_allowed_origins(mut self) -> Self {
        self.allowed_origins.clear();
        self
    }
    pub fn with_sse_keep_alive(mut self, duration: Option<Duration>) -> Self {
        self.sse_keep_alive = duration;
        self
    }

    pub fn with_sse_retry(mut self, duration: Option<Duration>) -> Self {
        self.sse_retry = duration;
        self
    }

    pub fn with_stateful_mode(mut self, stateful: bool) -> Self {
        self.stateful_mode = stateful;
        self
    }

    pub fn with_json_response(mut self, json_response: bool) -> Self {
        self.json_response = json_response;
        self
    }

    pub fn with_cancellation_token(mut self, token: CancellationToken) -> Self {
        self.cancellation_token = token;
        self
    }

    /// Set the maximum POST request body size in bytes.
    pub fn with_max_request_body_bytes(mut self, bytes: usize) -> Self {
        self.max_request_body_bytes = bytes;
        self
    }
}

#[expect(
    clippy::result_large_err,
    reason = "BoxResponse is intentionally large; matches other handlers in this file"
)]
/// Validates the `MCP-Protocol-Version` header on incoming HTTP requests.
///
/// Per the MCP 2025-06-18 spec:
/// - If the header is present but contains an unsupported version, return 400 Bad Request.
/// - If the header is absent, assume `2025-03-26` for backwards compatibility (no error).
fn validate_protocol_version_header(
    headers: &http::HeaderMap,
    allow_unknown: bool,
) -> Result<(), BoxResponse> {
    if let Some(value) = headers.get(HEADER_MCP_PROTOCOL_VERSION) {
        let version_str = value.to_str().map_err(|_| {
            Response::builder()
                .status(http::StatusCode::BAD_REQUEST)
                .body(
                    Full::new(Bytes::from(
                        "Bad Request: Invalid MCP-Protocol-Version header encoding",
                    ))
                    .boxed(),
                )
                .expect("valid response")
        })?;
        let is_known = ProtocolVersion::KNOWN_VERSIONS
            .iter()
            .any(|v| v.as_str() == version_str);
        if !allow_unknown && !is_known {
            return Err(Response::builder()
                .status(http::StatusCode::BAD_REQUEST)
                .body(
                    Full::new(Bytes::from(format!(
                        "Bad Request: Unsupported MCP-Protocol-Version: {version_str}"
                    )))
                    .boxed(),
                )
                .expect("valid response"));
        }
    }
    Ok(())
}

fn message_has_per_request_protocol_version(message: &ClientJsonRpcMessage) -> bool {
    match message {
        ClientJsonRpcMessage::Request(request) => {
            request.request.get_meta().protocol_version().is_some()
        }
        _ => false,
    }
}

fn invalid_request_jsonrpc_response(
    id: Option<RequestId>,
    message: impl Into<Cow<'static, str>>,
) -> BoxResponse {
    let err = JsonRpcError::new(id, ErrorData::invalid_request(message, None));
    let body = serde_json::to_vec(&err).expect("serialize JsonRpcError");
    Response::builder()
        .status(http::StatusCode::BAD_REQUEST)
        .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE)
        .body(Full::new(Bytes::from(body)).boxed())
        .expect("valid response")
}

fn invalid_params_jsonrpc_response(
    id: Option<RequestId>,
    message: impl Into<Cow<'static, str>>,
) -> BoxResponse {
    let err = JsonRpcError::new(id, ErrorData::invalid_params(message, None));
    let body = serde_json::to_vec(&err).expect("serialize JsonRpcError");
    Response::builder()
        .status(http::StatusCode::BAD_REQUEST)
        .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE)
        .body(Full::new(Bytes::from(body)).boxed())
        .expect("valid response")
}

#[expect(
    clippy::result_large_err,
    reason = "BoxResponse is intentionally large; matches other handlers in this file"
)]
/// Absent header is allowed; the first initialize round-trip may legitimately omit it.
fn validate_header_matches_init_body(
    headers: &http::HeaderMap,
    body_version: &str,
    request_id: Option<RequestId>,
) -> Result<(), BoxResponse> {
    let Some(header_value) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else {
        return Ok(());
    };
    let header_str = header_value.to_str().map_err(|_| {
        invalid_request_jsonrpc_response(
            request_id.clone(),
            "Invalid Request: MCP-Protocol-Version header is not valid UTF-8",
        )
    })?;
    if header_str != body_version {
        tracing::warn!(
            header = header_str,
            body = body_version,
            "rejecting initialize: MCP-Protocol-Version header does not match params.protocolVersion"
        );
        return Err(invalid_request_jsonrpc_response(
            request_id,
            format!(
                "Invalid Request: MCP-Protocol-Version header ({header_str}) does not match initialize params.protocolVersion ({body_version})"
            ),
        ));
    }
    Ok(())
}

#[expect(
    clippy::result_large_err,
    reason = "BoxResponse is intentionally large; matches other handlers in this file"
)]
fn validate_request_protocol_version_meta(
    headers: &HeaderMap,
    message: &ClientJsonRpcMessage,
) -> Result<(), BoxResponse> {
    let ClientJsonRpcMessage::Request(request) = message else {
        return Ok(());
    };
    if matches!(&request.request, ClientRequest::InitializeRequest(_)) {
        return Ok(());
    }
    let is_discover = matches!(&request.request, ClientRequest::DiscoverRequest(_));
    let Some(meta_version) = request.request.get_meta().protocol_version() else {
        if is_discover {
            return Err(invalid_params_jsonrpc_response(
                Some(request.id.clone()),
                "Invalid params: server/discover requires protocolVersion in request _meta",
            ));
        }
        return Ok(());
    };
    let Some(header_version) = headers
        .get(HEADER_MCP_PROTOCOL_VERSION)
        .and_then(|value| value.to_str().ok())
    else {
        return Err(invalid_request_jsonrpc_response(
            Some(request.id.clone()),
            "Invalid Request: request _meta protocolVersion requires MCP-Protocol-Version header",
        ));
    };
    if header_version != meta_version.as_str() {
        return Err(header_mismatch_jsonrpc_response(
            Some(request.id.clone()),
            format!(
                "MCP-Protocol-Version header ({header_version}) does not match request _meta protocolVersion ({meta_version})"
            ),
        ));
    }
    Ok(())
}

fn jsonrpc_http_status(message: &ServerJsonRpcMessage) -> http::StatusCode {
    let ServerJsonRpcMessage::Error(error) = message else {
        return http::StatusCode::OK;
    };
    // Modern per-request HTTP treats invalid params as a malformed request.
    // Legacy requests bypass this mapper and retain HTTP 200 JSON-RPC errors.
    match error.error.code {
        ErrorCode::UNSUPPORTED_PROTOCOL_VERSION
        | ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY
        | ErrorCode::INVALID_PARAMS => http::StatusCode::BAD_REQUEST,
        ErrorCode::METHOD_NOT_FOUND => http::StatusCode::NOT_FOUND,
        _ => http::StatusCode::OK,
    }
}

fn jsonrpc_message_response(
    message: ServerJsonRpcMessage,
    map_protocol_status: bool,
) -> Result<BoxResponse, BoxResponse> {
    let status = if map_protocol_status {
        jsonrpc_http_status(&message)
    } else {
        http::StatusCode::OK
    };
    let body =
        serde_json::to_vec(&message).map_err(internal_error_response("serialize json response"))?;
    Ok(Response::builder()
        .status(status)
        .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE)
        .body(Full::new(Bytes::from(body)).boxed())
        .expect("valid response"))
}

fn header_mismatch_jsonrpc_response(
    id: Option<RequestId>,
    message: impl Into<Cow<'static, str>>,
) -> BoxResponse {
    let err = JsonRpcError::new(id, ErrorData::header_mismatch(message, None));
    let body = serde_json::to_vec(&err).expect("serialize JsonRpcError");
    Response::builder()
        .status(http::StatusCode::BAD_REQUEST)
        .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE)
        .body(Full::new(Bytes::from(body)).boxed())
        .expect("valid response")
}

/// Validates SEP-2243 `Mcp-Method` / `Mcp-Name` / `Mcp-Param-*` headers against the body.
///
/// Only enforced when the request declares a protocol version `>= STANDARD_HEADERS`.
/// The `initialize` handshake is exempt: clients emit these headers only after the
/// version has been negotiated. `tool_schema` supplies the called tool's input schema
/// so annotated `Mcp-Param-*` headers can be checked (no schema => those are skipped).
#[expect(
    clippy::result_large_err,
    reason = "BoxResponse is intentionally large; matches other handlers in this file"
)]
fn validate_standard_headers(
    headers: &HeaderMap,
    message: &ClientJsonRpcMessage,
    tool_schema: impl Fn(&str) -> Option<Arc<JsonObject>>,
) -> Result<(), BoxResponse> {
    let version_requires_headers = headers
        .get(HEADER_MCP_PROTOCOL_VERSION)
        .and_then(|value| value.to_str().ok())
        .is_some_and(|version| version >= ProtocolVersion::STANDARD_HEADERS.as_str());
    if !version_requires_headers {
        return Ok(());
    }

    let request_id = match message {
        ClientJsonRpcMessage::Request(req) => {
            if matches!(&req.request, ClientRequest::InitializeRequest(_)) {
                return Ok(());
            }
            Some(req.id.clone())
        }
        ClientJsonRpcMessage::Notification(_) => None,
        _ => return Ok(()),
    };

    let Ok(value) = serde_json::to_value(message) else {
        return Ok(());
    };
    // For tools/call, look up the tool schema so Mcp-Param-* headers are validated.
    let schema = value
        .get("method")
        .and_then(|method| method.as_str())
        .filter(|method| *method == "tools/call")
        .and_then(|_| value.get("params"))
        .and_then(|params| params.get("name"))
        .and_then(|name| name.as_str())
        .and_then(tool_schema);
    if let Err(reason) = mcp_headers::validate_request_headers(headers, &value, schema.as_deref()) {
        return Err(header_mismatch_jsonrpc_response(request_id, reason));
    }
    Ok(())
}

fn forbidden_response(message: impl Into<String>) -> BoxResponse {
    Response::builder()
        .status(http::StatusCode::FORBIDDEN)
        .body(Full::new(Bytes::from(message.into())).boxed())
        .expect("valid response")
}

fn normalize_host(host: &str) -> String {
    host.trim_matches('[')
        .trim_matches(']')
        .to_ascii_lowercase()
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct NormalizedAuthority {
    host: String,
    port: Option<u16>,
}

fn normalize_authority(host: &str, port: Option<u16>) -> NormalizedAuthority {
    NormalizedAuthority {
        host: normalize_host(host),
        port,
    }
}

fn parse_allowed_authority(allowed: &str) -> Option<NormalizedAuthority> {
    let allowed = allowed.trim();
    if allowed.is_empty() {
        return None;
    }

    if let Ok(authority) = http::uri::Authority::try_from(allowed) {
        return Some(normalize_authority(authority.host(), authority.port_u16()));
    }

    Some(normalize_authority(allowed, None))
}

fn host_is_allowed(host: &NormalizedAuthority, allowed_hosts: &[String]) -> bool {
    if allowed_hosts.is_empty() {
        // If the allowed hosts list is empty, allow all hosts (not recommended).
        return true;
    }
    allowed_hosts
        .iter()
        .filter_map(|allowed| parse_allowed_authority(allowed))
        .any(|allowed| {
            allowed.host == host.host
                && match allowed.port {
                    Some(port) => host.port == Some(port),
                    None => true,
                }
        })
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum NormalizedOrigin {
    Null,
    Tuple {
        scheme: String,
        host: String,
        port: Option<u16>,
    },
}

fn parse_origin_value(value: &str) -> Option<NormalizedOrigin> {
    let value = value.trim();
    if value.is_empty() {
        return None;
    }
    if value.eq_ignore_ascii_case("null") {
        return Some(NormalizedOrigin::Null);
    }
    let uri = http::Uri::try_from(value).ok()?;
    let scheme = uri.scheme_str()?.to_ascii_lowercase();
    let authority = uri.authority()?;
    Some(NormalizedOrigin::Tuple {
        scheme,
        host: normalize_host(authority.host()),
        port: authority.port_u16(),
    })
}

fn origin_is_allowed(origin: &NormalizedOrigin, allowed_origins: &[String]) -> bool {
    if allowed_origins.is_empty() {
        return true;
    }
    allowed_origins
        .iter()
        .filter_map(|raw| parse_origin_value(raw))
        .any(|allowed| match (&allowed, origin) {
            (NormalizedOrigin::Null, NormalizedOrigin::Null) => true,
            (
                NormalizedOrigin::Tuple {
                    scheme: a_scheme,
                    host: a_host,
                    port: a_port,
                },
                NormalizedOrigin::Tuple {
                    scheme: o_scheme,
                    host: o_host,
                    port: o_port,
                },
            ) => a_scheme == o_scheme && a_host == o_host && (a_port.is_none() || a_port == o_port),
            _ => false,
        })
}

fn bad_request_response(message: &str) -> BoxResponse {
    let body = Full::from(message.to_string()).boxed();

    http::Response::builder()
        .status(http::StatusCode::BAD_REQUEST)
        .header(http::header::CONTENT_TYPE, "text/plain; charset=utf-8")
        .body(body)
        .expect("failed to build bad request response")
}

fn parse_host_header(
    uri: &http::Uri,
    headers: &HeaderMap,
) -> Result<NormalizedAuthority, BoxResponse> {
    if let Some(host) = headers.get(http::header::HOST) {
        let host_str = host
            .to_str()
            .inspect_err(|_| {
                tracing::warn!(host = ?host, "rejected request with non-UTF-8 Host header");
            })
            .map_err(|_| bad_request_response("Bad Request: Invalid Host header encoding"))?;
        let authority = http::uri::Authority::try_from(host_str)
            .inspect_err(|_| {
                tracing::warn!(
                    host = host_str,
                    "rejected request with malformed Host header"
                );
            })
            .map_err(|_| bad_request_response("Bad Request: Invalid Host header"))?;
        return Ok(normalize_authority(authority.host(), authority.port_u16()));
    }
    // HTTP/2 carries the host in `:authority`; middleware such as
    // `axum::Router::nest` can drop the `Host` header hyper synthesizes from it.
    let authority = uri.authority().ok_or_else(|| {
        tracing::warn!("rejected request with missing Host header and no :authority");
        bad_request_response("Bad Request: missing Host header")
    })?;
    Ok(normalize_authority(authority.host(), authority.port_u16()))
}

fn validate_dns_rebinding_headers(
    uri: &http::Uri,
    headers: &HeaderMap,
    config: &StreamableHttpServerConfig,
) -> Result<(), BoxResponse> {
    let host = parse_host_header(uri, headers)?;
    if !host_is_allowed(&host, &config.allowed_hosts) {
        tracing::warn!(
            host = ?host,
            "rejected request with disallowed Host header (possible DNS rebinding attempt)",
        );
        return Err(forbidden_response("Forbidden: Host header is not allowed"));
    }
    validate_origin_header(headers, &config.allowed_origins)?;
    Ok(())
}

fn validate_origin_header(
    headers: &HeaderMap,
    allowed_origins: &[String],
) -> Result<(), BoxResponse> {
    if allowed_origins.is_empty() {
        return Ok(());
    }
    let Some(origin_header) = headers.get(http::header::ORIGIN) else {
        return Ok(());
    };
    let origin_str = origin_header
        .to_str()
        .inspect_err(|_| {
            tracing::warn!(origin = ?origin_header, "rejected request with non-UTF-8 Origin header");
        })
        .map_err(|_| bad_request_response("Bad Request: Invalid Origin header encoding"))?;
    let origin = parse_origin_value(origin_str).ok_or_else(|| {
        tracing::warn!(
            origin = origin_str,
            "rejected request with malformed Origin header",
        );
        bad_request_response("Bad Request: Invalid Origin header")
    })?;
    if !origin_is_allowed(&origin, allowed_origins) {
        tracing::warn!(
            origin = ?origin,
            "rejected request with disallowed Origin header (possible cross-origin attack)",
        );
        return Err(forbidden_response(
            "Forbidden: Origin header is not allowed",
        ));
    }
    Ok(())
}

/// # Streamable HTTP server
///
/// An HTTP service that implements the
/// [Streamable HTTP transport](https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#streamable-http)
/// for MCP servers.
///
/// ## Session management
///
/// When [`StreamableHttpServerConfig::stateful_mode`] is `true` (the default),
/// the server creates a session for each client that sends an `initialize`
/// request. The session ID is returned in the `Mcp-Session-Id` response header
/// and the client must include it on all subsequent requests.
///
/// Two tool calls carrying the same `Mcp-Session-Id` come from the same logical
/// session (typically one conversation in an LLM client). Different session IDs
/// mean different sessions.
///
/// The [`SessionManager`] trait controls how sessions are stored and routed:
///
/// * [`LocalSessionManager`](super::session::local::LocalSessionManager) —
///   in-memory session store (default).
/// * [`NeverSessionManager`](super::session::never::NeverSessionManager) —
///   disables sessions entirely (stateless mode).
///
/// ## Accessing HTTP request data from tool handlers
///
/// The service consumes the request body but injects the remaining
/// [`http::request::Parts`] into [`crate::model::Extensions`], which is
/// accessible through [`crate::service::RequestContext`].
///
/// ### Reading the raw HTTP parts
///
/// ```rust
/// use rmcp::handler::server::tool::Extension;
/// use http::request::Parts;
/// async fn my_tool(Extension(parts): Extension<Parts>) {
///     tracing::info!("http parts:{parts:?}")
/// }
/// ```
///
/// ### Reading the session ID inside a tool handler
///
/// ```rust,ignore
/// use rmcp::handler::server::tool::Extension;
/// use rmcp::service::RequestContext;
/// use rmcp::model::RoleServer;
///
/// #[tool(description = "session-aware tool")]
/// async fn my_tool(
///     &self,
///     Extension(parts): Extension<http::request::Parts>,
/// ) -> Result<CallToolResult, rmcp::ErrorData> {
///     if let Some(session_id) = parts.headers.get("mcp-session-id") {
///         tracing::info!(?session_id, "called from session");
///     }
///     // ...
///     # todo!()
/// }
/// ```
///
/// ### Accessing custom axum/tower extension state
///
/// State added via axum's `Extension` layer is available inside
/// `Parts.extensions`:
///
/// ```rust,ignore
/// use rmcp::service::RequestContext;
/// use rmcp::model::RoleServer;
///
/// #[derive(Clone)]
/// struct AppState { /* ... */ }
///
/// #[tool(description = "example")]
/// async fn my_tool(
///     &self,
///     ctx: RequestContext<RoleServer>,
/// ) -> Result<CallToolResult, rmcp::ErrorData> {
///     let parts = ctx.extensions.get::<http::request::Parts>().unwrap();
///     let state = parts.extensions.get::<AppState>().unwrap();
///     // use state...
///     # todo!()
/// }
/// ```
pub struct StreamableHttpService<S, M> {
    pub config: StreamableHttpServerConfig,
    session_manager: Arc<M>,
    service_factory: Arc<dyn Fn() -> Result<S, std::io::Error> + Send + Sync>,
    /// Tracks in-progress session restores so that concurrent requests for the
    /// same unknown session ID wait for the first restore to complete rather
    /// than racing to replay the initialize handshake. `None` when no external
    /// session store is configured (avoids allocating the map).
    pending_restores: Option<
        Arc<tokio::sync::RwLock<HashMap<SessionId, tokio::sync::watch::Sender<Option<bool>>>>>,
    >,
    /// Caches tool input schemas by name for SEP-2243 `Mcp-Param-*` validation.
    /// Populated lazily via `get_tool` so the service factory runs at most once
    /// per tool name. `None` value means the tool exposes no schema.
    tool_schemas: Arc<std::sync::RwLock<HashMap<String, Option<Arc<JsonObject>>>>>,
}

impl<S, M> Clone for StreamableHttpService<S, M> {
    fn clone(&self) -> Self {
        Self {
            config: self.config.clone(),
            session_manager: self.session_manager.clone(),
            service_factory: self.service_factory.clone(),
            pending_restores: self.pending_restores.clone(),
            tool_schemas: self.tool_schemas.clone(),
        }
    }
}

impl<RequestBody, S, M> tower_service::Service<Request<RequestBody>> for StreamableHttpService<S, M>
where
    RequestBody: Body + Send + 'static,
    S: crate::ServerHandler + Send + 'static,
    M: SessionManager,
    RequestBody::Error: Display,
    RequestBody::Data: Send + 'static,
{
    type Response = BoxResponse;
    type Error = Infallible;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
    fn call(&mut self, req: http::Request<RequestBody>) -> Self::Future {
        let service = self.clone();
        Box::pin(async move {
            let response = service.handle(req).await;
            Ok(response)
        })
    }
    fn poll_ready(
        &mut self,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        std::task::Poll::Ready(Ok(()))
    }
}

/// Guard used inside [`StreamableHttpService::try_restore_from_store`].
///
/// Ensures the `pending_restores` map entry is always cleaned up — even when
/// the future is cancelled mid-await.
///
/// `result` defaults to `false` (failure / cancellation). Only the success path
/// needs to set it to `true` before returning.
struct PendingRestoreGuard {
    pending_restores:
        Arc<tokio::sync::RwLock<HashMap<SessionId, tokio::sync::watch::Sender<Option<bool>>>>>,
    session_id: SessionId,
    watch_tx: tokio::sync::watch::Sender<Option<bool>>,
    /// The value that will be broadcast to waiting tasks on drop.
    result: bool,
}

impl Drop for PendingRestoreGuard {
    fn drop(&mut self) {
        // `send` is synchronous — unblocks waiters immediately, no lock needed.
        let _ = self.watch_tx.send(Some(self.result));
        // Remove the map entry asynchronously (requires the async write lock).
        let pending_restores = self.pending_restores.clone();
        let session_id = self.session_id.clone();
        tokio::spawn(async move {
            pending_restores.write().await.remove(&session_id);
        });
    }
}

impl<S, M> StreamableHttpService<S, M>
where
    S: crate::ServerHandler + Send + 'static,
    M: SessionManager,
{
    pub fn new(
        service_factory: impl Fn() -> Result<S, std::io::Error> + Send + Sync + 'static,
        session_manager: Arc<M>,
        config: StreamableHttpServerConfig,
    ) -> Self {
        let pending_restores = config.session_store.is_some().then(|| {
            Arc::new(tokio::sync::RwLock::new(HashMap::<
                SessionId,
                tokio::sync::watch::Sender<Option<bool>>,
            >::new()))
        });
        Self {
            config,
            session_manager,
            service_factory: Arc::new(service_factory),
            pending_restores,
            tool_schemas: Arc::new(std::sync::RwLock::new(HashMap::new())),
        }
    }
    fn get_service(&self) -> Result<S, std::io::Error> {
        (self.service_factory)()
    }

    // The HTTP status must be known before opening an SSE stream.
    async fn serve_negotiated_request_directly(
        &self,
        service: S,
        mut request: crate::model::JsonRpcRequest<ClientRequest>,
        parts: http::request::Parts,
    ) -> Result<BoxResponse, BoxResponse> {
        let peer_info = Self::peer_info_for_stateless_request(&request, &parts.headers);
        request.request.extensions_mut().insert(parts);
        let (transport, mut receiver) =
            OneshotTransport::<RoleServer>::new(ClientJsonRpcMessage::Request(request));
        // Give this stateless request its own cancellation token so a client
        // disconnect can cancel the in-flight handler (#857), as in the
        // non-negotiated stateless path below.
        let request_ct = CancellationToken::new();
        let service = serve_directly_with_ct(service, transport, peer_info, request_ct.clone());
        tokio::spawn(async move {
            let _ = service.waiting().await;
        });

        let cancel = self.config.cancellation_token.child_token();
        // Cancel the handler if the client disconnects while it is still
        // producing its first message (this future is dropped before
        // `receiver.recv()` completes). Disarmed once the handler emits
        // anything, so a normal response is never cancelled.
        let mut disconnect_guard = Some(request_ct.clone().drop_guard());
        let first = tokio::select! {
            message = receiver.recv() => {
                if let Some(guard) = disconnect_guard.take() {
                    guard.disarm();
                }
                message
            }
            _ = cancel.cancelled() => None,
        }
        .ok_or_else(|| {
            internal_error_response("empty response")(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "no response message received from handler",
            ))
        })?;

        if self.config.json_response || jsonrpc_http_status(&first) != http::StatusCode::OK {
            // This message is the whole reply, so `receiver` is dropped here and
            // anything the handler emits afterwards is undeliverable. Cancel it so
            // a still-running handler stops instead of running on unobserved: its
            // terminal `send` would otherwise fail before adding the termination
            // permit, leaving the serve loop parked forever. A no-op when the
            // handler already completed.
            request_ct.cancel();
            return jsonrpc_message_response(first, true);
        }

        // The handler may still be streaming, so guard the response: dropping it
        // (client disconnect) must cancel the handler.
        let stream = futures::stream::once(async move { first })
            .chain(ReceiverStream::new(receiver))
            .map(|message| {
                tracing::trace!(?message);
                ServerSseMessage::from_message(message)
            });
        Ok(sse_stream_response(
            CancelOnDisconnect::new(stream, request_ct),
            self.config.sse_keep_alive,
            self.config.cancellation_token.child_token(),
        ))
    }

    /// Returns the cached input schema for `name`, constructing a service once
    /// per name to read its `ServerHandler::get_tool` definition. Used to
    /// validate SEP-2243 `Mcp-Param-*` headers against the request body.
    fn tool_schema(&self, name: &str) -> Option<Arc<JsonObject>> {
        if let Ok(cache) = self.tool_schemas.read() {
            if let Some(schema) = cache.get(name) {
                return schema.clone();
            }
        }
        let schema = self
            .get_service()
            .ok()
            .and_then(|service| service.get_tool(name))
            .map(|tool| tool.input_schema);
        if let Ok(mut cache) = self.tool_schemas.write() {
            cache.insert(name.to_owned(), schema.clone());
        }
        schema
    }

    /// Spawn a task that runs `serve_server` for the given session, waits for
    /// it to finish, and then calls `close_session`.
    ///
    /// `init_done_tx`: when `Some`, the sender is fired after `serve_server`
    /// returns successfully, signalling to the caller that the MCP handshake
    /// is complete. Used by `try_restore_from_store` to synchronise with the
    /// restore `initialize` replay; `handle_post` passes `None`.
    fn spawn_session_worker(
        session_manager: Arc<M>,
        session_id: SessionId,
        service: S,
        transport: M::Transport,
        init_done_tx: Option<tokio::sync::oneshot::Sender<()>>,
    ) where
        S: crate::ServerHandler + Send + 'static,
        M: SessionManager,
    {
        tokio::spawn(async move {
            let svc =
                serve_server::<S, M::Transport, _, TransportAdapterIdentity>(service, transport)
                    .await;
            match svc {
                Ok(svc) => {
                    if let Some(tx) = init_done_tx {
                        let _ = tx.send(());
                    }
                    let _ = svc.waiting().await;
                }
                Err(e) => {
                    tracing::error!("Failed to serve session: {e}");
                    // Dropping init_done_tx (if Some) signals failure to the caller.
                }
            }
            let _ = session_manager
                .close_session(&session_id)
                .await
                .inspect_err(|e| {
                    tracing::error!("Failed to close session {session_id}: {e}");
                });
        });
    }

    /// Attempt to restore a session from the external store.
    ///
    /// Returns `true` when the session is available and ready to serve the
    /// current request (either just restored or already in memory). Returns
    /// `false` when no store is configured or the session ID is unknown.
    ///
    /// Concurrent requests for the same unknown session ID are serialized: the
    /// first caller performs the full restore and handshake replay while others
    /// subscribe to a `watch` channel and wait, avoiding duplicate handshakes.
    async fn try_restore_from_store(
        &self,
        session_id: &SessionId,
        parts: &http::request::Parts,
    ) -> Result<bool, std::io::Error>
    where
        S: crate::ServerHandler + Send + 'static,
        M: SessionManager,
    {
        // Both fields are Some iff a session store is configured.
        let (Some(pending_restores), Some(store)) =
            (&self.pending_restores, &self.config.session_store)
        else {
            return Ok(false);
        };

        // Serialize concurrent restores for the same session ID.
        // Write-lock once: if another task is already restoring, subscribe and wait;
        // otherwise, register ourselves as the restoring task.
        // Channel value: None = in progress, Some(true) = restored, Some(false) = not found/failed.
        let (watch_tx, _watch_rx) = tokio::sync::watch::channel(None::<bool>);
        {
            let mut pending = pending_restores.write().await;
            if let Some(tx) = pending.get(session_id) {
                let mut rx = tx.subscribe();
                drop(pending);
                // Wait for the restore to finish, then propagate the outcome.
                let result = rx
                    .wait_for(|r| r.is_some())
                    .await
                    .map(|r| r.unwrap_or(false))
                    .unwrap_or(false);
                return Ok(result);
            }
            pending.insert(session_id.clone(), watch_tx.clone());
        }

        // Guard: signals waiters and cleans up the map entry on drop
        let mut guard = PendingRestoreGuard {
            pending_restores: pending_restores.clone(),
            session_id: session_id.clone(),
            watch_tx: watch_tx.clone(),
            result: false,
        };

        // --- Step 3: load from external store ---
        let state = match store.load(session_id.as_ref()).await {
            Ok(Some(s)) => s,
            Ok(None) => {
                return Ok(false);
            }
            Err(e) => {
                tracing::error!(
                    session_id = session_id.as_ref(),
                    error = %e,
                    "session store load failed during restore"
                );
                return Err(std::io::Error::other(e));
            }
        };

        // --- Step 4: ask the session manager to allocate an in-memory worker ---
        let transport = match self
            .session_manager
            .restore_session(session_id.clone())
            .await
            .map_err(|e| std::io::Error::other(e.to_string()))
        {
            Ok(RestoreOutcome::Restored(t)) => t,
            Ok(RestoreOutcome::AlreadyPresent) => {
                // Invariant violation: pending_restores ensures only one task can call
                // restore_session per session ID, so AlreadyPresent is impossible here.
                return Err(std::io::Error::other(
                    "restore_session returned AlreadyPresent unexpectedly; session manager might have modified the session store outside of the restore_session API",
                ));
            }
            Ok(RestoreOutcome::NotSupported) => {
                return Ok(false);
            }
            Err(e) => {
                return Err(e);
            }
        };

        // --- Step 5: replay the MCP initialize handshake ---
        let service = match self.get_service() {
            Ok(s) => s,
            Err(e) => {
                return Err(e);
            }
        };

        // `serve_server` requires both the `initialize` request and the
        // `notifications/initialized` notification before transitioning to
        // the running state — we must send both before returning.
        let mut restore_init = ClientJsonRpcMessage::request(
            ClientRequest::InitializeRequest(InitializeRequest {
                params: state.initialize_params,
                ..Default::default()
            }),
            crate::model::NumberOrString::Number(0),
        );
        restore_init.insert_extension(parts.clone());
        restore_init.insert_extension(SessionRestoreMarker {
            id: session_id.clone(),
        });
        let mut restore_initialized = ClientJsonRpcMessage::notification(
            ClientNotification::InitializedNotification(InitializedNotification {
                ..Default::default()
            }),
        );
        restore_initialized.insert_extension(parts.clone());
        restore_initialized.insert_extension(SessionRestoreMarker {
            id: session_id.clone(),
        });
        // Signal from the spawned task once serve_server finishes initialising.
        let (init_done_tx, init_done_rx) = tokio::sync::oneshot::channel::<()>();

        Self::spawn_session_worker(
            self.session_manager.clone(),
            session_id.clone(),
            service,
            transport,
            Some(init_done_tx),
        );

        if let Err(e) = self
            .session_manager
            .initialize_session(session_id, restore_init)
            .await
            .map_err(|e| std::io::Error::other(e.to_string()))
        {
            return Err(e);
        }

        if let Err(e) = self
            .session_manager
            .accept_message(session_id, restore_initialized)
            .await
            .map_err(|e| std::io::Error::other(e.to_string()))
        {
            return Err(e);
        }

        if init_done_rx.await.is_err() {
            return Err(std::io::Error::other(
                "serve_server initialization failed during restore",
            ));
        }

        // Restore complete — wake any waiting concurrent requests.
        guard.result = true;

        tracing::debug!(
            session_id = session_id.as_ref(),
            "session restored from external store"
        );
        Ok(true)
    }
    pub async fn handle<B>(&self, request: Request<B>) -> Response<BoxBody<Bytes, Infallible>>
    where
        B: Body + Send + 'static,
        B::Error: Display,
    {
        if let Err(response) =
            validate_dns_rebinding_headers(request.uri(), request.headers(), &self.config)
        {
            return response;
        }
        let method = request.method().clone();
        let allowed_methods = match self.config.stateful_mode {
            true => "GET, POST, DELETE",
            false => "POST",
        };
        let result = match (method, self.config.stateful_mode) {
            (Method::POST, _) => self.handle_post(request).await,
            // if we're not in stateful mode, we don't support GET or DELETE because there is no session
            (Method::GET, true) => self.handle_get(request).await,
            (Method::DELETE, true) => self.handle_delete(request).await,
            _ => {
                // Handle other methods or return an error
                let response = Response::builder()
                    .status(http::StatusCode::METHOD_NOT_ALLOWED)
                    .header(ALLOW, allowed_methods)
                    .body(Full::new(Bytes::from("Method Not Allowed")).boxed())
                    .expect("valid response");
                return response;
            }
        };
        match result {
            Ok(response) => response,
            Err(response) => response,
        }
    }
    async fn handle_get<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
    where
        B: Body + Send + 'static,
        B::Error: Display,
    {
        // check accept header
        if !request
            .headers()
            .get(http::header::ACCEPT)
            .and_then(|header| header.to_str().ok())
            .is_some_and(|header| header.contains(EVENT_STREAM_MIME_TYPE))
        {
            return Ok(Response::builder()
                .status(http::StatusCode::NOT_ACCEPTABLE)
                .body(
                    Full::new(Bytes::from(
                        "Not Acceptable: Client must accept text/event-stream",
                    ))
                    .boxed(),
                )
                .expect("valid response"));
        }
        // check session id
        let session_id = request
            .headers()
            .get(HEADER_SESSION_ID)
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_owned().into());
        let Some(session_id) = session_id else {
            // MCP spec: servers that require a session ID SHOULD respond with 400 Bad Request
            return Ok(Response::builder()
                .status(http::StatusCode::BAD_REQUEST)
                .body(Full::new(Bytes::from("Bad Request: Session ID is required")).boxed())
                .expect("valid response"));
        };
        // check if session exists
        let has_session = self
            .session_manager
            .has_session(&session_id)
            .await
            .map_err(internal_error_response("check session"))?;
        let (parts, _) = request.into_parts();
        if !has_session {
            // Attempt transparent cross-instance restore from external store.
            let restored = self
                .try_restore_from_store(&session_id, &parts)
                .await
                .map_err(internal_error_response("restore session"))?;
            if !restored {
                // MCP spec: server MUST respond with 404 Not Found for terminated/unknown sessions
                return Ok(Response::builder()
                    .status(http::StatusCode::NOT_FOUND)
                    .body(Full::new(Bytes::from("Not Found: Session not found")).boxed())
                    .expect("valid response"));
            }
        }
        // Validate MCP-Protocol-Version header (per 2025-06-18 spec)
        validate_protocol_version_header(&parts.headers, false)?;
        // check if last event id is provided
        let last_event_id = parts
            .headers
            .get(HEADER_LAST_EVENT_ID)
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_owned());
        if let Some(last_event_id) = last_event_id {
            match self
                .session_manager
                .resume(&session_id, last_event_id)
                .await
            {
                Ok(stream) => {
                    return Ok(sse_stream_response(
                        stream,
                        self.config.sse_keep_alive,
                        self.config.cancellation_token.child_token(),
                    ));
                }
                Err(e) => {
                    // Return 200 with an immediately-closed empty stream.
                    // Returning an HTTP error would cause EventSource to retry
                    // with the same Last-Event-ID in an infinite loop. An empty
                    // 200 cleanly terminates the EventSource without delivering
                    // events from a different stream.
                    tracing::warn!("Resume failed ({e}), returning empty stream");
                    return Ok(sse_stream_response(
                        futures::stream::empty(),
                        None,
                        self.config.cancellation_token.child_token(),
                    ));
                }
            }
        }
        // No Last-Event-ID — create standalone stream
        let stream = self
            .session_manager
            .create_standalone_stream(&session_id)
            .await
            .map_err(internal_error_response("create standalone stream"))?;
        let stream = if let Some(retry) = self.config.sse_retry {
            let priming = ServerSseMessage::priming("0", retry);
            futures::stream::once(async move { priming })
                .chain(stream)
                .left_stream()
        } else {
            stream.right_stream()
        };
        Ok(sse_stream_response(
            stream,
            self.config.sse_keep_alive,
            self.config.cancellation_token.child_token(),
        ))
    }

    async fn handle_post<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
    where
        B: Body + Send + 'static,
        B::Error: Display,
    {
        // check accept header
        if !request
            .headers()
            .get(http::header::ACCEPT)
            .and_then(|header| header.to_str().ok())
            .is_some_and(|header| {
                header.contains(JSON_MIME_TYPE) && header.contains(EVENT_STREAM_MIME_TYPE)
            })
        {
            return Ok(Response::builder()
                .status(http::StatusCode::NOT_ACCEPTABLE)
                .body(Full::new(Bytes::from("Not Acceptable: Client must accept both application/json and text/event-stream")).boxed())
                .expect("valid response"));
        }

        // check content type
        if !request
            .headers()
            .get(http::header::CONTENT_TYPE)
            .and_then(|header| header.to_str().ok())
            .is_some_and(|header| header.starts_with(JSON_MIME_TYPE))
        {
            return Ok(Response::builder()
                .status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
                .body(
                    Full::new(Bytes::from(
                        "Unsupported Media Type: Content-Type must be application/json",
                    ))
                    .boxed(),
                )
                .expect("valid response"));
        }

        // json deserialize request body
        let (part, body) = request.into_parts();
        let mut message = match expect_json(body, self.config.max_request_body_bytes).await {
            Ok(message) => message,
            Err(response) => return Ok(response),
        };

        if self.config.stateful_mode {
            // do we have a session id?
            let session_id = part
                .headers
                .get(HEADER_SESSION_ID)
                .and_then(|v| v.to_str().ok());
            if let Some(session_id) = session_id {
                let session_id = session_id.to_owned().into();
                let has_session = self
                    .session_manager
                    .has_session(&session_id)
                    .await
                    .map_err(internal_error_response("check session"))?;
                if !has_session {
                    // Attempt transparent cross-instance restore from external store.
                    let restored = self
                        .try_restore_from_store(&session_id, &part)
                        .await
                        .map_err(internal_error_response("restore session"))?;
                    if !restored {
                        // MCP spec: server MUST respond with 404 Not Found for terminated/unknown sessions
                        return Ok(Response::builder()
                            .status(http::StatusCode::NOT_FOUND)
                            .body(Full::new(Bytes::from("Not Found: Session not found")).boxed())
                            .expect("valid response"));
                    }
                }

                // Validate MCP-Protocol-Version header (per 2025-06-18 spec)
                let has_per_request_version = message_has_per_request_protocol_version(&message);
                validate_protocol_version_header(&part.headers, has_per_request_version)?;
                validate_request_protocol_version_meta(&part.headers, &message)?;
                // Validate SEP-2243 standard headers against the body
                validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?;

                // inject request part to extensions
                match &mut message {
                    ClientJsonRpcMessage::Request(req) => {
                        req.request.extensions_mut().insert(part);
                    }
                    ClientJsonRpcMessage::Notification(not) => {
                        not.notification.extensions_mut().insert(part);
                    }
                    _ => {
                        // skip
                    }
                }

                match message {
                    ClientJsonRpcMessage::Request(_) => {
                        // Priming for request-wise streams is handled by the
                        // session layer (SessionManager::create_stream) which
                        // has access to the http_request_id for correct event IDs.
                        let stream = self
                            .session_manager
                            .create_stream(&session_id, message)
                            .await
                            .map_err(internal_error_response("get session"))?;
                        Ok(sse_stream_response(
                            stream,
                            self.config.sse_keep_alive,
                            self.config.cancellation_token.child_token(),
                        ))
                    }
                    ClientJsonRpcMessage::Notification(_)
                    | ClientJsonRpcMessage::Response(_)
                    | ClientJsonRpcMessage::Error(_) => {
                        // handle notification
                        self.session_manager
                            .accept_message(&session_id, message)
                            .await
                            .map_err(internal_error_response("accept message"))?;
                        Ok(accepted_response())
                    }
                }
            } else {
                if matches!(
                    &message,
                    ClientJsonRpcMessage::Request(request)
                        if matches!(&request.request, ClientRequest::DiscoverRequest(_))
                ) {
                    validate_protocol_version_header(
                        &part.headers,
                        message_has_per_request_protocol_version(&message),
                    )?;
                    validate_standard_headers(&part.headers, &message, |name| {
                        self.tool_schema(name)
                    })?;
                    validate_request_protocol_version_meta(&part.headers, &message)?;
                    let ClientJsonRpcMessage::Request(request) = message else {
                        unreachable!("guarded as a request above");
                    };
                    let service = self
                        .get_service()
                        .map_err(internal_error_response("get service"))?;
                    return self
                        .serve_negotiated_request_directly(service, request, part)
                        .await;
                }
                // Capture init params for external store persistence before
                // extensions are injected (which would require Clone).
                let stored_init_params = match &mut message {
                    ClientJsonRpcMessage::Request(req) => {
                        let ClientRequest::InitializeRequest(init_req) = &req.request else {
                            return Err(unexpected_message_response("initialize request"));
                        };
                        // Reject mismatched MCP-Protocol-Version header before binding the session to anything.
                        validate_header_matches_init_body(
                            &part.headers,
                            init_req.params.protocol_version.as_str(),
                            Some(req.id.clone()),
                        )?;
                        let stored_init_params = self
                            .config
                            .session_store
                            .as_ref()
                            .map(|_| init_req.params.clone());
                        // inject request part to extensions
                        req.request.extensions_mut().insert(part);
                        stored_init_params
                    }
                    _ => {
                        return Err(unexpected_message_response("initialize request"));
                    }
                };
                let service = self
                    .get_service()
                    .map_err(internal_error_response("get service"))?;
                let (session_id, transport) = self
                    .session_manager
                    .create_session()
                    .await
                    .map_err(internal_error_response("create session"))?;
                // spawn a task to serve the session
                Self::spawn_session_worker(
                    self.session_manager.clone(),
                    session_id.clone(),
                    service,
                    transport,
                    None,
                );
                // get initialize response
                let response = self
                    .session_manager
                    .initialize_session(&session_id, message)
                    .await
                    .map_err(internal_error_response("create stream"))?;
                // Persist session state to external store after a successful handshake.
                if let (Some(store), Some(params)) =
                    (&self.config.session_store, stored_init_params)
                {
                    let state = SessionState {
                        initialize_params: params,
                    };
                    let _ = store
                        .store(session_id.as_ref(), &state)
                        .await
                        .inspect_err(|e| {
                            tracing::warn!(
                                "Failed to persist session {} to store: {e}",
                                session_id
                            );
                        });
                }
                let stream =
                    futures::stream::once(async move { ServerSseMessage::from_message(response) });
                // Prepend priming event if sse_retry configured
                let stream = if let Some(retry) = self.config.sse_retry {
                    let priming = ServerSseMessage::priming("0", retry);
                    futures::stream::once(async move { priming })
                        .chain(stream)
                        .left_stream()
                } else {
                    stream.right_stream()
                };
                let mut response = sse_stream_response(
                    stream,
                    self.config.sse_keep_alive,
                    self.config.cancellation_token.child_token(),
                );

                response.headers_mut().insert(
                    HEADER_SESSION_ID,
                    session_id
                        .parse()
                        .map_err(internal_error_response("create session id header"))?,
                );
                Ok(response)
            }
        } else {
            // Stateless mode:
            // - on initialize: the header (if present) must match `params.protocolVersion`
            // - on every other request: the header must name a known version.
            let has_per_request_version = message_has_per_request_protocol_version(&message);
            match &message {
                ClientJsonRpcMessage::Request(req) => {
                    if let ClientRequest::InitializeRequest(init_req) = &req.request {
                        validate_header_matches_init_body(
                            &part.headers,
                            init_req.params.protocol_version.as_str(),
                            Some(req.id.clone()),
                        )?;
                    } else {
                        validate_protocol_version_header(&part.headers, has_per_request_version)?;
                    }
                }
                _ => {
                    validate_protocol_version_header(&part.headers, has_per_request_version)?;
                }
            }
            // Validate SEP-2243 standard headers against the body
            validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?;
            validate_request_protocol_version_meta(&part.headers, &message)?;
            let service = self
                .get_service()
                .map_err(internal_error_response("get service"))?;
            match message {
                ClientJsonRpcMessage::Request(mut request) => {
                    let negotiates_per_request = has_per_request_version
                        || matches!(&request.request, ClientRequest::DiscoverRequest(_));
                    if negotiates_per_request {
                        return self
                            .serve_negotiated_request_directly(service, request, part)
                            .await;
                    }
                    // Build a peer_info so context.protocol_version() works inside handlers.
                    // serve_directly skips the handshake and receives None by default, making
                    // protocol_version() always return None in stateless mode. We reconstruct it:
                    // - initialize requests: version comes from the request body params
                    // - all other requests: version comes from the MCP-Protocol-Version header
                    //   (already validated above; absent header defaults to 2025-03-26)
                    let peer_info = Self::peer_info_for_stateless_request(&request, &part.headers);
                    request.request.extensions_mut().insert(part);
                    let (transport, mut receiver) =
                        OneshotTransport::<RoleServer>::new(ClientJsonRpcMessage::Request(request));
                    // Give this stateless request its own cancellation token so a
                    // client disconnect can cancel the in-flight handler (#857). A
                    // stateless request is one-shot (no session, no resumption), so a
                    // dropped response is terminal and safe to cancel.
                    let request_ct = CancellationToken::new();
                    let service =
                        serve_directly_with_ct(service, transport, peer_info, request_ct.clone());
                    tokio::spawn(async move {
                        // on service created
                        let _ = service.waiting().await;
                    });
                    if self.config.json_response {
                        // Prefer JSON for a terminal first message. If the handler
                        // emits an intermediate notification or request, preserve
                        // the complete message sequence by falling back to SSE.
                        let cancel = self.config.cancellation_token.child_token();
                        // Cancel the handler if the client disconnects while it is
                        // still producing its first message (this future is dropped
                        // before `receiver.recv()` completes). Disarmed once the
                        // handler emits anything, so a normal response is never
                        // cancelled.
                        let mut disconnect_guard = Some(request_ct.clone().drop_guard());
                        let Some(message) = (tokio::select! {
                            res = receiver.recv() => {
                                if let Some(guard) = disconnect_guard.take() {
                                    guard.disarm();
                                }
                                res
                            }
                            _ = cancel.cancelled() => None,
                        }) else {
                            return Err(internal_error_response("empty response")(
                                std::io::Error::new(
                                    std::io::ErrorKind::UnexpectedEof,
                                    "no response message received from handler",
                                ),
                            ));
                        };
                        tracing::trace!(?message);
                        if matches!(
                            message,
                            ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_)
                        ) {
                            let body = serde_json::to_vec(&message).map_err(|e| {
                                internal_error_response("serialize json response")(e)
                            })?;
                            Ok(Response::builder()
                                .status(http::StatusCode::OK)
                                .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE)
                                .body(Full::new(Bytes::from(body)).boxed())
                                .expect("valid response"))
                        } else {
                            // The handler emitted an intermediate message and is still
                            // running, so guard the streamed sequence too: dropping it
                            // (client disconnect) must cancel the handler.
                            let first = futures::stream::once(async move {
                                ServerSseMessage::from_message(message)
                            });
                            let remaining = ReceiverStream::new(receiver).map(|message| {
                                tracing::trace!(?message);
                                ServerSseMessage::from_message(message)
                            });
                            Ok(sse_stream_response(
                                CancelOnDisconnect::new(first.chain(remaining), request_ct),
                                self.config.sse_keep_alive,
                                self.config.cancellation_token.child_token(),
                            ))
                        }
                    } else {
                        // SSE mode (default): cancel the handler if the client
                        // disconnects (drops the response stream) before it completes.
                        let stream = ReceiverStream::new(receiver).map(|message| {
                            tracing::trace!(?message);
                            ServerSseMessage::from_message(message)
                        });
                        let stream = CancelOnDisconnect::new(stream, request_ct);
                        Ok(sse_stream_response(
                            stream,
                            self.config.sse_keep_alive,
                            self.config.cancellation_token.child_token(),
                        ))
                    }
                }
                ClientJsonRpcMessage::Notification(_notification) => {
                    // ignore
                    Ok(accepted_response())
                }
                ClientJsonRpcMessage::Response(_json_rpc_response) => Ok(accepted_response()),
                ClientJsonRpcMessage::Error(_json_rpc_error) => Ok(accepted_response()),
            }
        }
    }

    async fn handle_delete<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
    where
        B: Body + Send + 'static,
        B::Error: Display,
    {
        // check session id
        let session_id = request
            .headers()
            .get(HEADER_SESSION_ID)
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_owned().into());
        let Some(session_id) = session_id else {
            // MCP spec: servers that require a session ID SHOULD respond with 400 Bad Request
            return Ok(Response::builder()
                .status(http::StatusCode::BAD_REQUEST)
                .body(Full::new(Bytes::from("Bad Request: Session ID is required")).boxed())
                .expect("valid response"));
        };
        // Validate MCP-Protocol-Version header (per 2025-06-18 spec)
        validate_protocol_version_header(request.headers(), false)?;
        // close session
        self.session_manager
            .close_session(&session_id)
            .await
            .map_err(internal_error_response("close session"))?;
        // Remove from external store: a DELETE means the client intentionally
        // ends the session, so the store entry is no longer needed.
        if let Some(store) = &self.config.session_store {
            let _ = store.delete(session_id.as_ref()).await.inspect_err(|e| {
                tracing::warn!("Failed to delete session {} from store: {e}", session_id);
            });
        }
        Ok(accepted_response())
    }

    /// Build a `ClientInfo` (peer_info) for a stateless request so that
    /// `context.protocol_version()` returns the correct value inside handlers.
    ///
    /// `serve_directly` skips the MCP handshake and accepts `peer_info = None`,
    /// which means `context.protocol_version()` is always `None` in stateless mode.
    /// We reconstruct the protocol version from the available signal per request type:
    /// - initialize: version is in the request body params (authoritative)
    /// - all other requests: version is in the MCP-Protocol-Version header
    ///   (validated before this point; absent header defaults to 2025-03-26)
    fn peer_info_for_stateless_request(
        request: &crate::model::JsonRpcRequest<ClientRequest>,
        headers: &HeaderMap,
    ) -> Option<InitializeRequestParams> {
        let version = if let ClientRequest::InitializeRequest(ref init) = request.request {
            init.params.protocol_version.clone()
        } else {
            headers
                .get(HEADER_MCP_PROTOCOL_VERSION)
                .and_then(|v| v.to_str().ok())
                .and_then(|s| serde_json::from_value(serde_json::Value::String(s.to_owned())).ok())
                .unwrap_or(ProtocolVersion::V_2025_03_26)
        };
        Some(InitializeRequestParams {
            meta: None,
            protocol_version: version,
            capabilities: ClientCapabilities::default(),
            client_info: Implementation::default(),
        })
    }
}

pin_project! {
    /// Wraps a stateless SSE response stream so a client disconnect cancels the
    /// in-flight request.
    ///
    /// A stateless streamable-HTTP request is one-shot: it has no session and no
    /// resumption, so a dropped response stream means the client is gone for
    /// good. When the stream is dropped *before* it ends naturally, the request's
    /// cancellation token is fired, which stops the dedicated `serve_directly`
    /// loop and cancels the handler's `RequestContext::ct` (see #857). If the
    /// stream ends naturally (the request completed), the guard is disarmed so
    /// normal completion cancels nothing.
    struct CancelOnDisconnect<S> {
        #[pin]
        inner: S,
        ct: Option<CancellationToken>,
    }
    impl<S> PinnedDrop for CancelOnDisconnect<S> {
        fn drop(this: Pin<&mut Self>) {
            let this = this.project();
            if let Some(ct) = this.ct.take() {
                ct.cancel();
            }
        }
    }
}

impl<S> CancelOnDisconnect<S> {
    fn new(inner: S, ct: CancellationToken) -> Self {
        Self {
            inner,
            ct: Some(ct),
        }
    }
}

impl<S: Stream> Stream for CancelOnDisconnect<S> {
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.project();
        let polled = this.inner.poll_next(cx);
        if let Poll::Ready(None) = &polled {
            // Ended naturally: the request completed, so don't cancel on drop.
            *this.ct = None;
        }
        polled
    }
}