feat(transport): expose internal worker error in fatal

Author: 4t145

crates/rmcp/src/transport/streamable_http_client.rs

@@ -269,7 +269,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
     async fn run(
         self,
         mut context: super::worker::WorkerContext<Self>,
-    ) -> Result<(), WorkerQuitReason> {
+    ) -> Result<(), WorkerQuitReason<Self::Error>> {
         let channel_buffer_capacity = self.config.channel_buffer_capacity;
         let (sse_worker_tx, mut sse_worker_rx) =
             tokio::sync::mpsc::channel::<ServerJsonRpcMessage>(channel_buffer_capacity);
@@ -286,7 +286,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
             .post_message(config.uri.clone(), initialize_request, None, None)
             .await
             .map_err(WorkerQuitReason::fatal_context("send initialize request"))?
-            .expect_initialized::<Self::Error>()
+            .expect_initialized::<C::Error>()
             .await
             .map_err(WorkerQuitReason::fatal_context(
                 "process initialize response",
@@ -346,7 +346,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
             .map_err(WorkerQuitReason::fatal_context(
                 "send initialized notification",
             ))?
-            .expect_accepted::<Self::Error>()
+            .expect_accepted::<C::Error>()
             .map_err(WorkerQuitReason::fatal_context(
                 "process initialized notification response",
             ))?;

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

@@ -296,12 +296,8 @@ pub enum SessionError {
     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 {
@@ -696,24 +692,30 @@ impl LocalSessionHandle {
 pub type SessionTransport = WorkerTransport<LocalSessionWorker>;
 
 #[derive(Debug, Error)]
-pub enum LocalSessionError {
+pub enum LocalSessionWorkerError {
     #[error("transport terminated")]
     TransportTerminated,
     #[error("unexpected message: {0:?}")]
     UnexpectedEvent(SessionEvent),
     #[error("fail to send initialize request {0}")]
     FailToSendInitializeRequest(SessionError),
-    #[error("keep alive timeout")]
-    KeepAliveTimeout,
+    #[error("fail to handle message: {0}")]
+    FailToHandleMessage(SessionError),
+    #[error("keep alive timeout after {}ms", _0.as_millis())]
+    KeepAliveTimeout(Duration),
+    #[error("Transport closed")]
+    TransportClosed,
+    #[error("Tokio join error {0}")]
+    TokioJoinError(#[from] tokio::task::JoinError),
 }
 impl Worker for LocalSessionWorker {
-    type Error = SessionError;
+    type Error = LocalSessionWorkerError;
     type Role = RoleServer;
     fn err_closed() -> Self::Error {
-        SessionError::TransportClosed
+        LocalSessionWorkerError::TransportClosed
     }
     fn err_join(e: tokio::task::JoinError) -> Self::Error {
-        SessionError::TokioJoinError(e)
+        LocalSessionWorkerError::TokioJoinError(e)
     }
     fn config(&self) -> crate::transport::worker::WorkerConfig {
         crate::transport::worker::WorkerConfig {
@@ -722,21 +724,24 @@ impl Worker for LocalSessionWorker {
         }
     }
     #[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))]
-    async fn run(mut self, mut context: WorkerContext<Self>) -> Result<(), WorkerQuitReason> {
+    async fn run(
+        mut self,
+        mut context: WorkerContext<Self>,
+    ) -> Result<(), WorkerQuitReason<Self::Error>> {
         enum InnerEvent {
             FromHttpService(SessionEvent),
             FromHandler(WorkerSendRequest<LocalSessionWorker>),
         }
         // waiting for initialize request
         let evt = self.event_rx.recv().await.ok_or_else(|| {
             WorkerQuitReason::fatal(
-                LocalSessionError::TransportTerminated,
+                LocalSessionWorkerError::TransportTerminated,
                 "get initialize request",
             )
         })?;
         let SessionEvent::InitializeRequest { request, responder } = evt else {
             return Err(WorkerQuitReason::fatal(
-                LocalSessionError::UnexpectedEvent(evt),
+                LocalSessionWorkerError::UnexpectedEvent(evt),
                 "get initialize request",
             ));
         };
@@ -746,7 +751,7 @@ impl Worker for LocalSessionWorker {
             .send(Ok(send_initialize_response.message))
             .map_err(|_| {
                 WorkerQuitReason::fatal(
-                    LocalSessionError::FailToSendInitializeRequest(
+                    LocalSessionWorkerError::FailToSendInitializeRequest(
                         SessionError::SessionServiceTerminated,
                     ),
                     "send initialize response",
@@ -765,7 +770,7 @@ impl Worker for LocalSessionWorker {
                     if let Some(event) = event {
                         InnerEvent::FromHttpService(event)
                     } else {
-                        return Err(WorkerQuitReason::fatal(LocalSessionError::TransportTerminated, "waiting next session event"))
+                        return Err(WorkerQuitReason::fatal(LocalSessionWorkerError::TransportTerminated, "waiting next session event"))
                     }
                 },
                 from_handler = context.recv_from_handler() => {
@@ -775,7 +780,7 @@ impl Worker for LocalSessionWorker {
                     return Err(WorkerQuitReason::Cancelled)
                 }
                 _ = keep_alive_timeout => {
-                    return Err(WorkerQuitReason::fatal(LocalSessionError::KeepAliveTimeout, "poll next session event"))
+                    return Err(WorkerQuitReason::fatal(LocalSessionWorkerError::KeepAliveTimeout(keep_alive), "poll next session event"))
                 }
             };
             match event {
@@ -795,7 +800,10 @@ impl Worker for LocalSessionWorker {
                             // no need to unregister resource
                         }
                     };
-                    let handle_result = self.handle_server_message(message).await;
+                    let handle_result = self
+                        .handle_server_message(message)
+                        .await
+                        .map_err(LocalSessionWorkerError::FailToHandleMessage);
                     let _ = responder.send(handle_result).inspect_err(|error| {
                         tracing::warn!(?error, "failed to send message to http service handler");
                     });

crates/rmcp/src/transport/worker.rs

@@ -7,12 +7,12 @@ use super::{IntoTransport, Transport};
 use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
 
 #[derive(Debug, thiserror::Error)]
-pub enum WorkerQuitReason {
+pub enum WorkerQuitReason<E> {
     #[error("Join error {0}")]
     Join(#[from] tokio::task::JoinError),
     #[error("Transport fatal {error}, when {context}")]
     Fatal {
-        error: Box<dyn std::error::Error + Send>,
+        error: E,
         context: Cow<'static, str>,
     },
     #[error("Transport canncelled")]
@@ -23,21 +23,16 @@ pub enum WorkerQuitReason {
     HandlerTerminated,
 }
 
-impl WorkerQuitReason {
-    pub fn fatal(
-        error: impl std::error::Error + Send + 'static,
-        context: impl Into<Cow<'static, str>>,
- ) -> Self {
+impl<E: std::error::Error + Send + 'static> WorkerQuitReason<E> {
+    pub fn fatal(error: E, context: impl Into<Cow<'static, str>>) -> Self {
         Self::Fatal {
-            error: Box::new(error),
+            error,
             context: context.into(),
         }
     }
-    pub fn fatal_context<E: std::error::Error + Send + 'static>(
-        context: impl Into<Cow<'static, str>>,
-    ) -> impl FnOnce(E) -> Self {
+    pub fn fatal_context(context: impl Into<Cow<'static, str>>) -> impl FnOnce(E) -> Self {
         |e| Self::Fatal {
-            error: Box::new(e),
+            error: e,
             context: context.into(),
         }
     }
@@ -51,7 +46,7 @@ pub trait Worker: Sized + Send + 'static {
     fn run(
         self,
         context: WorkerContext<Self>,
-    ) -> impl Future<Output = Result<(), WorkerQuitReason>> + Send;
+    ) -> impl Future<Output = Result<(), WorkerQuitReason<Self::Error>>> + Send;
     fn config(&self) -> WorkerConfig {
         WorkerConfig::default()
     }
@@ -65,7 +60,7 @@ pub struct WorkerSendRequest<W: Worker> {
 pub struct WorkerTransport<W: Worker> {
     rx: tokio::sync::mpsc::Receiver<RxJsonRpcMessage<W::Role>>,
     send_service: tokio::sync::mpsc::Sender<WorkerSendRequest<W>>,
-    join_handle: Option<tokio::task::JoinHandle<Result<(), WorkerQuitReason>>>,
+    join_handle: Option<tokio::task::JoinHandle<Result<(), WorkerQuitReason<W::Error>>>>,
     _drop_guard: tokio_util::sync::DropGuard,
     ct: CancellationToken,
 }
@@ -162,14 +157,16 @@ impl<W: Worker> WorkerContext<W> {
     pub async fn send_to_handler(
         &mut self,
         item: RxJsonRpcMessage<W::Role>,
-    ) -> Result<(), WorkerQuitReason> {
+    ) -> Result<(), WorkerQuitReason<W::Error>> {
         self.to_handler_tx
             .send(item)
             .await
             .map_err(|_| WorkerQuitReason::HandlerTerminated)
     }
 
-    pub async fn recv_from_handler(&mut self) -> Result<WorkerSendRequest<W>, WorkerQuitReason> {
+    pub async fn recv_from_handler(
+        &mut self,
+    ) -> Result<WorkerSendRequest<W>, WorkerQuitReason<W::Error>> {
         self.from_handler_rx
             .recv()
             .await