feat(service): add close() method for graceful connection shutdown

crates/rmcp/src/service.rs

@@ -434,7 +434,7 @@ impl<R: ServiceRole> Peer<R> {
 pub struct RunningService<R: ServiceRole, S: Service<R>> {
     service: Arc<S>,
     peer: Peer<R>,
-    handle: tokio::task::JoinHandle<QuitReason>,
+    handle: Option<tokio::task::JoinHandle<QuitReason>>,
     cancellation_token: CancellationToken,
     dg: DropGuard,
 }
@@ -459,14 +459,104 @@ impl<R: ServiceRole, S: Service<R>> RunningService<R, S> {
     pub fn cancellation_token(&self) -> RunningServiceCancellationToken {
         RunningServiceCancellationToken(self.cancellation_token.clone())
     }
+
+    /// Returns true if the service has been closed or cancelled.
     #[inline]
-    pub async fn waiting(self) -> Result<QuitReason, tokio::task::JoinError> {
-        self.handle.await
+    pub fn is_closed(&self) -> bool {
+        self.handle.is_none() || self.cancellation_token.is_cancelled()
+    }
+
+    /// Wait for the service to complete.
+    ///
+    /// This will block until the service loop terminates (either due to
+    /// cancellation, transport closure, or an error).
+    #[inline]
+    pub async fn waiting(mut self) -> Result<QuitReason, tokio::task::JoinError> {
+        match self.handle.take() {
+            Some(handle) => handle.await,
+            None => Ok(QuitReason::Closed),
+        }
+    }
+
+    /// Gracefully close the connection and wait for cleanup to complete.
+    ///
+    /// This method cancels the service, waits for the background task to finish
+    /// (which includes calling `transport.close()`), and ensures all cleanup
+    /// operations complete before returning.
+    ///
+    /// Unlike [`cancel`](Self::cancel), this method takes `&mut self` and can be
+    /// called without consuming the `RunningService`. After calling this method,
+    /// the service is considered closed and subsequent operations will fail.
+    ///
+    /// # Example
+    ///
+    /// ```rust,ignore
+    /// let mut client = ().serve(transport).await?;
+    /// // ... use the client ...
+    /// client.close().await?;
+    /// ```
+    pub async fn close(&mut self) -> Result<QuitReason, tokio::task::JoinError> {
+        if let Some(handle) = self.handle.take() {
+            // Disarm the drop guard so it doesn't try to cancel again
+            // We need to cancel manually and wait for completion
+            self.cancellation_token.cancel();
+            handle.await
+        } else {
+            // Already closed
+            Ok(QuitReason::Closed)
+        }
     }
-    pub async fn cancel(self) -> Result<QuitReason, tokio::task::JoinError> {
-        let RunningService { dg, handle, .. } = self;
-        dg.disarm().cancel();
-        handle.await
+
+    /// Gracefully close the connection with a timeout.
+    ///
+    /// Similar to [`close`](Self::close), but returns after the specified timeout
+    /// if the cleanup doesn't complete in time. This is useful for ensuring
+    /// a bounded shutdown time.
+    ///
+    /// Returns `Ok(Some(reason))` if shutdown completed within the timeout,
+    /// `Ok(None)` if the timeout was reached, or `Err` if there was a join error.
+    pub async fn close_with_timeout(
+        &mut self,
+        timeout: Duration,
+    ) -> Result<Option<QuitReason>, tokio::task::JoinError> {
+        if let Some(handle) = self.handle.take() {
+            self.cancellation_token.cancel();
+            match tokio::time::timeout(timeout, handle).await {
+                Ok(result) => result.map(Some),
+                Err(_elapsed) => {
+                    tracing::warn!(
+                        "close_with_timeout: cleanup did not complete within {:?}",
+                        timeout
+                    );
+                    Ok(None)
+                }
+            }
+        } else {
+            Ok(Some(QuitReason::Closed))
+        }
+    }
+
+    /// Cancel the service and wait for cleanup to complete.
+    ///
+    /// This consumes the `RunningService` and ensures the connection is properly
+    /// closed. For a non-consuming alternative, see [`close`](Self::close).
+    pub async fn cancel(mut self) -> Result<QuitReason, tokio::task::JoinError> {
+        // Disarm the drop guard since we're handling cancellation explicitly
+        let _ = std::mem::replace(&mut self.dg, self.cancellation_token.clone().drop_guard());
+        self.close().await
+    }
+}
+
+impl<R: ServiceRole, S: Service<R>> Drop for RunningService<R, S> {
+    fn drop(&mut self) {
+        if self.handle.is_some() && !self.cancellation_token.is_cancelled() {
+            tracing::debug!(
+                "RunningService dropped without explicit close(). \
+                 The connection will be closed asynchronously. \
+                 For guaranteed cleanup, call close() or cancel() before dropping."
+            );
+        }
+        // The DropGuard will handle cancellation
     }
 }
 
@@ -847,7 +937,7 @@ where
     RunningService {
         service,
         peer: peer_return,
-        handle,
+        handle: Some(handle),
         cancellation_token: ct.clone(),
         dg: ct.drop_guard(),
     }

crates/rmcp/src/transport/streamable_http_client.rs

@@ -333,37 +333,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
             }
             None
         };
-        // delete session when drop guard is dropped
-        if let Some(session_id) = &session_id {
-            let ct = transport_task_ct.clone();
-            let client = self.client.clone();
-            let session_id = session_id.clone();
-            let url = config.uri.clone();
-            let auth_header = config.auth_header.clone();
-            tokio::spawn(async move {
-                ct.cancelled().await;
-                let delete_session_result = client
-                    .delete_session(url, session_id.clone(), auth_header.clone())
-                    .await;
-                match delete_session_result {
-                    Ok(_) => {
-                        tracing::info!(session_id = session_id.as_ref(), "delete session success")
-                    }
-                    Err(StreamableHttpError::ServerDoesNotSupportDeleteSession) => {
-                        tracing::info!(
-                            session_id = session_id.as_ref(),
-                            "server doesn't support delete session"
-                        )
-                    }
-                    Err(e) => {
-                        tracing::error!(
-                            session_id = session_id.as_ref(),
-                            "fail to delete session: {e}"
-                        );
-                    }
-                };
-            });
-        }
+        // Store session info for cleanup when run() exits (not spawned, so cleanup completes before close() returns)
+        let session_cleanup_info = session_id.as_ref().map(|sid| {
+            (self.client.clone(), config.uri.clone(), sid.clone(), config.auth_header.clone())
+        });
 
         context.send_to_handler(message).await?;
         let initialized_notification = context.recv_from_handler().await?;
@@ -437,20 +410,23 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                 }
             });
         }
-        loop {
+        // Main event loop - capture exit reason so we can do cleanup before returning
+        let loop_result: Result<(), WorkerQuitReason<Self::Error>> = 'main_loop: loop {
             let event = tokio::select! {
                 _ = transport_task_ct.cancelled() => {
                     tracing::debug!("cancelled");
-                    return Err(WorkerQuitReason::Cancelled);
+                    break 'main_loop Err(WorkerQuitReason::Cancelled);
                 }
                 message = context.recv_from_handler() => {
-                    let message = message?;
-                    Event::ClientMessage(message)
+                    match message {
+                        Ok(msg) => Event::ClientMessage(msg),
+                        Err(e) => break 'main_loop Err(e),
+                    }
                 },
                 message = sse_worker_rx.recv() => {
                     let Some(message) = message else {
                         tracing::trace!("transport dropped, exiting");
-                        return Err(WorkerQuitReason::HandlerTerminated);
+                        break 'main_loop Err(WorkerQuitReason::HandlerTerminated);
                     };
                     Event::ServerMessage(message)
                 },
@@ -525,7 +501,9 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                 }
                 Event::ServerMessage(json_rpc_message) => {
                     // send the message to the handler
-                    context.send_to_handler(json_rpc_message).await?;
+                    if let Err(e) = context.send_to_handler(json_rpc_message).await {
+                        break 'main_loop Err(e);
+                    }
                 }
                 Event::StreamResult(result) => {
                     if result.is_err() {
@@ -536,7 +514,44 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                     }
                 }
             }
+        };
+
+        // Cleanup session before returning (ensures close() waits for session deletion)
+        // Use a timeout to prevent indefinite hangs if the server is unresponsive
+        if let Some((client, url, session_id, auth_header)) = session_cleanup_info {
+            const SESSION_CLEANUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
+            match tokio::time::timeout(
+                SESSION_CLEANUP_TIMEOUT,
+                client.delete_session(url, session_id.clone(), auth_header),
+            )
+            .await
+            {
+                Ok(Ok(_)) => {
+                    tracing::info!(session_id = session_id.as_ref(), "delete session success")
+                }
+                Ok(Err(StreamableHttpError::ServerDoesNotSupportDeleteSession)) => {
+                    tracing::info!(
+                        session_id = session_id.as_ref(),
+                        "server doesn't support delete session"
+                    )
+                }
+                Ok(Err(e)) => {
+                    tracing::error!(
+                        session_id = session_id.as_ref(),
+                        "fail to delete session: {e}"
+                    );
+                }
+                Err(_elapsed) => {
+                    tracing::warn!(
+                        session_id = session_id.as_ref(),
+                        "session cleanup timed out after {:?}",
+                        SESSION_CLEANUP_TIMEOUT
+                    );
+                }
+            }
         }
+
+        loop_result
     }
 }
 

crates/rmcp/tests/test_close_connection.rs

@@ -0,0 +1,127 @@
+//cargo test --test test_close_connection --features "client server"
+
+mod common;
+use std::time::Duration;
+
+use common::handlers::{TestClientHandler, TestServer};
+use rmcp::{service::QuitReason, ServiceExt};
+
+/// Test that close() properly shuts down the connection
+#[tokio::test]
+async fn test_close_method() -> anyhow::Result<()> {
+    let (server_transport, client_transport) = tokio::io::duplex(4096);
+
+    // Start server
+    let server_handle = tokio::spawn(async move {
+        let server = TestServer::new().serve(server_transport).await?;
+        server.waiting().await?;
+        anyhow::Ok(())
+    });
+
+    // Start client
+    let handler = TestClientHandler::new(true, true);
+    let mut client = handler.serve(client_transport).await?;
+
+    // Verify client is not closed
+    assert!(!client.is_closed());
+
+    // Call close() and verify it returns
+    let result = client.close().await?;
+    assert!(matches!(result, QuitReason::Cancelled));
+
+    // Verify client is now closed
+    assert!(client.is_closed());
+
+    // Calling close() again should return Closed immediately
+    let result = client.close().await?;
+    assert!(matches!(result, QuitReason::Closed));
+
+    // Wait for server to finish
+    server_handle.await??;
+    Ok(())
+}
+
+/// Test that close_with_timeout() respects the timeout
+#[tokio::test]
+async fn test_close_with_timeout() -> anyhow::Result<()> {
+    let (server_transport, client_transport) = tokio::io::duplex(4096);
+
+    // Start server
+    let server_handle = tokio::spawn(async move {
+        let server = TestServer::new().serve(server_transport).await?;
+        server.waiting().await?;
+        anyhow::Ok(())
+    });
+
+    // Start client
+    let handler = TestClientHandler::new(true, true);
+    let mut client = handler.serve(client_transport).await?;
+
+    // Close with a reasonable timeout
+    let result = client.close_with_timeout(Duration::from_secs(5)).await?;
+    assert!(result.is_some());
+    assert!(matches!(result.unwrap(), QuitReason::Cancelled));
+
+    // Verify client is now closed
+    assert!(client.is_closed());
+
+    // Wait for server to finish
+    server_handle.await??;
+    Ok(())
+}
+
+/// Test that cancel() still works and consumes self
+#[tokio::test]
+async fn test_cancel_method() -> anyhow::Result<()> {
+    let (server_transport, client_transport) = tokio::io::duplex(4096);
+
+    // Start server
+    let server_handle = tokio::spawn(async move {
+        let server = TestServer::new().serve(server_transport).await?;
+        server.waiting().await?;
+        anyhow::Ok(())
+    });
+
+    // Start client
+    let handler = TestClientHandler::new(true, true);
+    let client = handler.serve(client_transport).await?;
+
+    // Cancel should work as before
+    let result = client.cancel().await?;
+    assert!(matches!(result, QuitReason::Cancelled));
+
+    // Wait for server to finish
+    server_handle.await??;
+    Ok(())
+}
+
+/// Test that dropping without close() logs a debug message (we can't easily test
+/// the log output, but we can verify the drop doesn't panic)
+#[tokio::test]
+async fn test_drop_without_close() -> anyhow::Result<()> {
+    let (server_transport, client_transport) = tokio::io::duplex(4096);
+
+    // Start server that will handle the drop
+    let server_handle = tokio::spawn(async move {
+        let server = TestServer::new().serve(server_transport).await?;
+        // The server should close when the client drops
+        let result = server.waiting().await?;
+        // Server should detect closure
+        assert!(matches!(result, QuitReason::Closed | QuitReason::Cancelled));
+        anyhow::Ok(())
+    });
+
+    // Create and immediately drop the client
+    {
+        let handler = TestClientHandler::new(true, true);
+        let _client = handler.serve(client_transport).await?;
+        // Client dropped here without calling close()
+    }
+
+    // Give the async cleanup a moment to run
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    // Wait for server to finish (it should detect the closure)
+    server_handle.await??;
+    Ok(())
+}