refactor: provide http server as tower service (#228)

* refactor: streamable http server as tower service

crates/rmcp/Cargo.toml

@@ -39,7 +39,9 @@ reqwest = { version = "0.12", default-features = false, features = [
     "json",
     "stream",
 ], optional = true }
-sse-stream = { version = "0.2.0", optional = true }
+
+sse-stream = { version = "0.2", optional = true }
+
 http = { version = "1", optional = true }
 url = { version = "2.4", optional = true }
 
@@ -57,7 +59,9 @@ axum = { version = "0.8", features = [], optional = true }
 rand = { version = "0.9", optional = true }
 tokio-stream = { version = "0.1", optional = true }
 uuid = { version = "1", features = ["v4"], optional = true }
-
+http-body = { version = "1", optional = true }
+http-body-util = { version = "0.1", optional = true }
+bytes = { version = "1", optional = true }
 # macro
 rmcp-macros = { version = "0.1", workspace = true, optional = true }
 
@@ -74,7 +78,17 @@ reqwest = ["__reqwest", "reqwest?/rustls-tls"]
 
 reqwest-tls-no-provider = ["__reqwest", "reqwest?/rustls-tls-no-provider"]
 
-axum = ["dep:axum"]
+server-side-http = [
+    "uuid",
+    "dep:rand",
+    "dep:tokio-stream",
+    "dep:http",
+    "dep:http-body",
+    "dep:http-body-util",
+    "dep:bytes",
+    "dep:sse-stream",
+    "tower",
+]
 # SSE client
 client-side-sse = ["dep:sse-stream", "dep:http"]
 
@@ -97,15 +111,12 @@ transport-child-process = [
 transport-sse-server = [
     "transport-async-rw",
     "transport-worker",
-    "axum",
+    "server-side-http",
-    "dep:rand",
+    "dep:axum",
-    "dep:tokio-stream",
-    "uuid",
 ]
 transport-streamable-http-server = [
     "transport-streamable-http-server-session",
-    "axum",
+    "server-side-http",
-    "uuid",
 ]
 transport-streamable-http-server-session = [
     "transport-async-rw",

crates/rmcp/src/service.rs

@@ -488,7 +488,7 @@ pub struct RequestContext<R: ServiceRole> {
 }
 
 /// Use this function to skip initialization process
-pub async fn serve_directly<R, S, T, E, A>(
+pub fn serve_directly<R, S, T, E, A>(
     service: S,
     transport: T,
     peer_info: Option<R::PeerInfo>,
@@ -499,11 +499,11 @@ where
     T: IntoTransport<R, E, A>,
     E: std::error::Error + Send + Sync + 'static,
 {
-    serve_directly_with_ct(service, transport, peer_info, Default::default()).await
+    serve_directly_with_ct(service, transport, peer_info, Default::default())
 }
 
 /// Use this function to skip initialization process
-pub async fn serve_directly_with_ct<R, S, T, E, A>(
+pub fn serve_directly_with_ct<R, S, T, E, A>(
     service: S,
     transport: T,
     peer_info: Option<R::PeerInfo>,
@@ -516,11 +516,11 @@ where
     E: std::error::Error + Send + Sync + 'static,
 {
     let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info);
-    serve_inner(service, transport, peer, peer_rx, ct).await
+    serve_inner(service, transport, peer, peer_rx, ct)
 }
 
 #[instrument(skip_all)]
-async fn serve_inner<R, S, T, E, A>(
+fn serve_inner<R, S, T, E, A>(
     service: S,
     transport: T,
     peer: Peer<R>,

crates/rmcp/src/service/client.rs

@@ -111,7 +111,7 @@ pub async fn serve_client<S, T, E, A>(
 where
     S: Service<RoleClient>,
     T: IntoTransport<RoleClient, E, A>,
-    E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
+    E: std::error::Error + Send + Sync + 'static,
 {
     serve_client_with_ct(service, transport, Default::default()).await
 }
@@ -124,7 +124,7 @@ pub async fn serve_client_with_ct<S, T, E, A>(
 where
     S: Service<RoleClient>,
     T: IntoTransport<RoleClient, E, A>,
-    E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
+    E: std::error::Error + Send + Sync + 'static,
 {
     let mut transport = transport.into_transport();
     let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
@@ -175,7 +175,7 @@ where
             context: "send initialized notification".into(),
         })?;
     let (peer, peer_rx) = Peer::new(id_provider, Some(initialize_result));
-    Ok(serve_inner(service, transport, peer, peer_rx, ct).await)
+    Ok(serve_inner(service, transport, peer, peer_rx, ct))
 }
 
 macro_rules! method {

crates/rmcp/src/service/server.rs

@@ -70,7 +70,7 @@ impl<S: Service<RoleServer>> ServiceExt<RoleServer> for S {
     ) -> impl Future<Output = Result<RunningService<RoleServer, Self>, ServerInitializeError<E>>> + Send
     where
         T: IntoTransport<RoleServer, E, A>,
-        E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
+        E: std::error::Error + Send + Sync + 'static,
         Self: Sized,
     {
         serve_server_with_ct(self, transport, ct)
@@ -84,7 +84,7 @@ pub async fn serve_server<S, T, E, A>(
 where
     S: Service<RoleServer>,
     T: IntoTransport<RoleServer, E, A>,
-    E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
+    E: std::error::Error + Send + Sync + 'static,
 {
     serve_server_with_ct(service, transport, CancellationToken::new()).await
 }
@@ -143,7 +143,7 @@ pub async fn serve_server_with_ct<S, T, E, A>(
 where
     S: Service<RoleServer>,
     T: IntoTransport<RoleServer, E, A>,
-    E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
+    E: std::error::Error + Send + Sync + 'static,
 {
     let mut transport = transport.into_transport();
     let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
@@ -212,7 +212,7 @@ where
     };
     let _ = service.handle_notification(notification).await;
     // Continue processing service
-    Ok(serve_inner(service, transport, peer, peer_rx, ct).await)
+    Ok(serve_inner(service, transport, peer, peer_rx, ct))
 }
 
 macro_rules! method {

crates/rmcp/src/transport.rs

@@ -7,7 +7,7 @@
 //! | transport         | client                                                    | server                                                |
 //! |:-:                |:-:                                                        |:-:                                                    |
 //! | std IO            | [`child_process::TokioChildProcess`]                      | [`io::stdio`]                                         |
-//! | streamable http   | [`streamable_http_client::StreamableHttpClientTransport`] | [`streamable_http_server::session::create_session`]   |
+//! | streamable http   | [`streamable_http_client::StreamableHttpClientTransport`] | [`streamable_http_server::StreamableHttpService`]     |
 //! | sse               | [`sse_client::SseClientTransport`]                        | [`sse_server::SseServer`]                             |
 //!
 //！## Helper Transport Types
@@ -64,6 +64,8 @@
 //! }
 //! ```
 
+use std::sync::Arc;
+
 use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
 
 pub mod sink_stream;
@@ -122,7 +124,7 @@ pub use auth::{AuthError, AuthorizationManager, AuthorizationSession, Authorized
 pub mod streamable_http_server;
 #[cfg(feature = "transport-streamable-http-server")]
 #[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
-pub use streamable_http_server::axum::StreamableHttpServer;
+pub use streamable_http_server::tower::{StreamableHttpServerConfig, StreamableHttpService};
 
 #[cfg(feature = "transport-streamable-http-client")]
 #[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-client")))]
@@ -138,7 +140,7 @@ pub trait Transport<R>: Send
 where
     R: ServiceRole,
 {
-    type Error;
+    type Error: std::error::Error + Send + Sync + 'static;
     /// Send a message to the transport
     ///
     /// Notice that the future returned by this function should be `Send` and `'static`.
@@ -169,9 +171,73 @@ impl<R, T, E> IntoTransport<R, E, TransportAdapterIdentity> for T
 where
     T: Transport<R, Error = E> + Send + 'static,
     R: ServiceRole,
-    E: std::error::Error + Send + 'static,
+    E: std::error::Error + Send + Sync + 'static,
 {
     fn into_transport(self) -> impl Transport<R, Error = E> + 'static {
         self
     }
 }
+
+/// A transport that can send a single message and then close itself
+pub struct OneshotTransport<R>
+where
+    R: ServiceRole,
+{
+    message: Option<RxJsonRpcMessage<R>>,
+    sender: tokio::sync::mpsc::Sender<TxJsonRpcMessage<R>>,
+    finished_signal: Arc<tokio::sync::Notify>,
+}
+
+impl<R> OneshotTransport<R>
+where
+    R: ServiceRole,
+{
+    pub fn new(
+        message: RxJsonRpcMessage<R>,
+    ) -> (Self, tokio::sync::mpsc::Receiver<TxJsonRpcMessage<R>>) {
+        let (sender, receiver) = tokio::sync::mpsc::channel(16);
+        (
+            Self {
+                message: Some(message),
+                sender,
+                finished_signal: Arc::new(tokio::sync::Notify::new()),
+            },
+            receiver,
+        )
+    }
+}
+
+impl<R> Transport<R> for OneshotTransport<R>
+where
+    R: ServiceRole,
+{
+    type Error = tokio::sync::mpsc::error::SendError<TxJsonRpcMessage<R>>;
+
+    fn send(
+        &mut self,
+        item: TxJsonRpcMessage<R>,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
+        let sender = self.sender.clone();
+        let terminate = matches!(item, TxJsonRpcMessage::<R>::Response(_));
+        let signal = self.finished_signal.clone();
+        async move {
+            sender.send(item).await?;
+            if terminate {
+                signal.notify_waiters();
+            }
+            Ok(())
+        }
+    }
+
+    async fn receive(&mut self) -> Option<RxJsonRpcMessage<R>> {
+        if self.message.is_none() {
+            self.finished_signal.notified().await;
+        }
+        self.message.take()
+    }
+
+    fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        self.message.take();
+        std::future::ready(Ok(()))
+    }
+}

crates/rmcp/src/transport/common.rs

@@ -2,7 +2,7 @@
     feature = "transport-streamable-http-server",
     feature = "transport-sse-server"
 ))]
-pub mod axum;
+pub mod server_side_http;
 
 pub mod http_header;
 

crates/rmcp/src/transport/common/axum.rs

@@ -1,9 +0,0 @@
-use std::{sync::Arc, time::Duration};
-
-pub type SessionId = Arc<str>;
-
-pub fn session_id() -> SessionId {
-    uuid::Uuid::new_v4().to_string().into()
-}
-
-pub const DEFAULT_AUTO_PING_INTERVAL: Duration = Duration::from_secs(15);

crates/rmcp/src/transport/common/server_side_http.rs

@@ -0,0 +1,140 @@
+use std::{convert::Infallible, fmt::Display, sync::Arc, time::Duration};
+
+use bytes::{Buf, Bytes};
+use http::Response;
+use http_body::Body;
+use http_body_util::{BodyExt, Empty, Full, combinators::UnsyncBoxBody};
+use sse_stream::{KeepAlive, Sse, SseBody};
+
+use super::http_header::EVENT_STREAM_MIME_TYPE;
+use crate::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
+
+pub type SessionId = Arc<str>;
+
+pub fn session_id() -> SessionId {
+    uuid::Uuid::new_v4().to_string().into()
+}
+
+pub const DEFAULT_AUTO_PING_INTERVAL: Duration = Duration::from_secs(15);
+
+pub(crate) type BoxResponse = Response<UnsyncBoxBody<Bytes, Infallible>>;
+
+pub(crate) fn accepted_response() -> Response<UnsyncBoxBody<Bytes, Infallible>> {
+    Response::builder()
+        .status(http::StatusCode::ACCEPTED)
+        .body(Empty::new().boxed_unsync())
+        .expect("valid response")
+}
+pin_project_lite::pin_project! {
+    struct TokioTimer {
+        #[pin]
+        sleep: tokio::time::Sleep,
+    }
+}
+impl Future for TokioTimer {
+    type Output = ();
+
+    fn poll(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        let this = self.project();
+        this.sleep.poll(cx)
+    }
+}
+impl sse_stream::Timer for TokioTimer {
+    fn from_duration(duration: Duration) -> Self {
+        Self {
+            sleep: tokio::time::sleep(duration),
+        }
+    }
+
+    fn reset(self: std::pin::Pin<&mut Self>, when: std::time::Instant) {
+        let this = self.project();
+        this.sleep.reset(tokio::time::Instant::from_std(when));
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ServerSseMessage {
+    pub event_id: Option<String>,
+    pub message: Arc<ServerJsonRpcMessage>,
+}
+
+pub(crate) fn sse_stream_response(
+    stream: impl futures::Stream<Item = ServerSseMessage> + Send + 'static,
+    keep_alive: Option<Duration>,
+) -> Response<UnsyncBoxBody<Bytes, Infallible>> {
+    use futures::StreamExt;
+    let stream = SseBody::new(stream.map(|message| {
+        let data = serde_json::to_string(&message.message).expect("valid message");
+        let mut sse = Sse::default().data(data);
+        sse.id = message.event_id;
+        Result::<Sse, Infallible>::Ok(sse)
+    }));
+    let stream = match keep_alive {
+        Some(duration) => stream
+            .with_keep_alive::<TokioTimer>(KeepAlive::new().interval(duration))
+            .boxed_unsync(),
+        None => stream.boxed_unsync(),
+    };
+    Response::builder()
+        .status(http::StatusCode::OK)
+        .header(http::header::CONTENT_TYPE, EVENT_STREAM_MIME_TYPE)
+        .header(http::header::CACHE_CONTROL, "no-cache")
+        .body(stream)
+        .expect("valid response")
+}
+
+pub(crate) const fn internal_error_response<E: Display>(
+    context: &str,
+) -> impl FnOnce(E) -> Response<UnsyncBoxBody<Bytes, Infallible>> {
+    move |error| {
+        tracing::error!("Internal server error when {context}: {error}");
+        Response::builder()
+            .status(http::StatusCode::INTERNAL_SERVER_ERROR)
+            .body(
+                Full::new(Bytes::from(format!(
+                    "Encounter an error when {context}: {error}"
+                )))
+                .boxed_unsync(),
+            )
+            .expect("valid response")
+    }
+}
+
+pub(crate) async fn expect_json<B>(
+    body: B,
+) -> Result<ClientJsonRpcMessage, Response<UnsyncBoxBody<Bytes, Infallible>>>
+where
+    B: Body + Send + 'static,
+    B::Error: Display,
+{
+    match body.collect().await {
+        Ok(bytes) => {
+            match serde_json::from_reader::<_, ClientJsonRpcMessage>(bytes.aggregate().reader()) {
+                Ok(message) => Ok(message),
+                Err(e) => {
+                    let response = Response::builder()
+                        .status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
+                        .body(
+                            Full::new(Bytes::from(format!("fail to deserialize request body {e}")))
+                                .boxed_unsync(),
+                        )
+                        .expect("valid response");
+                    Err(response)
+                }
+            }
+        }
+        Err(e) => {
+            let response = Response::builder()
+                .status(http::StatusCode::INTERNAL_SERVER_ERROR)
+                .body(
+                    Full::new(Bytes::from(format!("Failed to read request body: {e}")))
+                        .boxed_unsync(),
+                )
+                .expect("valid response");
+            Err(response)
+        }
+    }
+}

crates/rmcp/src/transport/sink_stream.rs

@@ -24,6 +24,7 @@ impl<Role: ServiceRole, Si, St> Transport<Role> for SinkStreamTransport<Si, St>
 where
     St: Send + Stream<Item = RxJsonRpcMessage<Role>> + Unpin,
     Si: Send + Sink<TxJsonRpcMessage<Role>> + Unpin + 'static,
+    Si::Error: std::error::Error + Send + Sync + 'static,
 {
     type Error = Si::Error;
 
@@ -56,7 +57,7 @@ where
     Role: ServiceRole,
     Si: Send + Sink<TxJsonRpcMessage<Role>> + Unpin + 'static,
     St: Send + Stream<Item = RxJsonRpcMessage<Role>> + Unpin + 'static,
-    Si::Error: std::error::Error + Send + 'static,
+    Si::Error: std::error::Error + Send + Sync + 'static,
 {
     fn into_transport(self) -> impl Transport<Role, Error = Si::Error> + 'static {
         SinkStreamTransport::new(self.0, self.1)
@@ -68,7 +69,7 @@ impl<Role, S> IntoTransport<Role, S::Error, TransportAdapterAsyncCombinedRW> for
 where
     Role: ServiceRole,
     S: Sink<TxJsonRpcMessage<Role>> + Stream<Item = RxJsonRpcMessage<Role>> + Send + 'static,
-    S::Error: std::error::Error + Send + 'static,
+    S::Error: std::error::Error + Send + Sync + 'static,
 {
     fn into_transport(self) -> impl Transport<Role, Error = S::Error> + 'static {
         use futures::StreamExt;

crates/rmcp/src/transport/sse_server.rs

@@ -19,7 +19,7 @@ use crate::{
     RoleServer, Service,
     model::ClientJsonRpcMessage,
     service::{RxJsonRpcMessage, TxJsonRpcMessage, serve_directly_with_ct},
-    transport::common::axum::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
+    transport::common::server_side_http::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
 };
 
 type TxStore =
@@ -313,7 +313,7 @@ impl SseServer {
                 let service = service_provider();
                 let ct = self.config.ct.child_token();
                 tokio::spawn(async move {
-                    let server = serve_directly_with_ct(service, transport, None, ct).await;
+                    let server = serve_directly_with_ct(service, transport, None, ct);
                     server.waiting().await?;
                     tokio::io::Result::Ok(())
                 });

crates/rmcp/src/transport/streamable_http_server.rs

@@ -1,5 +1,8 @@
+pub mod session;
 #[cfg(feature = "transport-streamable-http-server")]
 #[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
-pub mod axum;
+pub mod tower;
-pub mod session;
+pub use session::{SessionId, SessionManager};
-pub use session::{SessionConfig, create_session};
+#[cfg(feature = "transport-streamable-http-server")]
+#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
+pub use tower::{StreamableHttpServerConfig, StreamableHttpService};

crates/rmcp/src/transport/streamable_http_server/axum.rs

@@ -1,349 +0,0 @@
-use std::{
-    collections::HashMap,
-    io,
-    net::{Ipv6Addr, SocketAddr, SocketAddrV6},
-    sync::Arc,
-    time::Duration,
-};
-
-use axum::{
-    Json, Router,
-    extract::State,
-    http::{HeaderMap, HeaderValue, StatusCode, request::Parts},
-    response::{
-        IntoResponse, Response,
-        sse::{Event, KeepAlive, Sse},
-    },
-    routing::get,
-};
-use futures::Stream;
-use tokio_stream::wrappers::ReceiverStream;
-use tokio_util::sync::CancellationToken;
-use tracing::Instrument;
-
-use super::session::{EventId, SessionHandle, SessionWorker, StreamableHttpMessageReceiver};
-use crate::{
-    RoleServer, Service,
-    model::ClientJsonRpcMessage,
-    transport::common::{
-        axum::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
-        http_header::{HEADER_LAST_EVENT_ID, HEADER_SESSION_ID},
-    },
-};
-type SessionManager = Arc<tokio::sync::RwLock<HashMap<SessionId, SessionHandle>>>;
-
-#[derive(Clone)]
-struct App {
-    session_manager: SessionManager,
-    transport_tx: tokio::sync::mpsc::UnboundedSender<SessionWorker>,
-    sse_ping_interval: Duration,
-}
-
-impl App {
-    pub fn new(
-        sse_ping_interval: Duration,
-    ) -> (Self, tokio::sync::mpsc::UnboundedReceiver<SessionWorker>) {
-        let (transport_tx, transport_rx) = tokio::sync::mpsc::unbounded_channel();
-        (
-            Self {
-                session_manager: Default::default(),
-                transport_tx,
-                sse_ping_interval,
-            },
-            transport_rx,
-        )
-    }
-}
-
-fn receiver_as_stream(
-    receiver: StreamableHttpMessageReceiver,
-) -> impl Stream<Item = Result<Event, io::Error>> {
-    use futures::StreamExt;
-    ReceiverStream::new(receiver.inner).map(|message| {
-        match serde_json::to_string(&message.message) {
-            Ok(bytes) => Ok(Event::default()
-                .event("message")
-                .data(&bytes)
-                .id(message.event_id.to_string())),
-            Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)),
-        }
-    })
-}
-
-async fn post_handler(
-    State(app): State<App>,
-    parts: Parts,
-    Json(mut message): Json<ClientJsonRpcMessage>,
-) -> Result<Response, Response> {
-    use futures::StreamExt;
-    if let Some(session_id) = parts.headers.get(HEADER_SESSION_ID).cloned() {
-        let session_id = session_id
-            .to_str()
-            .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()).into_response())?;
-        tracing::debug!(session_id, ?message, "new client message");
-        let handle = {
-            let sm = app.session_manager.read().await;
-            let session = sm
-                .get(session_id)
-                .ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?;
-            session.clone()
-        };
-        // inject request part
-        message.insert_extension(parts);
-        match &message {
-            ClientJsonRpcMessage::Request(_) | ClientJsonRpcMessage::BatchRequest(_) => {
-                let receiver = handle.establish_request_wise_channel().await.map_err(|e| {
-                    (
-                        StatusCode::INTERNAL_SERVER_ERROR,
-                        format!("fail to to establish request channel: {e}"),
-                    )
-                        .into_response()
-                })?;
-                let http_request_id = receiver.http_request_id;
-                if let Err(push_err) = handle.push_message(message, http_request_id).await {
-                    tracing::error!(session_id, ?push_err, "push message error");
-                    return Err((
-                        StatusCode::INTERNAL_SERVER_ERROR,
-                        format!("fail to push message: {push_err}"),
-                    )
-                        .into_response());
-                }
-                let stream =
-                    ReceiverStream::new(receiver.inner).map(|message| match serde_json::to_string(
-                        &message.message,
-                    ) {
-                        Ok(bytes) => Ok(Event::default()
-                            .event("message")
-                            .data(&bytes)
-                            .id(message.event_id.to_string())),
-                        Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)),
-                    });
-                Ok(Sse::new(stream)
-                    .keep_alive(KeepAlive::new().interval(app.sse_ping_interval))
-                    .into_response())
-            }
-            _ => {
-                let result = handle.push_message(message, None).await;
-                if result.is_err() {
-                    Err((StatusCode::GONE, "session terminated").into_response())
-                } else {
-                    Ok(StatusCode::ACCEPTED.into_response())
-                }
-            }
-        }
-    } else {
-        // expect initialize message
-        let session_id = session_id();
-        // inject request part
-        message.insert_extension(parts);
-        let (session, transport) =
-            super::session::create_session(session_id.clone(), Default::default());
-        let Ok(_) = app.transport_tx.send(transport) else {
-            return Err((StatusCode::GONE, "session terminated").into_response());
-        };
-
-        let response = session.initialize(message).await.map_err(|e| {
-            (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                format!("fail to initialize: {e}"),
-            )
-                .into_response()
-        })?;
-        let mut response = Json(response).into_response();
-        response.headers_mut().insert(
-            HEADER_SESSION_ID,
-            HeaderValue::from_bytes(session_id.as_bytes()).expect("should be valid header value"),
-        );
-        app.session_manager
-            .write()
-            .await
-            .insert(session_id, session);
-        Ok(response)
-    }
-}
-
-async fn get_handler(
-    State(app): State<App>,
-    header_map: HeaderMap,
-) -> Result<Sse<impl Stream<Item = Result<Event, io::Error>>>, Response> {
-    let session_id = header_map
-        .get(HEADER_SESSION_ID)
-        .and_then(|v| v.to_str().ok());
-    if let Some(session_id) = session_id {
-        let last_event_id = header_map
-            .get(HEADER_LAST_EVENT_ID)
-            .and_then(|v| v.to_str().ok());
-        let session = {
-            let sm = app.session_manager.read().await;
-            sm.get(session_id)
-                .ok_or_else(|| {
-                    (
-                        StatusCode::NOT_FOUND,
-                        format!("session {session_id} not found"),
-                    )
-                        .into_response()
-                })?
-                .clone()
-        };
-        match last_event_id {
-            Some(last_event_id) => {
-                let last_event_id = last_event_id.parse::<EventId>().map_err(|e| {
-                    (StatusCode::BAD_REQUEST, format!("invalid event_id {e}")).into_response()
-                })?;
-                let receiver = session.resume(last_event_id).await.map_err(|e| {
-                    (
-                        StatusCode::INTERNAL_SERVER_ERROR,
-                        format!("resume error {e}"),
-                    )
-                        .into_response()
-                })?;
-                let stream = receiver_as_stream(receiver);
-                Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(app.sse_ping_interval)))
-            }
-            None => {
-                let receiver = session.establish_common_channel().await.map_err(|e| {
-                    (
-                        StatusCode::INTERNAL_SERVER_ERROR,
-                        format!("establish common channel error {e}"),
-                    )
-                        .into_response()
-                })?;
-                let stream = receiver_as_stream(receiver);
-                Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(app.sse_ping_interval)))
-            }
-        }
-    } else {
-        Err((StatusCode::BAD_REQUEST, "missing session id").into_response())
-    }
-}
-
-async fn delete_handler(
-    State(app): State<App>,
-    header_map: HeaderMap,
-) -> Result<StatusCode, Response> {
-    if let Some(session_id) = header_map.get(HEADER_SESSION_ID) {
-        let session_id = session_id
-            .to_str()
-            .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()).into_response())?;
-        let session = {
-            let mut sm = app.session_manager.write().await;
-            sm.remove(session_id)
-                .ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?
-        };
-        session.close().await.map_err(|e| {
-            (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                format!("fail to cancel session {session_id}: tokio join error: {e}"),
-            )
-                .into_response()
-        })?;
-        tracing::debug!(session_id, "session deleted");
-        Ok(StatusCode::ACCEPTED)
-    } else {
-        Err((StatusCode::BAD_REQUEST, "missing session id").into_response())
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct StreamableHttpServerConfig {
-    pub bind: SocketAddr,
-    pub path: String,
-    pub ct: CancellationToken,
-    pub sse_keep_alive: Option<Duration>,
-}
-impl Default for StreamableHttpServerConfig {
-    fn default() -> Self {
-        Self {
-            bind: SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 80, 0, 0)),
-            path: "/".to_string(),
-            ct: CancellationToken::new(),
-            sse_keep_alive: None,
-        }
-    }
-}
-
-#[derive(Debug)]
-pub struct StreamableHttpServer {
-    transport_rx: tokio::sync::mpsc::UnboundedReceiver<SessionWorker>,
-    pub config: StreamableHttpServerConfig,
-}
-
-impl StreamableHttpServer {
-    pub async fn serve(bind: SocketAddr) -> io::Result<Self> {
-        Self::serve_with_config(StreamableHttpServerConfig {
-            bind,
-            ..Default::default()
-        })
-        .await
-    }
-    pub async fn serve_with_config(config: StreamableHttpServerConfig) -> io::Result<Self> {
-        let (streamable_http_server, service) = Self::new(config);
-        let listener = tokio::net::TcpListener::bind(streamable_http_server.config.bind).await?;
-        let ct = streamable_http_server.config.ct.child_token();
-        let server = axum::serve(listener, service).with_graceful_shutdown(async move {
-            ct.cancelled().await;
-            tracing::info!("streamable http server cancelled");
-        });
-        tokio::spawn(
-            async move {
-                if let Err(e) = server.await {
-                    tracing::error!(error = %e, "streamable http server shutdown with error");
-                }
-            }
-            .instrument(tracing::info_span!("streamable-http-server", bind_address = %streamable_http_server.config.bind)),
-        );
-        Ok(streamable_http_server)
-    }
-
-    /// Warning: This function creates a new StreamableHttpServer instance with the provided configuration.
-    /// `App.post_path` may be incorrect if using `Router` as an embedded router.
-    pub fn new(config: StreamableHttpServerConfig) -> (StreamableHttpServer, Router) {
-        let (app, transport_rx) =
-            App::new(config.sse_keep_alive.unwrap_or(DEFAULT_AUTO_PING_INTERVAL));
-        let router = Router::new()
-            .route(
-                &config.path,
-                get(get_handler).post(post_handler).delete(delete_handler),
-            )
-            .with_state(app);
-
-        let server = StreamableHttpServer {
-            transport_rx,
-            config,
-        };
-
-        (server, router)
-    }
-
-    pub fn with_service<S, F>(mut self, service_provider: F) -> CancellationToken
-    where
-        S: Service<RoleServer>,
-        F: Fn() -> S + Send + 'static,
-    {
-        use crate::service::ServiceExt;
-        let ct = self.config.ct.clone();
-        tokio::spawn(async move {
-            while let Some(transport) = self.next_transport().await {
-                let service = service_provider();
-                let ct = self.config.ct.child_token();
-                tokio::spawn(async move {
-                    let server = service
-                        .serve_with_ct(transport, ct)
-                        .await
-                        .map_err(tokio::io::Error::other)?;
-                    server.waiting().await?;
-                    tokio::io::Result::Ok(())
-                });
-            }
-        });
-        ct
-    }
-
-    pub fn cancel(&self) {
-        self.config.ct.cancel();
-    }
-
-    pub async fn next_transport(&mut self) -> Option<SessionWorker> {
-        self.transport_rx.recv().await
-    }
-}

crates/rmcp/src/transport/streamable_http_server/session.rs

@@ -1,783 +1,54 @@
-use std::{
+use futures::Stream;
-    borrow::Cow,
-    collections::{HashMap, HashSet, VecDeque},
-    num::ParseIntError,
-    sync::Arc,
-};
-
-use thiserror::Error;
-use tokio::sync::{
-    mpsc::{Receiver, Sender},
-    oneshot,
-};
-use tracing::instrument;
 
+pub use crate::transport::common::server_side_http::SessionId;
 use crate::{
     RoleServer,
-    model::{
+    model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
-        CancelledNotificationParam, ClientJsonRpcMessage, ClientNotification, ClientRequest,
+    transport::common::server_side_http::ServerSseMessage,
-        JsonRpcNotification, JsonRpcRequest, Notification, ProgressNotificationParam,
-        ProgressToken, RequestId, ServerJsonRpcMessage, ServerNotification,
-    },
-    transport::{
-        WorkerTransport,
-        worker::{Worker, WorkerContext, WorkerQuitReason, WorkerSendRequest},
-    },
 };
 
-#[derive(Debug, Clone)]
+pub mod local;
-pub struct ServerSessionMessage {
+pub mod never;
-    pub event_id: EventId,
-    pub message: Arc<ServerJsonRpcMessage>,
-}
 
-/// `<index>/request_id>`
+pub trait SessionManager: Send + Sync + 'static {
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+    type Error: std::error::Error + Send + 'static;
-pub struct EventId {
+    type Transport: crate::transport::Transport<RoleServer>;
-    http_request_id: Option<HttpRequestId>,
+    /// Create a new session with the given id and configuration.
-    index: usize,
+    fn create_session(
-}
+        &self,
-
+    ) -> impl Future<Output = Result<(SessionId, Self::Transport), Self::Error>> + Send;
-impl std::fmt::Display for EventId {
+    fn initialize_session(
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        &self,
-        write!(f, "{}", self.index)?;
+        id: &SessionId,
-        match &self.http_request_id {
-            Some(http_request_id) => write!(f, "/{http_request_id}"),
-            None => write!(f, ""),
-        }
-    }
-}
-
-#[derive(Debug, Clone, Error)]
-pub enum EventIdParseError {
-    #[error("Invalid index: {0}")]
-    InvalidIndex(ParseIntError),
-    #[error("Invalid numeric request id: {0}")]
-    InvalidNumericRequestId(ParseIntError),
-    #[error("Missing request id type")]
-    InvalidRequestIdType,
-    #[error("Missing request id")]
-    MissingRequestId,
-}
-
-impl std::str::FromStr for EventId {
-    type Err = EventIdParseError;
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        if let Some((index, request_id)) = s.split_once("/") {
-            let index = usize::from_str(index).map_err(EventIdParseError::InvalidIndex)?;
-            let request_id = u64::from_str(request_id).map_err(EventIdParseError::InvalidIndex)?;
-            Ok(EventId {
-                http_request_id: Some(request_id),
-                index,
-            })
-        } else {
-            let index = usize::from_str(s).map_err(EventIdParseError::InvalidIndex)?;
-            Ok(EventId {
-                http_request_id: None,
-                index,
-            })
-        }
-    }
-}
-
-pub use crate::transport::common::axum::SessionId;
-
-struct CachedTx {
-    tx: Sender<ServerSessionMessage>,
-    cache: VecDeque<ServerSessionMessage>,
-    http_request_id: Option<HttpRequestId>,
-    capacity: usize,
-}
-
-impl CachedTx {
-    fn new(tx: Sender<ServerSessionMessage>, http_request_id: Option<HttpRequestId>) -> Self {
-        Self {
-            cache: VecDeque::with_capacity(tx.capacity()),
-            capacity: tx.capacity(),
-            tx,
-            http_request_id,
-        }
-    }
-    fn new_common(tx: Sender<ServerSessionMessage>) -> Self {
-        Self::new(tx, None)
-    }
-
-    async fn send(&mut self, message: ServerJsonRpcMessage) {
-        let index = self.cache.back().map_or(0, |m| m.event_id.index + 1);
-        let event_id = EventId {
-            http_request_id: self.http_request_id,
-            index,
-        };
-        let message = ServerSessionMessage {
-            event_id: event_id.clone(),
-            message: Arc::new(message),
-        };
-        if self.cache.len() >= self.capacity {
-            self.cache.pop_front();
-            self.cache.push_back(message.clone());
-        } else {
-            self.cache.push_back(message.clone());
-        }
-        let _ = self.tx.send(message).await.inspect_err(|e| {
-            let event_id = &e.0.event_id;
-            tracing::trace!(%event_id, "trying to send message in a closed session")
-        });
-    }
-
-    async fn sync(&mut self, index: usize) -> Result<(), SessionError> {
-        let Some(front) = self.cache.front() else {
-            return Ok(());
-        };
-        let sync_index = index.saturating_sub(front.event_id.index);
-        if sync_index > self.cache.len() {
-            // invalid index
-            return Err(SessionError::InvalidEventId);
-        }
-        for message in self.cache.iter().skip(sync_index) {
-            let send_result = self.tx.send(message.clone()).await;
-            if send_result.is_err() {
-                return Err(SessionError::ChannelClosed(
-                    message.event_id.http_request_id,
-                ));
-            }
-        }
-        Ok(())
-    }
-}
-
-struct HttpRequestWise {
-    resources: HashSet<ResourceKey>,
-    tx: CachedTx,
-}
-
-type HttpRequestId = u64;
-#[derive(Debug, Clone, Hash, PartialEq, Eq)]
-enum ResourceKey {
-    McpRequestId(RequestId),
-    ProgressToken(ProgressToken),
-}
-
-pub struct SessionWorker {
-    id: SessionId,
-    next_http_request_id: HttpRequestId,
-    tx_router: HashMap<HttpRequestId, HttpRequestWise>,
-    resource_router: HashMap<ResourceKey, HttpRequestId>,
-    common: CachedTx,
-    event_rx: Receiver<SessionEvent>,
-    session_config: SessionConfig,
-}
-
-impl SessionWorker {
-    pub fn id(&self) -> &SessionId {
-        &self.id
-    }
-}
-
-#[derive(Debug, Error)]
-pub enum SessionError {
-    #[error("Invalid request id: {0}")]
-    DuplicatedRequestId(HttpRequestId),
-    #[error("Channel closed: {0:?}")]
-    ChannelClosed(Option<HttpRequestId>),
-    #[error("Cannot parse event id: {0}")]
-    EventIdParseError(Cow<'static, str>),
-    #[error("Session service terminated")]
-    SessionServiceTerminated,
-    #[error("Invalid event id")]
-    InvalidEventId,
-    #[error("Transport closed")]
-    TransportClosed,
-    #[error("IO error: {0}")]
-    Io(#[from] std::io::Error),
-    #[error("Tokio join error {0}")]
-    TokioJoinError(#[from] tokio::task::JoinError),
-}
-
-impl From<SessionError> for std::io::Error {
-    fn from(value: SessionError) -> Self {
-        match value {
-            SessionError::Io(io) => io,
-            _ => std::io::Error::new(std::io::ErrorKind::Other, format!("Session error: {value}")),
-        }
-    }
-}
-
-enum OutboundChannel {
-    RequestWise { id: HttpRequestId, close: bool },
-    Common,
-}
-
-pub struct StreamableHttpMessageReceiver {
-    pub http_request_id: Option<HttpRequestId>,
-    pub inner: Receiver<ServerSessionMessage>,
-}
-
-impl SessionWorker {
-    fn unregister_resource(&mut self, resource: &ResourceKey) {
-        if let Some(http_request_id) = self.resource_router.remove(resource) {
-            tracing::trace!(?resource, http_request_id, "unregister resource");
-            if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
-                channel.resources.remove(resource);
-                if channel.resources.is_empty() {
-                    tracing::debug!(http_request_id, "close http request wise channel");
-                    self.tx_router.remove(&http_request_id);
-                }
-            }
-        }
-    }
-    fn register_resource(&mut self, resource: ResourceKey, http_request_id: HttpRequestId) {
-        tracing::trace!(?resource, http_request_id, "register resource");
-        if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
-            channel.resources.insert(resource.clone());
-            self.resource_router.insert(resource, http_request_id);
-        }
-    }
-    fn register_request(
-        &mut self,
-        request: &JsonRpcRequest<ClientRequest>,
-        http_request_id: HttpRequestId,
-    ) {
-        use crate::model::GetMeta;
-        self.register_resource(
-            ResourceKey::McpRequestId(request.id.clone()),
-            http_request_id,
-        );
-        if let Some(progress_token) = request.request.get_meta().get_progress_token() {
-            self.register_resource(
-                ResourceKey::ProgressToken(progress_token.clone()),
-                http_request_id,
-            );
-        }
-    }
-    fn catch_cancellation_notification(
-        &mut self,
-        notification: &JsonRpcNotification<ClientNotification>,
-    ) {
-        if let ClientNotification::CancelledNotification(n) = &notification.notification {
-            let request_id = n.params.request_id.clone();
-            let resource = ResourceKey::McpRequestId(request_id);
-            self.unregister_resource(&resource);
-        }
-    }
-    fn next_http_request_id(&mut self) -> HttpRequestId {
-        let id = self.next_http_request_id;
-        self.next_http_request_id = self.next_http_request_id.wrapping_add(1);
-        id
-    }
-    async fn establish_request_wise_channel(
-        &mut self,
-    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
-        let http_request_id = self.next_http_request_id();
-        let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
-        self.tx_router.insert(
-            http_request_id,
-            HttpRequestWise {
-                resources: Default::default(),
-                tx: CachedTx::new(tx, Some(http_request_id)),
-            },
-        );
-        tracing::debug!(http_request_id, "establish new request wise channel");
-        Ok(StreamableHttpMessageReceiver {
-            http_request_id: Some(http_request_id),
-            inner: rx,
-        })
-    }
-    fn resolve_outbound_channel(&self, message: &ServerJsonRpcMessage) -> OutboundChannel {
-        match &message {
-            ServerJsonRpcMessage::Request(_) => OutboundChannel::Common,
-            ServerJsonRpcMessage::Notification(JsonRpcNotification {
-                notification:
-                    ServerNotification::ProgressNotification(Notification {
-                        params: ProgressNotificationParam { progress_token, .. },
-                        ..
-                    }),
-                ..
-            }) => {
-                let id = self
-                    .resource_router
-                    .get(&ResourceKey::ProgressToken(progress_token.clone()));
-
-                if let Some(id) = id {
-                    OutboundChannel::RequestWise {
-                        id: *id,
-                        close: false,
-                    }
-                } else {
-                    OutboundChannel::Common
-                }
-            }
-            ServerJsonRpcMessage::Notification(JsonRpcNotification {
-                notification:
-                    ServerNotification::CancelledNotification(Notification {
-                        params: CancelledNotificationParam { request_id, .. },
-                        ..
-                    }),
-                ..
-            }) => {
-                if let Some(id) = self
-                    .resource_router
-                    .get(&ResourceKey::McpRequestId(request_id.clone()))
-                {
-                    OutboundChannel::RequestWise {
-                        id: *id,
-                        close: false,
-                    }
-                } else {
-                    OutboundChannel::Common
-                }
-            }
-            ServerJsonRpcMessage::Notification(_) => OutboundChannel::Common,
-            ServerJsonRpcMessage::Response(json_rpc_response) => {
-                if let Some(id) = self
-                    .resource_router
-                    .get(&ResourceKey::McpRequestId(json_rpc_response.id.clone()))
-                {
-                    OutboundChannel::RequestWise {
-                        id: *id,
-                        close: false,
-                    }
-                } else {
-                    OutboundChannel::Common
-                }
-            }
-            ServerJsonRpcMessage::Error(json_rpc_error) => {
-                if let Some(id) = self
-                    .resource_router
-                    .get(&ResourceKey::McpRequestId(json_rpc_error.id.clone()))
-                {
-                    OutboundChannel::RequestWise {
-                        id: *id,
-                        close: false,
-                    }
-                } else {
-                    OutboundChannel::Common
-                }
-            }
-            ServerJsonRpcMessage::BatchRequest(_) | ServerJsonRpcMessage::BatchResponse(_) => {
-                // the server side should never yield a batch request or response now
-                unreachable!("server side won't yield batch request or response")
-            }
-        }
-    }
-    async fn handle_server_message(
-        &mut self,
-        message: ServerJsonRpcMessage,
-    ) -> Result<(), SessionError> {
-        let outbound_channel = self.resolve_outbound_channel(&message);
-        match outbound_channel {
-            OutboundChannel::RequestWise { id, close } => {
-                if let Some(request_wise) = self.tx_router.get_mut(&id) {
-                    request_wise.tx.send(message).await;
-                    if close {
-                        self.tx_router.remove(&id);
-                    }
-                } else {
-                    return Err(SessionError::ChannelClosed(Some(id)));
-                }
-            }
-            OutboundChannel::Common => self.common.send(message).await,
-        }
-        Ok(())
-    }
-    async fn resume(
-        &mut self,
-        last_event_id: EventId,
-    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
-        match last_event_id.http_request_id {
-            Some(http_request_id) => {
-                let request_wise = self
-                    .tx_router
-                    .get_mut(&http_request_id)
-                    .ok_or(SessionError::ChannelClosed(Some(http_request_id)))?;
-                let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
-                let (tx, rx) = channel;
-                request_wise.tx.tx = tx;
-                let index = last_event_id.index;
-                // sync messages after index
-                request_wise.tx.sync(index).await?;
-                Ok(StreamableHttpMessageReceiver {
-                    http_request_id: Some(http_request_id),
-                    inner: rx,
-                })
-            }
-            None => {
-                let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
-                let (tx, rx) = channel;
-                self.common.tx = tx;
-                let index = last_event_id.index;
-                // sync messages after index
-                self.common.sync(index).await?;
-                Ok(StreamableHttpMessageReceiver {
-                    http_request_id: None,
-                    inner: rx,
-                })
-            }
-        }
-    }
-}
-
-enum SessionEvent {
-    ClientMessage {
         message: ClientJsonRpcMessage,
-        http_request_id: Option<HttpRequestId>,
+    ) -> impl Future<Output = Result<ServerJsonRpcMessage, Self::Error>> + Send;
-    },
+    fn has_session(&self, id: &SessionId)
-    EstablishRequestWiseChannel {
+    -> impl Future<Output = Result<bool, Self::Error>> + Send;
-        responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
+    fn close_session(&self, id: &SessionId)
-    },
+    -> impl Future<Output = Result<(), Self::Error>> + Send;
-    CloseRequestWiseChannel {
+    fn create_stream(
-        id: HttpRequestId,
-        responder: oneshot::Sender<Result<(), SessionError>>,
-    },
-    Resume {
-        last_event_id: EventId,
-        responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
-    },
-    InitializeRequest {
-        request: ClientJsonRpcMessage,
-        responder: oneshot::Sender<Result<ServerJsonRpcMessage, SessionError>>,
-    },
-    Close,
-}
-
-#[derive(Debug, Clone)]
-pub enum SessionQuitReason {
-    ServiceTerminated,
-    ClientTerminated,
-    ExpectInitializeRequest,
-    ExpectInitializeResponse,
-    Cancelled,
-}
-
-#[derive(Debug, Clone)]
-pub struct SessionHandle {
-    id: SessionId,
-    // after all event_tx drop, inner task will be terminated
-    event_tx: Sender<SessionEvent>,
-}
-
-impl SessionHandle {
-    /// Get the session id
-    pub fn id(&self) -> &SessionId {
-        &self.id
-    }
-
-    /// Close the session
-    pub async fn close(&self) -> Result<(), SessionError> {
-        self.event_tx
-            .send(SessionEvent::Close)
-            .await
-            .map_err(|_| SessionError::SessionServiceTerminated)?;
-        Ok(())
-    }
-
-    /// Send a message to the session
-    pub async fn push_message(
         &self,
+        id: &SessionId,
         message: ClientJsonRpcMessage,
-        http_request_id: Option<HttpRequestId>,
+    ) -> impl Future<
-    ) -> Result<(), SessionError> {
+        Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
-        self.event_tx
+    > + Send;
-            .send(SessionEvent::ClientMessage {
+    fn accept_message(
-                message,
-                http_request_id,
-            })
-            .await
-            .map_err(|_| SessionError::SessionServiceTerminated)?;
-        Ok(())
-    }
-
-    /// establish a channel for a http-request, the corresponded message from server will be
-    /// sent through this channel. The channel will be closed when the request is completed,
-    /// or you can close it manually by calling [`SessionHandle::close_request_wise_channel`].
-    pub async fn establish_request_wise_channel(
         &self,
-    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
+        id: &SessionId,
-        let (tx, rx) = tokio::sync::oneshot::channel();
+        message: ClientJsonRpcMessage,
-        self.event_tx
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send;
-            .send(SessionEvent::EstablishRequestWiseChannel { responder: tx })
+    fn create_standalone_stream(
-            .await
-            .map_err(|_| SessionError::SessionServiceTerminated)?;
-        rx.await
-            .map_err(|_| SessionError::SessionServiceTerminated)?
-    }
-
-    /// close the http-request wise channel.
-    pub async fn close_request_wise_channel(
         &self,
-        request_id: HttpRequestId,
+        id: &SessionId,
-    ) -> Result<(), SessionError> {
+    ) -> impl Future<
-        let (tx, rx) = tokio::sync::oneshot::channel();
+        Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
-        self.event_tx
+    > + Send;
-            .send(SessionEvent::CloseRequestWiseChannel {
+    fn resume(
-                id: request_id,
-                responder: tx,
-            })
-            .await
-            .map_err(|_| SessionError::SessionServiceTerminated)?;
-        rx.await
-            .map_err(|_| SessionError::SessionServiceTerminated)?
-    }
-
-    /// Establish a common channel for general purpose messages.
-    pub async fn establish_common_channel(
         &self,
-    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
+        id: &SessionId,
-        let (tx, rx) = tokio::sync::oneshot::channel();
+        last_event_id: String,
-        self.event_tx
+    ) -> impl Future<
-            .send(SessionEvent::Resume {
+        Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
-                last_event_id: EventId {
+    > + Send;
-                    http_request_id: None,
-                    index: 0,
-                },
-                responder: tx,
-            })
-            .await
-            .map_err(|_| SessionError::SessionServiceTerminated)?;
-        rx.await
-            .map_err(|_| SessionError::SessionServiceTerminated)?
-    }
-
-    /// Resume streaming response by the last event id. This is suitable for both request wise and common channel.
-    pub async fn resume(
-        &self,
-        last_event_id: EventId,
-    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
-        let (tx, rx) = tokio::sync::oneshot::channel();
-        self.event_tx
-            .send(SessionEvent::Resume {
-                last_event_id,
-                responder: tx,
-            })
-            .await
-            .map_err(|_| SessionError::SessionServiceTerminated)?;
-        rx.await
-            .map_err(|_| SessionError::SessionServiceTerminated)?
-    }
-
-    /// Send an initialize request to the session. And wait for the initialized response.
-    ///
-    /// This is used to establish a session with the server.
-    pub async fn initialize(
-        &self,
-        request: ClientJsonRpcMessage,
-    ) -> Result<ServerJsonRpcMessage, SessionError> {
-        let (tx, rx) = tokio::sync::oneshot::channel();
-        self.event_tx
-            .send(SessionEvent::InitializeRequest {
-                request,
-                responder: tx,
-            })
-            .await
-            .map_err(|_| SessionError::SessionServiceTerminated)?;
-        rx.await
-            .map_err(|_| SessionError::SessionServiceTerminated)?
-    }
-}
-
-pub type SessionTransport = WorkerTransport<SessionWorker>;
-
-impl Worker for SessionWorker {
-    type Error = SessionError;
-    type Role = RoleServer;
-    fn err_closed() -> Self::Error {
-        SessionError::TransportClosed
-    }
-    fn err_join(e: tokio::task::JoinError) -> Self::Error {
-        SessionError::TokioJoinError(e)
-    }
-    fn config(&self) -> crate::transport::worker::WorkerConfig {
-        crate::transport::worker::WorkerConfig {
-            name: Some(format!("streamable-http-session-{}", self.id)),
-            channel_buffer_capacity: self.session_config.channel_capacity,
-        }
-    }
-    #[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))]
-    async fn run(mut self, mut context: WorkerContext<Self>) -> Result<(), WorkerQuitReason> {
-        enum InnerEvent {
-            FromHttpService(SessionEvent),
-            FromHandler(WorkerSendRequest<SessionWorker>),
-        }
-        // waiting for initialize request
-        let evt = self.event_rx.recv().await.ok_or_else(|| {
-            WorkerQuitReason::fatal("transport terminated", "get initialize request")
-        })?;
-        let SessionEvent::InitializeRequest { request, responder } = evt else {
-            return Err(WorkerQuitReason::fatal(
-                "unexpected message",
-                "get initialize request",
-            ));
-        };
-        context.send_to_handler(request).await?;
-        let send_initialize_response = context.recv_from_handler().await?;
-        responder
-            .send(Ok(send_initialize_response.message))
-            .map_err(|_| {
-                WorkerQuitReason::fatal(
-                    "failed to send initialize response to http service",
-                    "send initialize response",
-                )
-            })?;
-        send_initialize_response
-            .responder
-            .send(Ok(()))
-            .map_err(|_| WorkerQuitReason::HandlerTerminated)?;
-        let ct = context.cancellation_token.clone();
-        loop {
-            let event = tokio::select! {
-                event = self.event_rx.recv() => {
-                    if let Some(event) = event {
-                        InnerEvent::FromHttpService(event)
-                    } else {
-                        return Err(WorkerQuitReason::fatal("session dropped", "waiting next session event"))
-                    }
-                },
-                from_handler = context.recv_from_handler() => {
-                    InnerEvent::FromHandler(from_handler?)
-                }
-                _ = ct.cancelled() => {
-                    return Err(WorkerQuitReason::Cancelled)
-                }
-            };
-            match event {
-                InnerEvent::FromHandler(WorkerSendRequest { message, responder }) => {
-                    // catch response
-                    match &message {
-                        crate::model::JsonRpcMessage::Response(json_rpc_response) => {
-                            let request_id = json_rpc_response.id.clone();
-                            self.unregister_resource(&ResourceKey::McpRequestId(request_id));
-                        }
-                        crate::model::JsonRpcMessage::Error(json_rpc_error) => {
-                            let request_id = json_rpc_error.id.clone();
-                            self.unregister_resource(&ResourceKey::McpRequestId(request_id));
-                        }
-                        // unlikely happen
-                        crate::model::JsonRpcMessage::BatchResponse(
-                            json_rpc_batch_response_items,
-                        ) => {
-                            for item in json_rpc_batch_response_items {
-                                let request_id = match item {
-                                    crate::model::JsonRpcBatchResponseItem::Response(
-                                        json_rpc_response,
-                                    ) => json_rpc_response.id.clone(),
-                                    crate::model::JsonRpcBatchResponseItem::Error(
-                                        json_rpc_error,
-                                    ) => json_rpc_error.id.clone(),
-                                };
-                                self.unregister_resource(&ResourceKey::McpRequestId(request_id));
-                            }
-                        }
-                        _ => {
-                            // no need to unregister resource
-                        }
-                    }
-                    let handle_result = self.handle_server_message(message).await;
-                    let _ = responder.send(handle_result);
-                }
-                InnerEvent::FromHttpService(SessionEvent::ClientMessage {
-                    message: json_rpc_message,
-                    http_request_id,
-                }) => {
-                    match &json_rpc_message {
-                        crate::model::JsonRpcMessage::Request(request) => {
-                            if let Some(http_request_id) = http_request_id {
-                                self.register_request(request, http_request_id)
-                            }
-                        }
-                        crate::model::JsonRpcMessage::Notification(notification) => {
-                            self.catch_cancellation_notification(notification)
-                        }
-                        crate::model::JsonRpcMessage::BatchRequest(items) => {
-                            for r in items {
-                                match r {
-                                    crate::model::JsonRpcBatchRequestItem::Request(request) => {
-                                        if let Some(http_request_id) = http_request_id {
-                                            self.register_request(request, http_request_id)
-                                        }
-                                    }
-                                    crate::model::JsonRpcBatchRequestItem::Notification(
-                                        notification,
-                                    ) => self.catch_cancellation_notification(notification),
-                                }
-                            }
-                        }
-                        _ => {}
-                    }
-                    context.send_to_handler(json_rpc_message).await?;
-                }
-                InnerEvent::FromHttpService(SessionEvent::EstablishRequestWiseChannel {
-                    responder,
-                }) => {
-                    let handle_result = self.establish_request_wise_channel().await;
-                    let _ = responder.send(handle_result);
-                }
-                InnerEvent::FromHttpService(SessionEvent::CloseRequestWiseChannel {
-                    id,
-                    responder,
-                }) => {
-                    let _handle_result = self.tx_router.remove(&id);
-                    let _ = responder.send(Ok(()));
-                }
-                InnerEvent::FromHttpService(SessionEvent::Resume {
-                    last_event_id,
-                    responder,
-                }) => {
-                    let handle_result = self.resume(last_event_id).await;
-                    let _ = responder.send(handle_result);
-                }
-                InnerEvent::FromHttpService(SessionEvent::Close) => {
-                    return Err(WorkerQuitReason::TransportClosed);
-                }
-                _ => {
-                    // ignore
-                }
-            }
-        }
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct SessionConfig {
-    channel_capacity: usize,
-}
-
-impl SessionConfig {
-    pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
-}
-
-impl Default for SessionConfig {
-    fn default() -> Self {
-        Self {
-            channel_capacity: Self::DEFAULT_CHANNEL_CAPACITY,
-        }
-    }
-}
-
-/// Create a new session with the given id and configuration.
-///
-/// This function will return a pair of [`SessionHandle`] and [`SessionWorker`].
-///
-/// You can run the [`SessionWorker`] as a transport for mcp server. And use the [`SessionHandle`] operate the session.
-pub fn create_session(
-    id: impl Into<SessionId>,
-    config: SessionConfig,
-) -> (SessionHandle, SessionWorker) {
-    let id = id.into();
-    let (event_tx, event_rx) = tokio::sync::mpsc::channel(config.channel_capacity);
-    let (common_tx, _) = tokio::sync::mpsc::channel(config.channel_capacity);
-    let common = CachedTx::new_common(common_tx);
-    tracing::info!(session_id = ?id, "create new session");
-    let handle = SessionHandle {
-        event_tx,
-        id: id.clone(),
-    };
-    let session_worker = SessionWorker {
-        next_http_request_id: 0,
-        id,
-        tx_router: HashMap::new(),
-        resource_router: HashMap::new(),
-        common,
-        event_rx,
-        session_config: config.clone(),
-    };
-    (handle, session_worker)
 }

crates/rmcp/src/transport/streamable_http_server/session/local.rs

@@ -0,0 +1,893 @@
+use std::{
+    collections::{HashMap, HashSet, VecDeque},
+    num::ParseIntError,
+    sync::Arc,
+    time::Duration,
+};
+
+use futures::Stream;
+use thiserror::Error;
+use tokio::sync::{
+    mpsc::{Receiver, Sender},
+    oneshot,
+};
+use tokio_stream::wrappers::ReceiverStream;
+use tracing::instrument;
+
+use crate::{
+    RoleServer,
+    model::{
+        CancelledNotificationParam, ClientJsonRpcMessage, ClientNotification, ClientRequest,
+        JsonRpcNotification, JsonRpcRequest, Notification, ProgressNotificationParam,
+        ProgressToken, RequestId, ServerJsonRpcMessage, ServerNotification,
+    },
+    transport::{
+        WorkerTransport,
+        common::server_side_http::{SessionId, session_id},
+        worker::{Worker, WorkerContext, WorkerQuitReason, WorkerSendRequest},
+    },
+};
+
+#[derive(Debug, Default)]
+pub struct LocalSessionManager {
+    pub sessions: tokio::sync::RwLock<HashMap<SessionId, LocalSessionHandle>>,
+    pub session_config: SessionConfig,
+}
+
+#[derive(Debug, Error)]
+pub enum LocalSessionManagerError {
+    #[error("Session not found: {0}")]
+    SessionNotFound(SessionId),
+    #[error("Session error: {0}")]
+    SessionError(#[from] SessionError),
+    #[error("Invalid event id: {0}")]
+    InvalidEventId(#[from] EventIdParseError),
+}
+impl SessionManager for LocalSessionManager {
+    type Error = LocalSessionManagerError;
+    type Transport = WorkerTransport<LocalSessionWorker>;
+    async fn create_session(&self) -> Result<(SessionId, Self::Transport), Self::Error> {
+        let id = session_id();
+        let (handle, worker) = create_local_session(id.clone(), self.session_config.clone());
+        self.sessions.write().await.insert(id.clone(), handle);
+        Ok((id, WorkerTransport::spawn(worker)))
+    }
+    async fn initialize_session(
+        &self,
+        id: &SessionId,
+        message: ClientJsonRpcMessage,
+    ) -> Result<ServerJsonRpcMessage, Self::Error> {
+        let sessions = self.sessions.read().await;
+        let handle = sessions
+            .get(id)
+            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
+        let response = handle.initialize(message).await?;
+        Ok(response)
+    }
+    async fn close_session(&self, id: &SessionId) -> Result<(), Self::Error> {
+        let mut sessions = self.sessions.write().await;
+        if let Some(handle) = sessions.remove(id) {
+            handle.close().await?;
+        }
+        Ok(())
+    }
+    async fn has_session(&self, id: &SessionId) -> Result<bool, Self::Error> {
+        let sessions = self.sessions.read().await;
+        Ok(sessions.contains_key(id))
+    }
+    async fn create_stream(
+        &self,
+        id: &SessionId,
+        message: ClientJsonRpcMessage,
+    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
+        let sessions = self.sessions.read().await;
+        let handle = sessions
+            .get(id)
+            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
+        let receiver = handle.establish_request_wise_channel().await?;
+        handle
+            .push_message(message, receiver.http_request_id)
+            .await?;
+        Ok(ReceiverStream::new(receiver.inner))
+    }
+
+    async fn create_standalone_stream(
+        &self,
+        id: &SessionId,
+    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
+        let sessions = self.sessions.read().await;
+        let handle = sessions
+            .get(id)
+            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
+        let receiver = handle.establish_common_channel().await?;
+        Ok(ReceiverStream::new(receiver.inner))
+    }
+
+    async fn resume(
+        &self,
+        id: &SessionId,
+        last_event_id: String,
+    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
+        let sessions = self.sessions.read().await;
+        let handle = sessions
+            .get(id)
+            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
+        let receiver = handle.resume(last_event_id.parse()?).await?;
+        Ok(ReceiverStream::new(receiver.inner))
+    }
+
+    async fn accept_message(
+        &self,
+        id: &SessionId,
+        message: ClientJsonRpcMessage,
+    ) -> Result<(), Self::Error> {
+        let sessions = self.sessions.read().await;
+        let handle = sessions
+            .get(id)
+            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
+        handle.push_message(message, None).await?;
+        Ok(())
+    }
+}
+
+/// `<index>/request_id>`
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct EventId {
+    http_request_id: Option<HttpRequestId>,
+    index: usize,
+}
+
+impl std::fmt::Display for EventId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.index)?;
+        match &self.http_request_id {
+            Some(http_request_id) => write!(f, "/{http_request_id}"),
+            None => write!(f, ""),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Error)]
+pub enum EventIdParseError {
+    #[error("Invalid index: {0}")]
+    InvalidIndex(ParseIntError),
+    #[error("Invalid numeric request id: {0}")]
+    InvalidNumericRequestId(ParseIntError),
+    #[error("Missing request id type")]
+    InvalidRequestIdType,
+    #[error("Missing request id")]
+    MissingRequestId,
+}
+
+impl std::str::FromStr for EventId {
+    type Err = EventIdParseError;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Some((index, request_id)) = s.split_once("/") {
+            let index = usize::from_str(index).map_err(EventIdParseError::InvalidIndex)?;
+            let request_id = u64::from_str(request_id).map_err(EventIdParseError::InvalidIndex)?;
+            Ok(EventId {
+                http_request_id: Some(request_id),
+                index,
+            })
+        } else {
+            let index = usize::from_str(s).map_err(EventIdParseError::InvalidIndex)?;
+            Ok(EventId {
+                http_request_id: None,
+                index,
+            })
+        }
+    }
+}
+
+use super::{ServerSseMessage, SessionManager};
+
+struct CachedTx {
+    tx: Sender<ServerSseMessage>,
+    cache: VecDeque<ServerSseMessage>,
+    http_request_id: Option<HttpRequestId>,
+    capacity: usize,
+}
+
+impl CachedTx {
+    fn new(tx: Sender<ServerSseMessage>, http_request_id: Option<HttpRequestId>) -> Self {
+        Self {
+            cache: VecDeque::with_capacity(tx.capacity()),
+            capacity: tx.capacity(),
+            tx,
+            http_request_id,
+        }
+    }
+    fn new_common(tx: Sender<ServerSseMessage>) -> Self {
+        Self::new(tx, None)
+    }
+
+    async fn send(&mut self, message: ServerJsonRpcMessage) {
+        let index = self.cache.back().map_or(0, |m| {
+            m.event_id
+                .as_deref()
+                .unwrap_or_default()
+                .parse::<EventId>()
+                .expect("valid event id")
+                .index
+                + 1
+        });
+        let event_id = EventId {
+            http_request_id: self.http_request_id,
+            index,
+        };
+        let message = ServerSseMessage {
+            event_id: Some(event_id.to_string()),
+            message: Arc::new(message),
+        };
+        if self.cache.len() >= self.capacity {
+            self.cache.pop_front();
+            self.cache.push_back(message.clone());
+        } else {
+            self.cache.push_back(message.clone());
+        }
+        let _ = self.tx.send(message).await.inspect_err(|e| {
+            let event_id = &e.0.event_id;
+            tracing::trace!(?event_id, "trying to send message in a closed session")
+        });
+    }
+
+    async fn sync(&mut self, index: usize) -> Result<(), SessionError> {
+        let Some(front) = self.cache.front() else {
+            return Ok(());
+        };
+        let front_event_id = front
+            .event_id
+            .as_deref()
+            .unwrap_or_default()
+            .parse::<EventId>()?;
+        let sync_index = index.saturating_sub(front_event_id.index);
+        if sync_index > self.cache.len() {
+            // invalid index
+            return Err(SessionError::InvalidEventId);
+        }
+        for message in self.cache.iter().skip(sync_index) {
+            let send_result = self.tx.send(message.clone()).await;
+            if send_result.is_err() {
+                let event_id: EventId = message.event_id.as_deref().unwrap_or_default().parse()?;
+                return Err(SessionError::ChannelClosed(Some(event_id.index as u64)));
+            }
+        }
+        Ok(())
+    }
+}
+
+struct HttpRequestWise {
+    resources: HashSet<ResourceKey>,
+    tx: CachedTx,
+}
+
+type HttpRequestId = u64;
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+enum ResourceKey {
+    McpRequestId(RequestId),
+    ProgressToken(ProgressToken),
+}
+
+pub struct LocalSessionWorker {
+    id: SessionId,
+    next_http_request_id: HttpRequestId,
+    tx_router: HashMap<HttpRequestId, HttpRequestWise>,
+    resource_router: HashMap<ResourceKey, HttpRequestId>,
+    common: CachedTx,
+    event_rx: Receiver<SessionEvent>,
+    session_config: SessionConfig,
+}
+
+impl LocalSessionWorker {
+    pub fn id(&self) -> &SessionId {
+        &self.id
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum SessionError {
+    #[error("Invalid request id: {0}")]
+    DuplicatedRequestId(HttpRequestId),
+    #[error("Channel closed: {0:?}")]
+    ChannelClosed(Option<HttpRequestId>),
+    #[error("Cannot parse event id: {0}")]
+    EventIdParseError(#[from] EventIdParseError),
+    #[error("Session service terminated")]
+    SessionServiceTerminated,
+    #[error("Invalid event id")]
+    InvalidEventId,
+    #[error("Transport closed")]
+    TransportClosed,
+    #[error("IO error: {0}")]
+    Io(#[from] std::io::Error),
+    #[error("Tokio join error {0}")]
+    TokioJoinError(#[from] tokio::task::JoinError),
+}
+
+impl From<SessionError> for std::io::Error {
+    fn from(value: SessionError) -> Self {
+        match value {
+            SessionError::Io(io) => io,
+            _ => std::io::Error::new(std::io::ErrorKind::Other, format!("Session error: {value}")),
+        }
+    }
+}
+
+enum OutboundChannel {
+    RequestWise { id: HttpRequestId, close: bool },
+    Common,
+}
+
+pub struct StreamableHttpMessageReceiver {
+    pub http_request_id: Option<HttpRequestId>,
+    pub inner: Receiver<ServerSseMessage>,
+}
+
+impl LocalSessionWorker {
+    fn unregister_resource(&mut self, resource: &ResourceKey) {
+        if let Some(http_request_id) = self.resource_router.remove(resource) {
+            tracing::trace!(?resource, http_request_id, "unregister resource");
+            if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
+                channel.resources.remove(resource);
+                if channel.resources.is_empty() {
+                    tracing::debug!(http_request_id, "close http request wise channel");
+                    self.tx_router.remove(&http_request_id);
+                }
+            }
+        }
+    }
+    fn register_resource(&mut self, resource: ResourceKey, http_request_id: HttpRequestId) {
+        tracing::trace!(?resource, http_request_id, "register resource");
+        if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
+            channel.resources.insert(resource.clone());
+            self.resource_router.insert(resource, http_request_id);
+        }
+    }
+    fn register_request(
+        &mut self,
+        request: &JsonRpcRequest<ClientRequest>,
+        http_request_id: HttpRequestId,
+    ) {
+        use crate::model::GetMeta;
+        self.register_resource(
+            ResourceKey::McpRequestId(request.id.clone()),
+            http_request_id,
+        );
+        if let Some(progress_token) = request.request.get_meta().get_progress_token() {
+            self.register_resource(
+                ResourceKey::ProgressToken(progress_token.clone()),
+                http_request_id,
+            );
+        }
+    }
+    fn catch_cancellation_notification(
+        &mut self,
+        notification: &JsonRpcNotification<ClientNotification>,
+    ) {
+        if let ClientNotification::CancelledNotification(n) = &notification.notification {
+            let request_id = n.params.request_id.clone();
+            let resource = ResourceKey::McpRequestId(request_id);
+            self.unregister_resource(&resource);
+        }
+    }
+    fn next_http_request_id(&mut self) -> HttpRequestId {
+        let id = self.next_http_request_id;
+        self.next_http_request_id = self.next_http_request_id.wrapping_add(1);
+        id
+    }
+    async fn establish_request_wise_channel(
+        &mut self,
+    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
+        let http_request_id = self.next_http_request_id();
+        let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
+        self.tx_router.insert(
+            http_request_id,
+            HttpRequestWise {
+                resources: Default::default(),
+                tx: CachedTx::new(tx, Some(http_request_id)),
+            },
+        );
+        tracing::debug!(http_request_id, "establish new request wise channel");
+        Ok(StreamableHttpMessageReceiver {
+            http_request_id: Some(http_request_id),
+            inner: rx,
+        })
+    }
+    fn resolve_outbound_channel(&self, message: &ServerJsonRpcMessage) -> OutboundChannel {
+        match &message {
+            ServerJsonRpcMessage::Request(_) => OutboundChannel::Common,
+            ServerJsonRpcMessage::Notification(JsonRpcNotification {
+                notification:
+                    ServerNotification::ProgressNotification(Notification {
+                        params: ProgressNotificationParam { progress_token, .. },
+                        ..
+                    }),
+                ..
+            }) => {
+                let id = self
+                    .resource_router
+                    .get(&ResourceKey::ProgressToken(progress_token.clone()));
+
+                if let Some(id) = id {
+                    OutboundChannel::RequestWise {
+                        id: *id,
+                        close: false,
+                    }
+                } else {
+                    OutboundChannel::Common
+                }
+            }
+            ServerJsonRpcMessage::Notification(JsonRpcNotification {
+                notification:
+                    ServerNotification::CancelledNotification(Notification {
+                        params: CancelledNotificationParam { request_id, .. },
+                        ..
+                    }),
+                ..
+            }) => {
+                if let Some(id) = self
+                    .resource_router
+                    .get(&ResourceKey::McpRequestId(request_id.clone()))
+                {
+                    OutboundChannel::RequestWise {
+                        id: *id,
+                        close: false,
+                    }
+                } else {
+                    OutboundChannel::Common
+                }
+            }
+            ServerJsonRpcMessage::Notification(_) => OutboundChannel::Common,
+            ServerJsonRpcMessage::Response(json_rpc_response) => {
+                if let Some(id) = self
+                    .resource_router
+                    .get(&ResourceKey::McpRequestId(json_rpc_response.id.clone()))
+                {
+                    OutboundChannel::RequestWise {
+                        id: *id,
+                        close: false,
+                    }
+                } else {
+                    OutboundChannel::Common
+                }
+            }
+            ServerJsonRpcMessage::Error(json_rpc_error) => {
+                if let Some(id) = self
+                    .resource_router
+                    .get(&ResourceKey::McpRequestId(json_rpc_error.id.clone()))
+                {
+                    OutboundChannel::RequestWise {
+                        id: *id,
+                        close: false,
+                    }
+                } else {
+                    OutboundChannel::Common
+                }
+            }
+            ServerJsonRpcMessage::BatchRequest(_) | ServerJsonRpcMessage::BatchResponse(_) => {
+                // the server side should never yield a batch request or response now
+                unreachable!("server side won't yield batch request or response")
+            }
+        }
+    }
+    async fn handle_server_message(
+        &mut self,
+        message: ServerJsonRpcMessage,
+    ) -> Result<(), SessionError> {
+        let outbound_channel = self.resolve_outbound_channel(&message);
+        match outbound_channel {
+            OutboundChannel::RequestWise { id, close } => {
+                if let Some(request_wise) = self.tx_router.get_mut(&id) {
+                    request_wise.tx.send(message).await;
+                    if close {
+                        self.tx_router.remove(&id);
+                    }
+                } else {
+                    return Err(SessionError::ChannelClosed(Some(id)));
+                }
+            }
+            OutboundChannel::Common => self.common.send(message).await,
+        }
+        Ok(())
+    }
+    async fn resume(
+        &mut self,
+        last_event_id: EventId,
+    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
+        match last_event_id.http_request_id {
+            Some(http_request_id) => {
+                let request_wise = self
+                    .tx_router
+                    .get_mut(&http_request_id)
+                    .ok_or(SessionError::ChannelClosed(Some(http_request_id)))?;
+                let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
+                let (tx, rx) = channel;
+                request_wise.tx.tx = tx;
+                let index = last_event_id.index;
+                // sync messages after index
+                request_wise.tx.sync(index).await?;
+                Ok(StreamableHttpMessageReceiver {
+                    http_request_id: Some(http_request_id),
+                    inner: rx,
+                })
+            }
+            None => {
+                let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
+                let (tx, rx) = channel;
+                self.common.tx = tx;
+                let index = last_event_id.index;
+                // sync messages after index
+                self.common.sync(index).await?;
+                Ok(StreamableHttpMessageReceiver {
+                    http_request_id: None,
+                    inner: rx,
+                })
+            }
+        }
+    }
+}
+
+enum SessionEvent {
+    ClientMessage {
+        message: ClientJsonRpcMessage,
+        http_request_id: Option<HttpRequestId>,
+    },
+    EstablishRequestWiseChannel {
+        responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
+    },
+    CloseRequestWiseChannel {
+        id: HttpRequestId,
+        responder: oneshot::Sender<Result<(), SessionError>>,
+    },
+    Resume {
+        last_event_id: EventId,
+        responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
+    },
+    InitializeRequest {
+        request: ClientJsonRpcMessage,
+        responder: oneshot::Sender<Result<ServerJsonRpcMessage, SessionError>>,
+    },
+    Close,
+}
+
+#[derive(Debug, Clone)]
+pub enum SessionQuitReason {
+    ServiceTerminated,
+    ClientTerminated,
+    ExpectInitializeRequest,
+    ExpectInitializeResponse,
+    Cancelled,
+}
+
+#[derive(Debug, Clone)]
+pub struct LocalSessionHandle {
+    id: SessionId,
+    // after all event_tx drop, inner task will be terminated
+    event_tx: Sender<SessionEvent>,
+}
+
+impl LocalSessionHandle {
+    /// Get the session id
+    pub fn id(&self) -> &SessionId {
+        &self.id
+    }
+
+    /// Close the session
+    pub async fn close(&self) -> Result<(), SessionError> {
+        self.event_tx
+            .send(SessionEvent::Close)
+            .await
+            .map_err(|_| SessionError::SessionServiceTerminated)?;
+        Ok(())
+    }
+
+    /// Send a message to the session
+    pub async fn push_message(
+        &self,
+        message: ClientJsonRpcMessage,
+        http_request_id: Option<HttpRequestId>,
+    ) -> Result<(), SessionError> {
+        self.event_tx
+            .send(SessionEvent::ClientMessage {
+                message,
+                http_request_id,
+            })
+            .await
+            .map_err(|_| SessionError::SessionServiceTerminated)?;
+        Ok(())
+    }
+
+    /// establish a channel for a http-request, the corresponded message from server will be
+    /// sent through this channel. The channel will be closed when the request is completed,
+    /// or you can close it manually by calling [`LocalSessionHandle::close_request_wise_channel`].
+    pub async fn establish_request_wise_channel(
+        &self,
+    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.event_tx
+            .send(SessionEvent::EstablishRequestWiseChannel { responder: tx })
+            .await
+            .map_err(|_| SessionError::SessionServiceTerminated)?;
+        rx.await
+            .map_err(|_| SessionError::SessionServiceTerminated)?
+    }
+
+    /// close the http-request wise channel.
+    pub async fn close_request_wise_channel(
+        &self,
+        request_id: HttpRequestId,
+    ) -> Result<(), SessionError> {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.event_tx
+            .send(SessionEvent::CloseRequestWiseChannel {
+                id: request_id,
+                responder: tx,
+            })
+            .await
+            .map_err(|_| SessionError::SessionServiceTerminated)?;
+        rx.await
+            .map_err(|_| SessionError::SessionServiceTerminated)?
+    }
+
+    /// Establish a common channel for general purpose messages.
+    pub async fn establish_common_channel(
+        &self,
+    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.event_tx
+            .send(SessionEvent::Resume {
+                last_event_id: EventId {
+                    http_request_id: None,
+                    index: 0,
+                },
+                responder: tx,
+            })
+            .await
+            .map_err(|_| SessionError::SessionServiceTerminated)?;
+        rx.await
+            .map_err(|_| SessionError::SessionServiceTerminated)?
+    }
+
+    /// Resume streaming response by the last event id. This is suitable for both request wise and common channel.
+    pub async fn resume(
+        &self,
+        last_event_id: EventId,
+    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.event_tx
+            .send(SessionEvent::Resume {
+                last_event_id,
+                responder: tx,
+            })
+            .await
+            .map_err(|_| SessionError::SessionServiceTerminated)?;
+        rx.await
+            .map_err(|_| SessionError::SessionServiceTerminated)?
+    }
+
+    /// Send an initialize request to the session. And wait for the initialized response.
+    ///
+    /// This is used to establish a session with the server.
+    pub async fn initialize(
+        &self,
+        request: ClientJsonRpcMessage,
+    ) -> Result<ServerJsonRpcMessage, SessionError> {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.event_tx
+            .send(SessionEvent::InitializeRequest {
+                request,
+                responder: tx,
+            })
+            .await
+            .map_err(|_| SessionError::SessionServiceTerminated)?;
+        rx.await
+            .map_err(|_| SessionError::SessionServiceTerminated)?
+    }
+}
+
+pub type SessionTransport = WorkerTransport<LocalSessionWorker>;
+
+impl Worker for LocalSessionWorker {
+    type Error = SessionError;
+    type Role = RoleServer;
+    fn err_closed() -> Self::Error {
+        SessionError::TransportClosed
+    }
+    fn err_join(e: tokio::task::JoinError) -> Self::Error {
+        SessionError::TokioJoinError(e)
+    }
+    fn config(&self) -> crate::transport::worker::WorkerConfig {
+        crate::transport::worker::WorkerConfig {
+            name: Some(format!("streamable-http-session-{}", self.id)),
+            channel_buffer_capacity: self.session_config.channel_capacity,
+        }
+    }
+    #[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))]
+    async fn run(mut self, mut context: WorkerContext<Self>) -> Result<(), WorkerQuitReason> {
+        enum InnerEvent {
+            FromHttpService(SessionEvent),
+            FromHandler(WorkerSendRequest<LocalSessionWorker>),
+        }
+        // waiting for initialize request
+        let evt = self.event_rx.recv().await.ok_or_else(|| {
+            WorkerQuitReason::fatal("transport terminated", "get initialize request")
+        })?;
+        let SessionEvent::InitializeRequest { request, responder } = evt else {
+            return Err(WorkerQuitReason::fatal(
+                "unexpected message",
+                "get initialize request",
+            ));
+        };
+        context.send_to_handler(request).await?;
+        let send_initialize_response = context.recv_from_handler().await?;
+        responder
+            .send(Ok(send_initialize_response.message))
+            .map_err(|_| {
+                WorkerQuitReason::fatal(
+                    "failed to send initialize response to http service",
+                    "send initialize response",
+                )
+            })?;
+        send_initialize_response
+            .responder
+            .send(Ok(()))
+            .map_err(|_| WorkerQuitReason::HandlerTerminated)?;
+        let ct = context.cancellation_token.clone();
+        let keep_alive = self.session_config.keep_alive.unwrap_or(Duration::MAX);
+        loop {
+            let keep_alive_timeout = tokio::time::sleep(keep_alive);
+            let event = tokio::select! {
+                event = self.event_rx.recv() => {
+                    if let Some(event) = event {
+                        InnerEvent::FromHttpService(event)
+                    } else {
+                        return Err(WorkerQuitReason::fatal("session dropped", "waiting next session event"))
+                    }
+                },
+                from_handler = context.recv_from_handler() => {
+                    InnerEvent::FromHandler(from_handler?)
+                }
+                _ = ct.cancelled() => {
+                    return Err(WorkerQuitReason::Cancelled)
+                }
+                _ = keep_alive_timeout => {
+                    return Err(WorkerQuitReason::fatal("keep live timeout", "poll next session event"))
+                }
+            };
+            match event {
+                InnerEvent::FromHandler(WorkerSendRequest { message, responder }) => {
+                    // catch response
+                    let to_unregister = match &message {
+                        crate::model::JsonRpcMessage::Response(json_rpc_response) => {
+                            let request_id = json_rpc_response.id.clone();
+                            Some(ResourceKey::McpRequestId(request_id))
+                        }
+                        crate::model::JsonRpcMessage::Error(json_rpc_error) => {
+                            let request_id = json_rpc_error.id.clone();
+                            Some(ResourceKey::McpRequestId(request_id))
+                        }
+                        _ => {
+                            None
+                            // no need to unregister resource
+                        }
+                    };
+                    let handle_result = self.handle_server_message(message).await;
+                    let _ = responder.send(handle_result).inspect_err(|error| {
+                        tracing::warn!(?error, "failed to send message to http service handler");
+                    });
+                    if let Some(to_unregister) = to_unregister {
+                        self.unregister_resource(&to_unregister);
+                    }
+                }
+                InnerEvent::FromHttpService(SessionEvent::ClientMessage {
+                    message: json_rpc_message,
+                    http_request_id,
+                }) => {
+                    match &json_rpc_message {
+                        crate::model::JsonRpcMessage::Request(request) => {
+                            if let Some(http_request_id) = http_request_id {
+                                self.register_request(request, http_request_id)
+                            }
+                        }
+                        crate::model::JsonRpcMessage::Notification(notification) => {
+                            self.catch_cancellation_notification(notification)
+                        }
+                        crate::model::JsonRpcMessage::BatchRequest(items) => {
+                            for r in items {
+                                match r {
+                                    crate::model::JsonRpcBatchRequestItem::Request(request) => {
+                                        if let Some(http_request_id) = http_request_id {
+                                            self.register_request(request, http_request_id)
+                                        }
+                                    }
+                                    crate::model::JsonRpcBatchRequestItem::Notification(
+                                        notification,
+                                    ) => self.catch_cancellation_notification(notification),
+                                }
+                            }
+                        }
+                        _ => {}
+                    }
+                    context.send_to_handler(json_rpc_message).await?;
+                }
+                InnerEvent::FromHttpService(SessionEvent::EstablishRequestWiseChannel {
+                    responder,
+                }) => {
+                    let handle_result = self.establish_request_wise_channel().await;
+                    let _ = responder.send(handle_result);
+                }
+                InnerEvent::FromHttpService(SessionEvent::CloseRequestWiseChannel {
+                    id,
+                    responder,
+                }) => {
+                    let _handle_result = self.tx_router.remove(&id);
+                    let _ = responder.send(Ok(()));
+                }
+                InnerEvent::FromHttpService(SessionEvent::Resume {
+                    last_event_id,
+                    responder,
+                }) => {
+                    let handle_result = self.resume(last_event_id).await;
+                    let _ = responder.send(handle_result);
+                }
+                InnerEvent::FromHttpService(SessionEvent::Close) => {
+                    return Err(WorkerQuitReason::TransportClosed);
+                }
+                _ => {
+                    // ignore
+                }
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct SessionConfig {
+    /// the capacity of the channel for the session. Default is 16.
+    pub channel_capacity: usize,
+    /// if set, the session will be closed after this duration of inactivity.
+    pub keep_alive: Option<Duration>,
+}
+
+impl SessionConfig {
+    pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
+}
+
+impl Default for SessionConfig {
+    fn default() -> Self {
+        Self {
+            channel_capacity: Self::DEFAULT_CHANNEL_CAPACITY,
+            keep_alive: None,
+        }
+    }
+}
+
+/// Create a new session with the given id and configuration.
+///
+/// This function will return a pair of [`LocalSessionHandle`] and [`LocalSessionWorker`].
+///
+/// You can run the [`LocalSessionWorker`] as a transport for mcp server. And use the [`LocalSessionHandle`] operate the session.
+pub fn create_local_session(
+    id: impl Into<SessionId>,
+    config: SessionConfig,
+) -> (LocalSessionHandle, LocalSessionWorker) {
+    let id = id.into();
+    let (event_tx, event_rx) = tokio::sync::mpsc::channel(config.channel_capacity);
+    let (common_tx, _) = tokio::sync::mpsc::channel(config.channel_capacity);
+    let common = CachedTx::new_common(common_tx);
+    tracing::info!(session_id = ?id, "create new session");
+    let handle = LocalSessionHandle {
+        event_tx,
+        id: id.clone(),
+    };
+    let session_worker = LocalSessionWorker {
+        next_http_request_id: 0,
+        id,
+        tx_router: HashMap::new(),
+        resource_router: HashMap::new(),
+        common,
+        event_rx,
+        session_config: config.clone(),
+    };
+    (handle, session_worker)
+}

crates/rmcp/src/transport/streamable_http_server/session/never.rs

@@ -0,0 +1,107 @@
+use futures::Stream;
+use thiserror::Error;
+
+use super::{ServerSseMessage, SessionId, SessionManager};
+use crate::{
+    RoleServer,
+    model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
+    transport::Transport,
+};
+
+#[derive(Debug, Clone, Error)]
+#[error("Session management is not supported")]
+pub struct ErrorSessionManagementNotSupported;
+#[derive(Debug, Clone, Default)]
+pub struct NeverSessionManager {}
+pub enum NeverTransport {}
+impl Transport<RoleServer> for NeverTransport {
+    type Error = ErrorSessionManagementNotSupported;
+
+    fn send(
+        &mut self,
+        _item: ServerJsonRpcMessage,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
+        futures::future::ready(Err(ErrorSessionManagementNotSupported))
+    }
+
+    fn receive(&mut self) -> impl Future<Output = Option<ClientJsonRpcMessage>> {
+        futures::future::ready(None)
+    }
+
+    async fn close(&mut self) -> Result<(), Self::Error> {
+        Err(ErrorSessionManagementNotSupported)
+    }
+}
+
+impl SessionManager for NeverSessionManager {
+    type Error = ErrorSessionManagementNotSupported;
+    type Transport = NeverTransport;
+
+    fn create_session(
+        &self,
+    ) -> impl Future<Output = Result<(SessionId, Self::Transport), Self::Error>> + Send {
+        futures::future::ready(Err(ErrorSessionManagementNotSupported))
+    }
+
+    fn initialize_session(
+        &self,
+        _id: &SessionId,
+        _message: ClientJsonRpcMessage,
+    ) -> impl Future<Output = Result<ServerJsonRpcMessage, Self::Error>> + Send {
+        futures::future::ready(Err(ErrorSessionManagementNotSupported))
+    }
+
+    fn has_session(
+        &self,
+        _id: &SessionId,
+    ) -> impl Future<Output = Result<bool, Self::Error>> + Send {
+        futures::future::ready(Err(ErrorSessionManagementNotSupported))
+    }
+
+    fn close_session(
+        &self,
+        _id: &SessionId,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        futures::future::ready(Err(ErrorSessionManagementNotSupported))
+    }
+
+    fn create_stream(
+        &self,
+        _id: &SessionId,
+        _message: ClientJsonRpcMessage,
+    ) -> impl Future<
+        Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
+    > + Send {
+        futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
+            ErrorSessionManagementNotSupported,
+        ))
+    }
+    fn create_standalone_stream(
+        &self,
+        _id: &SessionId,
+    ) -> impl Future<
+        Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
+    > + Send {
+        futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
+            ErrorSessionManagementNotSupported,
+        ))
+    }
+    fn resume(
+        &self,
+        _id: &SessionId,
+        _last_event_id: String,
+    ) -> impl Future<
+        Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
+    > + Send {
+        futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
+            ErrorSessionManagementNotSupported,
+        ))
+    }
+    fn accept_message(
+        &self,
+        _id: &SessionId,
+        _message: ClientJsonRpcMessage,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        futures::future::ready(Err(ErrorSessionManagementNotSupported))
+    }
+}

crates/rmcp/src/transport/streamable_http_server/tower.rs

@@ -0,0 +1,427 @@
+use std::{convert::Infallible, fmt::Display, sync::Arc, time::Duration};
+
+use bytes::Bytes;
+use futures::{StreamExt, future::BoxFuture};
+use http::{Method, Request, Response, header::ALLOW};
+use http_body::Body;
+use http_body_util::{BodyExt, Full, combinators::UnsyncBoxBody};
+use tokio_stream::wrappers::ReceiverStream;
+
+use super::session::SessionManager;
+use crate::{
+    RoleServer,
+    model::ClientJsonRpcMessage,
+    serve_server,
+    service::serve_directly,
+    transport::{
+        OneshotTransport, TransportAdapterIdentity,
+        common::{
+            http_header::{
+                EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
+            },
+            server_side_http::{
+                BoxResponse, ServerSseMessage, accepted_response, expect_json,
+                internal_error_response, sse_stream_response,
+            },
+        },
+    },
+};
+
+#[derive(Debug, Clone)]
+pub struct StreamableHttpServerConfig {
+    /// The ping message duration for SSE connections.
+    pub sse_keep_alive: Option<Duration>,
+    /// If true, the server will create a session for each request and keep it alive.
+    pub stateful_mode: bool,
+}
+
+impl Default for StreamableHttpServerConfig {
+    fn default() -> Self {
+        Self {
+            sse_keep_alive: Some(Duration::from_secs(15)),
+            stateful_mode: true,
+        }
+    }
+}
+
+pub struct StreamableHttpService<S, M = super::session::local::LocalSessionManager> {
+    pub config: StreamableHttpServerConfig,
+    session_manager: Arc<M>,
+    service_factory: Arc<dyn Fn() -> S + Send + Sync>,
+}
+
+impl<S, M> Clone for StreamableHttpService<S, M> {
+    fn clone(&self) -> Self {
+        Self {
+            config: self.config.clone(),
+            session_manager: self.session_manager.clone(),
+            service_factory: self.service_factory.clone(),
+        }
+    }
+}
+
+impl<RequestBody, S, M> tower_service::Service<Request<RequestBody>> for StreamableHttpService<S, M>
+where
+    RequestBody: Body + Send + 'static,
+    S: crate::Service<RoleServer>,
+    M: SessionManager,
+    RequestBody::Error: Display,
+    RequestBody::Data: Send + 'static,
+{
+    type Response = BoxResponse;
+    type Error = Infallible;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+    fn call(&mut self, req: http::Request<RequestBody>) -> Self::Future {
+        let service = self.clone();
+        Box::pin(async move {
+            let response = service.handle(req).await;
+            Ok(response)
+        })
+    }
+    fn poll_ready(
+        &mut self,
+        _cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), Self::Error>> {
+        std::task::Poll::Ready(Ok(()))
+    }
+}
+
+impl<S, M> StreamableHttpService<S, M>
+where
+    S: crate::Service<RoleServer> + Send + 'static,
+    M: SessionManager,
+{
+    pub fn new(
+        service_factory: impl Fn() -> S + Send + Sync + 'static,
+        session_manager: Arc<M>,
+        config: StreamableHttpServerConfig,
+    ) -> Self {
+        Self {
+            config,
+            session_manager,
+            service_factory: Arc::new(service_factory),
+        }
+    }
+    fn get_service(&self) -> S {
+        (self.service_factory)()
+    }
+    pub async fn handle<B>(&self, request: Request<B>) -> Response<UnsyncBoxBody<Bytes, Infallible>>
+    where
+        B: Body + Send + 'static,
+        B::Error: Display,
+    {
+        let method = request.method().clone();
+        let result = match method {
+            Method::GET => self.handle_get(request).await,
+            Method::POST => self.handle_post(request).await,
+            Method::DELETE => self.handle_delete(request).await,
+            _ => {
+                // Handle other methods or return an error
+                let response = Response::builder()
+                    .status(http::StatusCode::METHOD_NOT_ALLOWED)
+                    .header(ALLOW, "GET, POST, DELETE")
+                    .body(Full::new(Bytes::from("Method Not Allowed")).boxed_unsync())
+                    .expect("valid response");
+                return response;
+            }
+        };
+        match result {
+            Ok(response) => response,
+            Err(response) => response,
+        }
+    }
+    async fn handle_get<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
+    where
+        B: Body + Send + 'static,
+        B::Error: Display,
+    {
+        // check accept header
+        if !request
+            .headers()
+            .get(http::header::ACCEPT)
+            .and_then(|header| header.to_str().ok())
+            .is_some_and(|header| header.contains(EVENT_STREAM_MIME_TYPE))
+        {
+            return Ok(Response::builder()
+                .status(http::StatusCode::NOT_ACCEPTABLE)
+                .body(
+                    Full::new(Bytes::from(
+                        "Not Acceptable: Client must accept text/event-stream",
+                    ))
+                    .boxed_unsync(),
+                )
+                .expect("valid response"));
+        }
+        // check session id
+        let session_id = request
+            .headers()
+            .get(HEADER_SESSION_ID)
+            .and_then(|v| v.to_str().ok())
+            .map(|s| s.to_owned().into());
+        let Some(session_id) = session_id else {
+            // unauthorized
+            return Ok(Response::builder()
+                .status(http::StatusCode::UNAUTHORIZED)
+                .body(Full::new(Bytes::from("Unauthorized: Session ID is required")).boxed_unsync())
+                .expect("valid response"));
+        };
+        // check if session exists
+        let has_session = self
+            .session_manager
+            .has_session(&session_id)
+            .await
+            .map_err(internal_error_response("check session"))?;
+        if !has_session {
+            // unauthorized
+            return Ok(Response::builder()
+                .status(http::StatusCode::UNAUTHORIZED)
+                .body(Full::new(Bytes::from("Unauthorized: Session not found")).boxed_unsync())
+                .expect("valid response"));
+        }
+        // check if last event id is provided
+        let last_event_id = request
+            .headers()
+            .get(HEADER_LAST_EVENT_ID)
+            .and_then(|v| v.to_str().ok())
+            .map(|s| s.to_owned());
+        if let Some(last_event_id) = last_event_id {
+            // check if session has this event id
+            let stream = self
+                .session_manager
+                .resume(&session_id, last_event_id)
+                .await
+                .map_err(internal_error_response("resume session"))?;
+            Ok(sse_stream_response(stream, self.config.sse_keep_alive))
+        } else {
+            // create standalone stream
+            let stream = self
+                .session_manager
+                .create_standalone_stream(&session_id)
+                .await
+                .map_err(internal_error_response("create standalone stream"))?;
+            Ok(sse_stream_response(stream, self.config.sse_keep_alive))
+        }
+    }
+
+    async fn handle_post<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
+    where
+        B: Body + Send + 'static,
+        B::Error: Display,
+    {
+        // check accept header
+        if !request
+            .headers()
+            .get(http::header::ACCEPT)
+            .and_then(|header| header.to_str().ok())
+            .is_some_and(|header| {
+                header.contains(JSON_MIME_TYPE) && header.contains(EVENT_STREAM_MIME_TYPE)
+            })
+        {
+            return Ok(Response::builder()
+                .status(http::StatusCode::NOT_ACCEPTABLE)
+                .body(Full::new(Bytes::from("Not Acceptable: Client must accept both application/json and text/event-stream")).boxed_unsync())
+                .expect("valid response"));
+        }
+
+        // check content type
+        if !request
+            .headers()
+            .get(http::header::CONTENT_TYPE)
+            .and_then(|header| header.to_str().ok())
+            .is_some_and(|header| header.starts_with(JSON_MIME_TYPE))
+        {
+            return Ok(Response::builder()
+                .status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
+                .body(
+                    Full::new(Bytes::from(
+                        "Unsupported Media Type: Content-Type must be application/json",
+                    ))
+                    .boxed_unsync(),
+                )
+                .expect("valid response"));
+        }
+
+        // json deserialize request body
+        let (part, body) = request.into_parts();
+        let message = match expect_json(body).await {
+            Ok(message) => message,
+            Err(response) => return Ok(response),
+        };
+
+        if self.config.stateful_mode {
+            // do we have a session id?
+            let session_id = part
+                .headers
+                .get(HEADER_SESSION_ID)
+                .and_then(|v| v.to_str().ok());
+            if let Some(session_id) = session_id {
+                let session_id = session_id.to_owned().into();
+                let has_session = self
+                    .session_manager
+                    .has_session(&session_id)
+                    .await
+                    .map_err(internal_error_response("check session"))?;
+                if !has_session {
+                    // unauthorized
+                    return Ok(Response::builder()
+                        .status(http::StatusCode::UNAUTHORIZED)
+                        .body(
+                            Full::new(Bytes::from("Unauthorized: Session not found"))
+                                .boxed_unsync(),
+                        )
+                        .expect("valid response"));
+                }
+                match message {
+                    ClientJsonRpcMessage::Request(_) => {
+                        let stream = self
+                            .session_manager
+                            .create_stream(&session_id, message)
+                            .await
+                            .map_err(internal_error_response("get session"))?;
+                        Ok(sse_stream_response(stream, self.config.sse_keep_alive))
+                    }
+                    ClientJsonRpcMessage::Notification(_)
+                    | ClientJsonRpcMessage::Response(_)
+                    | ClientJsonRpcMessage::Error(_) => {
+                        // handle notification
+                        self.session_manager
+                            .accept_message(&session_id, message)
+                            .await
+                            .map_err(internal_error_response("accept message"))?;
+                        Ok(accepted_response())
+                    }
+                    _ => Ok(Response::builder()
+                        .status(http::StatusCode::NOT_IMPLEMENTED)
+                        .body(
+                            Full::new(Bytes::from("Batch requests are not supported yet"))
+                                .boxed_unsync(),
+                        )
+                        .expect("valid response")),
+                }
+            } else {
+                let (session_id, transport) = self
+                    .session_manager
+                    .create_session()
+                    .await
+                    .map_err(internal_error_response("create session"))?;
+                let service = self.get_service();
+                // spawn a task to serve the session
+                tokio::spawn({
+                    let session_manager = self.session_manager.clone();
+                    let session_id = session_id.clone();
+                    async move {
+                        let service = serve_server::<S, M::Transport, _, TransportAdapterIdentity>(
+                            service, transport,
+                        )
+                        .await;
+                        match service {
+                            Ok(service) => {
+                                // on service created
+                                let _ = service.waiting().await;
+                            }
+                            Err(e) => {
+                                tracing::error!("Failed to create service: {e}");
+                            }
+                        }
+                        let _ = session_manager
+                            .close_session(&session_id)
+                            .await
+                            .inspect_err(|e| {
+                                tracing::error!("Failed to close session {session_id}: {e}");
+                            });
+                    }
+                });
+                // get initialize response
+                let response = self
+                    .session_manager
+                    .initialize_session(&session_id, message)
+                    .await
+                    .map_err(internal_error_response("create stream"))?;
+                let mut response = sse_stream_response(
+                    futures::stream::once({
+                        async move {
+                            ServerSseMessage {
+                                event_id: None,
+                                message: response.into(),
+                            }
+                        }
+                    }),
+                    self.config.sse_keep_alive,
+                );
+
+                response.headers_mut().insert(
+                    HEADER_SESSION_ID,
+                    session_id
+                        .parse()
+                        .map_err(internal_error_response("create session id header"))?,
+                );
+                Ok(response)
+            }
+        } else {
+            let service = self.get_service();
+            match message {
+                ClientJsonRpcMessage::Request(request) => {
+                    let (transport, receiver) =
+                        OneshotTransport::<RoleServer>::new(ClientJsonRpcMessage::Request(request));
+                    let service = serve_directly(service, transport, None);
+                    tokio::spawn(async move {
+                        // on service created
+                        let _ = service.waiting().await;
+                    });
+                    Ok(sse_stream_response(
+                        ReceiverStream::new(receiver).map(|message| {
+                            tracing::info!(?message);
+                            ServerSseMessage {
+                                event_id: None,
+                                message: message.into(),
+                            }
+                        }),
+                        self.config.sse_keep_alive,
+                    ))
+                }
+                ClientJsonRpcMessage::Notification(notification) => {
+                    service
+                        .handle_notification(notification.notification)
+                        .await
+                        .map_err(internal_error_response("handle notification"))?;
+                    Ok(accepted_response())
+                }
+                ClientJsonRpcMessage::Response(_json_rpc_response) => Ok(accepted_response()),
+                ClientJsonRpcMessage::Error(_json_rpc_error) => Ok(accepted_response()),
+                _ => Ok(Response::builder()
+                    .status(http::StatusCode::NOT_IMPLEMENTED)
+                    .body(
+                        Full::new(Bytes::from("Batch requests are not supported yet"))
+                            .boxed_unsync(),
+                    )
+                    .expect("valid response")),
+            }
+        }
+    }
+
+    async fn handle_delete<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
+    where
+        B: Body + Send + 'static,
+        B::Error: Display,
+    {
+        // check session id
+        let session_id = request
+            .headers()
+            .get(HEADER_SESSION_ID)
+            .and_then(|v| v.to_str().ok())
+            .map(|s| s.to_owned().into());
+        let Some(session_id) = session_id else {
+            // unauthorized
+            return Ok(Response::builder()
+                .status(http::StatusCode::UNAUTHORIZED)
+                .body(Full::new(Bytes::from("Unauthorized: Session ID is required")).boxed_unsync())
+                .expect("valid response"));
+        };
+        // close session
+        self.session_manager
+            .close_session(&session_id)
+            .await
+            .map_err(internal_error_response("close session"))?;
+        Ok(accepted_response())
+    }
+}

crates/rmcp/tests/test_with_js.rs

@@ -2,10 +2,14 @@ use rmcp::{
     ServiceExt,
     service::QuitReason,
     transport::{
-        ConfigureCommandExt, SseServer, StreamableHttpClientTransport, TokioChildProcess,
+        ConfigureCommandExt, SseServer, StreamableHttpClientTransport, StreamableHttpServerConfig,
-        streamable_http_server::axum::StreamableHttpServer,
+        TokioChildProcess,
+        streamable_http_server::{
+            session::local::LocalSessionManager, tower::StreamableHttpService,
+        },
     },
 };
+use tokio_util::sync::CancellationToken;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 mod common;
 use common::calculator::Calculator;
@@ -90,10 +94,26 @@ async fn test_with_js_streamable_http_client() -> anyhow::Result<()> {
         .wait()
         .await?;
 
-    let ct = StreamableHttpServer::serve(STREAMABLE_HTTP_BIND_ADDRESS.parse()?)
+    let service: StreamableHttpService<Calculator, LocalSessionManager> =
-        .await?
+        StreamableHttpService::new(
-        .with_service(Calculator::default);
+            Calculator::default,
-
+            Default::default(),
+            StreamableHttpServerConfig {
+                stateful_mode: true,
+                sse_keep_alive: None,
+            },
+        );
+    let router = axum::Router::new().nest_service("/mcp", service);
+    let tcp_listener = tokio::net::TcpListener::bind(STREAMABLE_HTTP_BIND_ADDRESS).await?;
+    let ct = CancellationToken::new();
+    let handle = tokio::spawn({
+        let ct = ct.clone();
+        async move {
+            let _ = axum::serve(tcp_listener, router)
+                .with_graceful_shutdown(async move { ct.cancelled_owned().await })
+                .await;
+        }
+    });
     let exit_status = tokio::process::Command::new("node")
         .arg("tests/test_with_js/streamable_client.js")
         .spawn()?
@@ -101,6 +121,7 @@ async fn test_with_js_streamable_http_client() -> anyhow::Result<()> {
         .await?;
     assert!(exit_status.success());
     ct.cancel();
+    handle.await?;
     Ok(())
 }
 

crates/rmcp/tests/test_with_js/streamable_client.js

@@ -1,7 +1,7 @@
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 
-const transport = new StreamableHTTPClientTransport(new URL(`http://127.0.0.1:8001/`));
+const transport = new StreamableHTTPClientTransport(new URL(`http://127.0.0.1:8001/mcp/`));
 
 const client = new Client(
   {

examples/servers/Cargo.toml

@@ -7,8 +7,20 @@ edition = "2024"
 publish = false
 
 [dependencies]
-rmcp= { path = "../../crates/rmcp", features = ["server", "transport-sse-server", "transport-io", "transport-streamable-http-server", "auth"] }
+rmcp = { path = "../../crates/rmcp", features = [
-tokio = { version = "1", features = ["macros", "rt", "rt-multi-thread", "io-std", "signal"] }
+    "server",
+    "transport-sse-server",
+    "transport-io",
+    "transport-streamable-http-server",
+    "auth",
+] }
+tokio = { version = "1", features = [
+    "macros",
+    "rt",
+    "rt-multi-thread",
+    "io-std",
+    "signal",
+] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 anyhow = "1.0"
@@ -26,9 +38,10 @@ reqwest = { version = "0.12", features = ["json"] }
 chrono = "0.4"
 uuid = { version = "1.6", features = ["v4", "serde"] }
 serde_urlencoded = "0.7"
-askama = { version = "0.14"}
+askama = { version = "0.14" }
 tower-http = { version = "0.6", features = ["cors"] }
-
+hyper = { version = "1" }
+hyper-util = { version = "0", features = ["server"] }
 
 [dev-dependencies]
 tokio-stream = { version = "0.1" }
@@ -60,4 +73,8 @@ path = "src/complex_auth_sse.rs"
 
 [[example]]
 name = "servers_simple_auth_sse"
-path = "src/simple_auth_sse.rs"
+path = "src/simple_auth_sse.rs"
+
+[[example]]
+name = "counter_hyper_streamable_http"
+path = "src/counter_hyper_streamable_http.rs"

examples/servers/README.md

@@ -38,12 +38,19 @@ A minimal server example using stdio transport.
 
 ### Counter Streamable HTTP Server (`counter_streamhttp.rs`)
 
-A server using streamable HTTP transport for MCP communication.
+A server using streamable HTTP transport for MCP communication, with axum.
 
 - Runs on HTTP with streaming capabilities
 - Provides counter tools via HTTP streaming
 - Demonstrates streamable HTTP transport configuration
 
+### Counter Streamable HTTP Server with Hyper (`counter_hyper_streamable_http.rs`)
+
+A server using streamable HTTP transport for MCP communication, with hyper.
+- Runs on HTTP with streaming capabilities
+- Provides counter tools via HTTP streaming
+- Demonstrates streamable HTTP transport configuration
+
 ### Complex OAuth SSE Server (`complex_auth_sse.rs`)
 
 A comprehensive example demonstrating OAuth 2.0 integration with MCP servers.