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        // Don't save empty messages - they cause problems with LLM API
111        if content.is_empty() {
112            tracing::debug!("Not saving empty assistant message to DB");
113            return;
114        }
115        match robit_agent::storage::insert_message(
116            &db,
117            &self.session_id,
118            "assistant",
119            content,
120            None,
121            None,
122            None,
123        ) {
124            Ok(_) => {
125                tracing::debug!("Saved assistant message to DB: session_id={}", self.session_id);
126                let _ = robit_agent::storage::touch_session(&db, &self.session_id);
127            }
128            Err(e) => tracing::warn!("Failed to save assistant message: {}", e),
129        }
130    }
131
132    /// Save a user message to the database (called from manager).
133    pub async fn save_user_message(&self, content: &str) {
134        let db = self.db.lock().await;
135        match robit_agent::storage::insert_message(
136            &db,
137            &self.session_id,
138            "user",
139            content,
140            None,
141            None,
142            None,
143        ) {
144            Ok(_) => {
145                tracing::debug!("Saved user message to DB: session_id={}", self.session_id);
146                let _ = robit_agent::storage::touch_session(&db, &self.session_id);
147            }
148            Err(e) => tracing::warn!("Failed to save user message: {}", e),
149        }
150    }
151
152    /// Append a delta to the buffer (no streaming send, just accumulate).
153    /// For QQ Bot, we send the full sanitized message at TurnComplete to
154    /// avoid duplicate messages and ensure proper Markdown handling.
155    async fn append_delta(&self, delta: &str) {
156        let mut buffer = self.buffer.lock().await;
157        buffer.push_str(delta);
158    }
159
160    /// Flush the buffer: take all accumulated text, sanitize it, send it, and save to DB.
161    /// This ensures Markdown is parsed as a whole and we only send once per turn.
162    async fn flush_buffer(&self) {
163        let mut buffer = self.buffer.lock().await;
164        if buffer.is_empty() {
165            return;
166        }
167        let text = std::mem::take(&mut *buffer);
168        drop(buffer);
169
170        let caps = self.platform_sender.capabilities();
171
172        let prepared = if caps.supports_markdown {
173            prepare_markdown_for_platform(&text, &caps.markdown_features)
174        } else {
175            text.clone()
176        };
177        let prepared = truncate_to_max(&prepared, caps.max_message_length);
178
179        // Just send once, no edit tricks - simple and reliable
180        let mut last_msg_id = self.last_msg_id.lock().await;
181        if caps.supports_edit && last_msg_id.is_some() {
182            // Edit if we already sent something this turn (e.g., progress hint)
183            let msg_id = last_msg_id.clone().unwrap();
184            if self.platform_sender.edit(&self.chat_id, &msg_id, &prepared).await.is_err() {
185                // Edit failed, fall back to send
186                if let Ok(res) = self.platform_sender.send(&self.chat_id, &prepared).await {
187                    *last_msg_id = Some(res.msg_id);
188                }
189            }
190        } else {
191            // No previous message this turn, just send
192            if let Ok(res) = self.platform_sender.send(&self.chat_id, &prepared).await {
193                *last_msg_id = Some(res.msg_id);
194            }
195        }
196
197        // Save the original (un-truncated) message to database
198        if !text.is_empty() {
199            self.save_assistant_message(&text).await;
200        }
201    }
202
203    /// Send a brief progress hint. Only sends once per turn to avoid spam.
204    async fn send_progress_hint(&self, tool_name: &str) {
205        let mut sent = self.progress_hint_sent.lock().await;
206        if *sent {
207            return;
208        }
209        *sent = true;
210        drop(sent);
211
212        let hint = match tool_name {
213            "bash" => "🔧 正在执行命令...".to_string(),
214            "read" => "📖 正在读取文件...".to_string(),
215            "write" => "✏️ 正在写入文件...".to_string(),
216            "edit" => "✏️ 正在编辑文件...".to_string(),
217            "grep" => "🔍 正在搜索...".to_string(),
218            "find" => "🔍 正在查找...".to_string(),
219            _ => "🔧 正在处理...".to_string(),
220        };
221        if let Ok(res) = self.platform_sender.send(&self.chat_id, &hint).await {
222            *self.last_msg_id.lock().await = Some(res.msg_id);
223        }
224    }
225}
226
227#[async_trait]
228impl Frontend for ChatbotFrontend {
229    async fn on_event(&self, event: AgentEvent) -> Result<()> {
230        match event {
231            AgentEvent::TextDelta(delta) => {
232                self.append_delta(&delta).await;
233            }
234            AgentEvent::ToolCallRequested { name, .. } => {
235                // Flush any buffered text before showing tool progress.
236                self.flush_buffer().await;
237                // In auto-approve mode, send a progress hint so the user knows
238                // the bot is working. In manual mode, the Confirmer already
239                // sends a confirm prompt — no extra hint needed.
240                if self.auto_approve {
241                    self.send_progress_hint(&name).await;
242                }
243            }
244            AgentEvent::ToolCallResult { .. } => {
245                // Silent: tool outputs are internal; the user only sees the
246                // final text reply. Any progress hint is replaced by the reply
247                // on TurnComplete.
248            }
249            AgentEvent::TurnComplete => {
250                self.flush_buffer().await;
251                // Reset per-turn state.
252                *self.progress_hint_sent.lock().await = false;
253                *self.last_msg_id.lock().await = None;
254            }
255            AgentEvent::Error(e) => {
256                self.flush_buffer().await;
257                let msg = format!("❌ Error: {}", e);
258                let _ = self.platform_sender.send(&self.chat_id, &msg).await;
259            }
260            AgentEvent::SkillTriggered { .. } => {
261                // Silent: skill trigger is internal; the skill's own output
262                // arrives as TextDelta events.
263            }
264        }
265        Ok(())
266    }
267
268    async fn request_tool_confirmation(&self, info: &ToolCallInfo) -> Result<bool> {
269        self.confirmer
270            .request(&self.chat_id, info, self.auto_approve)
271            .await
272    }
273}
274
275#[async_trait]
276impl PlatformExt for ChatbotFrontend {
277    async fn upload_file(&self, file_path: &str, media_type: &str) -> Result<UploadResult> {
278        self.platform_sender
279            .upload_file(&self.chat_id, file_path, media_type)
280            .await
281    }
282
283    async fn send_media_message(
284        &self,
285        file_url: &str,
286        file_name: &str,
287        media_type: &str,
288    ) -> Result<SendResult> {
289        self.platform_sender
290            .send_media_message(&self.chat_id, file_url, file_name, media_type)
291            .await
292    }
293}
294
295/// Truncate `text` to `max` characters, appending an ellipsis if cut.
296fn truncate_to_max(text: &str, max: usize) -> String {
297    if max == 0 {
298        return text.to_string();
299    }
300    if text.chars().count() <= max {
301        return text.to_string();
302    }
303    let mut out: String = text.chars().take(max.saturating_sub(1)).collect();
304    out.push('…');
305    out
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311    use crate::adapter::MarkdownFeatures;
312    use std::sync::Mutex as StdMutex;
313
314    struct MockSender {
315        sent: StdMutex<Vec<(String, String)>>,
316        edits: StdMutex<Vec<(String, String, String)>>,
317        caps: PlatformCaps,
318    }
319
320    impl MockSender {
321        fn new_with_edit(edit: bool) -> Arc<Self> {
322            Arc::new(Self {
323                sent: StdMutex::new(Vec::new()),
324                edits: StdMutex::new(Vec::new()),
325                caps: PlatformCaps {
326                    supports_edit: edit,
327                    returns_msg_id: true,
328                    supports_markdown: true,
329                    markdown_features: MarkdownFeatures::qq(),
330                    max_message_length: 2000,
331                    supports_images: true,
332                    supports_files: true,
333                    max_upload_size: 20 * 1024 * 1024,
334                },
335            })
336        }
337        fn sent_texts(&self) -> Vec<String> {
338            self.sent.lock().unwrap().iter().map(|(_, t)| t.clone()).collect()
339        }
340    }
341
342    #[async_trait]
343    impl PlatformSender for MockSender {
344        async fn send(&self, chat_id: &str, text: &str) -> Result<SendResult> {
345            self.sent
346                .lock()
347                .unwrap()
348                .push((chat_id.to_string(), text.to_string()));
349            Ok(SendResult {
350                msg_id: format!("msg-{}", self.sent.lock().unwrap().len()),
351            })
352        }
353        async fn edit(&self, chat_id: &str, msg_id: &str, text: &str) -> Result<()> {
354            self.edits.lock().unwrap().push((
355                chat_id.to_string(),
356                msg_id.to_string(),
357                text.to_string(),
358            ));
359            Ok(())
360        }
361        async fn upload_file(&self, _chat_id: &str, _file_path: &str, _media_type: &str) -> Result<UploadResult> {
362            Ok(UploadResult {
363                file_id: "mock-file-id".into(),
364                url: "/mock/upload.png".into(),
365            })
366        }
367        async fn send_media_message(
368            &self,
369            chat_id: &str,
370            _file_url: &str,
371            file_name: &str,
372            media_type: &str,
373        ) -> Result<SendResult> {
374            self.sent
375                .lock()
376                .unwrap()
377                .push((chat_id.to_string(), format!("[media:{}] {}", media_type, file_name)));
378            Ok(SendResult {
379                msg_id: format!("msg-{}", self.sent.lock().unwrap().len()),
380            })
381        }
382        fn capabilities(&self) -> PlatformCaps {
383            self.caps.clone()
384        }
385    }
386
387    async fn make_frontend(sender: Arc<dyn PlatformSender>, auto_approve: bool) -> ChatbotFrontend {
388        let confirmer = Arc::new(Confirmer::new(sender.clone(), std::time::Duration::from_secs(60)));
389        // 创建一个内存中的 SQLite 连接用于测试
390        let db = Arc::new(Mutex::new(rusqlite::Connection::open_in_memory().unwrap()));
391        // 初始化数据库 schema
392        {
393            let db = db.lock().await;
394            robit_agent::storage::init_db(&db).unwrap();
395        }
396        ChatbotFrontend::new(
397            "group:1".to_string(),
398            "test-session-1".to_string(),
399            sender,
400            confirmer,
401            db,
402            auto_approve
403        )
404    }
405
406    #[tokio::test]
407    async fn textdelta_accumulates_until_turn_complete() {
408        let sender = MockSender::new_with_edit(false);
409        let fe = make_frontend(sender.clone(), false).await;
410        // Text deltas only accumulate, nothing is sent until TurnComplete.
411        fe.on_event(AgentEvent::TextDelta("你好".to_string())).await.unwrap();
412        fe.on_event(AgentEvent::TextDelta("世界".to_string())).await.unwrap();
413        assert!(sender.sent_texts().is_empty());
414        // Nothing sent until TurnComplete.
415        fe.on_event(AgentEvent::TurnComplete).await.unwrap();
416        assert_eq!(sender.sent_texts().len(), 1);
417        assert!(sender.sent_texts()[0].contains("你好世界"));
418    }
419
420    #[tokio::test]
421    async fn turn_complete_flushes_accumulated_text() {
422        let sender = MockSender::new_with_edit(false);
423        let fe = make_frontend(sender.clone(), false).await;
424        fe.on_event(AgentEvent::TextDelta("一段未被刷新的文本".to_string())).await.unwrap();
425        assert!(sender.sent_texts().is_empty());
426        fe.on_event(AgentEvent::TurnComplete).await.unwrap();
427        assert_eq!(sender.sent_texts().len(), 1);
428        assert!(sender.sent_texts()[0].contains("一段未被刷新的文本"));
429    }
430
431    #[tokio::test]
432    async fn progress_hint_rate_limited_per_turn() {
433        let sender = MockSender::new_with_edit(false);
434        let fe = make_frontend(sender.clone(), true).await; // auto_approve
435        // Two tool calls in one turn — only one hint should be sent.
436        fe.on_event(AgentEvent::ToolCallRequested {
437            tool_call_id: "tc1".into(),
438            name: "bash".into(),
439            arguments: "{}".into(),
440        }).await.unwrap();
441        fe.on_event(AgentEvent::ToolCallRequested {
442            tool_call_id: "tc2".into(),
443            name: "read".into(),
444            arguments: "{}".into(),
445        }).await.unwrap();
446        let hints: Vec<_> = sender
447            .sent_texts()
448            .into_iter()
449            .filter(|t| t.contains("正在"))
450            .collect();
451        assert_eq!(hints.len(), 1);
452    }
453
454    #[tokio::test]
455    async fn progress_hint_resets_on_turn_complete() {
456        let sender = MockSender::new_with_edit(false);
457        let fe = make_frontend(sender.clone(), true).await;
458        fe.on_event(AgentEvent::ToolCallRequested {
459            tool_call_id: "tc1".into(),
460            name: "bash".into(),
461            arguments: "{}".into(),
462        }).await.unwrap();
463        fe.on_event(AgentEvent::TurnComplete).await.unwrap();
464        // After TurnComplete, a new tool call should send another hint.
465        fe.on_event(AgentEvent::ToolCallRequested {
466            tool_call_id: "tc2".into(),
467            name: "bash".into(),
468            arguments: "{}".into(),
469        }).await.unwrap();
470        let hints: Vec<_> = sender
471            .sent_texts()
472            .into_iter()
473            .filter(|t| t.contains("正在"))
474            .collect();
475        assert_eq!(hints.len(), 2);
476    }
477
478    #[tokio::test]
479    async fn no_hint_in_manual_mode() {
480        let sender = MockSender::new_with_edit(false);
481        let fe = make_frontend(sender.clone(), false).await; // manual confirm
482        fe.on_event(AgentEvent::ToolCallRequested {
483            tool_call_id: "tc1".into(),
484            name: "bash".into(),
485            arguments: "{}".into(),
486        }).await.unwrap();
487        // No progress hint in manual mode (Confirmer sends the prompt instead).
488        assert!(!sender.sent_texts().iter().any(|t| t.contains("正在")));
489    }
490
491    #[tokio::test]
492    async fn error_sends_error_message() {
493        let sender = MockSender::new_with_edit(false);
494        let fe = make_frontend(sender.clone(), false).await;
495        fe.on_event(AgentEvent::Error(robit_agent::AgentError::ToolError("boom".into())))
496            .await
497            .unwrap();
498        assert!(sender.sent_texts().iter().any(|t| t.contains("Error") && t.contains("boom")));
499    }
500
501}
robit_chatbot/frontend.rs

robit_chatbot/
frontend.rs
