chore: update dispatcher to use Duration timeout param

Author: hashemix

Makefile.toml

@@ -29,6 +29,10 @@ args = [
     "rust-mcp-transport",
 ]
 
+
+[tasks.check]
+dependencies = ["fmt", "clippy", "test", "doc-test"]
+
 [tasks.clippy-fix]
 command = "cargo"
 args = ["clippy", "--fix", "--allow-dirty"]

crates/rust-mcp-transport/src/mcp_stream.rs

@@ -9,6 +9,7 @@ use std::{
     collections::HashMap,
     pin::Pin,
     sync::{atomic::AtomicI64, Arc},
+    time::Duration,
 };
 use tokio::{
     io::{AsyncBufReadExt, BufReader},
@@ -34,7 +35,7 @@ impl MCPStream {
         writable: Mutex<Pin<Box<dyn tokio::io::AsyncWrite + Send + Sync>>>,
         error_io: IoStream,
         pending_requests: Arc<Mutex<HashMap<RequestId, tokio::sync::oneshot::Sender<R>>>>,
-        timeout_msec: u64,
+        request_timeout: Duration,
         shutdown_rx: Receiver<bool>,
     ) -> (
         Pin<Box<dyn Stream<Item = R> + Send>>,
@@ -62,7 +63,7 @@ impl MCPStream {
             pending_requests,
             writable,
             Arc::new(AtomicI64::new(0)),
-            timeout_msec,
+            request_timeout,
         );
 
         (stream, sender, error_io)

crates/rust-mcp-transport/src/message_dispatcher.rs

@@ -28,7 +28,7 @@ pub struct MessageDispatcher<R> {
     pending_requests: Arc<Mutex<HashMap<RequestId, oneshot::Sender<R>>>>,
     writable_std: Mutex<Pin<Box<dyn tokio::io::AsyncWrite + Send + Sync>>>,
     message_id_counter: Arc<AtomicI64>,
-    timeout_msec: u64,
+    request_timeout: Duration,
 }
 
 impl<R> MessageDispatcher<R> {
@@ -38,21 +38,21 @@ impl<R> MessageDispatcher<R> {
     /// * `pending_requests` - A thread-safe map for storing pending request IDs and their response channels.
     /// * `writable_std` - A mutex-protected, pinned writer (e.g., stdout) for sending serialized messages.
     /// * `message_id_counter` - An atomic counter for generating unique request IDs.
-    /// * `timeout_msec` - The timeout duration in milliseconds for awaiting responses.
+    /// * `request_timeout` - The timeout duration in milliseconds for awaiting responses.
     ///
     /// # Returns
     /// A new `MessageDispatcher` instance configured for MCP message handling.
     pub fn new(
         pending_requests: Arc<Mutex<HashMap<RequestId, oneshot::Sender<R>>>>,
         writable_std: Mutex<Pin<Box<dyn tokio::io::AsyncWrite + Send + Sync>>>,
         message_id_counter: Arc<AtomicI64>,
-        timeout_msec: u64,
+        request_timeout: Duration,
     ) -> Self {
         Self {
             pending_requests,
             writable_std,
             message_id_counter,
-            timeout_msec,
+            request_timeout,
         }
     }
 
@@ -148,7 +148,7 @@ impl McpDispatch<ServerMessage, MessageFromClient> for MessageDispatcher<ServerM
 
         if let Some(rx) = rx_response {
             // Wait for the response with timeout
-            match await_timeout(rx, Duration::from_millis(self.timeout_msec)).await {
+            match await_timeout(rx, self.request_timeout).await {
                 Ok(response) => Ok(Some(response)),
                 Err(error) => match error {
                     TransportError::OneshotRecvError(_) => {
@@ -220,7 +220,7 @@ impl McpDispatch<ClientMessage, MessageFromServer> for MessageDispatcher<ClientM
         writable_std.flush().await?;
 
         if let Some(rx) = rx_response {
-            match await_timeout(rx, Duration::from_millis(self.timeout_msec)).await {
+            match await_timeout(rx, self.request_timeout).await {
                 Ok(response) => Ok(Some(response)),
                 Err(error) => Err(error),
             }

crates/rust-mcp-transport/src/transport.rs

@@ -1,4 +1,4 @@
-use std::pin::Pin;
+use std::{pin::Pin, time::Duration};
 
 use async_trait::async_trait;
 use rust_mcp_schema::{schema_utils::McpMessage, RequestId};
@@ -29,12 +29,12 @@ pub struct TransportOptions {
     ///
     /// This value defines the maximum amount of time to wait for a response before
     /// considering the request as timed out.
-    pub timeout: u64,
+    pub timeout: Duration,
 }
 impl Default for TransportOptions {
     fn default() -> Self {
         Self {
-            timeout: DEFAULT_TIMEOUT_MSEC,
+            timeout: Duration::from_millis(DEFAULT_TIMEOUT_MSEC),
         }
     }
 }