add an error type for the sub-classes

Author: EItanya

crates/rmcp/src/transport/sse.rs

@@ -14,13 +14,13 @@ const MIME_TYPE: &str = "text/event-stream";
 const HEADER_LAST_EVENT_ID: &str = "Last-Event-ID";
 
 #[derive(Error, Debug)]
-pub enum SseTransportError {
+pub enum SseTransportError<E: std::error::Error + Send + Sync + 'static> {
     #[error("SSE error: {0}")]
     Sse(#[from] SseError),
     #[error("IO error: {0}")]
     Io(#[from] std::io::Error),
-    #[error("Reqwest error: {0}")]
-    Reqwest(#[from] reqwest::Error),
+    #[error("Transport error: {0}")]
+    Transport(E),
     #[error("unexpected end of stream")]
     UnexpectedEndOfStream,
     #[error("Url error: {0}")]
@@ -29,13 +29,13 @@ pub enum SseTransportError {
     UnexpectedContentType(Option<HeaderValue>),
 }
 
-enum SseTransportState {
+enum SseTransportState<E: std::error::Error + Send + Sync + 'static> {
     Connected(BoxStream<'static, Result<Sse, SseError>>),
     Retrying {
         times: usize,
         fut: BoxFuture<
             'static,
-            Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError>,
+            Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError<E>>,
         >,
     },
     Fatal {
@@ -60,17 +60,23 @@ impl Default for SseTransportRetryConfig {
     }
 }
 
-pub trait SseClient: Clone + Send + Sync + 'static {
+impl From<reqwest::Error> for SseTransportError<reqwest::Error> {
+    fn from(e: reqwest::Error) -> Self {
+        SseTransportError::Transport(e)
+    }
+}
+
+pub trait SseClient<E: std::error::Error + Send + Sync>: Clone + Send + Sync + 'static {
     fn connect(
         &self,
         last_event_id: Option<String>,
-    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError>>;
+    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError<E>>>;
 
     fn post(
         &self,
         endpoint: &str,
         message: ClientJsonRpcMessage,
-    ) -> BoxFuture<'static, Result<(), SseTransportError>>;
+    ) -> BoxFuture<'static, Result<(), SseTransportError<E>>>;
 }
 
 pub struct RetryConfig {
@@ -84,7 +90,7 @@ pub struct ReqwestSseClient {
     sse_url: Url,
 }
 impl ReqwestSseClient {
-    pub fn new<U>(url: U) -> Result<Self, SseTransportError>
+    pub fn new<U>(url: U) -> Result<Self, SseTransportError<reqwest::Error>>
     where
         U: IntoUrl,
     {
@@ -95,7 +101,10 @@ impl ReqwestSseClient {
         })
     }
 
-    pub async fn new_with_timeout<U>(url: U, timeout: Duration) -> Result<Self, SseTransportError>
+    pub async fn new_with_timeout<U>(
+        url: U,
+        timeout: Duration,
+    ) -> Result<Self, SseTransportError<reqwest::Error>>
     where
         U: IntoUrl,
     {
@@ -109,7 +118,10 @@ impl ReqwestSseClient {
         })
     }
 
-    pub async fn new_with_client<U>(url: U, client: HttpClient) -> Result<Self, SseTransportError>
+    pub async fn new_with_client<U>(
+        url: U,
+        client: HttpClient,
+    ) -> Result<Self, SseTransportError<reqwest::Error>>
     where
         U: IntoUrl,
     {
@@ -121,12 +133,14 @@ impl ReqwestSseClient {
     }
 }
 
-impl SseClient for ReqwestSseClient {
+impl SseClient<reqwest::Error> for ReqwestSseClient {
     fn connect(
         &self,
         last_event_id: Option<String>,
-    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError>>
-    {
+    ) -> BoxFuture<
+        'static,
+        Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError<reqwest::Error>>,
+    > {
         let client = self.http_client.clone();
         let sse_url = self.sse_url.as_ref().to_string();
         let last_event_id = last_event_id.clone();
@@ -157,7 +171,7 @@ impl SseClient for ReqwestSseClient {
         &self,
         session_id: &str,
         message: ClientJsonRpcMessage,
-    ) -> BoxFuture<'static, Result<(), SseTransportError>> {
+    ) -> BoxFuture<'static, Result<(), SseTransportError<reqwest::Error>>> {
         let client = self.http_client.clone();
         let sse_url = self.sse_url.clone();
         let session_id = session_id.to_string();
@@ -179,19 +193,21 @@ impl SseClient for ReqwestSseClient {
 /// Call [`SseTransport::start`] to create a new transport from url.
 ///
 /// Call [`SseTransport::start_with_client`] to create a new transport with a customized reqwest client.
-pub struct SseTransport<C: SseClient> {
+pub struct SseTransport<C: SseClient<E>, E: std::error::Error + Send + Sync + 'static> {
     client: Arc<C>,
-    state: SseTransportState,
+    state: SseTransportState<E>,
     last_event_id: Option<String>,
     recommended_retry_duration_ms: Option<u64>,
     session_id: String,
     #[allow(clippy::type_complexity)]
-    request_queue: VecDeque<tokio::sync::oneshot::Receiver<Result<(), SseTransportError>>>,
+    request_queue: VecDeque<tokio::sync::oneshot::Receiver<Result<(), SseTransportError<E>>>>,
     pub retry_config: SseTransportRetryConfig,
 }
 
-impl SseTransport<ReqwestSseClient> {
-    pub async fn start<U>(url: U) -> Result<SseTransport<ReqwestSseClient>, SseTransportError>
+impl SseTransport<ReqwestSseClient, reqwest::Error> {
+    pub async fn start<U>(
+        url: U,
+    ) -> Result<SseTransport<ReqwestSseClient, reqwest::Error>, SseTransportError<reqwest::Error>>
     where
         U: IntoUrl,
     {
@@ -200,8 +216,8 @@ impl SseTransport<ReqwestSseClient> {
     }
 }
 
-impl<C: SseClient> SseTransport<C> {
-    pub async fn start_with_client(client: C) -> Result<Self, SseTransportError> {
+impl<C: SseClient<E>, E: std::error::Error + Send + Sync + 'static> SseTransport<C, E> {
+    pub async fn start_with_client(client: C) -> Result<Self, SseTransportError<E>> {
         let mut event_stream = client.connect(None).await?;
         let mut last_event_id = None;
         let mut retry = None;
@@ -233,7 +249,7 @@ impl<C: SseClient> SseTransport<C> {
 
     fn retry_connection(
         &self,
-    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError>>
+    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError<E>>>
     {
         let retry_duration = {
             let recommended_retry_duration = self
@@ -250,7 +266,7 @@ impl<C: SseClient> SseTransport<C> {
     }
 }
 
-impl<C: SseClient> Stream for SseTransport<C> {
+impl<C: SseClient<E>, E: std::error::Error + Send + Sync + 'static> Stream for SseTransport<C, E> {
     type Item = ServerJsonRpcMessage;
 
     fn poll_next(
@@ -323,8 +339,10 @@ impl<C: SseClient> Stream for SseTransport<C> {
     }
 }
 
-impl<C: SseClient> Sink<ClientJsonRpcMessage> for SseTransport<C> {
-    type Error = SseTransportError;
+impl<C: SseClient<E>, E: std::error::Error + Send + Sync + 'static> Sink<ClientJsonRpcMessage>
+    for SseTransport<C, E>
+{
+    type Error = SseTransportError<E>;
 
     fn poll_ready(
         mut self: std::pin::Pin<&mut Self>,