robit_chatbot/
frontend.rs

1//! Per-session Frontend implementation for Bot platforms.
2//!
3//! [`ChatbotFrontend`] implements [`robit_agent::frontend::Frontend`] for a
4//! single chat (group or private). It buffers streaming `TextDelta` events and
5//! flushes them in natural-boundary segments (so Markdown and code blocks
6//! aren't cut mid-construct), sends a rate-limited progress hint when tools
7//! run, and delegates tool confirmation to the shared [`Confirmer`].
8//!
9//! On platforms with edit support, the first message is sent then edited in
10//! place to create a "growing message" effect; otherwise segments are sent as
11//! separate messages.
12
13use std::sync::Arc;
14
15use async_trait::async_trait;
16use robit_agent::error::Result;
17use robit_agent::event::AgentEvent;
18use robit_agent::frontend::Frontend;
19use robit_agent::tool::ToolCallInfo;
20use tokio::sync::Mutex;
21
22use crate::adapter::{PlatformCaps, SendResult, UploadResult};
23use crate::confirmer::Confirmer;
24use crate::markdown::prepare_markdown_for_platform;
25/// Abstracted message sending capability (platform-agnostic).
26///
27/// `ChatbotFrontend` talks to the platform through this trait rather than
28/// `PlatformAdapter` directly, so the manager can supply a bridge that wraps
29/// the concrete adapter.
30#[async_trait]
31pub trait PlatformSender: Send + Sync {
32    /// Send a text message to a chat; returns the platform message ID.
33    async fn send(&self, chat_id: &str, text: &str) -> Result<SendResult>;
34    /// Edit a previously-sent message in place.
35    async fn edit(&self, chat_id: &str, msg_id: &str, text: &str) -> Result<()>;
36    /// Upload a file to the platform. Returns the platform file URL/ID.
37    async fn upload_file(&self, chat_id: &str, file_path: &str, media_type: &str) -> Result<UploadResult>;
38    /// Send a media message (image/file) to a chat.
39    async fn send_media_message(&self, chat_id: &str, file_url: &str, file_name: &str, media_type: &str) -> Result<SendResult>;
40    /// Platform capabilities (drives streaming strategy).
41    fn capabilities(&self) -> PlatformCaps;
42}
43
44/// Platform extension for file/media operations.
45///
46/// Exposed to tools via `ToolContext.extensions` under key `"chatbot.platform_ext"`.
47/// `ChatbotFrontend` implements this by delegating to its `PlatformSender`,
48/// using the frontend's own `chat_id` internally.
49#[async_trait]
50pub trait PlatformExt: Send + Sync {
51    /// Upload a local file to the platform. Returns the platform identifier.
52    async fn upload_file(&self, file_path: &str, media_type: &str) -> Result<UploadResult>;
53    /// Send an already-uploaded media file to the chat.
54    async fn send_media_message(&self, file_url: &str, file_name: &str, media_type: &str) -> Result<SendResult>;
55}
56
57/// Per-session Frontend trait implementation for Bot platforms.
58///
59/// Each chat (group or private) gets its own `ChatbotFrontend` instance.
60/// `TextDelta` events are buffered and flushed once at TurnComplete with
61/// full Markdown sanitization to ensure clean output on platforms like QQ.
62pub struct ChatbotFrontend {
63    /// The chat this frontend belongs to (`group:{id}` or `private:{id}`).
64    pub chat_id: String,
65    /// The session ID for database storage.
66    pub session_id: String,
67    /// Platform message sender (shared across all frontends).
68    pub platform_sender: Arc<dyn PlatformSender>,
69    /// Tool confirmation coordinator (shared).
70    pub confirmer: Arc<Confirmer>,
71    /// Database connection for persisting messages.
72    pub db: Arc<Mutex<rusqlite::Connection>>,
73    /// Buffer to accumulate text until TurnComplete.
74    pub buffer: Mutex<String>,
75    /// ID of the last message sent (for edit-based updates, e.g., replacing a
76    /// progress hint with the actual response).
77    pub last_msg_id: Mutex<Option<String>>,
78    /// Whether a progress hint has already been sent this turn (rate limit).
79    pub progress_hint_sent: Mutex<bool>,
80    /// Auto-approve all tool calls.
81    pub auto_approve: bool,
82}
83
84impl ChatbotFrontend {
85    /// Create a new `ChatbotFrontend` for `chat_id`.
86    pub fn new(
87        chat_id: String,
88        session_id: String,
89        platform_sender: Arc<dyn PlatformSender>,
90        confirmer: Arc<Confirmer>,
91        db: Arc<Mutex<rusqlite::Connection>>,
92        auto_approve: bool,
93    ) -> Self {
94        Self {
95            chat_id,
96            session_id,
97            platform_sender,
98            confirmer,
99            db,
100            buffer: Mutex::new(String::new()),
101            last_msg_id: Mutex::new(None),
102            progress_hint_sent: Mutex::new(false),
103            auto_approve,
104        }
105    }
106
107    /// Save an assistant message to the database.
108    async fn save_assistant_message(&self, content: &str) {
109        let db = self.db.lock().await;
110        match robit_agent::storage::insert_message(
111            &db,
112            &self.session_id,
113            "assistant",
114            content,
115            None,
116            None,
117            None,
118        ) {
119            Ok(_) => {
120                tracing::debug!("Saved assistant message to DB: session_id={}", self.session_id);
121                let _ = robit_agent::storage::touch_session(&db, &self.session_id);
122            }
123            Err(e) => tracing::warn!("Failed to save assistant message: {}", e),
124        }
125    }
126
127    /// Save a user message to the database (called from manager).
128    pub async fn save_user_message(&self, content: &str) {
129        let db = self.db.lock().await;
130        match robit_agent::storage::insert_message(
131            &db,
132            &self.session_id,
133            "user",
134            content,
135            None,
136            None,
137            None,
138        ) {
139            Ok(_) => {
140                tracing::debug!("Saved user message to DB: session_id={}", self.session_id);
141                let _ = robit_agent::storage::touch_session(&db, &self.session_id);
142            }
143            Err(e) => tracing::warn!("Failed to save user message: {}", e),
144        }
145    }
146
147    /// Append a delta to the buffer (no streaming send, just accumulate).
148    /// For QQ Bot, we send the full sanitized message at TurnComplete to
149    /// avoid duplicate messages and ensure proper Markdown handling.
150    async fn append_delta(&self, delta: &str) {
151        let mut buffer = self.buffer.lock().await;
152        buffer.push_str(delta);
153    }
154
155    /// Flush the buffer: take all accumulated text, sanitize it, send it, and save to DB.
156    /// This ensures Markdown is parsed as a whole and we only send once per turn.
157    async fn flush_buffer(&self) {
158        let mut buffer = self.buffer.lock().await;
159        if buffer.is_empty() {
160            return;
161        }
162        let text = std::mem::take(&mut *buffer);
163        drop(buffer);
164
165        let caps = self.platform_sender.capabilities();
166
167        let prepared = if caps.supports_markdown {
168            prepare_markdown_for_platform(&text, &caps.markdown_features)
169        } else {
170            text.clone()
171        };
172        let prepared = truncate_to_max(&prepared, caps.max_message_length);
173
174        // Just send once, no edit tricks - simple and reliable
175        let mut last_msg_id = self.last_msg_id.lock().await;
176        if caps.supports_edit && last_msg_id.is_some() {
177            // Edit if we already sent something this turn (e.g., progress hint)
178            let msg_id = last_msg_id.clone().unwrap();
179            if self.platform_sender.edit(&self.chat_id, &msg_id, &prepared).await.is_err() {
180                // Edit failed, fall back to send
181                if let Ok(res) = self.platform_sender.send(&self.chat_id, &prepared).await {
182                    *last_msg_id = Some(res.msg_id);
183                }
184            }
185        } else {
186            // No previous message this turn, just send
187            if let Ok(res) = self.platform_sender.send(&self.chat_id, &prepared).await {
188                *last_msg_id = Some(res.msg_id);
189            }
190        }
191
192        // Save the original (un-truncated) message to database
193        if !text.is_empty() {
194            self.save_assistant_message(&text).await;
195        }
196    }
197
198    /// Send a brief progress hint. Only sends once per turn to avoid spam.
199    async fn send_progress_hint(&self, tool_name: &str) {
200        let mut sent = self.progress_hint_sent.lock().await;
201        if *sent {
202            return;
203        }
204        *sent = true;
205        drop(sent);
206
207        let hint = match tool_name {
208            "bash" => "🔧 正在执行命令...".to_string(),
209            "read" => "📖 正在读取文件...".to_string(),
210            "write" => "✏️ 正在写入文件...".to_string(),
211            "edit" => "✏️ 正在编辑文件...".to_string(),
212            "grep" => "🔍 正在搜索...".to_string(),
213            "find" => "🔍 正在查找...".to_string(),
214            _ => "🔧 正在处理...".to_string(),
215        };
216        if let Ok(res) = self.platform_sender.send(&self.chat_id, &hint).await {
217            *self.last_msg_id.lock().await = Some(res.msg_id);
218        }
219    }
220}
221
222#[async_trait]
223impl Frontend for ChatbotFrontend {
224    async fn on_event(&self, event: AgentEvent) -> Result<()> {
225        match event {
226            AgentEvent::TextDelta(delta) => {
227                self.append_delta(&delta).await;
228            }
229            AgentEvent::ToolCallRequested { name, .. } => {
230                // Flush any buffered text before showing tool progress.
231                self.flush_buffer().await;
232                // In auto-approve mode, send a progress hint so the user knows
233                // the bot is working. In manual mode, the Confirmer already
234                // sends a confirm prompt — no extra hint needed.
235                if self.auto_approve {
236                    self.send_progress_hint(&name).await;
237                }
238            }
239            AgentEvent::ToolCallResult { .. } => {
240                // Silent: tool outputs are internal; the user only sees the
241                // final text reply. Any progress hint is replaced by the reply
242                // on TurnComplete.
243            }
244            AgentEvent::TurnComplete => {
245                self.flush_buffer().await;
246                // Reset per-turn state.
247                *self.progress_hint_sent.lock().await = false;
248                *self.last_msg_id.lock().await = None;
249            }
250            AgentEvent::Error(e) => {
251                self.flush_buffer().await;
252                let msg = format!("❌ Error: {}", e);
253                let _ = self.platform_sender.send(&self.chat_id, &msg).await;
254            }
255            AgentEvent::SkillTriggered { .. } => {
256                // Silent: skill trigger is internal; the skill's own output
257                // arrives as TextDelta events.
258            }
259        }
260        Ok(())
261    }
262
263    async fn request_tool_confirmation(&self, info: &ToolCallInfo) -> Result<bool> {
264        self.confirmer
265            .request(&self.chat_id, info, self.auto_approve)
266            .await
267    }
268}
269
270#[async_trait]
271impl PlatformExt for ChatbotFrontend {
272    async fn upload_file(&self, file_path: &str, media_type: &str) -> Result<UploadResult> {
273        self.platform_sender
274            .upload_file(&self.chat_id, file_path, media_type)
275            .await
276    }
277
278    async fn send_media_message(
279        &self,
280        file_url: &str,
281        file_name: &str,
282        media_type: &str,
283    ) -> Result<SendResult> {
284        self.platform_sender
285            .send_media_message(&self.chat_id, file_url, file_name, media_type)
286            .await
287    }
288}
289
290/// Truncate `text` to `max` characters, appending an ellipsis if cut.
291fn truncate_to_max(text: &str, max: usize) -> String {
292    if max == 0 {
293        return text.to_string();
294    }
295    if text.chars().count() <= max {
296        return text.to_string();
297    }
298    let mut out: String = text.chars().take(max.saturating_sub(1)).collect();
299    out.push('…');
300    out
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306    use crate::adapter::MarkdownFeatures;
307    use std::sync::Mutex as StdMutex;
308
309    struct MockSender {
310        sent: StdMutex<Vec<(String, String)>>,
311        edits: StdMutex<Vec<(String, String, String)>>,
312        caps: PlatformCaps,
313    }
314
315    impl MockSender {
316        fn new_with_edit(edit: bool) -> Arc<Self> {
317            Arc::new(Self {
318                sent: StdMutex::new(Vec::new()),
319                edits: StdMutex::new(Vec::new()),
320                caps: PlatformCaps {
321                    supports_edit: edit,
322                    returns_msg_id: true,
323                    supports_markdown: true,
324                    markdown_features: MarkdownFeatures::qq(),
325                    max_message_length: 2000,
326                    supports_images: true,
327                    supports_files: true,
328                    max_upload_size: 20 * 1024 * 1024,
329                },
330            })
331        }
332        fn sent_texts(&self) -> Vec<String> {
333            self.sent.lock().unwrap().iter().map(|(_, t)| t.clone()).collect()
334        }
335    }
336
337    #[async_trait]
338    impl PlatformSender for MockSender {
339        async fn send(&self, chat_id: &str, text: &str) -> Result<SendResult> {
340            self.sent
341                .lock()
342                .unwrap()
343                .push((chat_id.to_string(), text.to_string()));
344            Ok(SendResult {
345                msg_id: format!("msg-{}", self.sent.lock().unwrap().len()),
346            })
347        }
348        async fn edit(&self, chat_id: &str, msg_id: &str, text: &str) -> Result<()> {
349            self.edits.lock().unwrap().push((
350                chat_id.to_string(),
351                msg_id.to_string(),
352                text.to_string(),
353            ));
354            Ok(())
355        }
356        async fn upload_file(&self, _chat_id: &str, _file_path: &str, _media_type: &str) -> Result<UploadResult> {
357            Ok(UploadResult {
358                file_id: "mock-file-id".into(),
359                url: "/mock/upload.png".into(),
360            })
361        }
362        async fn send_media_message(
363            &self,
364            chat_id: &str,
365            _file_url: &str,
366            file_name: &str,
367            media_type: &str,
368        ) -> Result<SendResult> {
369            self.sent
370                .lock()
371                .unwrap()
372                .push((chat_id.to_string(), format!("[media:{}] {}", media_type, file_name)));
373            Ok(SendResult {
374                msg_id: format!("msg-{}", self.sent.lock().unwrap().len()),
375            })
376        }
377        fn capabilities(&self) -> PlatformCaps {
378            self.caps.clone()
379        }
380    }
381
382    fn make_frontend(sender: Arc<dyn PlatformSender>, auto_approve: bool) -> ChatbotFrontend {
383        let confirmer = Arc::new(Confirmer::new(sender.clone(), std::time::Duration::from_secs(60)));
384        // 创建一个内存中的 SQLite 连接用于测试
385        let db = Arc::new(Mutex::new(rusqlite::Connection::open_in_memory().unwrap()));
386        // 初始化数据库 schema
387        {
388            let db = db.blocking_lock();
389            robit_agent::storage::init_db(&db).unwrap();
390        }
391        ChatbotFrontend::new(
392            "group:1".to_string(),
393            "test-session-1".to_string(),
394            sender,
395            confirmer,
396            db,
397            auto_approve
398        )
399    }
400
401    #[tokio::test]
402    async fn textdelta_accumulates_until_turn_complete() {
403        let sender = MockSender::new_with_edit(false);
404        let fe = make_frontend(sender.clone(), false);
405        // Text deltas only accumulate, nothing is sent until TurnComplete.
406        fe.on_event(AgentEvent::TextDelta("你好".to_string())).await.unwrap();
407        fe.on_event(AgentEvent::TextDelta("世界".to_string())).await.unwrap();
408        assert!(sender.sent_texts().is_empty());
409        // Nothing sent until TurnComplete.
410        fe.on_event(AgentEvent::TurnComplete).await.unwrap();
411        assert_eq!(sender.sent_texts().len(), 1);
412        assert!(sender.sent_texts()[0].contains("你好世界"));
413    }
414
415    #[tokio::test]
416    async fn turn_complete_flushes_accumulated_text() {
417        let sender = MockSender::new_with_edit(false);
418        let fe = make_frontend(sender.clone(), false);
419        fe.on_event(AgentEvent::TextDelta("一段未被刷新的文本".to_string())).await.unwrap();
420        assert!(sender.sent_texts().is_empty());
421        fe.on_event(AgentEvent::TurnComplete).await.unwrap();
422        assert_eq!(sender.sent_texts().len(), 1);
423        assert!(sender.sent_texts()[0].contains("一段未被刷新的文本"));
424    }
425
426    #[tokio::test]
427    async fn progress_hint_rate_limited_per_turn() {
428        let sender = MockSender::new_with_edit(false);
429        let fe = make_frontend(sender.clone(), true); // auto_approve
430        // Two tool calls in one turn — only one hint should be sent.
431        fe.on_event(AgentEvent::ToolCallRequested {
432            tool_call_id: "tc1".into(),
433            name: "bash".into(),
434            arguments: "{}".into(),
435        }).await.unwrap();
436        fe.on_event(AgentEvent::ToolCallRequested {
437            tool_call_id: "tc2".into(),
438            name: "read".into(),
439            arguments: "{}".into(),
440        }).await.unwrap();
441        let hints: Vec<_> = sender
442            .sent_texts()
443            .into_iter()
444            .filter(|t| t.contains("正在"))
445            .collect();
446        assert_eq!(hints.len(), 1);
447    }
448
449    #[tokio::test]
450    async fn progress_hint_resets_on_turn_complete() {
451        let sender = MockSender::new_with_edit(false);
452        let fe = make_frontend(sender.clone(), true);
453        fe.on_event(AgentEvent::ToolCallRequested {
454            tool_call_id: "tc1".into(),
455            name: "bash".into(),
456            arguments: "{}".into(),
457        }).await.unwrap();
458        fe.on_event(AgentEvent::TurnComplete).await.unwrap();
459        // After TurnComplete, a new tool call should send another hint.
460        fe.on_event(AgentEvent::ToolCallRequested {
461            tool_call_id: "tc2".into(),
462            name: "bash".into(),
463            arguments: "{}".into(),
464        }).await.unwrap();
465        let hints: Vec<_> = sender
466            .sent_texts()
467            .into_iter()
468            .filter(|t| t.contains("正在"))
469            .collect();
470        assert_eq!(hints.len(), 2);
471    }
472
473    #[tokio::test]
474    async fn no_hint_in_manual_mode() {
475        let sender = MockSender::new_with_edit(false);
476        let fe = make_frontend(sender.clone(), false); // manual confirm
477        fe.on_event(AgentEvent::ToolCallRequested {
478            tool_call_id: "tc1".into(),
479            name: "bash".into(),
480            arguments: "{}".into(),
481        }).await.unwrap();
482        // No progress hint in manual mode (Confirmer sends the prompt instead).
483        assert!(!sender.sent_texts().iter().any(|t| t.contains("正在")));
484    }
485
486    #[tokio::test]
487    async fn error_sends_error_message() {
488        let sender = MockSender::new_with_edit(false);
489        let fe = make_frontend(sender.clone(), false);
490        fe.on_event(AgentEvent::Error(robit_agent::AgentError::ToolError("boom".into())))
491            .await
492            .unwrap();
493        assert!(sender.sent_texts().iter().any(|t| t.contains("Error") && t.contains("boom")));
494    }
495
496}
robit_chatbot/frontend.rs

robit_chatbot/
frontend.rs
