import logging import logging.handlers import queue as queue_module import signal import sys import threading import time from concurrent.futures import ThreadPoolExecutor from src.config import load_config, Config, UserConfig from src.summarizer import poll_accounts, ai_process, tg_send_and_mark, _mail_label from src.tg_bot import poll_telegram, send_thinking from src.database import init_db, is_uid_processed, mark_fetched, mark_processing, mark_sent, mark_failed # Thread-safe logging _log_queue = queue_module.Queue(-1) _queue_handler = logging.handlers.QueueHandler(_log_queue) _console_handler = logging.StreamHandler() _console_handler.setFormatter(logging.Formatter( "%(asctime)s [%(levelname)s] [%(threadName)s] %(name)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S", )) _log_listener = logging.handlers.QueueListener(_log_queue, _console_handler) root = logging.getLogger() root.addHandler(_queue_handler) root.setLevel(logging.INFO) _log_listener.start() logger = logging.getLogger("main") _running = True _shutdown_event = threading.Event() # Queue 1: raw emails → AI processor (user, acct_idx, mail, tg_msg_id) _email_queue: queue_module.Queue = queue_module.Queue() # Queue 2: AI summaries → TG sender (user, info_dict) _tg_queue: queue_module.Queue = queue_module.Queue() def _signal_handler(signum, frame): global _running logger.info("收到退出信号,正在停止...") _running = False _shutdown_event.set() def _email_poller(cfg: Config): logger.info("邮件轮询线程已启动") while _running: try: for user in cfg.users: if not _running: return for acct_idx, mail in poll_accounts(user): if not _running: return uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid) if not is_uid_processed(uid_str): mark_fetched(uid_str, mail.account_email, mail.folder, mail.subject, mail.sender) try: tg_msg_id = send_thinking(cfg.bot_token, user.chat_id, f"📥 *获取邮件中…*\n{_mail_label(mail)}") logger.info(" [%s] 获取邮件中消息已发送: msg_id=%d", user.chat_id, tg_msg_id) except Exception as e: logger.error(" [%s] 获取邮件中消息发送失败: %s", user.chat_id, e) tg_msg_id = 0 _email_queue.put((user, acct_idx, mail, tg_msg_id)) except Exception as e: logger.error(f"邮件轮询线程异常: {e}", exc_info=True) time.sleep(5) if _running: _shutdown_event.wait(cfg.polling_interval) def _ai_processor(cfg: Config): logger.info("AI 处理线程已启动 (并行 workers=10)") max_workers = 10 with ThreadPoolExecutor(max_workers=max_workers) as executor: pending_futures = {} while _running: while len(pending_futures) < max_workers: try: user, acct_idx, mail, tg_msg_id = _email_queue.get_nowait() except queue_module.Empty: break uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid) mark_processing(uid_str) future = executor.submit(ai_process, cfg, user, acct_idx, mail, tg_msg_id) pending_futures[future] = (user, acct_idx, mail) completed = [f for f in pending_futures if f.done()] for future in completed: user, acct_idx, mail = pending_futures.pop(future) try: info = future.result() if info: _tg_queue.put((user, info)) except Exception as e: logger.error(f"AI 处理邮件失败: {e}", exc_info=True) uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid) mark_failed(uid_str) time.sleep(0.1 if pending_futures else 0.5) def _tg_poller(cfg: Config): logger.info("TG 轮询线程已启动(按钮回调专用)") last_update_id = 0 while _running: try: last_update_id = poll_telegram(cfg.bot_token, cfg, last_update_id) except Exception as e: logger.error(f"TG 轮询错误: {e}", exc_info=True) if _running: _shutdown_event.wait(0.05) def _tg_sender(cfg: Config): logger.info("TG 发送线程已启动(邮件摘要专用)") while _running: try: user, info = _tg_queue.get_nowait() tg_send_and_mark(cfg, user, info) uid_str = info["uid"].decode("utf-8", errors="replace") if isinstance(info["uid"], bytes) else str(info["uid"]) mark_sent(uid_str) except queue_module.Empty: pass except Exception as e: logger.error(f"TG 发送失败: {e}", exc_info=True) try: uid_str = info["uid"].decode("utf-8", errors="replace") if isinstance(info["uid"], bytes) else str(info["uid"]) mark_failed(uid_str) except Exception: pass if _running: _shutdown_event.wait(0.2) def main(): global _running signal.signal(signal.SIGINT, _signal_handler) signal.signal(signal.SIGTERM, _signal_handler) cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.json" cfg = load_config(cfg_path) init_db() user_count = len(cfg.users) acct_count = sum(len(u.email_accounts) for u in cfg.users) logger.info(f"AI邮件摘要机器人已启动: {user_count} 用户, {acct_count} 邮箱, 轮询间隔: {cfg.polling_interval}s") threads = [ threading.Thread(target=_email_poller, args=(cfg,), daemon=True), threading.Thread(target=_ai_processor, args=(cfg,), daemon=True), threading.Thread(target=_tg_poller, args=(cfg,), daemon=True), threading.Thread(target=_tg_sender, args=(cfg,), daemon=True), ] for t in threads: t.start() logger.info("所有线程已启动,进入主循环") tick = 0 try: while _running: tick += 1 if tick % 30 == 0: logger.info("队列状态: email_queue=%d tg_queue=%d", _email_queue.qsize(), _tg_queue.qsize()) _shutdown_event.wait(1) except KeyboardInterrupt: _running = False _shutdown_event.set() for t in threads: t.join(timeout=5) _log_listener.stop() logger.info("机器人已停止") if __name__ == "__main__": main()