From 0512d9172972016c9bbda7a17fb4db6beaffb4d1 Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Sun, 19 Jul 2026 15:47:27 +0800 Subject: [PATCH] fix(feishu): harden websocket event processing --- docs/ARCHITECTURE.md | 1 + src/channels/feishu.rs | 381 ++++++++++++++++++++++++++++++++++------- 2 files changed, 320 insertions(+), 62 deletions(-) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index b26ba7e..66613a1 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -139,6 +139,7 @@ sequenceDiagram - `LivePolicy::Snapshot` 按渠道间隔发送最新运行态;`FinalOnly` 忽略运行态,只处理终态。终态绕过节流并只对明确的瞬态错误重试。 - `TurnDeliveryService` 返回可等待的终态句柄;sink 生命周期启动不等于终态已送达。Session 在终态重试最终失败时通过普通出站路径兜底一次。 - `cli_chat` 将同一 `turn_updated` 快照发给 TUI 和 WebUI。客户端只保留当前 session 中 revision 更新的 `active_turn`,终态随后由持久化历史校准。 +- 飞书对每个 DATA 帧先在 2 秒硬期限内 ACK,再进行有界分片重组,并把完整事件交给容量 32 的连接内处理队列;媒体下载和引用查询不占用正常的 WebSocket 读循环。队列饱和时当前事件在连接任务中同步处理而不丢弃。连接异常采用有上限的指数退避持续重连,不因累计故障永久停止。 - 飞书在协议解析阶段按 `allow_from` 拒绝未授权用户;群聊默认必须明确 @ 运行时解析出的机器人身份,身份解析失败时安全地忽略群消息。飞书默认 `FinalOnly`;开启 `live_updates` 后,第一个可见快照创建卡片,后续编辑同一卡片,终态编辑失败则发送完整结果兜底。reaction 清理在 finish、abort 和 Gateway shutdown 中幂等执行。 - DeliveryCoordinator 与 OutboundDispatcher 共享 `(channel, chat_id)` 写锁,避免活动 Turn 终态与独立消息并发写入同一目标。 diff --git a/src/channels/feishu.rs b/src/channels/feishu.rs index 56dd706..ed737d3 100644 --- a/src/channels/feishu.rs +++ b/src/channels/feishu.rs @@ -33,6 +33,12 @@ const TOKEN_REFRESH_SKEW: Duration = Duration::from_secs(120); const DEFAULT_TOKEN_TTL: Duration = Duration::from_secs(7200); /// Dedup cache TTL (30 minutes). const DEDUP_CACHE_TTL: Duration = Duration::from_secs(30 * 60); +const WS_EVENT_QUEUE_CAPACITY: usize = 32; +const WS_ACK_TIMEOUT: Duration = Duration::from_secs(2); +const WS_FRAGMENT_TTL: Duration = Duration::from_secs(5 * 60); +const MAX_WS_FRAGMENTS: usize = 256; +const MAX_WS_EVENT_BYTES: usize = 8 * 1024 * 1024; +const STABLE_CONNECTION_WINDOW: Duration = Duration::from_secs(30); // ───────────────────────────────────────────────────────────────────────────── // Protobuf types for Feishu WebSocket protocol (pbbp2.proto) @@ -64,6 +70,21 @@ struct PbFrame { pub payload: Option>, } +impl PbFrame { + fn header_value(&self, key: &str) -> Option<&str> { + self.headers + .iter() + .find(|header| header.key == key) + .map(|header| header.value.as_str()) + } +} + +struct FragmentAssembly { + parts: Vec>>, + total_bytes: usize, + created_at: Instant, +} + /// POST /callback/ws/endpoint response #[derive(Deserialize)] struct WsEndpointResp { @@ -974,7 +995,7 @@ impl FeishuChannel { .unwrap_or(0) } - /// Handle incoming binary PbFrame - returns Some(ParsedMessage) if we need to ack + /// Parse one complete event frame. Transport ACK and fragment handling happen earlier. async fn handle_frame(&self, frame: &PbFrame) -> Result, ChannelError> { // method 0 = CONTROL (ping/pong) if frame.method == 0 { @@ -1202,6 +1223,54 @@ impl FeishuChannel { Ok((text, media)) } + async fn process_event_frame(&self, bus: &Arc, frame: &PbFrame) { + let parsed = match self.handle_frame(frame).await { + Ok(Some(parsed)) => parsed, + Ok(None) => return, + Err(error) => { + tracing::warn!(error = %error, "Failed to parse Feishu frame"); + return; + } + }; + + let message_id = parsed.message_id.clone(); + let reaction_id = match self.add_reaction(&message_id).await { + Ok(Some(reaction_id)) => Some(reaction_id), + Ok(None) => None, + Err(error) => { + tracing::debug!(error = %error, message_id = %message_id, "Failed to add reaction"); + None + } + }; + + let mut private_context = HashMap::new(); + private_context.insert("feishu.message_id".to_string(), message_id); + if let Some(reaction_id) = reaction_id { + private_context.insert("feishu.reaction_id".to_string(), reaction_id); + } + + let msg = crate::bus::InboundMessage { + channel: "feishu".to_string(), + sender_id: parsed.open_id.clone(), + chat_id: parsed.chat_id.clone(), + content: parsed.content, + received_at: crate::bus::message::current_timestamp(), + media: parsed.media, + channel_context: crate::bus::ChannelContext { + reply_to: parsed.parent_id, + private: private_context, + }, + }; + if let Err(error) = self.handle_and_publish(bus, &msg).await { + tracing::error!( + error = %error, + open_id = %parsed.open_id, + chat_id = %parsed.chat_id, + "Failed to publish Feishu message to bus" + ); + } + } + /// Send acknowledgment for a message async fn send_ack( frame: &PbFrame, @@ -1231,21 +1300,22 @@ impl FeishuChannel { &self, bus: Arc, shutdown: CancellationToken, - ) -> Result<(), ChannelError> { + ) -> Result { let (wss_url, client_config) = tokio::select! { result = self.get_ws_endpoint(&self.http_client) => result?, - _ = shutdown.cancelled() => return Ok(()), + _ = shutdown.cancelled() => return Ok(Duration::ZERO), }; let service_id = Self::extract_service_id(&wss_url); - tracing::info!(url = %wss_url, "Connecting to Feishu WebSocket"); + tracing::info!(service_id, "Connecting to Feishu WebSocket"); 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(()), + _ = shutdown.cancelled() => return Ok(Duration::ZERO), }; + let connected_at = Instant::now(); *self.connected.write().await = true; tracing::info!("Feishu WebSocket connected"); @@ -1270,15 +1340,27 @@ impl FeishuChannel { )) => result.map_err(|e| { ChannelError::ConnectionError(format!("Failed to send initial ping: {}", e)) })?, - _ = shutdown.cancelled() => return Ok(()), + _ = shutdown.cancelled() => return Ok(connected_at.elapsed()), }; + let (event_tx, mut event_rx) = tokio::sync::mpsc::channel(WS_EVENT_QUEUE_CAPACITY); + let worker_channel = self.clone(); + let worker_bus = bus.clone(); + let mut event_worker = tokio::spawn(async move { + while let Some(frame) = event_rx.recv().await { + worker_channel + .process_event_frame(&worker_bus, &frame) + .await; + } + }); + let ping_interval = client_config.ping_interval.unwrap_or(120).max(10); let mut ping_interval_tok = tokio::time::interval(tokio::time::Duration::from_secs(ping_interval)); let mut timeout_check = tokio::time::interval(tokio::time::Duration::from_secs(10)); let mut seq: u64 = 1; let mut last_recv = Instant::now(); + let mut fragment_cache: HashMap = HashMap::new(); // Consume the immediate tick ping_interval_tok.tick().await; @@ -1292,54 +1374,46 @@ impl FeishuChannel { last_recv = Instant::now(); let bytes: Bytes = data; if let Ok(frame) = PbFrame::decode(bytes.as_ref()) { - match self.handle_frame(&frame).await { - Ok(Some(parsed)) => { - // Send ACK immediately (Feishu requires within 3 s) - if let Err(e) = Self::send_ack(&frame, &mut write).await { - tracing::error!(error = %e, "Failed to send ACK to Feishu"); - } + if frame.method != 1 { + continue; + } + match tokio::time::timeout( + WS_ACK_TIMEOUT, + Self::send_ack(&frame, &mut write), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "Failed to ACK Feishu DATA frame"); + break; + } + Err(_) => { + tracing::warn!("Timed out ACKing Feishu DATA frame"); + break; + } + } - // Add reaction emoji (await so we get the reaction_id for later removal) - let message_id = parsed.message_id.clone(); - let reaction_id = match self.add_reaction(&message_id).await { - Ok(Some(rid)) => Some(rid), - Ok(None) => None, - Err(e) => { - tracing::debug!(error = %e, message_id = %message_id, "Failed to add reaction"); - None + match reassemble_data_frame(&frame, &mut fragment_cache) { + Ok(Some(complete_frame)) => { + match event_tx.try_send(complete_frame) { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(frame)) => { + tracing::warn!( + capacity = WS_EVENT_QUEUE_CAPACITY, + "Feishu event queue is full; processing one event on the connection task" + ); + self.process_event_frame(&bus, &frame).await; + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + tracing::warn!("Feishu event worker stopped unexpectedly"); + break; } - }; - - let mut private_context = std::collections::HashMap::new(); - private_context.insert("feishu.message_id".to_string(), message_id.clone()); - if let Some(ref rid) = reaction_id { - private_context.insert("feishu.reaction_id".to_string(), rid.clone()); - } - - #[cfg(debug_assertions)] - tracing::debug!(open_id = %parsed.open_id, chat_id = %parsed.chat_id, content_len = %parsed.content.len(), media_count = %parsed.media.len(), "Publishing message to bus"); - let msg = crate::bus::InboundMessage { - channel: "feishu".to_string(), - sender_id: parsed.open_id.clone(), - chat_id: parsed.chat_id.clone(), - content: parsed.content.clone(), - received_at: crate::bus::message::current_timestamp(), - media: parsed.media.clone(), - channel_context: crate::bus::ChannelContext { - reply_to: parsed.parent_id.clone(), - private: private_context, - }, - }; - if let Err(e) = self.handle_and_publish(&bus, &msg).await { - tracing::error!(error = %e, open_id = %parsed.open_id, chat_id = %parsed.chat_id, "Failed to publish Feishu message to bus"); - } else { - #[cfg(debug_assertions)] - tracing::debug!(open_id = %parsed.open_id, chat_id = %parsed.chat_id, "Message published to bus successfully"); } } Ok(None) => {} - Err(e) => { - tracing::warn!(error = %e, "Failed to parse Feishu frame"); + Err(error) => { + tracing::warn!(error = %error, "Rejected invalid Feishu fragment sequence"); } } } @@ -1401,6 +1475,9 @@ impl FeishuChannel { let now = Instant::now(); let mut seen = self.seen_message_ids.write().await; seen.retain(|_, ts| now.duration_since(*ts) < DEDUP_CACHE_TTL); + fragment_cache.retain(|_, assembly| { + now.duration_since(assembly.created_at) < WS_FRAGMENT_TTL + }); } _ = shutdown.cancelled() => { tracing::info!("Feishu channel shutdown signal received"); @@ -1410,10 +1487,106 @@ impl FeishuChannel { } *self.connected.write().await = false; - Ok(()) + drop(event_tx); + if tokio::time::timeout(CHANNEL_STOP_GRACE, &mut event_worker) + .await + .is_err() + { + event_worker.abort(); + let _ = event_worker.await; + } + Ok(connected_at.elapsed()) } } +fn reassemble_data_frame( + frame: &PbFrame, + cache: &mut HashMap, +) -> Result, ChannelError> { + let sum = frame + .header_value("sum") + .and_then(|value| value.parse::().ok()) + .unwrap_or(1) + .max(1); + if sum == 1 { + if frame + .payload + .as_ref() + .is_some_and(|payload| payload.len() > MAX_WS_EVENT_BYTES) + { + return Err(ChannelError::Other(format!( + "Feishu event exceeds {MAX_WS_EVENT_BYTES} byte limit" + ))); + } + return Ok(Some(frame.clone())); + } + if sum > MAX_WS_FRAGMENTS { + return Err(ChannelError::Other(format!( + "Feishu event declares too many fragments: {sum}" + ))); + } + + let message_id = frame.header_value("message_id").unwrap_or_default(); + if message_id.is_empty() { + return Err(ChannelError::Other( + "Fragmented Feishu event has no message_id".to_string(), + )); + } + let seq = frame + .header_value("seq") + .and_then(|value| value.parse::().ok()) + .unwrap_or(0); + if seq >= sum { + return Err(ChannelError::Other(format!( + "Feishu fragment index {seq} is outside declared count {sum}" + ))); + } + + let assembly = cache + .entry(message_id.to_string()) + .or_insert_with(|| FragmentAssembly { + parts: vec![None; sum], + total_bytes: 0, + created_at: Instant::now(), + }); + if assembly.parts.len() != sum { + *assembly = FragmentAssembly { + parts: vec![None; sum], + total_bytes: 0, + created_at: Instant::now(), + }; + } + let payload = frame.payload.clone().unwrap_or_default(); + let previous_len = assembly.parts[seq].as_ref().map_or(0, Vec::len); + assembly.total_bytes = assembly.total_bytes - previous_len + payload.len(); + if assembly.total_bytes > MAX_WS_EVENT_BYTES { + cache.remove(message_id); + return Err(ChannelError::Other(format!( + "Feishu fragmented event exceeds {MAX_WS_EVENT_BYTES} byte limit" + ))); + } + assembly.parts[seq] = Some(payload); + if assembly.parts.iter().any(Option::is_none) { + return Ok(None); + } + + let complete_payload = assembly + .parts + .iter() + .flat_map(|part| part.as_deref().unwrap_or_default()) + .copied() + .collect(); + cache.remove(message_id); + let mut complete = frame.clone(); + complete.payload = Some(complete_payload); + Ok(Some(complete)) +} + +fn reconnect_delay(attempt: u32) -> Duration { + let exponent = attempt.saturating_sub(1).min(6); + Duration::from_secs(1_u64 << exponent) +} + fn parse_post_content(content: &str) -> String { /// Extract text from a single post element (text, link, at-mention). fn extract_element(el: &serde_json::Value, out: &mut Vec) { @@ -2317,8 +2490,7 @@ impl Channel for FeishuChannel { let channel = self.clone(); let bus = bus.clone(); *run_task = Some(tokio::spawn(async move { - let mut consecutive_failures = 0; - let max_failures = 3; + let mut retry_attempt = 0_u32; loop { if !*channel.running.read().await { @@ -2326,16 +2498,20 @@ impl Channel for FeishuChannel { } match channel.run_ws_loop(bus.clone(), shutdown.clone()).await { - Ok(_) => { - tracing::info!("Feishu WebSocket disconnected"); + Ok(connected_for) => { + if connected_for >= STABLE_CONNECTION_WINDOW { + retry_attempt = 0; + } else { + retry_attempt = retry_attempt.saturating_add(1).min(7); + } + tracing::info!( + connected_secs = connected_for.as_secs(), + "Feishu WebSocket disconnected" + ); } Err(e) => { - consecutive_failures += 1; - tracing::error!(attempt = consecutive_failures, error = %e, "Feishu WebSocket error"); - if consecutive_failures >= max_failures { - tracing::error!("Feishu channel: max failures reached, stopping"); - break; - } + retry_attempt = retry_attempt.saturating_add(1).min(6); + tracing::error!(attempt = retry_attempt, error = %e, "Feishu WebSocket error"); } } @@ -2343,9 +2519,13 @@ impl Channel for FeishuChannel { break; } - tracing::info!("Feishu channel retrying in 5s..."); + let retry_delay = reconnect_delay(retry_attempt); + tracing::info!( + delay_secs = retry_delay.as_secs(), + "Feishu channel reconnect scheduled" + ); tokio::select! { - _ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {} + _ = tokio::time::sleep(retry_delay) => {} _ = shutdown.cancelled() => break, } } @@ -2799,6 +2979,83 @@ mod tests { assert!(!message_mentions_bot(&message, Some("ou_other"))); } + fn fragment_frame(message_id: &str, seq: usize, sum: usize, payload: &[u8]) -> PbFrame { + PbFrame { + seq_id: seq as u64, + log_id: 1, + service: 1, + method: 1, + headers: vec![ + PbHeader { + key: "message_id".to_string(), + value: message_id.to_string(), + }, + PbHeader { + key: "seq".to_string(), + value: seq.to_string(), + }, + PbHeader { + key: "sum".to_string(), + value: sum.to_string(), + }, + ], + payload: Some(payload.to_vec()), + } + } + + #[test] + fn fragmented_frames_are_reassembled_in_sequence_order() { + let mut cache = HashMap::new(); + assert!( + reassemble_data_frame(&fragment_frame("event", 1, 3, b"world"), &mut cache) + .unwrap() + .is_none() + ); + assert!( + reassemble_data_frame(&fragment_frame("event", 0, 3, b"hello "), &mut cache) + .unwrap() + .is_none() + ); + let complete = reassemble_data_frame(&fragment_frame("event", 2, 3, b"!"), &mut cache) + .unwrap() + .expect("last fragment should complete the event"); + + assert_eq!( + complete.payload.as_deref(), + Some(b"hello world!".as_slice()) + ); + assert!(cache.is_empty()); + } + + #[test] + fn fragmented_frames_reject_invalid_shape_and_bound_total_size() { + let mut cache = HashMap::new(); + assert!( + reassemble_data_frame( + &fragment_frame("event", MAX_WS_FRAGMENTS, MAX_WS_FRAGMENTS, b"x"), + &mut cache, + ) + .is_err() + ); + assert!( + reassemble_data_frame( + &fragment_frame("event", 0, 2, &vec![0; MAX_WS_EVENT_BYTES + 1]), + &mut cache, + ) + .is_err() + ); + assert!(cache.is_empty()); + } + + #[test] + fn reconnect_backoff_is_bounded_and_resets_after_success() { + assert_eq!(reconnect_delay(0), Duration::from_secs(1)); + assert_eq!(reconnect_delay(1), Duration::from_secs(1)); + assert_eq!(reconnect_delay(2), Duration::from_secs(2)); + assert_eq!(reconnect_delay(7), Duration::from_secs(64)); + assert_eq!(reconnect_delay(100), Duration::from_secs(64)); + } + #[tokio::test] async fn turn_sink_creates_once_updates_same_card_and_cleans_up_at_finish() { let state = Arc::new(Mutex::new(MockTurnState::default()));