rust-mcp-axum 2.0.0

use crate::error::TransportServerResult;
use axum::body::Bytes;
use axum::routing::get;
use axum::Extension;
use axum::{
    extract::{Query, State},
    response::IntoResponse,
    routing::{delete, post},
    Router,
};
use http::{HeaderMap, Method, StatusCode, Uri};
use rust_mcp_sdk::mcp_http::{McpAppState, McpHttpHandler};
use std::{collections::HashMap, sync::Arc};

pub fn routes(streamable_http_endpoint: &str) -> Router<Arc<McpAppState>> {
    Router::new()
        .route(streamable_http_endpoint, get(handle_streamable_http_get))
        .route(streamable_http_endpoint, post(handle_streamable_http_post))
        .route(
            streamable_http_endpoint,
            delete(handle_streamable_http_delete),
        )
}

pub async fn handle_streamable_http_get(
    headers: HeaderMap,
    uri: Uri,
    State(state): State<Arc<McpAppState>>,
    Extension(http_handler): Extension<Arc<McpHttpHandler>>,
) -> TransportServerResult<impl IntoResponse> {
    let request = McpHttpHandler::create_request(Method::GET, uri, headers, None);
    let generic_res = http_handler.handle_streamable_http(request, state).await?;
    let (parts, body) = generic_res.into_parts();
    let resp = axum::response::Response::from_parts(parts, axum::body::Body::new(body));
    Ok(resp)
}

pub async fn handle_streamable_http_post(
    headers: HeaderMap,
    uri: Uri,
    State(state): State<Arc<McpAppState>>,
    Extension(http_handler): Extension<Arc<McpHttpHandler>>,
    Query(_params): Query<HashMap<String, String>>,
    payload: Bytes,
) -> TransportServerResult<impl IntoResponse> {
    // Borrow the raw body as UTF-8 instead of extracting an owned `String`,
    // avoiding an allocation and copy per request. JSON-RPC payloads are UTF-8;
    // anything else is rejected up front.
    let payload = match std::str::from_utf8(&payload) {
        Ok(payload) => payload,
        Err(_) => {
            return Ok(axum::response::Response::builder()
                .status(StatusCode::BAD_REQUEST)
                .body(axum::body::Body::from("Request body must be valid UTF-8"))
                .expect("static 400 response is always valid"));
        }
    };

    let request = McpHttpHandler::create_request(Method::POST, uri, headers, Some(payload));
    let generic_res = http_handler.handle_streamable_http(request, state).await?;
    let (parts, body) = generic_res.into_parts();
    let resp = axum::response::Response::from_parts(parts, axum::body::Body::new(body));
    Ok(resp)
}

pub async fn handle_streamable_http_delete(
    headers: HeaderMap,
    uri: Uri,
    State(state): State<Arc<McpAppState>>,
    Extension(http_handler): Extension<Arc<McpHttpHandler>>,
) -> TransportServerResult<impl IntoResponse> {
    let request = McpHttpHandler::create_request(Method::DELETE, uri, headers, None);
    let generic_res = http_handler.handle_streamable_http(request, state).await?;
    let (parts, body) = generic_res.into_parts();
    let resp = axum::response::Response::from_parts(parts, axum::body::Body::new(body));
    Ok(resp)
}