fix(feishu): bound shutdown during connection

src/channels/feishu.rs

@@ -8,8 +8,9 @@ use futures_util::{SinkExt, StreamExt};
 use prost::{Message as ProstMessage, bytes::Bytes};
 use regex::Regex;
 use serde::Deserialize;
-use tokio::sync::{Mutex, RwLock, broadcast};
+use tokio::sync::{Mutex, RwLock};
 use tokio::task::JoinHandle;
+use tokio_util::sync::CancellationToken;
 
 use crate::bus::{MediaItem, MessageBus, OutboundMessage};
 use crate::channels::base::{Channel, ChannelError};
@@ -21,6 +22,11 @@ const FEISHU_WS_BASE: &str = "https://open.feishu.cn";
 /// Heartbeat timeout for WS connection — must be larger than ping_interval (default 120 s).
 /// If no binary frame (pong or event) is received within this window, reconnect.
 const WS_HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(300);
+const CHANNEL_STOP_GRACE: Duration = if cfg!(test) {
+    Duration::from_millis(100)
+} else {
+    Duration::from_secs(5)
+};
 /// Refresh tenant token this many seconds before the announced expiry.
 const TOKEN_REFRESH_SKEW: Duration = Duration::from_secs(120);
 /// Default tenant token TTL when `expire`/`expires_in` is absent.
@@ -145,7 +151,7 @@ pub struct FeishuChannel {
     config: FeishuChannelConfig,
     http_client: reqwest::Client,
     running: Arc<RwLock<bool>>,
-    shutdown_tx: Arc<RwLock<Option<broadcast::Sender<()>>>>,
+    shutdown: Arc<RwLock<Option<CancellationToken>>>,
     run_task: Arc<Mutex<Option<JoinHandle<()>>>>,
     connected: Arc<RwLock<bool>>,
     /// Cached tenant access token with proactive refresh.
@@ -179,7 +185,7 @@ impl FeishuChannel {
             config,
             http_client: reqwest::Client::new(),
             running: Arc::new(RwLock::new(false)),
-            shutdown_tx: Arc::new(RwLock::new(None)),
+            shutdown: Arc::new(RwLock::new(None)),
             run_task: Arc::new(Mutex::new(None)),
             connected: Arc::new(RwLock::new(false)),
             tenant_token: Arc::new(RwLock::new(None)),
@@ -1166,18 +1172,22 @@ impl FeishuChannel {
     async fn run_ws_loop(
         &self,
         bus: Arc<MessageBus>,
-        mut shutdown_rx: broadcast::Receiver<()>,
+        shutdown: CancellationToken,
     ) -> Result<(), ChannelError> {
-        let (wss_url, client_config) = self.get_ws_endpoint(&self.http_client).await?;
+        let (wss_url, client_config) = tokio::select! {
+            result = self.get_ws_endpoint(&self.http_client) => result?,
+            _ = shutdown.cancelled() => return Ok(()),
+        };
 
         let service_id = Self::extract_service_id(&wss_url);
         tracing::info!(url = %wss_url, "Connecting to Feishu WebSocket");
 
-        let (ws_stream, _) = tokio_tungstenite::connect_async(&wss_url)
-            .await
-            .map_err(|e| {
+        let (ws_stream, _) = tokio::select! {
+            result = tokio_tungstenite::connect_async(&wss_url) => result.map_err(|e| {
                 ChannelError::ConnectionError(format!("WebSocket connection failed: {}", e))
-            })?;
+            })?,
+            _ = shutdown.cancelled() => return Ok(()),
+        };
 
         *self.connected.write().await = true;
         tracing::info!("Feishu WebSocket connected");
@@ -1196,14 +1206,14 @@ impl FeishuChannel {
             }],
             payload: None,
         };
-        write
-            .send(tokio_tungstenite::tungstenite::Message::Binary(
+        tokio::select! {
+            result = write.send(tokio_tungstenite::tungstenite::Message::Binary(
                 ping_frame.encode_to_vec().into(),
-            ))
-            .await
-            .map_err(|e| {
+            )) => result.map_err(|e| {
                 ChannelError::ConnectionError(format!("Failed to send initial ping: {}", e))
-            })?;
+            })?,
+            _ = shutdown.cancelled() => return Ok(()),
+        };
 
         let ping_interval = client_config.ping_interval.unwrap_or(120).max(10);
         let mut ping_interval_tok =
@@ -1336,7 +1346,7 @@ impl FeishuChannel {
                     let mut seen = self.seen_message_ids.write().await;
                     seen.retain(|_, ts| now.duration_since(*ts) < DEDUP_CACHE_TTL);
                 }
-                _ = shutdown_rx.recv() => {
+                _ = shutdown.cancelled() => {
                     tracing::info!("Feishu channel shutdown signal received");
                     break;
                 }
@@ -1954,8 +1964,8 @@ impl Channel for FeishuChannel {
 
         *self.running.write().await = true;
 
-        let (shutdown_tx, _) = broadcast::channel(1);
-        *self.shutdown_tx.write().await = Some(shutdown_tx.clone());
+        let shutdown = CancellationToken::new();
+        *self.shutdown.write().await = Some(shutdown.clone());
 
         let channel = self.clone();
         let bus = bus.clone();
@@ -1968,8 +1978,7 @@ impl Channel for FeishuChannel {
                     break;
                 }
 
-                let shutdown_rx = shutdown_tx.subscribe();
-                match channel.run_ws_loop(bus.clone(), shutdown_rx).await {
+                match channel.run_ws_loop(bus.clone(), shutdown.clone()).await {
                     Ok(_) => {
                         tracing::info!("Feishu WebSocket disconnected");
                     }
@@ -1983,12 +1992,15 @@ impl Channel for FeishuChannel {
                     }
                 }
 
-                if !*channel.running.read().await {
+                if !*channel.running.read().await || shutdown.is_cancelled() {
                     break;
                 }
 
                 tracing::info!("Feishu channel retrying in 5s...");
-                tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+                tokio::select! {
+                    _ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {}
+                    _ = shutdown.cancelled() => break,
+                }
             }
 
             *channel.running.write().await = false;
@@ -2001,15 +2013,27 @@ impl Channel for FeishuChannel {
 
     async fn stop(&self) -> Result<(), ChannelError> {
         *self.running.write().await = false;
+        *self.connected.write().await = false;
 
-        if let Some(tx) = self.shutdown_tx.write().await.take() {
-            let _ = tx.send(());
+        if let Some(shutdown) = self.shutdown.write().await.take() {
+            shutdown.cancel();
         }
 
-        if let Some(task) = self.run_task.lock().await.take() {
-            task.await.map_err(|error| {
+        let task = { self.run_task.lock().await.take() };
+        if let Some(mut task) = task {
+            match tokio::time::timeout(CHANNEL_STOP_GRACE, &mut task).await {
+                Ok(result) => result.map_err(|error| {
                     ChannelError::Other(format!("Feishu channel task failed to join: {error}"))
-            })?;
+                })?,
+                Err(_) => {
+                    tracing::warn!(
+                        grace_ms = CHANNEL_STOP_GRACE.as_millis(),
+                        "Feishu channel did not stop in time; aborting connection task"
+                    );
+                    task.abort();
+                    let _ = task.await;
+                }
+            }
         }
 
         Ok(())
@@ -2229,6 +2253,43 @@ impl Channel for FeishuChannel {
 mod tests {
     use super::*;
 
+    fn test_channel() -> FeishuChannel {
+        FeishuChannel::new(
+            FeishuChannelConfig {
+                enabled: true,
+                app_id: "test-app".to_string(),
+                app_secret: "test-secret".to_string(),
+                allow_from: vec!["*".to_string()],
+                agent: String::new(),
+                media_dir: String::new(),
+                reaction_emoji: "THUMBSUP".to_string(),
+            },
+            Path::new("/tmp"),
+        )
+        .expect("test channel should be valid")
+    }
+
+    #[tokio::test]
+    async fn stop_aborts_connection_task_that_ignores_cancellation() {
+        let channel = test_channel();
+        let shutdown = CancellationToken::new();
+
+        *channel.running.write().await = true;
+        *channel.connected.write().await = true;
+        *channel.shutdown.write().await = Some(shutdown.clone());
+        *channel.run_task.lock().await = Some(tokio::spawn(std::future::pending()));
+
+        tokio::time::timeout(Duration::from_secs(1), channel.stop())
+            .await
+            .expect("stop must have a hard deadline")
+            .expect("stop should succeed after aborting the stuck task");
+
+        assert!(shutdown.is_cancelled());
+        assert!(!channel.is_running());
+        assert!(!*channel.connected.read().await);
+        assert!(channel.run_task.lock().await.is_none());
+    }
+
     #[test]
     fn collect_post_image_keys_finds_nested_images() {
         let content = serde_json::json!({