From e326e12039f40281839dcf90d49a0a59b8f76d24 Mon Sep 17 00:00:00 2001
From: Chummy <chumyin0912@gmail.com>
Date: Wed, 18 Feb 2026 16:29:40 +0800
Subject: [PATCH] test(telegram): cover draft streaming paths and simplify
 stream modes

---
 src/channels/telegram.rs | 75 +++++++++++++++++++++++++++++++++++++---
 src/config/mod.rs        |  3 +-
 src/config/schema.rs     |  3 --
 3 files changed, 72 insertions(+), 9 deletions(-)

diff --git a/src/channels/telegram.rs b/src/channels/telegram.rs
index 32c91b2..6e67b3b 100644
--- a/src/channels/telegram.rs
+++ b/src/channels/telegram.rs
@@ -1251,7 +1251,9 @@ impl Channel for TelegramChannel {
             .and_then(|id| id.as_i64())
             .map(|id| id.to_string());
 
-        self.last_draft_edit.lock().insert(chat_id.to_string(), std::time::Instant::now());
+        self.last_draft_edit
+            .lock()
+            .insert(chat_id.to_string(), std::time::Instant::now());
 
         Ok(message_id)
     }
@@ -1313,7 +1315,9 @@ impl Channel for TelegramChannel {
             .await?;
 
         if resp.status().is_success() {
-            self.last_draft_edit.lock().insert(chat_id.clone(), std::time::Instant::now());
+            self.last_draft_edit
+                .lock()
+                .insert(chat_id.clone(), std::time::Instant::now());
         } else {
             let status = resp.status();
             let err = resp.text().await.unwrap_or_default();
@@ -1337,7 +1341,9 @@ impl Channel for TelegramChannel {
                 Ok(id) => id,
                 Err(e) => {
                     tracing::warn!("Invalid Telegram message_id '{message_id}': {e}");
-                    return self.send_text_chunks(text, &chat_id, thread_id.as_deref()).await;
+                    return self
+                        .send_text_chunks(text, &chat_id, thread_id.as_deref())
+                        .await;
                 }
             };
 
@@ -1362,7 +1368,9 @@ impl Channel for TelegramChannel {
             Ok(id) => id,
             Err(e) => {
                 tracing::warn!("Invalid Telegram message_id '{message_id}': {e}");
-                return self.send_text_chunks(text, &chat_id, thread_id.as_deref()).await;
+                return self
+                    .send_text_chunks(text, &chat_id, thread_id.as_deref())
+                    .await;
             }
         };
 
@@ -1641,6 +1649,65 @@ mod tests {
         assert!(guard.is_some());
     }
 
+    #[test]
+    fn supports_draft_updates_respects_stream_mode() {
+        let off = TelegramChannel::new("fake-token".into(), vec!["*".into()]);
+        assert!(!off.supports_draft_updates());
+
+        let partial = TelegramChannel::new("fake-token".into(), vec!["*".into()])
+            .with_streaming(StreamMode::Partial, 750);
+        assert!(partial.supports_draft_updates());
+        assert_eq!(partial.draft_update_interval_ms, 750);
+    }
+
+    #[tokio::test]
+    async fn send_draft_returns_none_when_stream_mode_off() {
+        let ch = TelegramChannel::new("fake-token".into(), vec!["*".into()]);
+        let id = ch
+            .send_draft(&SendMessage::new("draft", "123"))
+            .await
+            .unwrap();
+        assert!(id.is_none());
+    }
+
+    #[tokio::test]
+    async fn update_draft_rate_limit_short_circuits_network() {
+        let ch = TelegramChannel::new("fake-token".into(), vec!["*".into()])
+            .with_streaming(StreamMode::Partial, 60_000);
+        ch.last_draft_edit
+            .lock()
+            .insert("123".to_string(), std::time::Instant::now());
+
+        let result = ch.update_draft("123", "42", "delta text").await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn update_draft_utf8_truncation_is_safe_for_multibyte_text() {
+        let ch = TelegramChannel::new("fake-token".into(), vec!["*".into()])
+            .with_streaming(StreamMode::Partial, 0);
+        let long_emoji_text = "😀".repeat(TELEGRAM_MAX_MESSAGE_LENGTH + 20);
+
+        // Invalid message_id returns early after building display_text.
+        // This asserts truncation never panics on UTF-8 boundaries.
+        let result = ch
+            .update_draft("123", "not-a-number", &long_emoji_text)
+            .await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn finalize_draft_invalid_message_id_falls_back_to_chunk_send() {
+        let ch = TelegramChannel::new("fake-token".into(), vec!["*".into()])
+            .with_streaming(StreamMode::Partial, 0);
+        let long_text = "a".repeat(TELEGRAM_MAX_MESSAGE_LENGTH + 64);
+
+        // For oversized text + invalid draft message_id, finalize_draft should
+        // fall back to chunked send instead of returning early.
+        let result = ch.finalize_draft("123", "not-a-number", &long_text).await;
+        assert!(result.is_err());
+    }
+
     #[test]
     fn telegram_api_url() {
         let ch = TelegramChannel::new("123:ABC".into(), vec![]);
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 5c8d3d5..227789f 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -10,8 +10,7 @@ pub use schema::{
     PeripheralBoardConfig, PeripheralsConfig, QueryClassificationConfig, ReliabilityConfig,
     ResourceLimitsConfig, RuntimeConfig, SandboxBackend, SandboxConfig, SchedulerConfig,
     SecretsConfig, SecurityConfig, SlackConfig, StreamMode, TelegramConfig, TunnelConfig,
-    WebSearchConfig,
-    WebhookConfig,
+    WebSearchConfig, WebhookConfig,
 };
 
 #[cfg(test)]
diff --git a/src/config/schema.rs b/src/config/schema.rs
index db6d44a..e341b32 100644
--- a/src/config/schema.rs
+++ b/src/config/schema.rs
@@ -1449,9 +1449,6 @@ pub enum StreamMode {
     Off,
     /// Update a draft message with every flush interval.
     Partial,
-    /// Update a draft message in larger chunks (reserved for future use;
-    /// currently behaves the same as `Partial`).
-    Block,
 }
 
 fn default_draft_update_interval_ms() -> u64 {