turul_http_mcp_server/

streamable_http.rs

//! Streamable HTTP Transport for MCP 2025-11-25
//!
//! This module implements the "Streamable HTTP" transport mechanism introduced
//! in MCP 2025-03-26, which replaces the previous HTTP+SSE approach from 2024-11-05.
//!
//! ## Key Improvements over HTTP+SSE
//! - **Serverless Compatibility**: Enables deployment on AWS Lambda, Google Cloud Run
//! - **Improved Scalability**: Supports chunked transfer encoding and progressive delivery
//! - **Session Management**: Cryptographically secure session IDs for connection tracking
//! - **Enterprise Network Friendly**: No long-lived connections or polling requirements

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;

use bytes::Bytes;
use futures::Stream;
use http_body::Body;
use http_body_util::{BodyExt, Full};
use hyper::header::{ACCEPT, CONTENT_TYPE};
use hyper::{HeaderMap, Method, Request, Response, StatusCode};
use serde_json::Value;
use tracing::{debug, error, warn};
use turul_mcp_session_storage::SessionView;

use crate::ServerConfig;
use crate::protocol::normalize_header_value;

/// MCP Protocol versions
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum McpProtocolVersion {
    /// Original protocol without streamable HTTP (2024-11-05)
    V2024_11_05,
    /// Protocol including streamable HTTP (2025-03-26)
    V2025_03_26,
    /// Protocol with structured _meta, cursor, progressToken, and elicitation (2025-06-18)
    V2025_06_18,
    /// Protocol with tasks, icons, URL elicitation, and sampling tools (2025-11-25)
    #[default]
    V2025_11_25,
}

impl McpProtocolVersion {
    /// Parse from header string
    pub fn parse_version(s: &str) -> Option<Self> {
        match s {
            "2024-11-05" => Some(Self::V2024_11_05),
            "2025-03-26" => Some(Self::V2025_03_26),
            "2025-06-18" => Some(Self::V2025_06_18),
            "2025-11-25" => Some(Self::V2025_11_25),
            _ => None,
        }
    }

    /// Convert to string representation
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::V2024_11_05 => "2024-11-05",
            Self::V2025_03_26 => "2025-03-26",
            Self::V2025_06_18 => "2025-06-18",
            Self::V2025_11_25 => "2025-11-25",
        }
    }

    /// Returns whether this version supports streamable HTTP
    pub fn supports_streamable_http(&self) -> bool {
        matches!(
            self,
            Self::V2025_03_26 | Self::V2025_06_18 | Self::V2025_11_25
        )
    }

    /// Returns whether this version supports _meta fields
    pub fn supports_meta_fields(&self) -> bool {
        matches!(self, Self::V2025_06_18 | Self::V2025_11_25)
    }

    /// Returns whether this version supports cursor-based pagination
    pub fn supports_cursors(&self) -> bool {
        matches!(self, Self::V2025_06_18 | Self::V2025_11_25)
    }

    /// Returns whether this version supports progress tokens
    pub fn supports_progress_tokens(&self) -> bool {
        matches!(self, Self::V2025_06_18 | Self::V2025_11_25)
    }

    /// Returns whether this version supports elicitation
    pub fn supports_elicitation(&self) -> bool {
        matches!(self, Self::V2025_06_18 | Self::V2025_11_25)
    }

    /// Returns whether this version supports the task system (experimental)
    pub fn supports_tasks(&self) -> bool {
        matches!(self, Self::V2025_11_25)
    }

    /// Returns whether this version supports icons
    pub fn supports_icons(&self) -> bool {
        matches!(self, Self::V2025_11_25)
    }

    /// Get list of supported features for this version
    pub fn supported_features(&self) -> Vec<&'static str> {
        let mut features = vec![];
        if self.supports_streamable_http() {
            features.push("streamable-http");
        }
        if self.supports_meta_fields() {
            features.push("_meta-fields");
        }
        if self.supports_cursors() {
            features.push("cursor-pagination");
        }
        if self.supports_progress_tokens() {
            features.push("progress-tokens");
        }
        if self.supports_elicitation() {
            features.push("elicitation");
        }
        if self.supports_tasks() {
            features.push("tasks");
        }
        if self.supports_icons() {
            features.push("icons");
        }
        features
    }
}

/// Streamable HTTP request context
#[derive(Debug, Clone)]
pub struct StreamableHttpContext {
    /// Protocol version negotiated
    pub protocol_version: McpProtocolVersion,
    /// Session ID if provided
    pub session_id: Option<String>,
    /// Whether client wants SSE stream (text/event-stream)
    pub wants_sse_stream: bool,
    /// Whether client accepts stream frames (application/json, text/event-stream, or */*)
    pub accepts_stream_frames: bool,
    /// Additional request headers
    pub headers: HashMap<String, String>,
}

impl StreamableHttpContext {
    /// Parse context from HTTP request headers
    pub fn from_request<T>(req: &Request<T>) -> Self {
        let headers = req.headers();

        // Parse protocol version from MCP-Protocol-Version header
        let protocol_version = headers
            .get("MCP-Protocol-Version")
            .and_then(|h| h.to_str().ok())
            .and_then(McpProtocolVersion::parse_version)
            .unwrap_or_default();

        // Extract session ID from Mcp-Session-Id header (note capitalization)
        let session_id = headers
            .get("Mcp-Session-Id")
            .and_then(|h| h.to_str().ok())
            .map(|s| s.to_string());

        // Check Accept header for streaming and JSON support
        let accept_header = headers
            .get(ACCEPT)
            .and_then(|h| h.to_str().ok())
            .map(normalize_header_value)
            .unwrap_or_default();

        let wants_sse_stream = accept_header.contains("text/event-stream");
        let accepts_stream_frames = accept_header.contains("application/json")
            || accept_header.contains("text/event-stream")
            || accept_header.contains("*/*");

        // Collect additional headers for debugging/logging
        let mut header_map = HashMap::new();
        for (name, value) in headers.iter() {
            if let Ok(value_str) = value.to_str() {
                header_map.insert(name.to_string(), value_str.to_string());
            }
        }

        Self {
            protocol_version,
            session_id,
            wants_sse_stream,
            accepts_stream_frames,
            headers: header_map,
        }
    }

    /// Whether client wants SSE stream
    pub fn wants_sse_stream(&self) -> bool {
        self.wants_sse_stream
    }

    /// Whether client wants streaming POST responses
    pub fn wants_streaming_post(&self) -> bool {
        self.accepts_stream_frames && self.wants_sse_stream
    }

    /// Check if request is compatible with streamable HTTP
    pub fn is_streamable_compatible(&self) -> bool {
        self.protocol_version.supports_streamable_http() && self.accepts_stream_frames
    }

    /// Validate request for MCP compliance
    pub fn validate(&self, method: &Method) -> std::result::Result<(), String> {
        if !self.accepts_stream_frames {
            return Err(
                "Accept header must include application/json, text/event-stream, or */*"
                    .to_string(),
            );
        }

        if self.wants_sse_stream && !self.protocol_version.supports_streamable_http() {
            return Err(format!(
                "Protocol version {} does not support streamable HTTP",
                self.protocol_version.as_str()
            ));
        }

        // Only enforce session_id for GET requests with SSE streams
        // POST requests will validate session based on the JSON-RPC method (initialize vs others)
        if *method == Method::GET && self.wants_sse_stream && self.session_id.is_none() {
            return Err("Mcp-Session-Id header required for SSE streaming connections".to_string());
        }

        Ok(())
    }

    /// Create response headers for this context
    pub fn response_headers(&self) -> HeaderMap {
        let mut headers = HeaderMap::new();

        // Always include protocol version in response
        headers.insert(
            "MCP-Protocol-Version",
            self.protocol_version.as_str().parse().unwrap(),
        );

        // Include session ID if present
        if let Some(session_id) = &self.session_id {
            headers.insert("Mcp-Session-Id", session_id.parse().unwrap());
        }

        // Add capabilities header showing supported features
        let features = self.protocol_version.supported_features();
        if !features.is_empty() {
            headers.insert("MCP-Capabilities", features.join(",").parse().unwrap());
        }

        headers
    }
}

/// Why session validation failed — determines the HTTP status code.
///
/// MCP 2025-11-25 spec: terminated or unknown sessions MUST return 404 Not Found.
/// Missing `Mcp-Session-Id` header (a different code path) stays 401 Unauthorized.
enum SessionValidationError {
    /// Session ID not found or session has been terminated → 404 Not Found
    NotFound(String),
    /// Storage backend error → 500 Internal Server Error
    StorageError(String),
}

impl SessionValidationError {
    fn status_code(&self) -> StatusCode {
        match self {
            Self::NotFound(_) => StatusCode::NOT_FOUND,
            Self::StorageError(_) => StatusCode::INTERNAL_SERVER_ERROR,
        }
    }

    fn message(&self) -> &str {
        match self {
            Self::NotFound(msg) | Self::StorageError(msg) => msg,
        }
    }
}

/// Streamable HTTP response types
pub enum StreamableResponse {
    /// Single JSON response
    Json(Value),
    /// Streaming response with multiple JSON messages
    Stream(Pin<Box<dyn Stream<Item = std::result::Result<Value, String>> + Send>>),
    /// Error response
    Error { status: StatusCode, message: String },
}

impl std::fmt::Debug for StreamableResponse {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Json(value) => f.debug_tuple("Json").field(value).finish(),
            Self::Stream(_) => f.debug_tuple("Stream").field(&"<stream>").finish(),
            Self::Error { status, message } => f
                .debug_struct("Error")
                .field("status", status)
                .field("message", message)
                .finish(),
        }
    }
}

impl StreamableResponse {
    /// Convert to HTTP response
    pub fn into_response(self, context: &StreamableHttpContext) -> Response<Full<Bytes>> {
        let mut response_headers = context.response_headers();

        match self {
            StreamableResponse::Json(json) => {
                response_headers.insert(CONTENT_TYPE, "application/json".parse().unwrap());

                let body = serde_json::to_string(&json)
                    .unwrap_or_else(|_| r#"{"error": "Failed to serialize response"}"#.to_string());

                Response::builder()
                    .status(StatusCode::OK)
                    .body(Full::new(Bytes::from(body)))
                    .unwrap()
            }

            StreamableResponse::Stream(_stream) => {
                // For streaming responses, set appropriate headers
                response_headers.insert(CONTENT_TYPE, "text/event-stream".parse().unwrap());
                response_headers.insert("Cache-Control", "no-cache, no-transform".parse().unwrap());
                response_headers.insert("Connection", "keep-alive".parse().unwrap());

                // TODO: Implement actual streaming body with chunked transfer encoding
                // Should stream JSON messages over HTTP with proper Content-Type: text/event-stream
                // For now, return 202 Accepted to indicate streaming would happen
                Response::builder()
                    .status(StatusCode::ACCEPTED)
                    .body(Full::new(Bytes::from("Streaming response accepted")))
                    .unwrap()
            }

            StreamableResponse::Error { status, message } => {
                response_headers.insert(CONTENT_TYPE, "application/json".parse().unwrap());

                let error_json = serde_json::json!({
                    "error": {
                        "code": status.as_u16(),
                        "message": message
                    }
                });

                let body = serde_json::to_string(&error_json).unwrap_or_else(|_| {
                    r#"{"error": {"code": 500, "message": "Internal server error"}}"#.to_string()
                });

                Response::builder()
                    .status(status)
                    .body(Full::new(Bytes::from(body)))
                    .unwrap()
            }
        }
    }

    /// Convert to HTTP response with UnsyncBoxBody for streaming compatibility
    pub fn into_boxed_response(
        self,
        context: &StreamableHttpContext,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>> {
        self.into_response(context)
            .map(|body| body.map_err(|never| match never {}).boxed_unsync())
    }
}

/// Streamable HTTP transport handler
#[derive(Clone)]
pub struct StreamableHttpHandler {
    config: Arc<ServerConfig>,
    dispatcher: Arc<turul_mcp_json_rpc_server::JsonRpcDispatcher<turul_mcp_protocol::McpError>>,
    session_storage: Arc<turul_mcp_session_storage::BoxedSessionStorage>,
    stream_manager: Arc<crate::StreamManager>,
    server_capabilities: turul_mcp_protocol::ServerCapabilities,
    pub(crate) middleware_stack: Arc<crate::middleware::MiddlewareStack>,
}

impl StreamableHttpHandler {
    pub fn new(
        config: Arc<ServerConfig>,
        dispatcher: Arc<turul_mcp_json_rpc_server::JsonRpcDispatcher<turul_mcp_protocol::McpError>>,
        session_storage: Arc<turul_mcp_session_storage::BoxedSessionStorage>,
        stream_manager: Arc<crate::StreamManager>,
        server_capabilities: turul_mcp_protocol::ServerCapabilities,
        middleware_stack: Arc<crate::middleware::MiddlewareStack>,
    ) -> Self {
        Self {
            config,
            dispatcher,
            session_storage,
            stream_manager,
            server_capabilities,
            middleware_stack,
        }
    }

    /// Handle incoming HTTP request with streamable HTTP support
    pub async fn handle_request<T>(
        &self,
        req: Request<T>,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>>
    where
        T: Body + Send + 'static,
        T::Data: Send,
        T::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    {
        debug!(
            "Streamable handler request: method={}, uri={}",
            req.method(),
            req.uri()
        );
        // Parse streamable HTTP context from request
        let context = StreamableHttpContext::from_request(&req);

        debug!(
            "Streamable handler entry: method={}, protocol={}, session={:?}, accepts_stream_frames={}, wants_sse_stream={}",
            req.method(),
            context.protocol_version.as_str(),
            context.session_id,
            context.accepts_stream_frames,
            context.wants_sse_stream()
        );

        // Handle OPTIONS preflight before validation (no Accept header required).
        // CORS headers are added by CorsLayer in server.rs — not here.
        if *req.method() == Method::OPTIONS {
            return Response::builder()
                .status(StatusCode::OK)
                .body(Full::new(Bytes::new()))
                .unwrap()
                .map(|body| body.map_err(|never| match never {}).boxed_unsync());
        }

        // Validate request
        if let Err(error) = context.validate(req.method()) {
            warn!("Invalid streamable HTTP request: {}", error);
            return StreamableResponse::Error {
                status: StatusCode::BAD_REQUEST,
                message: error,
            }
            .into_boxed_response(&context);
        }

        // Route based on MCP 2025-11-25 specification
        match *req.method() {
            Method::POST => {
                // ALL client messages (requests, notifications, responses) come via POST
                // Server decides whether to respond with JSON or SSE stream
                self.handle_client_message(req, context).await
            }
            Method::GET => {
                // Optional SSE stream for server-initiated messages
                self.handle_get_sse_notifications(req, context).await
            }
            Method::DELETE => {
                // Optional session cleanup
                self.handle_session_delete(req, context).await
            }
            _ => StreamableResponse::Error {
                status: StatusCode::METHOD_NOT_ALLOWED,
                message: "Method not allowed for this endpoint".to_string(),
            }
            .into_boxed_response(&context),
        }
    }

    /// Handle GET request for long-lived server-initiated notifications (GET SSE)
    ///
    /// This is traditional Server-Sent Events - a long-lived GET connection for
    /// server-initiated notifications unrelated to specific client requests.
    /// NOT used for tool progress (that's POST Streamable HTTP).
    async fn handle_get_sse_notifications<T>(
        &self,
        req: Request<T>,
        context: StreamableHttpContext,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>>
    where
        T: Body + Send + 'static,
    {
        debug!(
            "Opening streaming connection for session: {:?}",
            context.session_id
        );

        // 1. Validate session exists and is authorized
        let session_id = match context.session_id {
            Some(ref id) => id.clone(),
            None => {
                warn!("Missing session ID for streaming GET request");
                return StreamableResponse::Error {
                    status: StatusCode::BAD_REQUEST,
                    message: "Mcp-Session-Id header required for streaming connection".to_string(),
                }
                .into_boxed_response(&context);
            }
        };

        // Validate session exists (do NOT create if missing)
        match self.validate_session_exists(&session_id).await {
            Ok(_) => {
                debug!(
                    "Session validation successful for streaming GET: {}",
                    session_id
                );
            }
            Err(err) => {
                error!(
                    "Session validation failed for streaming GET {}: {}",
                    session_id,
                    err.message()
                );
                return StreamableResponse::Error {
                    status: err.status_code(),
                    message: format!("Session validation failed: {}", err.message()),
                }
                .into_boxed_response(&context);
            }
        }

        // 2. Create bi-directional stream with chunked transfer encoding
        // For MCP Streamable HTTP, we create a stream that can handle bidirectional JSON-RPC
        // Unlike SSE which is unidirectional server->client, this supports client->server and server->client

        // Extract Last-Event-ID for resumability (if client supports it)
        let last_event_id = req
            .headers()
            .get("Last-Event-ID")
            .and_then(|h| h.to_str().ok())
            .and_then(|s| s.parse::<u64>().ok());

        // Generate unique connection ID for tracking this stream
        let connection_id = uuid::Uuid::now_v7().as_simple().to_string();

        debug!(
            "Creating streamable HTTP connection: session={}, connection={}, last_event_id={:?}",
            session_id, connection_id, last_event_id
        );

        // 3. Return streaming response supporting progressive message delivery
        // ✅ FIXED: Return the actual streaming response from StreamManager
        // This preserves event replay, resumability, and live streaming capabilities
        match self
            .stream_manager
            .handle_sse_connection(session_id.clone(), connection_id.clone(), last_event_id)
            .await
        {
            Ok(mut streaming_response) => {
                debug!(
                    "Streamable HTTP connection established: session={}, connection={}",
                    session_id, connection_id
                );

                // Merge MCP headers from context.response_headers()
                let mcp_headers = context.response_headers();
                for (key, value) in mcp_headers.iter() {
                    streaming_response.headers_mut().insert(key, value.clone());
                }

                // ✅ PRESERVE STREAMING: Return the streaming response with MCP headers
                // This maintains event replay from session storage and live streaming
                streaming_response
            }
            Err(err) => {
                error!("Failed to create streamable HTTP connection: {}", err);
                StreamableResponse::Error {
                    status: StatusCode::INTERNAL_SERVER_ERROR,
                    message: format!("Streaming connection failed: {}", err),
                }
                .into_boxed_response(&context)
            }
        }
    }

    /// Validate that a session exists and is not terminated — do NOT create if missing.
    ///
    /// Returns typed errors so callers can map to the correct HTTP status:
    /// - `NotFound` → 404 (MCP spec: terminated sessions MUST return 404)
    /// - `StorageError` → 500
    async fn validate_session_exists(
        &self,
        session_id: &str,
    ) -> std::result::Result<(), SessionValidationError> {
        match self.session_storage.get_session(session_id).await {
            Ok(Some(session_info)) => {
                if session_info.is_terminated() {
                    error!("Session '{}' has been terminated", session_id);
                    return Err(SessionValidationError::NotFound(format!(
                        "Session '{}' has been terminated. Create a new session to continue.",
                        session_id
                    )));
                }
                debug!("Session validation successful: {}", session_id);
                Ok(())
            }
            Ok(None) => {
                error!("Session not found: {}", session_id);
                Err(SessionValidationError::NotFound(format!(
                    "Session '{}' not found. Sessions must be created via initialize request first.",
                    session_id
                )))
            }
            Err(err) => {
                error!("Failed to validate session {}: {}", session_id, err);
                Err(SessionValidationError::StorageError(format!(
                    "Session validation failed: {}",
                    err
                )))
            }
        }
    }

    /// Handle POST request with JSON response (legacy compatibility)
    #[allow(dead_code)]
    async fn handle_json_post<T>(
        &self,
        req: Request<T>,
        context: StreamableHttpContext,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>>
    where
        T: Body + Send + 'static,
    {
        debug!("Handling JSON POST (non-streaming/legacy)");

        // 1. Parse JSON-RPC request(s) from request body (legacy clients don't require sessions)

        // Check content type
        let content_type = req
            .headers()
            .get(CONTENT_TYPE)
            .and_then(|ct| ct.to_str().ok())
            .map(normalize_header_value)
            .unwrap_or_default();

        if !content_type.starts_with("application/json") {
            warn!("Invalid content type for legacy POST: {}", content_type);
            return StreamableResponse::Error {
                status: StatusCode::BAD_REQUEST,
                message: "Content-Type must be application/json".to_string(),
            }
            .into_boxed_response(&context);
        }

        // Read request body
        let body_bytes = match req.into_body().collect().await {
            Ok(collected) => collected.to_bytes(),
            Err(_err) => {
                error!("Failed to read legacy POST request body");
                return StreamableResponse::Error {
                    status: StatusCode::BAD_REQUEST,
                    message: "Failed to read request body".to_string(),
                }
                .into_boxed_response(&context);
            }
        };

        // Check body size
        if body_bytes.len() > self.config.max_body_size {
            warn!(
                "Legacy POST request body too large: {} bytes",
                body_bytes.len()
            );
            return StreamableResponse::Error {
                status: StatusCode::PAYLOAD_TOO_LARGE,
                message: "Request body too large".to_string(),
            }
            .into_boxed_response(&context);
        }

        // Parse as UTF-8
        let body_str = match std::str::from_utf8(&body_bytes) {
            Ok(s) => s,
            Err(err) => {
                error!("Invalid UTF-8 in legacy POST request body: {}", err);
                return StreamableResponse::Error {
                    status: StatusCode::BAD_REQUEST,
                    message: "Request body must be valid UTF-8".to_string(),
                }
                .into_boxed_response(&context);
            }
        };

        debug!("Received legacy POST JSON-RPC request: {}", body_str);

        // Parse JSON-RPC message
        use turul_mcp_json_rpc_server::dispatch::{
            JsonRpcMessage, JsonRpcMessageResult, parse_json_rpc_message,
        };

        let message = match parse_json_rpc_message(body_str) {
            Ok(msg) => msg,
            Err(rpc_err) => {
                error!("JSON-RPC parse error in legacy POST: {}", rpc_err);
                let error_json =
                    serde_json::to_string(&rpc_err).unwrap_or_else(|_| "{}".to_string());
                return Response::builder()
                    .status(StatusCode::OK) // JSON-RPC parse errors still use 200 OK
                    .header(CONTENT_TYPE, "application/json")
                    .header("MCP-Protocol-Version", context.protocol_version.as_str())
                    .body(Full::new(Bytes::from(error_json)))
                    .unwrap()
                    .map(|body| body.map_err(|never| match never {}).boxed_unsync());
            }
        };

        // 2. Process via dispatcher (no session context for legacy clients)
        // Legacy clients (MCP 2024-11-05) don't use sessions, so no session context
        let message_result = match message {
            JsonRpcMessage::Request(request) => {
                debug!(
                    "Processing legacy POST JSON-RPC request: method={}",
                    request.method
                );

                // Special handling for initialize requests - legacy clients can create sessions too
                let response = if request.method == "initialize" {
                    debug!("Handling legacy initialize request - creating new session");

                    // Let session storage create the session and generate the ID
                    match self
                        .session_storage
                        .create_session(self.server_capabilities.clone())
                        .await
                    {
                        Ok(session_info) => {
                            debug!(
                                "Created new session for legacy client: {}",
                                session_info.session_id
                            );

                            // Create session context for initialize response
                            use crate::notification_bridge::StreamManagerNotificationBroadcaster;
                            use turul_mcp_json_rpc_server::r#async::SessionContext;

                            let broadcaster = Arc::new(StreamManagerNotificationBroadcaster::new(
                                Arc::clone(&self.stream_manager),
                            ));
                            let broadcaster_any =
                                Arc::new(broadcaster) as Arc<dyn std::any::Any + Send + Sync>;

                            let session_context = SessionContext {
                                session_id: session_info.session_id.clone(),
                                metadata: std::collections::HashMap::new(),
                                broadcaster: Some(broadcaster_any),
                                timestamp: chrono::Utc::now().timestamp_millis() as u64,
                                extensions: std::collections::HashMap::new(),
                            };

                            self.dispatcher
                                .handle_request_with_context(request, session_context)
                                .await
                        }
                        Err(err) => {
                            error!("Failed to create session during legacy initialize: {}", err);
                            let error_msg = format!("Session creation failed: {}", err);
                            turul_mcp_json_rpc_server::JsonRpcMessage::error(
                                turul_mcp_json_rpc_server::JsonRpcError::internal_error(
                                    Some(request.id),
                                    Some(error_msg),
                                ),
                            )
                        }
                    }
                } else {
                    // For non-initialize requests, process without session context (legacy mode)
                    self.dispatcher.handle_request(request).await
                };

                // Convert JsonRpcMessage to JsonRpcMessageResult
                match response {
                    turul_mcp_json_rpc_server::JsonRpcMessage::Response(resp) => {
                        JsonRpcMessageResult::Response(resp)
                    }
                    turul_mcp_json_rpc_server::JsonRpcMessage::Error(err) => {
                        JsonRpcMessageResult::Error(err)
                    }
                }
            }
            JsonRpcMessage::Notification(notification) => {
                debug!(
                    "Processing legacy POST JSON-RPC notification: method={}",
                    notification.method
                );

                // Process notification without session context (legacy mode)
                let result = self
                    .dispatcher
                    .handle_notification_with_context(notification, None)
                    .await;

                if let Err(err) = result {
                    error!("Legacy POST notification handling error: {}", err);
                }
                JsonRpcMessageResult::NoResponse
            }
        };

        // 3. Return single JSON response (no streaming) - legacy compatibility
        match message_result {
            JsonRpcMessageResult::Response(response) => {
                let response_json = serde_json::to_string(&response)
                    .unwrap_or_else(|_| r#"{"error": "Failed to serialize response"}"#.to_string());

                Response::builder()
                    .status(StatusCode::OK)
                    .header(CONTENT_TYPE, "application/json")
                    .header("MCP-Protocol-Version", context.protocol_version.as_str())
                    .body(Full::new(Bytes::from(response_json)))
                    .unwrap()
                    .map(|body| body.map_err(|never| match never {}).boxed_unsync())
            }
            JsonRpcMessageResult::Error(error) => {
                let error_json = serde_json::to_string(&error)
                    .unwrap_or_else(|_| r#"{"error": "Internal error"}"#.to_string());

                Response::builder()
                    .status(StatusCode::OK) // JSON-RPC errors still return 200 OK
                    .header(CONTENT_TYPE, "application/json")
                    .header("MCP-Protocol-Version", context.protocol_version.as_str())
                    .body(Full::new(Bytes::from(error_json)))
                    .unwrap()
                    .map(|body| body.map_err(|never| match never {}).boxed_unsync())
            }
            JsonRpcMessageResult::NoResponse => {
                // Notifications return 202 Accepted per MCP spec
                Response::builder()
                    .status(StatusCode::ACCEPTED)
                    .header("MCP-Protocol-Version", context.protocol_version.as_str())
                    .body(Full::new(Bytes::new()))
                    .unwrap()
                    .map(|body| body.map_err(|never| match never {}).boxed_unsync())
            }
        }
    }

    /// Handle DELETE request for session cleanup
    async fn handle_session_delete<T>(
        &self,
        _req: Request<T>,
        context: StreamableHttpContext,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>>
    where
        T: Body + Send + 'static,
    {
        if let Some(session_id) = &context.session_id {
            debug!("Deleting session: {}", session_id);

            // Implement proper session cleanup for Streamable HTTP
            // 1. Close any active streaming connections for this session
            let closed_connections = self
                .stream_manager
                .close_session_connections(session_id)
                .await;
            debug!(
                "Closed {} streaming connections for session: {}",
                closed_connections, session_id
            );

            // 2. Mark session as terminated instead of immediate deletion (for proper lifecycle management)
            match self.session_storage.get_session(session_id).await {
                Ok(Some(mut session_info)) => {
                    // Mark session as terminated in state
                    session_info
                        .state
                        .insert("terminated".to_string(), serde_json::Value::Bool(true));
                    session_info.state.insert(
                        "terminated_at".to_string(),
                        serde_json::Value::Number(serde_json::Number::from(
                            chrono::Utc::now().timestamp_millis(),
                        )),
                    );
                    session_info.touch();

                    // 3. Update session with termination markers
                    match self.session_storage.update_session(session_info).await {
                        Ok(()) => {
                            debug!(
                                "Session {} marked as terminated (TTL will handle cleanup)",
                                session_id
                            );

                            // Return success response with proper headers
                            Response::builder()
                                .status(StatusCode::OK)
                                .header(CONTENT_TYPE, "application/json")
                                .header("MCP-Protocol-Version", context.protocol_version.as_str())
                                .header("Mcp-Session-Id", session_id)
                                .body(Full::new(Bytes::from(
                                    serde_json::to_string(&serde_json::json!({
                                        "status": "session_terminated",
                                        "session_id": session_id,
                                        "closed_connections": closed_connections,
                                        "message": "Session marked for cleanup"
                                    }))
                                    .unwrap_or_else(|_| {
                                        r#"{"status":"session_terminated"}"#.to_string()
                                    }),
                                )))
                                .unwrap()
                                .map(|body| body.map_err(|never| match never {}).boxed_unsync())
                        }
                        Err(err) => {
                            error!(
                                "Error marking session {} as terminated: {}",
                                session_id, err
                            );
                            // Fallback to deletion if update fails
                            match self.session_storage.delete_session(session_id).await {
                                Ok(_) => {
                                    debug!("Session {} deleted as fallback", session_id);
                                    Response::builder()
                                        .status(StatusCode::OK)
                                        .header(CONTENT_TYPE, "application/json")
                                        .header(
                                            "MCP-Protocol-Version",
                                            context.protocol_version.as_str(),
                                        )
                                        .body(Full::new(Bytes::from(
                                            serde_json::to_string(&serde_json::json!({
                                                "status": "session_deleted",
                                                "session_id": session_id,
                                                "closed_connections": closed_connections,
                                                "message": "Session removed"
                                            }))
                                            .unwrap_or_else(|_| {
                                                r#"{"status":"session_deleted"}"#.to_string()
                                            }),
                                        )))
                                        .unwrap()
                                        .map(|body| {
                                            body.map_err(|never| match never {}).boxed_unsync()
                                        })
                                }
                                Err(delete_err) => {
                                    error!(
                                        "Error deleting session {} as fallback: {}",
                                        session_id, delete_err
                                    );
                                    StreamableResponse::Error {
                                        status: StatusCode::INTERNAL_SERVER_ERROR,
                                        message: "Session termination error".to_string(),
                                    }
                                    .into_boxed_response(&context)
                                }
                            }
                        }
                    }
                }
                Ok(None) => {
                    // Session not found
                    Response::builder()
                        .status(StatusCode::NOT_FOUND)
                        .header(CONTENT_TYPE, "application/json")
                        .header("MCP-Protocol-Version", context.protocol_version.as_str())
                        .body(Full::new(Bytes::from(
                            serde_json::to_string(&serde_json::json!({
                                "status": "session_not_found",
                                "session_id": session_id,
                                "message": "Session not found"
                            }))
                            .unwrap_or_else(|_| r#"{"status":"session_not_found"}"#.to_string()),
                        )))
                        .unwrap()
                        .map(|body| body.map_err(|never| match never {}).boxed_unsync())
                }
                Err(err) => {
                    error!(
                        "Error retrieving session {} for termination: {}",
                        session_id, err
                    );
                    StreamableResponse::Error {
                        status: StatusCode::INTERNAL_SERVER_ERROR,
                        message: "Session lookup error".to_string(),
                    }
                    .into_boxed_response(&context)
                }
            }
        } else {
            StreamableResponse::Error {
                status: StatusCode::BAD_REQUEST,
                message: "Mcp-Session-Id header required for session deletion".to_string(),
            }
            .into_boxed_response(&context)
        }
    }

    /// Handle POST with tool progress notifications (POST Streamable HTTP)
    ///
    /// Implements MCP 2025-11-25 Streamable HTTP where POST requests receive
    /// chunked responses containing both progress notifications AND final result.
    /// This is request-response, NOT a long-lived connection (that's GET SSE).
    ///
    /// Uses hyper::Body::channel() for chunked transfer encoding with background
    /// task forwarding (works for long-running servers, BROKEN in Lambda).
    async fn handle_post_streamable_http<T>(
        &self,
        req: Request<T>,
        mut context: StreamableHttpContext,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>>
    where
        T: Body + Send + 'static,
    {
        debug!("Streaming handler called - using true streaming POST");

        // Parse request body (still need to collect for JSON-RPC parsing)
        let body_bytes = match req.into_body().collect().await {
            Ok(collected) => collected.to_bytes(),
            Err(_err) => {
                error!("Failed to read streaming POST request body");
                return StreamableResponse::Error {
                    status: StatusCode::BAD_REQUEST,
                    message: "Failed to read request body".to_string(),
                }
                .into_boxed_response(&context);
            }
        };

        // Check body size
        if body_bytes.len() > self.config.max_body_size {
            warn!(
                "Streaming POST request body too large: {} bytes",
                body_bytes.len()
            );
            return StreamableResponse::Error {
                status: StatusCode::PAYLOAD_TOO_LARGE,
                message: "Request body too large".to_string(),
            }
            .into_boxed_response(&context);
        }

        // Parse as UTF-8
        let body_str = match std::str::from_utf8(&body_bytes) {
            Ok(s) => s,
            Err(err) => {
                error!("Invalid UTF-8 in streaming POST request body: {}", err);
                return StreamableResponse::Error {
                    status: StatusCode::BAD_REQUEST,
                    message: "Request body must be valid UTF-8".to_string(),
                }
                .into_boxed_response(&context);
            }
        };

        debug!("Streaming POST received JSON-RPC request: {}", body_str);

        // Parse JSON-RPC message
        use turul_mcp_json_rpc_server::dispatch::{JsonRpcMessage, parse_json_rpc_message};
        use turul_mcp_json_rpc_server::error::JsonRpcErrorObject;

        let message = match parse_json_rpc_message(body_str) {
            Ok(msg) => msg,
            Err(rpc_err) => {
                error!("JSON-RPC parse error in streaming POST: {}", rpc_err);
                let error_json =
                    serde_json::to_string(&rpc_err).unwrap_or_else(|_| "{}".to_string());

                // Return error with MCP headers (no session header for parse errors)
                return Response::builder()
                    .status(StatusCode::OK) // JSON-RPC parse errors still use 200 OK
                    .header(CONTENT_TYPE, "application/json")
                    .header("MCP-Protocol-Version", context.protocol_version.as_str())
                    .body(
                        Full::new(Bytes::from(error_json))
                            .map_err(|never| match never {})
                            .boxed_unsync(),
                    )
                    .unwrap();
            }
        };

        // Handle sessionless ping (pre-init ping support per MCP 2025-11-25)
        // Clients are permitted to send pings before initialization completes.
        // When allowed by config, dispatch through the shared middleware + dispatch pipeline
        // with session=None, so rate-limiting middleware can still block abuse.
        let is_sessionless_ping = match &message {
            JsonRpcMessage::Request(req) => req.method == "ping",
            JsonRpcMessage::Notification(notif) => notif.method == "ping",
        } && context.session_id.is_none()
            && self.config.allow_unauthenticated_ping;

        if is_sessionless_ping {
            return match message {
                JsonRpcMessage::Request(request) => {
                    // Sessionless ping request: shared dispatch path
                    let (response, _) = self
                        .run_middleware_and_dispatch(request, context.headers.clone(), None, None)
                        .await;
                    let response_value =
                        serde_json::to_value(&response).unwrap_or(serde_json::json!({}));
                    StreamableResponse::Json(response_value).into_boxed_response(&context)
                }
                JsonRpcMessage::Notification(notification) => {
                    // Sessionless ping notification: dispatch and return 202 Accepted
                    let dispatcher = Arc::clone(&self.dispatcher);
                    tokio::spawn(async move {
                        if let Err(e) = dispatcher
                            .handle_notification_with_context(notification, None)
                            .await
                        {
                            error!("Failed to process sessionless ping notification: {}", e);
                        }
                    });

                    Response::builder()
                        .status(StatusCode::ACCEPTED)
                        .header("MCP-Protocol-Version", context.protocol_version.as_str())
                        .body(
                            Full::new(Bytes::new())
                                .map_err(|never| match never {})
                                .boxed_unsync(),
                        )
                        .unwrap()
                }
            };
        }

        // --- Pre-session auth phase (D4) ---
        // Extract Bearer token using hardened parser (D6)
        let bearer_token = context
            .headers
            .get("authorization")
            .and_then(|v| extract_bearer_token(v));

        // Run pre-session middleware if any are registered
        let pre_session_extensions = if self.middleware_stack.has_pre_session_middleware() {
            let method_name = match &message {
                JsonRpcMessage::Request(req) => req.method.as_str(),
                JsonRpcMessage::Notification(notif) => notif.method.as_str(),
            };
            let mut pre_ctx = crate::middleware::RequestContext::new(method_name, None);
            if let Some(ref token) = bearer_token {
                pre_ctx.set_bearer_token(token.clone());
            }
            // Copy headers to metadata, excluding Bearer authorization (D5)
            for (k, v) in &context.headers {
                if k.eq_ignore_ascii_case("authorization") && is_bearer_scheme(v) {
                    continue;
                }
                pre_ctx.add_metadata(k.clone(), serde_json::json!(v));
            }
            match self
                .middleware_stack
                .execute_before_session(&mut pre_ctx)
                .await
            {
                Ok(()) => Some(pre_ctx.take_extensions()),
                Err(crate::middleware::MiddlewareError::HttpChallenge {
                    status,
                    www_authenticate,
                    body,
                }) => {
                    return build_http_challenge_response(
                        status,
                        &www_authenticate,
                        body.as_deref(),
                        &context,
                    );
                }
                Err(other_err) => {
                    // Non-challenge pre-session errors → JSON-RPC error
                    if let JsonRpcMessage::Request(ref req) = message {
                        let response =
                            Self::map_middleware_error_to_jsonrpc(other_err, req.id.clone());
                        let response_value =
                            serde_json::to_value(&response).unwrap_or(serde_json::json!({}));
                        return StreamableResponse::Json(response_value)
                            .into_boxed_response(&context);
                    } else {
                        // Notification — can't return JSON-RPC error, just reject
                        return Response::builder()
                            .status(StatusCode::FORBIDDEN)
                            .body(
                                Full::new(Bytes::from(other_err.to_string()))
                                    .map_err(|never| match never {})
                                    .boxed_unsync(),
                            )
                            .unwrap();
                    }
                }
            }
        } else {
            None
        };

        // Validate session requirements based on method
        let session_id = match &message {
            JsonRpcMessage::Request(req) if req.method == "initialize" => {
                // Initialize can create session if none exists
                if let Some(existing_id) = &context.session_id {
                    // Validate existing session for initialize
                    if let Err(err) = self.validate_session_exists(existing_id).await {
                        warn!(
                            "Invalid session ID {} during initialize: {}",
                            existing_id,
                            err.message()
                        );
                        return StreamableResponse::Error {
                            status: err.status_code(),
                            message: format!("Invalid or expired session: {}", err.message()),
                        }
                        .into_boxed_response(&context);
                    }
                    existing_id.clone()
                } else {
                    // Create new session for initialize
                    match self
                        .session_storage
                        .create_session(self.server_capabilities.clone())
                        .await
                    {
                        Ok(session_info) => {
                            debug!(
                                "Created new session for initialize: {}",
                                session_info.session_id
                            );
                            context.session_id = Some(session_info.session_id.clone());
                            session_info.session_id
                        }
                        Err(err) => {
                            error!("Failed to create session during initialize: {}", err);
                            return StreamableResponse::Error {
                                status: StatusCode::INTERNAL_SERVER_ERROR,
                                message: "Failed to create session".to_string(),
                            }
                            .into_boxed_response(&context);
                        }
                    }
                }
            }
            JsonRpcMessage::Request(_) | JsonRpcMessage::Notification(_) => {
                // All other methods REQUIRE session ID
                if let Some(existing_id) = &context.session_id {
                    // Validate existing session
                    if let Err(err) = self.validate_session_exists(existing_id).await {
                        warn!("Invalid session ID {}: {}", existing_id, err.message());
                        return StreamableResponse::Error {
                            status: err.status_code(),
                            message: format!("Invalid or expired session: {}", err.message()),
                        }
                        .into_boxed_response(&context);
                    }
                    existing_id.clone()
                } else {
                    // Return 401 for missing header — this is NOT a stale session (404),
                    // it's "no session ID provided at all"
                    let method_name = match &message {
                        JsonRpcMessage::Request(req) => &req.method,
                        JsonRpcMessage::Notification(notif) => &notif.method,
                    };
                    let request_id = match &message {
                        JsonRpcMessage::Request(req) => Some(req.id.clone()),
                        JsonRpcMessage::Notification(_) => None,
                    };

                    warn!("Missing session ID for method: {}", method_name);

                    let error_response = turul_mcp_json_rpc_server::JsonRpcError::new(
                        request_id,
                        JsonRpcErrorObject::server_error(
                            -32001,
                            "Missing Mcp-Session-Id header. Call initialize first.",
                            None::<serde_json::Value>,
                        ),
                    );

                    let error_json =
                        serde_json::to_string(&error_response).unwrap_or_else(|_| "{}".to_string());

                    return Response::builder()
                        .status(StatusCode::UNAUTHORIZED)
                        .header(CONTENT_TYPE, "application/json")
                        .header("MCP-Protocol-Version", context.protocol_version.as_str())
                        .body(
                            Full::new(Bytes::from(error_json))
                                .map_err(|never| match never {})
                                .boxed_unsync(),
                        )
                        .unwrap();
                }
            }
        };

        debug!("Processing streaming request with session: {}", session_id);

        // Create streaming response using hyper::Body::channel()
        match message {
            JsonRpcMessage::Request(request) => {
                debug!(
                    "Processing streaming JSON-RPC request: method={}",
                    request.method
                );
                self.create_streaming_response(
                    request,
                    session_id,
                    context,
                    pre_session_extensions.clone(),
                )
                .await
            }
            JsonRpcMessage::Notification(notification) => {
                debug!(
                    "Processing streaming JSON-RPC notification: method={}",
                    notification.method
                );

                // Create session context with notification broadcaster for notifications
                use crate::notification_bridge::StreamManagerNotificationBroadcaster;
                use turul_mcp_json_rpc_server::SessionContext;

                let broadcaster = Arc::new(StreamManagerNotificationBroadcaster::new(Arc::clone(
                    &self.stream_manager,
                )));
                let broadcaster_any = Arc::new(broadcaster) as Arc<dyn std::any::Any + Send + Sync>;

                let session_context = SessionContext {
                    session_id: session_id.clone(),
                    metadata: std::collections::HashMap::new(),
                    broadcaster: Some(broadcaster_any),
                    timestamp: chrono::Utc::now().timestamp_millis() as u64,
                    extensions: std::collections::HashMap::new(),
                };

                // Process notification through dispatcher (notifications don't return responses)
                let dispatcher = Arc::clone(&self.dispatcher);
                let notification_clone = notification.clone();

                // notifications/initialized MUST be processed synchronously before
                // returning 202 — otherwise the next request (tools/list) can race
                // ahead of the is_initialized state write. Other notifications are
                // fire-and-forget and can be processed asynchronously.
                //
                // Per MCP spec, notifications always return 202 even on failure.
                // If processing fails, we log the error — the next request will
                // fail with "session not initialized" which points operators to
                // the initialization failure in logs.
                if notification_clone.method == "notifications/initialized" {
                    if let Err(e) = dispatcher
                        .handle_notification_with_context(
                            notification_clone,
                            Some(session_context),
                        )
                        .await
                    {
                        error!(
                            "Failed to process notifications/initialized: {}. \
                             Session will remain uninitialized — subsequent requests will fail.",
                            e
                        );
                    }
                } else {
                    tokio::spawn(async move {
                        if let Err(e) = dispatcher
                            .handle_notification_with_context(
                                notification_clone,
                                Some(session_context),
                            )
                            .await
                        {
                            error!("Failed to process notification: {}", e);
                        }
                    });
                }

                // Return 202 Accepted with MCP headers
                Response::builder()
                    .status(StatusCode::ACCEPTED)
                    .header("MCP-Protocol-Version", context.protocol_version.as_str())
                    .header("Mcp-Session-Id", &session_id)
                    .body(
                        Full::new(Bytes::new())
                            .map_err(|never| match never {})
                            .boxed_unsync(),
                    )
                    .unwrap()
            }
        }
    }

    /// Create a streaming response using hyper::Body::channel()
    /// This enables true progressive responses with Transfer-Encoding: chunked
    async fn create_streaming_response(
        &self,
        request: turul_mcp_json_rpc_server::JsonRpcRequest,
        session_id: String,
        context: StreamableHttpContext,
        pre_session_extensions: Option<HashMap<String, serde_json::Value>>,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>> {
        debug!(
            "Creating streaming response for method: {}, session: {}",
            request.method, session_id
        );
        // Create channel for streaming response
        use http_body_util::StreamBody;
        use tokio_stream::StreamExt;
        use tokio_stream::wrappers::UnboundedReceiverStream; // Add StreamExt for map method

        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Result<bytes::Bytes, hyper::Error>>();
        let body_stream =
            UnboundedReceiverStream::new(rx).map(|item| item.map(http_body::Frame::data));
        let body = StreamBody::new(body_stream);

        // Create session context with notification broadcaster (same pattern as SessionMcpHandler)
        use crate::notification_bridge::{
            SharedNotificationBroadcaster, StreamManagerNotificationBroadcaster,
        };
        use turul_mcp_json_rpc_server::SessionContext;

        let broadcaster: SharedNotificationBroadcaster = Arc::new(
            StreamManagerNotificationBroadcaster::new(Arc::clone(&self.stream_manager)),
        );
        let broadcaster_any = Arc::new(broadcaster) as Arc<dyn std::any::Any + Send + Sync>;

        let session_context = SessionContext {
            session_id: session_id.clone(),
            metadata: std::collections::HashMap::new(),
            broadcaster: Some(broadcaster_any),
            timestamp: chrono::Utc::now().timestamp_millis() as u64,
            extensions: std::collections::HashMap::new(),
        };

        // Register streaming POST connection with StreamManager for progress events
        let wants_sse = context.wants_sse_stream();
        let connection_id = format!("post-{}", uuid::Uuid::now_v7().as_simple());

        // Progress forwarding only for SSE clients
        let (shutdown_tx, completion_rx) = if wants_sse {
            // Create shutdown signal for progress task (critical for no-progress-events case)
            let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
            let (completion_tx, completion_rx) = tokio::sync::oneshot::channel::<()>();
            let (progress_tx, mut progress_rx) = tokio::sync::mpsc::channel(100);

            // Register with StreamManager to receive progress events
            let registration_result = self
                .stream_manager
                .register_streaming_connection(&session_id, connection_id.clone(), progress_tx)
                .await;

            if let Err(e) = registration_result {
                error!("Failed to register POST streaming connection: {}", e);
                // Continue without streaming - will still work as regular POST
                (None, None)
            } else {
                debug!(
                    "Registered SSE streaming connection for session: {}",
                    session_id
                );

                // Spawn task to forward progress events to HTTP response
                let sender_clone = tx.clone();
                let session_id_clone = session_id.clone();
                let connection_id_clone = connection_id.clone();
                let stream_manager_clone = Arc::clone(&self.stream_manager);

                tokio::spawn(async move {
                    debug!(
                        "Starting progress forwarding task for session: {}",
                        session_id_clone
                    );

                    // CRITICAL: Use select to handle both progress events AND explicit shutdown
                    loop {
                        debug!(
                            "🔍 Progress task entering select loop for session: {}",
                            session_id_clone
                        );
                        tokio::select! {
                            // Handle progress events if they arrive
                            maybe_event = progress_rx.recv() => {
                                debug!("🔍 Progress task: progress_rx.recv() branch fired for session: {}", session_id_clone);
                                match maybe_event {
                                    Some(sse_event) => {
                                        debug!("🔍 Forwarding progress event to POST response: session={}, event={:?}", session_id_clone, sse_event.event_type);

                                        // Convert SSE event to fully-formatted SSE chunk with event metadata
                                        let sse_chunk = sse_event.format();

                                        if let Err(e) = sender_clone.send(Ok(Bytes::from(sse_chunk))) {
                                            error!("Failed to send progress event to POST response: {}", e);
                                            break;
                                        }
                                    }
                                    None => {
                                        // Progress channel closed naturally
                                        debug!("🔍 Progress channel closed naturally for session: {}", session_id_clone);
                                        break;
                                    }
                                }
                            }
                            // Handle explicit shutdown signal from main task
                            _ = &mut shutdown_rx => {
                                debug!("🔍 Progress task: shutdown_rx branch fired! Received explicit shutdown signal for session: {}", session_id_clone);
                                break;
                            }
                        }
                    }

                    // Clean up: Unregister from StreamManager to close progress_tx
                    debug!(
                        "Progress task unregistering connection for session: {}",
                        session_id_clone
                    );
                    stream_manager_clone
                        .unregister_connection(&session_id_clone, &connection_id_clone)
                        .await;

                    // CRITICAL: Drop the sender to ensure stream can close
                    debug!(
                        "🔍 Progress task: dropping sender_clone for session: {}",
                        session_id_clone
                    );
                    drop(sender_clone);

                    // Signal completion to main task
                    debug!(
                        "🔍 Progress task: signaling completion for session: {}",
                        session_id_clone
                    );
                    if completion_tx.send(()).is_err() {
                        debug!(
                            "🔍 Progress task: main task already dropped completion_rx for session: {}",
                            session_id_clone
                        );
                    }

                    debug!(
                        "🔍 Progress forwarding task completed for session: {}",
                        session_id_clone
                    );
                });

                // Return shutdown_tx and completion_rx for later use
                (Some(shutdown_tx), Some(completion_rx))
            }
        } else {
            // No SSE, no shutdown signal needed
            (None, None)
        };

        // Spawn task to handle streaming dispatch
        let request_id = request.id.clone();
        let sender = tx; // Rename for clarity

        // Capture headers for middleware (clone before move into spawn)
        let headers = context.headers.clone();
        let self_clone = self.clone();

        tokio::spawn(async move {
            debug!(
                "Spawning streaming task for request ID: {:?}, wants_sse: {}",
                request_id, wants_sse
            );

            // Process actual request through middleware pipeline
            // Injection is applied immediately inside run_middleware_and_dispatch
            let (response, _) = self_clone
                .run_middleware_and_dispatch(
                    request,
                    headers,
                    Some(session_context),
                    pre_session_extensions,
                )
                .await;

            // Send final result - format depends on client type
            if wants_sse {
                // For SSE clients, send as streaming frame with SSE framing
                let final_frame = match response {
                    turul_mcp_json_rpc_server::JsonRpcMessage::Response(resp) => {
                        turul_mcp_json_rpc_server::JsonRpcFrame::FinalResult {
                            request_id: request_id.clone(),
                            result: match resp.result {
                                turul_mcp_json_rpc_server::response::ResponseResult::Success(
                                    val,
                                ) => val,
                                turul_mcp_json_rpc_server::response::ResponseResult::Null => {
                                    serde_json::Value::Null
                                }
                            },
                        }
                    }
                    turul_mcp_json_rpc_server::JsonRpcMessage::Error(err) => {
                        turul_mcp_json_rpc_server::JsonRpcFrame::Error {
                            request_id: request_id.clone(),
                            error: turul_mcp_json_rpc_server::error::JsonRpcErrorObject {
                                code: err.error.code,
                                message: err.error.message,
                                data: err.error.data,
                            },
                        }
                    }
                };

                let final_json = final_frame.to_json();
                // SSE framing: data: {json}\n\n
                let final_chunk =
                    format!("data: {}\n\n", serde_json::to_string(&final_json).unwrap());

                if let Err(err) = sender.send(Ok(Bytes::from(final_chunk))) {
                    error!("Failed to send SSE final chunk: {}", err);
                }

                // CRITICAL: Send explicit shutdown signal to progress forwarding task (SSE only)
                // This breaks it out of the progress_rx.recv().await loop immediately
                if let Some(shutdown_tx) = shutdown_tx {
                    debug!(
                        "🔍 Main task sending shutdown signal to progress task for request: {:?}",
                        request_id
                    );
                    match shutdown_tx.send(()) {
                        Ok(()) => {
                            debug!(
                                "🔍 Main task: shutdown signal sent successfully for request: {:?}",
                                request_id
                            );

                            // CRITICAL: Wait for progress task to complete and drop its sender_clone
                            // This ensures both senders are dropped before the stream tries to close
                            if let Some(completion_rx) = completion_rx {
                                match tokio::time::timeout(
                                    tokio::time::Duration::from_millis(100),
                                    completion_rx,
                                )
                                .await
                                {
                                    Ok(Ok(())) => {
                                        debug!(
                                            "🔍 Main task: progress task completed successfully for request: {:?}",
                                            request_id
                                        );
                                    }
                                    Ok(Err(_)) => {
                                        debug!(
                                            "🔍 Main task: progress task completion signal dropped for request: {:?}",
                                            request_id
                                        );
                                    }
                                    Err(_) => {
                                        debug!(
                                            "🔍 Main task: progress task completion timeout for request: {:?}",
                                            request_id
                                        );
                                    }
                                }
                            }
                        }
                        Err(_) => {
                            debug!(
                                "🔍 Main task: progress task already completed (shutdown_rx dropped) for request: {:?}",
                                request_id
                            );
                        }
                    }
                } else {
                    debug!(
                        "🔍 Main task: no shutdown_tx available (not SSE client) for request: {:?}",
                        request_id
                    );
                }
            } else {
                // For JSON-only clients, send as regular JSON-RPC response (no streaming frames)
                let final_json = serde_json::to_string(&response).unwrap();

                if let Err(err) = sender.send(Ok(Bytes::from(final_json))) {
                    error!("Failed to send final JSON response: {}", err);
                }
            }

            debug!(
                "🔍 Main task: streaming task completed for request ID: {:?}",
                request_id
            );

            // CRITICAL: Drop the sender to close the stream and signal completion to client
            debug!(
                "🔍 Main task: dropping main sender for request ID: {:?}",
                request_id
            );
            drop(sender);
        });

        // Build response with MCP headers merged from context
        // Set content type based on client preference
        let content_type = if context.wants_sse_stream() {
            "text/event-stream"
        } else {
            "application/json"
        };

        let mut response = Response::builder()
            .status(StatusCode::OK)
            .header(CONTENT_TYPE, content_type)
            .header("Transfer-Encoding", "chunked") // Key: Enable chunked encoding!
            .header("Cache-Control", "no-cache")
            .body(http_body_util::BodyExt::boxed_unsync(body))
            .unwrap();

        // Merge MCP headers from context.response_headers()
        let mcp_headers = context.response_headers();
        for (key, value) in mcp_headers.iter() {
            response.headers_mut().insert(key, value.clone());
        }

        response
    }

    /// Handle POST with buffered response (fallback for legacy clients)
    #[allow(dead_code)]
    async fn handle_buffered_post<T>(
        &self,
        _req: Request<T>,
        context: StreamableHttpContext,
        session_id: String,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>>
    where
        T: Body + Send + 'static,
    {
        debug!(
            "Using buffered POST for legacy client, session: {}",
            session_id
        );

        // Use the existing logic (simplified version)
        // TODO: Extract common logic into helper method

        Response::builder()
            .status(StatusCode::OK)
            .header(CONTENT_TYPE, "application/json")
            .header("MCP-Protocol-Version", context.protocol_version.as_str())
            .header("Mcp-Session-Id", &session_id)
            .body(
                Full::new(Bytes::from(
                    r#"{"jsonrpc":"2.0","id":1,"result":"buffered"}"#,
                ))
                .map_err(|never| match never {})
                .boxed_unsync(),
            )
            .unwrap()
    }

    /// Handle POST request - unified handler for all client messages (MCP 2025-11-25 compliant)
    /// Processes JSON-RPC requests, notifications, and responses
    /// Server decides whether to respond with JSON or SSE stream based on message type
    async fn handle_client_message<T>(
        &self,
        req: Request<T>,
        context: StreamableHttpContext,
    ) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>>
    where
        T: Body + Send + 'static,
    {
        debug!("Handling client message via POST (MCP 2025-11-25)");

        // Reject POST if accepts_stream_frames is false
        // Per MCP spec: "Include Accept header with application/json and text/event-stream"
        if !context.accepts_stream_frames {
            warn!("Client POST missing application/json in Accept header");
            return StreamableResponse::Error {
                status: StatusCode::BAD_REQUEST,
                message: "Accept header must include application/json, text/event-stream, or */*"
                    .to_string(),
            }
            .into_boxed_response(&context);
        }

        // Check content type
        let content_type = req
            .headers()
            .get(CONTENT_TYPE)
            .and_then(|ct| ct.to_str().ok())
            .map(normalize_header_value)
            .unwrap_or_default();
        if !content_type.starts_with("application/json") {
            warn!("Invalid content type for POST: {}", content_type);
            return StreamableResponse::Error {
                status: StatusCode::BAD_REQUEST,
                message: "Content-Type must be application/json".to_string(),
            }
            .into_boxed_response(&context);
        }

        // Use streaming for all POST requests, but adapt based on client needs
        // For simple JSON clients, streaming will send only the final result (no progress frames)
        debug!("Using streaming POST handler for all requests");
        return self.handle_post_streamable_http(req, context).await;
    }

    /// Run middleware stack around dispatcher call
    ///
    /// This helper:
    /// 1. Fast-paths when middleware stack is empty
    /// 2. Normalizes headers to lowercase String keys
    /// 3. Builds RequestContext with method + headers
    /// 4. Executes before_dispatch
    /// 5. Applies or stashes SessionInjection
    /// 6. Calls dispatcher
    /// 7. Executes after_dispatch
    ///
    /// Returns (JsonRpcMessage, Option<SessionInjection>) where the injection
    /// is Some when session was None (initialize case) and needs to be applied
    /// after session creation.
    async fn run_middleware_and_dispatch(
        &self,
        request: turul_mcp_json_rpc_server::JsonRpcRequest,
        headers: HashMap<String, String>,
        session: Option<turul_mcp_json_rpc_server::SessionContext>,
        pre_session_extensions: Option<HashMap<String, serde_json::Value>>,
    ) -> (
        turul_mcp_json_rpc_server::JsonRpcMessage,
        Option<crate::middleware::SessionInjection>,
    ) {
        // Fast path: if middleware stack is empty, dispatch directly
        if self.middleware_stack.is_empty() {
            let result = if let Some(session_ctx) = session {
                self.dispatcher
                    .handle_request_with_context(request, session_ctx)
                    .await
            } else {
                self.dispatcher.handle_request(request).await
            };
            return (result, None);
        }

        // Normalize headers: lowercase String keys
        let normalized_headers: HashMap<String, String> = headers
            .iter()
            .map(|(k, v)| (k.to_lowercase(), v.clone()))
            .collect();

        // Build RequestContext with method and headers
        let method = request.method.clone();

        // Convert params to Option<Value>
        let params = request.params.clone().map(|p| match p {
            turul_mcp_json_rpc_server::RequestParams::Object(map) => {
                serde_json::Value::Object(map.into_iter().collect())
            }
            turul_mcp_json_rpc_server::RequestParams::Array(arr) => serde_json::Value::Array(arr),
        });
        let mut ctx = crate::middleware::RequestContext::new(&method, params);

        // Seed extensions from pre-session phase (auth claims etc.)
        if let Some(ext) = pre_session_extensions {
            for (k, v) in ext {
                ctx.set_extension(k, v);
            }
        }

        for (k, v) in normalized_headers {
            // D5: Bearer scheme NEVER enters metadata, even if token is malformed
            if k == "authorization" && is_bearer_scheme(&v) {
                continue;
            }
            ctx.add_metadata(k, serde_json::json!(v));
        }

        // Create SessionView adapter if session is available
        let session_view = session.as_ref().map(|s| {
            crate::middleware::StorageBackedSessionView::new(
                s.session_id.clone(),
                Arc::clone(&self.session_storage),
            )
        });

        // Execute before_dispatch with SessionView (None if sessionless)
        let injection = match self
            .middleware_stack
            .execute_before(
                &mut ctx,
                session_view.as_ref().map(|v| v as &dyn SessionView),
            )
            .await
        {
            Ok(inj) => inj,
            Err(err) => {
                return (Self::map_middleware_error_to_jsonrpc(err, request.id), None);
            }
        };

        // Apply injection to session storage (only if session exists)
        if !injection.is_empty()
            && let Some(ref sv) = session_view
        {
            for (key, value) in injection.state() {
                if let Err(e) = sv.set_state(key, value.clone()).await {
                    tracing::warn!("Failed to apply injection state '{}': {}", key, e);
                }
            }
            for (key, value) in injection.metadata() {
                if let Err(e) = sv.set_metadata(key, value.clone()).await {
                    tracing::warn!("Failed to apply injection metadata '{}': {}", key, e);
                }
            }
        }

        // Thread extensions from RequestContext → JSON-RPC SessionContext (D3 canonical flow)
        let session = session.map(|mut s| {
            s.extensions = ctx.extensions().clone();
            s
        });

        // Dispatch the request
        let result = if let Some(session_ctx) = session {
            self.dispatcher
                .handle_request_with_context(request, session_ctx)
                .await
        } else {
            self.dispatcher.handle_request(request).await
        };

        // Execute after_dispatch
        let mut dispatcher_result = match &result {
            turul_mcp_json_rpc_server::JsonRpcMessage::Response(resp) => match &resp.result {
                turul_mcp_json_rpc_server::response::ResponseResult::Success(val) => {
                    crate::middleware::DispatcherResult::Success(val.clone())
                }
                turul_mcp_json_rpc_server::response::ResponseResult::Null => {
                    crate::middleware::DispatcherResult::Success(serde_json::Value::Null)
                }
            },
            turul_mcp_json_rpc_server::JsonRpcMessage::Error(err) => {
                crate::middleware::DispatcherResult::Error(err.error.message.clone())
            }
        };

        let _ = self
            .middleware_stack
            .execute_after(&ctx, &mut dispatcher_result)
            .await;

        (result, None)
    }

    /// Map MiddlewareError to JSON-RPC error with semantic error codes
    fn map_middleware_error_to_jsonrpc(
        err: crate::middleware::MiddlewareError,
        request_id: turul_mcp_json_rpc_server::RequestId,
    ) -> turul_mcp_json_rpc_server::JsonRpcMessage {
        use crate::middleware::MiddlewareError;
        use crate::middleware::error::error_codes;

        let (code, message, data) = match err {
            MiddlewareError::Unauthenticated(msg) => (error_codes::UNAUTHENTICATED, msg, None),
            MiddlewareError::Unauthorized(msg) => (error_codes::UNAUTHORIZED, msg, None),
            MiddlewareError::RateLimitExceeded {
                message,
                retry_after,
            } => {
                let data = retry_after.map(|s| serde_json::json!({"retryAfter": s}));
                (error_codes::RATE_LIMIT_EXCEEDED, message, data)
            }
            MiddlewareError::InvalidRequest(msg) => (error_codes::INVALID_REQUEST, msg, None),
            MiddlewareError::Internal(msg) => (error_codes::INTERNAL_ERROR, msg, None),
            MiddlewareError::Custom { message, .. } => (error_codes::INTERNAL_ERROR, message, None),
            MiddlewareError::HttpChallenge { .. } => {
                unreachable!(
                    "HttpChallenge must be caught at transport level before JSON-RPC dispatch"
                )
            }
        };

        let error_obj = if let Some(d) = data {
            turul_mcp_json_rpc_server::error::JsonRpcErrorObject::server_error(
                code,
                &message,
                Some(d),
            )
        } else {
            turul_mcp_json_rpc_server::error::JsonRpcErrorObject::server_error(
                code,
                &message,
                None::<serde_json::Value>,
            )
        };

        turul_mcp_json_rpc_server::JsonRpcMessage::Error(
            turul_mcp_json_rpc_server::JsonRpcError::new(Some(request_id), error_obj),
        )
    }
}

use crate::middleware::bearer::{extract_bearer_token, is_bearer_scheme};

/// Build an HTTP challenge response (401/403 with WWW-Authenticate header).
///
/// Returns a raw HTTP response — never enters the JSON-RPC layer.
fn build_http_challenge_response(
    status: u16,
    www_authenticate: &str,
    body: Option<&str>,
    context: &StreamableHttpContext,
) -> Response<http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>> {
    let status_code = StatusCode::from_u16(status).unwrap_or(StatusCode::UNAUTHORIZED);
    let body_bytes = body.unwrap_or("").to_string();

    Response::builder()
        .status(status_code)
        .header("WWW-Authenticate", www_authenticate)
        .header("Cache-Control", "no-store")
        .header("Content-Type", "application/json")
        .header("MCP-Protocol-Version", context.protocol_version.as_str())
        .body(
            http_body_util::Full::new(Bytes::from(body_bytes))
                .map_err(|never| match never {})
                .boxed_unsync(),
        )
        .unwrap()
}

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

    #[test]
    fn test_version_parsing() {
        assert_eq!(
            McpProtocolVersion::parse_version("2024-11-05"),
            Some(McpProtocolVersion::V2024_11_05)
        );
        assert_eq!(
            McpProtocolVersion::parse_version("2025-03-26"),
            Some(McpProtocolVersion::V2025_03_26)
        );
        assert_eq!(
            McpProtocolVersion::parse_version("2025-06-18"),
            Some(McpProtocolVersion::V2025_06_18)
        );
        assert_eq!(McpProtocolVersion::parse_version("invalid"), None);
    }

    #[test]
    fn test_version_capabilities() {
        let v1 = McpProtocolVersion::V2024_11_05;
        assert!(!v1.supports_streamable_http());
        assert!(!v1.supports_meta_fields());

        let v2 = McpProtocolVersion::V2025_03_26;
        assert!(v2.supports_streamable_http());
        assert!(!v2.supports_meta_fields());

        let v3 = McpProtocolVersion::V2025_06_18;
        assert!(v3.supports_streamable_http());
        assert!(v3.supports_meta_fields());
        assert!(v3.supports_cursors());
        assert!(v3.supports_progress_tokens());
        assert!(v3.supports_elicitation());
    }

    #[test]
    fn test_context_validation() {
        let mut context = StreamableHttpContext {
            protocol_version: McpProtocolVersion::V2025_06_18,
            session_id: Some("test-session".to_string()),
            wants_sse_stream: true,
            accepts_stream_frames: true,
            headers: HashMap::new(),
        };

        // POST with session should be valid
        assert!(context.validate(&Method::POST).is_ok());
        // GET with session should be valid
        assert!(context.validate(&Method::GET).is_ok());

        // Test invalid cases
        context.accepts_stream_frames = false;
        assert!(context.validate(&Method::POST).is_err());

        context.accepts_stream_frames = true;
        context.protocol_version = McpProtocolVersion::V2024_11_05;
        context.wants_sse_stream = true;
        assert!(context.validate(&Method::POST).is_err());

        context.protocol_version = McpProtocolVersion::V2025_06_18;
        context.session_id = None;
        // POST without session should be OK (for initialize)
        assert!(context.validate(&Method::POST).is_ok());
        // GET without session should fail
        assert!(context.validate(&Method::GET).is_err());
    }

    // Bearer extraction tests are in middleware::bearer::tests
    // Transport-level integration tests:

    #[test]
    fn test_non_bearer_preserved_in_metadata() {
        // T5: Non-bearer auth headers pass through, Bearer excluded
        let mut ctx = crate::middleware::RequestContext::new("test/method", None);

        let headers = vec![
            (
                "authorization".to_string(),
                "Basic dXNlcjpwYXNz".to_string(),
            ),
            ("x-custom".to_string(), "value".to_string()),
        ];

        for (k, v) in &headers {
            if k == "authorization" && is_bearer_scheme(v) {
                continue;
            }
            ctx.add_metadata(k.clone(), serde_json::json!(v));
        }

        assert!(ctx.metadata().contains_key("authorization"));
        assert!(ctx.metadata().contains_key("x-custom"));

        // Bearer exclusion
        let mut ctx2 = crate::middleware::RequestContext::new("test/method", None);
        let bearer_headers = vec![
            ("authorization".to_string(), "Bearer abc123".to_string()),
            ("x-custom".to_string(), "value".to_string()),
        ];

        for (k, v) in &bearer_headers {
            if k == "authorization" && is_bearer_scheme(v) {
                continue;
            }
            ctx2.add_metadata(k.clone(), serde_json::json!(v));
        }

        assert!(!ctx2.metadata().contains_key("authorization"));
        assert!(ctx2.metadata().contains_key("x-custom"));
    }

    #[test]
    fn test_malformed_bearer_excluded_from_metadata() {
        // T47: Even malformed Bearer tokens are excluded from metadata
        let mut ctx = crate::middleware::RequestContext::new("test/method", None);

        let headers = vec![("authorization".to_string(), "Bearer ".to_string())];

        for (k, v) in &headers {
            if k == "authorization" && is_bearer_scheme(v) {
                continue;
            }
            ctx.add_metadata(k.clone(), serde_json::json!(v));
        }

        assert!(!ctx.metadata().contains_key("authorization"));
    }
}