refactor: Transport trait and worker transport, and streamable http client with those new features. (#167)

* refactor: a transport trait and streamable http client

* test: test with js streamable http server

* refactor: separate sse stream connection

* fix(test): streamable http client test with js

* fix(test): wait longer time for js server startup

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

@@ -1,4 +1,10 @@
-use std::{collections::HashMap, io, net::SocketAddr, sync::Arc, time::Duration};
+use std::{
+    collections::HashMap,
+    io,
+    net::{Ipv6Addr, SocketAddr, SocketAddrV6},
+    sync::Arc,
+    time::Duration,
+};
 
 use axum::{
     Json, Router,
@@ -15,30 +21,28 @@ use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::sync::CancellationToken;
 use tracing::Instrument;
 
-use super::session::{
-    EventId, HEADER_LAST_EVENT_ID, Session, SessionTransport, StreamableHttpMessageReceiver,
-};
+use super::session::{EventId, SessionHandle, SessionWorker, StreamableHttpMessageReceiver};
 use crate::{
     RoleServer, Service,
     model::ClientJsonRpcMessage,
-    transport::{
-        common::axum::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
-        streamable_http_server::session::HEADER_SESSION_ID,
+    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, Session>>>;
+type SessionManager = Arc<tokio::sync::RwLock<HashMap<SessionId, SessionHandle>>>;
 
 #[derive(Clone)]
 struct App {
     session_manager: SessionManager,
-    transport_tx: tokio::sync::mpsc::UnboundedSender<SessionTransport>,
+    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<SessionTransport>) {
+    ) -> (Self, tokio::sync::mpsc::UnboundedReceiver<SessionWorker>) {
         let (transport_tx, transport_rx) = tokio::sync::mpsc::unbounded_channel();
         (
             Self {
@@ -82,7 +86,7 @@ async fn post_handler(
             let session = sm
                 .get(session_id)
                 .ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?;
-            session.handle().clone()
+            session.clone()
         };
         // inject request part
         message.insert_extension(parts);
@@ -138,7 +142,7 @@ async fn post_handler(
             return Err((StatusCode::GONE, "session terminated").into_response());
         };
 
-        let response = session.handle().initialize(message).await.map_err(|e| {
+        let response = session.initialize(message).await.map_err(|e| {
             (
                 StatusCode::INTERNAL_SERVER_ERROR,
                 format!("fail to initialize: {e}"),
@@ -169,21 +173,24 @@ async fn get_handler(
         let last_event_id = header_map
             .get(HEADER_LAST_EVENT_ID)
             .and_then(|v| v.to_str().ok());
-        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 sm = app.session_manager.read().await;
-                let session = sm.get(session_id).ok_or_else(|| {
+        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 handle = session.handle();
-                let receiver = handle.resume(last_event_id).await.map_err(|e| {
+                let receiver = session.resume(last_event_id).await.map_err(|e| {
                     (
                         StatusCode::INTERNAL_SERVER_ERROR,
                         format!("resume error {e}"),
@@ -194,16 +201,7 @@ async fn get_handler(
                 Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(app.sse_ping_interval)))
             }
             None => {
-                let sm = app.session_manager.read().await;
-                let session = sm.get(session_id).ok_or_else(|| {
-                    (
-                        StatusCode::NOT_FOUND,
-                        format!("session {session_id} not found"),
-                    )
-                        .into_response()
-                })?;
-                let handle = session.handle();
-                let receiver = handle.establish_common_channel().await.map_err(|e| {
+                let receiver = session.establish_common_channel().await.map_err(|e| {
                     (
                         StatusCode::INTERNAL_SERVER_ERROR,
                         format!("establish common channel error {e}"),
@@ -227,18 +225,19 @@ async fn delete_handler(
         let session_id = session_id
             .to_str()
             .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()).into_response())?;
-        let mut sm = app.session_manager.write().await;
-        let session = sm
-            .remove(session_id)
-            .ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?;
-        let cancel_result = session.cancel().await.map_err(|e| {
+        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::info!(session_id, ?cancel_result, "session deleted");
+        tracing::debug!(session_id, "session deleted");
         Ok(StatusCode::ACCEPTED)
     } else {
         Err((StatusCode::BAD_REQUEST, "missing session id").into_response())
@@ -252,10 +251,20 @@ pub struct StreamableHttpServerConfig {
     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<SessionTransport>,
+    transport_rx: tokio::sync::mpsc::UnboundedReceiver<SessionWorker>,
     pub config: StreamableHttpServerConfig,
 }
 
@@ -263,9 +272,7 @@ impl StreamableHttpServer {
     pub async fn serve(bind: SocketAddr) -> io::Result<Self> {
         Self::serve_with_config(StreamableHttpServerConfig {
             bind,
-            path: "/".to_string(),
-            ct: CancellationToken::new(),
-            sse_keep_alive: None,
+            ..Default::default()
         })
         .await
     }
@@ -333,7 +340,7 @@ impl StreamableHttpServer {
         self.config.ct.cancel();
     }
 
-    pub async fn next_transport(&mut self) -> Option<SessionTransport> {
+    pub async fn next_transport(&mut self) -> Option<SessionWorker> {
         self.transport_rx.recv().await
     }
 }

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

@@ -5,14 +5,11 @@ use std::{
     sync::Arc,
 };
 
-use futures::{Sink, SinkExt, Stream};
 use thiserror::Error;
 use tokio::sync::{
     mpsc::{Receiver, Sender},
     oneshot,
 };
-use tokio_stream::wrappers::ReceiverStream;
-use tokio_util::sync::{CancellationToken, DropGuard, PollSender};
 use tracing::instrument;
 
 use crate::{
@@ -22,18 +19,19 @@ use crate::{
         JsonRpcNotification, JsonRpcRequest, Notification, ProgressNotificationParam,
         ProgressToken, RequestId, ServerJsonRpcMessage, ServerNotification,
     },
-    transport::IntoTransport,
+    transport::{
+        WorkerTransport,
+        worker::{Worker, WorkerContext, WorkerQuitReason, WorkerSendRequest},
+    },
 };
 
-pub const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
-pub const HEADER_LAST_EVENT_ID: &str = "Last-Event-Id";
 #[derive(Debug, Clone)]
 pub struct ServerSessionMessage {
     pub event_id: EventId,
     pub message: Arc<ServerJsonRpcMessage>,
 }
 
-/// <index>-<request_id>
+/// `<index>/request_id>`
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct EventId {
     http_request_id: Option<HttpRequestId>,
@@ -159,17 +157,22 @@ enum ResourceKey {
     ProgressToken(ProgressToken),
 }
 
-struct SessionContext {
+pub struct SessionWorker {
     id: SessionId,
     next_http_request_id: HttpRequestId,
     tx_router: HashMap<HttpRequestId, HttpRequestWise>,
     resource_router: HashMap<ResourceKey, HttpRequestId>,
     common: CachedTx,
-    to_service_tx: Sender<ClientJsonRpcMessage>,
     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}")]
@@ -185,7 +188,9 @@ pub enum SessionError {
     #[error("Transport closed")]
     TransportClosed,
     #[error("IO error: {0}")]
-    Io(std::io::Error),
+    Io(#[from] std::io::Error),
+    #[error("Tokio join error {0}")]
+    TokioJoinError(#[from] tokio::task::JoinError),
 }
 
 impl From<SessionError> for std::io::Error {
@@ -196,11 +201,6 @@ impl From<SessionError> for std::io::Error {
         }
     }
 }
-impl From<std::io::Error> for SessionError {
-    fn from(value: std::io::Error) -> Self {
-        SessionError::Io(value)
-    }
-}
 
 enum OutboundChannel {
     RequestWise { id: HttpRequestId, close: bool },
@@ -212,7 +212,7 @@ pub struct StreamableHttpMessageReceiver {
     pub inner: Receiver<ServerSessionMessage>,
 }
 
-impl SessionContext {
+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");
@@ -264,12 +264,6 @@ impl SessionContext {
         self.next_http_request_id = self.next_http_request_id.wrapping_add(1);
         id
     }
-    async fn send_to_service(&self, message: ClientJsonRpcMessage) -> Result<(), SessionError> {
-        if self.to_service_tx.send(message).await.is_err() {
-            return Err(SessionError::TransportClosed);
-        }
-        Ok(())
-    }
     async fn establish_request_wise_channel(
         &mut self,
     ) -> Result<StreamableHttpMessageReceiver, SessionError> {
@@ -423,7 +417,6 @@ impl SessionContext {
 }
 
 enum SessionEvent {
-    ServiceMessage(ServerJsonRpcMessage),
     ClientMessage {
         message: ClientJsonRpcMessage,
         http_request_id: Option<HttpRequestId>,
@@ -443,6 +436,7 @@ enum SessionEvent {
         request: ClientJsonRpcMessage,
         responder: oneshot::Sender<Result<ServerJsonRpcMessage, SessionError>>,
     },
+    Close,
 }
 
 #[derive(Debug, Clone)]
@@ -454,49 +448,200 @@ pub enum SessionQuitReason {
     Cancelled,
 }
 
-impl SessionContext {
+#[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,
+        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 [`SessionHandle::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<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()))]
-    pub async fn run(mut self, ct: CancellationToken) -> SessionQuitReason {
+    async fn run(mut self, mut context: WorkerContext<Self>) -> Result<(), WorkerQuitReason> {
+        enum InnerEvent {
+            FromHttpService(SessionEvent),
+            FromHandler(WorkerSendRequest<SessionWorker>),
+        }
         // waiting for initialize request
-        let Some(evt) = self.event_rx.recv().await else {
-            return SessionQuitReason::ServiceTerminated;
-        };
+        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 SessionQuitReason::ExpectInitializeRequest;
+            return Err(WorkerQuitReason::fatal(
+                "unexpected message",
+                "get initialize request",
+            ));
         };
-        let send_result = self.send_to_service(request).await;
-        if let Err(e) = send_result {
-            let _ = responder.send(Err(e));
-            return SessionQuitReason::ServiceTerminated;
-        }
-        let Some(evt) = self.event_rx.recv().await else {
-            return SessionQuitReason::ServiceTerminated;
-        };
-        let SessionEvent::ServiceMessage(response) = evt else {
-            return SessionQuitReason::ExpectInitializeResponse;
-        };
-        let response_result = responder.send(Ok(response));
-        if response_result.is_err() {
-            return SessionQuitReason::ClientTerminated;
-        }
-        let quit_reason = loop {
+        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 {
-                        event
+                        InnerEvent::FromHttpService(event)
                     } else {
-                        break SessionQuitReason::ServiceTerminated;
+                        return Err(WorkerQuitReason::fatal("session dropped", "waiting next session event"))
                     }
                 },
-
+                from_handler = context.recv_from_handler() => {
+                    InnerEvent::FromHandler(from_handler?)
+                }
                 _ = ct.cancelled() => {
-                    break SessionQuitReason::Cancelled;
+                    return Err(WorkerQuitReason::Cancelled)
                 }
             };
             match event {
-                SessionEvent::ServiceMessage(json_rpc_message) => {
+                InnerEvent::FromHandler(WorkerSendRequest { message, responder }) => {
                     // catch response
-                    match &json_rpc_message {
+                    match &message {
                         crate::model::JsonRpcMessage::Response(json_rpc_response) => {
                             let request_id = json_rpc_response.id.clone();
                             self.unregister_resource(&ResourceKey::McpRequestId(request_id));
@@ -525,12 +670,13 @@ impl SessionContext {
                             // no need to unregister resource
                         }
                     }
-                    let _handle_result = self.handle_server_message(json_rpc_message).await;
+                    let handle_result = self.handle_server_message(message).await;
+                    let _ = responder.send(handle_result);
                 }
-                SessionEvent::ClientMessage {
+                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 {
@@ -556,165 +702,36 @@ impl SessionContext {
                         }
                         _ => {}
                     }
-                    let _handle_result = self.send_to_service(json_rpc_message).await;
+                    context.send_to_handler(json_rpc_message).await?;
                 }
-                SessionEvent::EstablishRequestWiseChannel { responder } => {
+                InnerEvent::FromHttpService(SessionEvent::EstablishRequestWiseChannel {
+                    responder,
+                }) => {
                     let handle_result = self.establish_request_wise_channel().await;
                     let _ = responder.send(handle_result);
                 }
-                SessionEvent::CloseRequestWiseChannel { id, responder } => {
+                InnerEvent::FromHttpService(SessionEvent::CloseRequestWiseChannel {
+                    id,
+                    responder,
+                }) => {
                     let _handle_result = self.tx_router.remove(&id);
                     let _ = responder.send(Ok(()));
                 }
-                SessionEvent::Resume {
+                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
                 }
             }
-        };
-        tracing::debug!("session terminated: {:?}", quit_reason);
-        quit_reason
-    }
-}
-
-pub struct Session {
-    handle: SessionHandle,
-    guard: DropGuard,
-    task_handle: tokio::task::JoinHandle<SessionQuitReason>,
-}
-
-impl Session {
-    pub fn handle(&self) -> &SessionHandle {
-        &self.handle
-    }
-    pub async fn cancel(self) -> Result<SessionQuitReason, tokio::task::JoinError> {
-        self.guard.disarm().cancel();
-        self.task_handle.await
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct SessionHandle {
-    // after all event_tx drop, inner task will be terminated
-    event_tx: Sender<SessionEvent>,
-}
-
-impl SessionHandle {
-    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(())
-    }
-
-    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)?
-    }
-    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)?
-    }
-    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)?
-    }
-
-    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)?
-    }
-
-    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 struct SessionTransport {
-    session_handle: SessionHandle,
-    to_service_rx: Receiver<ClientJsonRpcMessage>,
-}
-
-impl IntoTransport<RoleServer, SessionError, ()> for SessionTransport {
-    fn into_transport(
-        self,
-    ) -> (
-        impl Sink<ServerJsonRpcMessage, Error = SessionError> + Send + 'static,
-        impl Stream<Item = ClientJsonRpcMessage> + Send + 'static,
-    ) {
-        let stream = ReceiverStream::new(self.to_service_rx);
-        let sink = PollSender::new(self.session_handle.event_tx.clone())
-            .sink_map_err(|_| SessionError::SessionServiceTerminated)
-            .with(async |m| Ok(SessionEvent::ServiceMessage(m)));
-        (sink, stream)
+        }
     }
 }
 
@@ -735,33 +752,32 @@ impl Default for SessionConfig {
     }
 }
 
-pub fn create_session(id: SessionId, config: SessionConfig) -> (Session, SessionTransport) {
-    let (to_service_tx, to_service_rx) = tokio::sync::mpsc::channel(config.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 session_context = SessionContext {
+    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,
-        to_service_tx,
         event_rx,
         session_config: config.clone(),
     };
-    let ct = CancellationToken::new();
-    let handle = SessionHandle { event_tx };
-    let task_handle = tokio::spawn(session_context.run(ct.child_token()));
-    let session = Session {
-        handle: handle.clone(),
-        task_handle,
-        guard: ct.drop_guard(),
-    };
-    let session_transport = SessionTransport {
-        to_service_rx,
-        session_handle: handle,
-    };
-    (session, session_transport)
+    (handle, session_worker)
 }