1use std::collections::HashMap;
12use std::path::PathBuf;
13use std::sync::Arc;
14use std::time::{Duration, Instant};
15
16use async_trait::async_trait;
17use robit_agent::event::{FrontendMessage, SessionId};
18use robit_agent::frontend::Frontend;
19use robit_agent::storage::{self, resolve_db_path};
20use robit_agent::tool::ToolCallInfo;
21use robit_agent::{Agent, AgentError, SkillRegistry, ToolRegistry};
22use robit_ai::config::RobitConfig;
23use robit_ai::LlmClient;
24use rusqlite::Connection;
25use tokio::sync::{mpsc, Mutex};
26use uuid::Uuid;
27
28use crate::adapter::{ChatMessage, PlatformAdapter, PlatformCaps, PlatformEvent, SendResult, UploadResult};
29use crate::confirmer::{ConfirmKeywords, Confirmer};
30use crate::extensions::PlatformExtWrapper;
31use crate::frontend::{ChatbotFrontend, PlatformSender, PlatformExt};
32
33const CLEANUP_INTERVAL: Duration = Duration::from_secs(300); pub struct AgentHandle {
38 pub message_tx: mpsc::Sender<FrontendMessage>,
40 pub session_id: String,
41 pub last_active_at: Instant,
42 pub frontend: Arc<ChatbotFrontend>,
44}
45
46struct PlatformSenderBridge<T: PlatformAdapter> {
49 platform: Arc<T>,
50 caps: PlatformCaps,
51}
52
53#[async_trait]
54impl<T: PlatformAdapter> PlatformSender for PlatformSenderBridge<T> {
55 async fn send(&self, chat_id: &str, text: &str) -> robit_agent::error::Result<SendResult> {
56 self.platform.send_message(chat_id, text).await
57 }
58 async fn edit(&self, chat_id: &str, msg_id: &str, text: &str) -> robit_agent::error::Result<()> {
59 self.platform.edit_message(chat_id, msg_id, text).await
60 }
61 async fn upload_file(
62 &self,
63 chat_id: &str,
64 file_path: &str,
65 media_type: &str,
66 ) -> robit_agent::error::Result<UploadResult> {
67 self.platform.upload_file(chat_id, file_path, media_type).await
68 }
69 async fn send_media_message(
70 &self,
71 chat_id: &str,
72 file_url: &str,
73 file_name: &str,
74 media_type: &str,
75 ) -> robit_agent::error::Result<SendResult> {
76 self.platform
77 .send_media_message(chat_id, file_url, file_name, media_type)
78 .await
79 }
80 fn capabilities(&self) -> PlatformCaps {
81 self.caps.clone()
82 }
83}
84
85pub struct ChatbotManager<T: PlatformAdapter> {
87 platform: Arc<T>,
89 agents: Mutex<HashMap<String, AgentHandle>>,
91 db: Arc<Mutex<Connection>>,
93 config: RobitConfig,
94 working_dir: PathBuf,
95 llm_client: Arc<LlmClient>,
96 tool_registry: Arc<ToolRegistry>,
97 skill_registry: Arc<SkillRegistry>,
98 platform_sender: Arc<dyn PlatformSender>,
100 confirmer: Arc<Confirmer>,
102 auto_approve: bool,
103 context_window: Option<u64>,
104 session_timeout: Duration,
106}
107
108impl<T: PlatformAdapter> ChatbotManager<T> {
109 #[allow(clippy::too_many_arguments)]
115 pub fn new(
116 platform: Arc<T>,
117 config: RobitConfig,
118 working_dir: PathBuf,
119 llm_client: Arc<LlmClient>,
120 tool_registry: Arc<ToolRegistry>,
121 skill_registry: Arc<SkillRegistry>,
122 ) -> Result<Self, ManagerError> {
123 let caps = T::capabilities();
124 let platform_sender: Arc<dyn PlatformSender> = Arc::new(PlatformSenderBridge {
125 platform: Arc::clone(&platform),
126 caps: caps.clone(),
127 });
128
129 let bot = config.app.as_ref().and_then(|a| a.bot.as_ref());
131 let auto_approve = config
132 .app
133 .as_ref()
134 .and_then(|a| a.auto_approve)
135 .unwrap_or(false);
136 let confirm_timeout = Duration::from_secs(
137 bot.and_then(|b| b.confirm_timeout_secs).unwrap_or(60),
138 );
139 let session_timeout = Duration::from_secs(
140 bot.and_then(|b| b.session_timeout_minutes).unwrap_or(30) * 60,
141 );
142 let global_storage = config
143 .app
144 .as_ref()
145 .and_then(|a| a.global_storage)
146 .unwrap_or(false);
147 let context_window = llm_client.resolved().context_window;
148
149 let confirmer = match bot.and_then(|b| b.confirm_keywords.as_ref()) {
151 Some(kw) => Confirmer::with_keywords(
152 Arc::clone(&platform_sender),
153 confirm_timeout,
154 ConfirmKeywords {
155 approve: kw.approve.clone().unwrap_or_default(),
156 reject: kw.reject.clone().unwrap_or_default(),
157 },
158 ),
159 None => Confirmer::new(Arc::clone(&platform_sender), confirm_timeout),
160 };
161 let confirmer = Arc::new(confirmer);
162
163 let db_path = resolve_db_path(&working_dir, global_storage)?;
165 if let Some(parent) = db_path.parent() {
166 let _ = std::fs::create_dir_all(parent);
167 }
168 let conn = Connection::open(&db_path).map_err(ManagerError::DbOpen)?;
169 storage::init_db(&conn).map_err(ManagerError::DbInit)?;
170 let db = Arc::new(Mutex::new(conn));
171
172 Ok(Self {
173 platform,
174 agents: Mutex::new(HashMap::new()),
175 db,
176 config,
177 working_dir,
178 llm_client,
179 tool_registry,
180 skill_registry,
181 platform_sender,
182 confirmer,
183 auto_approve,
184 context_window,
185 session_timeout,
186 })
187 }
188
189 pub async fn run(
193 &self,
194 shutdown: Arc<tokio::sync::Notify>,
195 ) -> Result<(), AgentError> {
196 let cleanup_db = Arc::clone(&self.db);
198 let session_timeout = self.session_timeout;
199 let cleanup_shutdown = shutdown.clone();
200 tokio::spawn(async move {
201 cleanup_loop(cleanup_db, session_timeout, cleanup_shutdown).await;
202 });
203
204 loop {
205 tokio::select! {
206 event = self.platform.recv_event() => {
207 match event {
208 Ok(PlatformEvent::Message(msg)) => {
209 self.handle_message(msg).await;
210 }
211 Ok(PlatformEvent::Disconnected) => {
212 tracing::warn!("Platform disconnected");
213 return Ok(());
214 }
215 Ok(PlatformEvent::Other(v)) => {
216 tracing::debug!("Ignoring platform event: {}", v);
217 }
218 Err(e) => {
219 tracing::error!("Platform recv error: {}", e);
220 return Err(e);
221 }
222 }
223 }
224 _ = shutdown.notified() => {
225 tracing::info!("Shutdown signal received, stopping event loop...");
226 return Ok(());
227 }
228 }
229 }
230 }
231
232 async fn handle_message(&self, msg: ChatMessage) {
234 let chat_id = msg.sender.chat_id.clone();
235 let text = msg.text.trim().to_lowercase();
236
237 if self
239 .confirmer
240 .check_confirmation_response(&chat_id, &text)
241 .is_some()
242 {
243 return;
244 }
245
246 let trimmed_text = msg.text.trim();
248 if trimmed_text.eq_ignore_ascii_case("/clear") {
249 self.handle_clear_command(&chat_id).await;
250 return;
251 }
252 if trimmed_text.eq_ignore_ascii_case("/stop") {
253 self.handle_stop_command(&chat_id).await;
254 return;
255 }
256 if trimmed_text.eq_ignore_ascii_case("/new") {
257 self.handle_new_command(&chat_id).await;
258 return;
259 }
260 if trimmed_text.eq_ignore_ascii_case("/list") {
261 self.handle_list_command(&chat_id).await;
262 return;
263 }
264 if trimmed_text.to_lowercase().starts_with("/switch ") {
265 self.handle_switch_command(&chat_id, &trimmed_text["/switch ".len()..]).await;
266 return;
267 }
268 if trimmed_text.eq_ignore_ascii_case("/help") {
269 self.handle_help_command(&chat_id).await;
270 return;
271 }
272
273 let media_dir = self.working_dir.join("media");
275 for attachment in &msg.attachments {
276 if let Err(e) = robit_agent::media::download_media(
277 &attachment.url,
278 attachment.filename.as_deref(),
279 &media_dir,
280 )
281 .await
282 {
283 tracing::warn!("Failed to download media: {}", e);
284 }
285 }
286
287 let attachments: Vec<robit_agent::event::MediaAttachment> =
289 msg.attachments.into_iter().map(|a| a.into()).collect();
290
291 match self.get_or_create_session(&chat_id, &msg.text).await {
293 Ok((tx, frontend)) => {
294 frontend.save_user_message(&msg.text).await;
296
297 if let Err(e) = tx
298 .send(robit_agent::event::FrontendMessage::UserInput {
299 text: msg.text,
300 attachments,
301 })
302 .await
303 {
304 tracing::warn!("Failed to send user message to agent for {}: {}", chat_id, e);
305 }
306 }
307 Err(e) => {
308 tracing::error!("Failed to get/create session for {}: {}", chat_id, e);
309 let _ = self
310 .platform_sender
311 .send(&chat_id, &format!("❌ 内部错误,无法处理消息:{}", e))
312 .await;
313 }
314 }
315 }
316
317 async fn handle_clear_command(&self, chat_id: &str) {
319 match self.get_or_create_session(chat_id, "clear command").await {
321 Ok((tx, _)) => {
322 if let Err(e) = tx.send("/clear".into()).await {
323 tracing::warn!("Failed to send /clear to agent: {}", e);
324 let _ = self.platform_sender.send(chat_id, "❌ 清空失败").await;
325 }
326 }
327 Err(e) => {
328 tracing::error!("Failed to get session for /clear: {}", e);
329 let _ = self.platform_sender.send(chat_id, &format!("❌ 无法执行清空:{}", e)).await;
330 }
331 }
332 }
333
334 async fn handle_stop_command(&self, chat_id: &str) {
336 let agents = self.agents.lock().await;
337 if let Some(handle) = agents.get(chat_id) {
338 if let Err(e) = handle.message_tx.send(robit_agent::event::FrontendMessage::Cancel).await {
340 tracing::warn!("Failed to send Cancel to agent: {}", e);
341 let _ = self.platform_sender.send(chat_id, "❌ 停止失败").await;
342 return;
343 }
344 let _ = self.platform_sender.send(chat_id, "⏹️ 已发送停止信号").await;
345 } else {
346 let _ = self.platform_sender.send(chat_id, "ℹ️ 当前没有活动的会话").await;
347 }
348 }
349
350 async fn handle_new_command(&self, chat_id: &str) {
352 let mut agents = self.agents.lock().await;
353
354 if let Some(old_handle) = agents.remove(chat_id) {
356 let db = self.db.lock().await;
358 if let Err(e) = robit_agent::storage::delete_session(&db, &old_handle.session_id) {
359 tracing::warn!("Failed to deactivate old session: {}", e);
360 }
361 drop(db);
362
363 drop(old_handle);
365 }
366 drop(agents);
367
368 match self.get_or_create_session(chat_id, "新会话").await {
370 Ok((_, frontend)) => {
371 let msg = format!(
372 "✨ 已创建新会话\n会话ID: {}\n旧会话已归档,使用 /list 查看历史",
373 frontend.session_id
374 );
375 let _ = self.platform_sender.send(chat_id, &msg).await;
376 }
377 Err(e) => {
378 tracing::error!("Failed to create new session: {}", e);
379 let _ = self.platform_sender.send(chat_id, &format!("❌ 创建新会话失败:{}", e)).await;
380 }
381 }
382 }
383
384 async fn handle_list_command(&self, chat_id: &str) {
386 let db = self.db.lock().await;
387 match robit_agent::storage::list_all_sessions_by_chat_id(&db, chat_id) {
388 Ok(sessions) if sessions.is_empty() => {
389 let _ = self.platform_sender.send(chat_id, "ℹ️ 暂无历史会话").await;
390 }
391 Ok(sessions) => {
392 let mut list_text = String::from("📜 历史会话列表\n\n");
393 for (i, session) in sessions.iter().enumerate() {
394 let indicator = if i == 0 { "👉" } else { " " };
395 let current_mark = if i == 0 { " [当前]" } else { "" };
396 list_text.push_str(&format!(
397 "{} {}. {}{}\n",
398 indicator,
399 i + 1,
400 session.title,
401 current_mark
402 ));
403 list_text.push_str(&format!(
404 " ID: {} | 创建: {}\n\n",
405 session.id,
406 session.created_at
407 ));
408 }
409 list_text.push_str("💡 使用 /switch <序号> 切换到对应会话");
410 let _ = self.platform_sender.send(chat_id, &list_text).await;
411 }
412 Err(e) => {
413 tracing::error!("Failed to list sessions: {}", e);
414 let _ = self.platform_sender.send(chat_id, &format!("❌ 获取会话列表失败:{}", e)).await;
415 }
416 }
417 }
418
419 async fn handle_switch_command(&self, chat_id: &str, arg: &str) {
421 let trimmed_arg = arg.trim();
422
423 let session_index = match trimmed_arg.parse::<usize>() {
425 Ok(n) if n > 0 => n - 1, _ => {
427 let _ = self.platform_sender.send(chat_id, "❌ 请输入有效的会话序号,如 /switch 1").await;
428 return;
429 }
430 };
431
432 let db = self.db.lock().await;
434 let sessions = match robit_agent::storage::list_all_sessions_by_chat_id(&db, chat_id) {
435 Ok(s) => s,
436 Err(e) => {
437 tracing::error!("Failed to list sessions: {}", e);
438 let _ = self.platform_sender.send(chat_id, &format!("❌ 获取会话列表失败:{}", e)).await;
439 return;
440 }
441 };
442 drop(db);
443
444 if session_index >= sessions.len() {
446 let _ = self.platform_sender.send(
447 chat_id,
448 &format!("❌ 会话序号无效,共有 {} 个会话", sessions.len())
449 ).await;
450 return;
451 }
452
453 let target_session = &sessions[session_index];
454 let target_id = target_session.id.clone();
455
456 {
458 let agents = self.agents.lock().await;
459 if let Some(current) = agents.get(chat_id) {
460 if current.session_id == target_id {
461 let _ = self.platform_sender.send(chat_id, "ℹ️ 已经是当前会话").await;
462 return;
463 }
464 }
465 }
466
467 let db = self.db.lock().await;
469 if let Err(e) = robit_agent::storage::activate_session(&db, &target_id, chat_id) {
470 tracing::error!("Failed to activate session: {}", e);
471 let _ = self.platform_sender.send(chat_id, &format!("❌ 切换失败:{}", e)).await;
472 return;
473 }
474 drop(db);
475
476 let mut agents = self.agents.lock().await;
478 agents.remove(chat_id);
480 drop(agents);
481
482 match self.get_or_create_session(chat_id, "切换会话").await {
484 Ok((_, _frontend)) => {
485 let _ = self.platform_sender.send(
486 chat_id,
487 &format!("✅ 已切换到会话:{}", target_session.title)
488 ).await;
489 }
490 Err(e) => {
491 tracing::error!("Failed to create agent after switch: {}", e);
492 let _ = self.platform_sender.send(chat_id, &format!("❌ 会话加载失败:{}", e)).await;
493 }
494 }
495 }
496
497 async fn handle_help_command(&self, chat_id: &str) {
499 let help_text = r#"🤖 Robit 帮助
500
501可用指令:
502- /clear - 清空当前对话上下文(仅内存中)
503- /stop - 停止当前执行
504- /new - 创建新会话(旧会话归档)
505- /list - 列出所有历史会话
506- /switch <序号> - 切换到指定会话
507- /help - 显示此帮助
508
509提示:直接发送消息与机器人对话即可。"#;
510 let _ = self.platform_sender.send(chat_id, help_text).await;
511 }
512
513 async fn get_or_create_session(
515 &self,
516 chat_id: &str,
517 first_message: &str,
518 ) -> Result<(mpsc::Sender<FrontendMessage>, Arc<ChatbotFrontend>), AgentError> {
519 let mut agents = self.agents.lock().await;
520 if let Some(handle) = agents.get_mut(chat_id) {
521 handle.last_active_at = Instant::now();
522 tracing::debug!("get_or_create_session: found active agent in memory for chat_id={}, session_id={}", chat_id, handle.session_id);
523 return Ok((handle.message_tx.clone(), handle.frontend.clone()));
524 }
525 drop(agents);
526
527 let session_id = {
529 let db = self.db.lock().await;
530 match storage::find_session_by_chat_id(&db, chat_id)
531 .map_err(|e| AgentError::InternalError(format!("DB lookup failed: {}", e)))?
532 {
533 Some(info) => {
534 tracing::info!("get_or_create_session: found existing session in DB for chat_id={}, session_id={}, title={}", chat_id, info.id, info.title);
535 info.id
536 }
537 None => {
538 let id = Uuid::new_v4().to_string();
540 let title = generate_title(first_message);
541 let model = self
542 .config
543 .default_model
544 .clone()
545 .unwrap_or_else(|| self.llm_client.model().to_string());
546 tracing::info!("get_or_create_session: creating new session in DB for chat_id={}, session_id={}, title={}", chat_id, id, title);
547 storage::insert_session(&db, &id, Some(chat_id), &title, &model, "qq")
548 .map_err(|e| {
549 AgentError::InternalError(format!("DB insert failed: {}", e))
550 })?;
551 id
552 }
553 }
554 };
555
556 let (tx, frontend) = self.spawn_session_agent(chat_id, &session_id).await?;
557
558 let mut agents = self.agents.lock().await;
559 agents.insert(
560 chat_id.to_string(),
561 AgentHandle {
562 message_tx: tx.clone(),
563 session_id,
564 last_active_at: Instant::now(),
565 frontend: frontend.clone(),
566 },
567 );
568 Ok((tx, frontend))
569 }
570
571 async fn spawn_session_agent(
573 &self,
574 chat_id: &str,
575 session_id: &str,
576 ) -> Result<(mpsc::Sender<FrontendMessage>, Arc<ChatbotFrontend>), AgentError> {
577 tracing::info!("spawn_session_agent: chat_id={}, session_id={}", chat_id, session_id);
578
579 let frontend = Arc::new(ChatbotFrontend::new(
580 chat_id.to_string(),
581 session_id.to_string(),
582 Arc::clone(&self.platform_sender),
583 Arc::clone(&self.confirmer),
584 Arc::clone(&self.db),
585 self.auto_approve,
586 ));
587
588 let (message_tx, message_rx) = mpsc::channel::<FrontendMessage>(16);
589
590 tracing::debug!("spawn_session_agent: loading history messages from DB...");
592 let db = self.db.lock().await;
593 let history_messages = robit_agent::storage::load_chat_messages(&db, session_id)
594 .unwrap_or_default();
595 drop(db);
596
597 tracing::info!("spawn_session_agent: loaded {} history messages", history_messages.len());
598
599 let session_id_obj = SessionId::from(session_id.to_string());
601
602 tracing::debug!("spawn_session_agent: creating Agent with history...");
603 let agent = Agent::with_history(
604 Arc::clone(&self.llm_client),
605 Arc::clone(&self.tool_registry),
606 Arc::clone(&self.skill_registry),
607 Arc::clone(&frontend) as Arc<dyn Frontend>,
608 self.config.app.as_ref().and_then(|a| a.context.as_ref()),
609 self.context_window,
610 self.working_dir.clone(),
611 self.auto_approve,
612 {
613 let mut exts = HashMap::new();
614 let platform_ext: Arc<dyn PlatformExt> = frontend.clone();
615 exts.insert(
616 crate::extensions::keys::PLATFORM_EXT.to_string(),
617 PlatformExtWrapper::new(platform_ext),
618 );
619 exts
620 },
621 session_id_obj,
622 history_messages,
623 );
624
625 let sid = session_id.to_string();
626 let cid = chat_id.to_string();
627 tokio::spawn(async move {
628 agent.run(message_rx).await;
629 tracing::info!("Agent task ended for chat {} (session {})", cid, sid);
630 });
631
632 Ok((message_tx, frontend))
633 }
634
635 pub async fn active_session_count(&self) -> usize {
637 self.agents.lock().await.len()
638 }
639}
640
641#[derive(Debug, thiserror::Error)]
643pub enum ManagerError {
644 #[error("Failed to resolve DB path: {0}")]
645 DbPath(#[from] robit_agent::AgentError),
646 #[error("Failed to open database: {0}")]
647 DbOpen(#[from] rusqlite::Error),
648 #[error("Failed to initialize database: {0}")]
649 DbInit(rusqlite::Error),
650}
651
652fn generate_title(message: &str) -> String {
654 let trimmed = message.trim();
655 const MAX: usize = 30;
656 let chars: Vec<char> = trimmed.chars().take(MAX).collect();
657 let mut title: String = chars.into_iter().collect();
658 if trimmed.chars().count() > MAX {
659 title.push('…');
660 }
661 if title.is_empty() {
662 "QQ 会话".to_string()
663 } else {
664 title
665 }
666}
667
668async fn cleanup_loop(
674 _db: Arc<Mutex<Connection>>,
675 _timeout: Duration,
676 shutdown: Arc<tokio::sync::Notify>,
677) {
678 loop {
679 tokio::select! {
680 _ = tokio::time::sleep(CLEANUP_INTERVAL) => {
681 tracing::debug!("cleanup tick (no-op in MVP)");
682 }
683 _ = shutdown.notified() => {
684 tracing::debug!("cleanup loop received shutdown signal");
685 return;
686 }
687 }
688 }
689}
690
691#[allow(dead_code)]
693fn _tool_call_info_used(_i: &ToolCallInfo) {}
694
695#[cfg(test)]
696mod tests {
697 use super::*;
698 use crate::adapter::{ChatType, SenderInfo};
699 use std::collections::VecDeque;
700
701 #[allow(dead_code)]
703 struct MockPlatform {
704 events: Mutex<VecDeque<PlatformEvent>>,
705 sent: std::sync::Mutex<Vec<(String, String)>>,
706 }
707
708 #[allow(dead_code)]
709 impl MockPlatform {
710 fn new() -> Arc<Self> {
711 Arc::new(Self {
712 events: Mutex::new(VecDeque::new()),
713 sent: std::sync::Mutex::new(Vec::new()),
714 })
715 }
716
717 async fn push_message(&self, chat_id: &str, text: &str) {
718 self.events.lock().await.push_back(PlatformEvent::Message(ChatMessage {
719 text: text.to_string(),
720 sender: SenderInfo {
721 user_id: "u1".into(),
722 chat_id: chat_id.to_string(),
723 chat_type: ChatType::Group,
724 },
725 attachments: vec![],
726 }));
727 }
728
729 fn sent(&self) -> Vec<(String, String)> {
730 self.sent.lock().unwrap().clone()
731 }
732 }
733
734 #[async_trait]
735 impl PlatformAdapter for MockPlatform {
736 fn capabilities() -> PlatformCaps {
737 PlatformCaps::qq()
738 }
739 async fn send_message(&self, chat_id: &str, text: &str) -> robit_agent::error::Result<SendResult> {
740 self.sent
741 .lock()
742 .unwrap()
743 .push((chat_id.to_string(), text.to_string()));
744 Ok(SendResult { msg_id: "m1".into() })
745 }
746 async fn recv_event(&self) -> robit_agent::error::Result<PlatformEvent> {
747 loop {
749 if let Some(ev) = self.events.lock().await.pop_front() {
750 return Ok(ev);
751 }
752 tokio::time::sleep(Duration::from_millis(10)).await;
753 }
754 }
755 }
756
757 #[test]
758 fn generate_title_truncates_long_messages() {
759 let long = "x".repeat(100);
760 let title = generate_title(&long);
761 assert!(title.ends_with('…'));
762 assert!(title.chars().count() <= 31);
763 }
764
765 #[test]
766 fn generate_title_short_message() {
767 assert_eq!(generate_title("hello"), "hello");
768 }
769
770 #[test]
771 fn generate_title_empty_message() {
772 assert_eq!(generate_title(" "), "QQ 会话");
773 }
774
775 }