refactor: replace all in-memory state with SQLite database

main.py

@@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor
 from src.config import load_config
 from src.summarizer import poll_accounts, ai_process, tg_send_and_mark
 from src.tg_bot import poll_telegram, send_thinking
+from src.database import init_db, is_uids_processed, mark_fetched, mark_processing, mark_sent, mark_failed
 
 # Thread-safe logging: all log records go through a single QueueListener thread
 _log_queue = queue_module.Queue(-1)
@@ -29,9 +30,6 @@ logger = logging.getLogger("main")
 
 _running = True
 _shutdown_event = threading.Event()
-_pending_lock = threading.Lock()
-_pending_uids: set[bytes] = set()
-_processing_uids: set[bytes] = set()  # 正在处理（含等待TG发送）的邮件UID
 
 # Queue 1: raw emails → AI processor
 _email_queue: queue_module.Queue = queue_module.Queue()
@@ -50,16 +48,15 @@ def _email_poller(cfg):
     logger.info("邮件轮询线程已启动")
     while _running:
         try:
-            for acct_idx, mail in poll_accounts(cfg, skip_uids=_processing_uids):
+            for acct_idx, mail in poll_accounts(cfg):
                 if not _running:
                     return
-                with _pending_lock:
+                uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
-                    if mail.uid not in _processing_uids:
+                if not is_uid_processed(uid_str):
-                        _processing_uids.add(mail.uid)
+                    mark_fetched(uid_str, mail.account_email, mail.folder,
-                        _pending_uids.add(mail.uid)
+                                 mail.subject, mail.sender)
-                        # 发送"获取邮件中"提示
                     try:
-                            tg_msg_id = send_thinking(cfg.telegram, cfg.telegram.chat_id, "📥 获取邮件中...")
+                        tg_msg_id = send_thinking(cfg.telegram, cfg.telegram.chat_id, "📥 获取邮件中…")
                     except Exception:
                         tg_msg_id = 0
                     _email_queue.put((acct_idx, mail, tg_msg_id))
@@ -81,6 +78,8 @@ def _ai_processor(cfg):
                     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, acct_idx, mail, tg_msg_id)
                 pending_futures[future] = (acct_idx, mail)
 
@@ -93,10 +92,8 @@ def _ai_processor(cfg):
                         _tg_queue.put(info)
                 except Exception as e:
                     logger.error(f"AI 处理邮件失败: {e}", exc_info=True)
-                    # AI 失败也要移除 pending，否则永远卡住
+                    uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
-                    with _pending_lock:
+                    mark_failed(uid_str)
-                        _pending_uids.discard(mail.uid)
-                        _processing_uids.discard(mail.uid)
 
             time.sleep(0.1 if pending_futures else 0.5)
 
@@ -114,16 +111,15 @@ def _tg_worker(cfg):
             try:
                 info = _tg_queue.get_nowait()
                 tg_send_and_mark(cfg, info)
-                # 发送并标记已读成功后移除 pending（_processing_uids 保留直到重启）
+                uid_str = info["uid"].decode("utf-8", errors="replace") if isinstance(info["uid"], bytes) else str(info["uid"])
-                with _pending_lock:
+                mark_sent(uid_str)
-                    _pending_uids.discard(info["uid"])
             except queue_module.Empty:
                 break
             except Exception as e:
                 logger.error(f"TG 发送失败: {e}", exc_info=True)
                 try:
-                    with _pending_lock:
+                    uid_str = info["uid"].decode("utf-8", errors="replace") if isinstance(info["uid"], bytes) else str(info["uid"])
-                        _pending_uids.discard(info["uid"])
+                    mark_failed(uid_str)
                 except Exception:
                     pass
 
@@ -138,6 +134,7 @@ def main():
 
     cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
     cfg = load_config(cfg_path)
+    init_db()
     logger.info(f"AI邮件摘要机器人已启动，轮询间隔: {cfg.polling.interval_seconds}s")
 
     threads = [
@@ -155,8 +152,8 @@ def main():
         while _running:
             tick += 1
             if tick % 30 == 0:
-                logger.info("队列状态: email_queue=%d tg_queue=%d pending_uids=%d",
+                logger.info("队列状态: email_queue=%d tg_queue=%d",
-                            _email_queue.qsize(), _tg_queue.qsize(), len(_pending_uids))
+                            _email_queue.qsize(), _tg_queue.qsize())
             _shutdown_event.wait(1)
     except KeyboardInterrupt:
         _running = False

src/database.py

@@ -0,0 +1,173 @@
+import json
+import logging
+import sqlite3
+import threading
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+DB_PATH = Path(__file__).parent.parent / "data" / "mail_bot.db"
+_local = threading.local()
+
+
+def init_db():
+    """初始化数据库表"""
+    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+    conn = _get_conn()
+    conn.executescript("""
+        CREATE TABLE IF NOT EXISTS processed_emails (
+            uid TEXT PRIMARY KEY,
+            account TEXT NOT NULL,
+            folder TEXT NOT NULL DEFAULT 'INBOX',
+            subject TEXT,
+            sender TEXT,
+            status TEXT NOT NULL DEFAULT 'fetched',
+            created_at TEXT NOT NULL,
+            sent_at TEXT
+        );
+
+        CREATE TABLE IF NOT EXISTS email_contexts (
+            msg_id INTEGER PRIMARY KEY,
+            data TEXT NOT NULL,
+            created_at TEXT NOT NULL
+        );
+
+        CREATE TABLE IF NOT EXISTS conversations (
+            chat_id INTEGER PRIMARY KEY,
+            data TEXT NOT NULL,
+            updated_at TEXT NOT NULL
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_processed_status ON processed_emails(status);
+        CREATE INDEX IF NOT EXISTS idx_processed_account ON processed_emails(account);
+    """)
+    conn.commit()
+    logger.info("数据库初始化完成: %s", DB_PATH)
+
+
+def _get_conn() -> sqlite3.Connection:
+    if not hasattr(_local, "conn") or _local.conn is None:
+        _local.conn = sqlite3.connect(str(DB_PATH), timeout=10)
+        _local.conn.execute("PRAGMA journal_mode=WAL")
+        _local.conn.execute("PRAGMA busy_timeout=5000")
+        _local.conn.row_factory = sqlite3.Row
+    return _local.conn
+
+
+@contextmanager
+def get_db():
+    conn = _get_conn()
+    try:
+        yield conn
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+
+
+# ── processed_emails ────────────────────────────────────
+
+def is_uid_processed(uid: str) -> bool:
+    with get_db() as conn:
+        row = conn.execute(
+            "SELECT 1 FROM processed_emails WHERE uid=? AND status != 'failed'",
+            (uid,)
+        ).fetchone()
+        return row is not None
+
+
+def is_uids_processed(uids: list[str]) -> set[str]:
+    """批量检查，返回已处理的 uid 集合"""
+    if not uids:
+        return set()
+    with get_db() as conn:
+        placeholders = ",".join("?" * len(uids))
+        rows = conn.execute(
+            f"SELECT uid FROM processed_emails WHERE uid IN ({placeholders}) AND status != 'failed'",
+            uids
+        ).fetchall()
+        return {row["uid"] for row in rows}
+
+
+def mark_fetched(uid: str, account: str, folder: str = "INBOX",
+                 subject: str = "", sender: str = ""):
+    with get_db() as conn:
+        conn.execute(
+            """INSERT OR IGNORE INTO processed_emails
+               (uid, account, folder, subject, sender, status, created_at)
+               VALUES (?, ?, ?, ?, ?, 'fetched', ?)""",
+            (uid, account, folder, subject, sender, datetime.now().isoformat())
+        )
+
+
+def mark_processing(uid: str):
+    with get_db() as conn:
+        conn.execute(
+            "UPDATE processed_emails SET status='processing' WHERE uid=?",
+            (uid,)
+        )
+
+
+def mark_sent(uid: str):
+    with get_db() as conn:
+        conn.execute(
+            "UPDATE processed_emails SET status='sent', sent_at=? WHERE uid=?",
+            (datetime.now().isoformat(), uid)
+        )
+
+
+def mark_failed(uid: str):
+    with get_db() as conn:
+        conn.execute(
+            "UPDATE processed_emails SET status='failed' WHERE uid=?",
+            (uid,)
+        )
+
+
+# ── email_contexts ──────────────────────────────────────
+
+def save_email_context(msg_id: int, data: dict):
+    with get_db() as conn:
+        conn.execute(
+            """INSERT OR REPLACE INTO email_contexts (msg_id, data, created_at)
+               VALUES (?, ?, ?)""",
+            (msg_id, json.dumps(data, ensure_ascii=False), datetime.now().isoformat())
+        )
+
+
+def load_email_context(msg_id: int) -> dict | None:
+    with get_db() as conn:
+        row = conn.execute(
+            "SELECT data FROM email_contexts WHERE msg_id=?", (msg_id,)
+        ).fetchone()
+        if row:
+            return json.loads(row["data"])
+        return None
+
+
+# ── conversations ───────────────────────────────────────
+
+def save_conversation(chat_id: int, data: dict):
+    with get_db() as conn:
+        conn.execute(
+            """INSERT OR REPLACE INTO conversations (chat_id, data, updated_at)
+               VALUES (?, ?, ?)""",
+            (chat_id, json.dumps(data, ensure_ascii=False), datetime.now().isoformat())
+        )
+
+
+def load_conversation(chat_id: int) -> dict | None:
+    with get_db() as conn:
+        row = conn.execute(
+            "SELECT data FROM conversations WHERE chat_id=?", (chat_id,)
+        ).fetchone()
+        if row:
+            return json.loads(row["data"])
+        return None
+
+
+def delete_conversation(chat_id: int):
+    with get_db() as conn:
+        conn.execute("DELETE FROM conversations WHERE chat_id=?", (chat_id,))

src/summarizer.py

@@ -58,14 +58,14 @@ def tg_send_and_mark(cfg: Config, info: dict):
     if thinking_msg_id:
         # 编辑思考中消息为实际摘要
         from src.tg_bot import _summary_keyboard
+        from src.database import save_email_context
         can_reply = info["data"].get("can_reply", True)
         is_promotion = info["data"].get("category") == "promotion"
         links = info["data"].get("links", [])
         edit_message(cfg.telegram, cfg.telegram.chat_id, thinking_msg_id,
                      info["text"], _summary_keyboard(can_reply, is_promotion, links))
         # 保存上下文用于回调
-        from src.tg_bot import _email_contexts
+        save_email_context(thinking_msg_id, {
-        _email_contexts[thinking_msg_id] = {
             "summary_text": info["text"],
             "summary_data": info["data"],
             "original_body": info["original_body"],
@@ -77,7 +77,7 @@ def tg_send_and_mark(cfg: Config, info: dict):
             "sender": info["data"].get("sender", ""),
             "subject": info["data"].get("subject", ""),
             "can_reply": can_reply,
-        }
+        })
     else:
         # 兜底：没有思考中消息就正常发送
         send_summary(cfg.telegram, cfg.telegram.chat_id,

src/tg_bot.py

@@ -6,13 +6,10 @@ from src.config import TelegramConfig, Config
 from src.ai_client import generate_more_replies, expand_promo_detail
 from src.smtp_client import send_reply
 from src.retry import retry
+from src.database import save_email_context, load_email_context, save_conversation, load_conversation, delete_conversation
 
 logger = logging.getLogger(__name__)
 
-_email_contexts: dict[int, dict] = {}
-# chat_id -> conversation state
-_conversations: dict[int, dict] = {}
-
 
 # ── Public API ──────────────────────────────────────────
 
@@ -31,7 +28,7 @@ def send_summary(tg_cfg: TelegramConfig, chat_id: str, summary_text: str,
         "reply_markup": _summary_keyboard(can_reply, is_promotion, links),
     })
     msg_id = result["result"]["message_id"]
-    _email_contexts[msg_id] = {
+    save_email_context(msg_id, {
         "summary_text": summary_text,
         "summary_data": summary_data,
         "original_body": original_body,
@@ -43,7 +40,7 @@ def send_summary(tg_cfg: TelegramConfig, chat_id: str, summary_text: str,
         "sender": summary_data.get("sender", ""),
         "subject": summary_data.get("subject", ""),
         "can_reply": can_reply,
-    }
+    })
     return msg_id
 
 
@@ -270,7 +267,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
     parts = data.split("|")
     action = parts[0]
 
-    ctx = _email_contexts.get(msg_id)
+    ctx = load_email_context(msg_id)
     subject = ctx["subject"] if ctx else "?"
     logger.info("TG 回调: action=%s msg_id=%d subject=%s", action, msg_id, subject)
 
@@ -298,10 +295,10 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
             tg_cfg, str(chat_id), "请输入你的回复内容：",
             reply_markup=_cancel_keyboard(),
         )
-        _conversations[chat_id] = {
+        save_conversation(chat_id, {
             "state": "awaiting_reply", "summary_msg_id": msg_id,
             "prompt_msg_id": prompt_msg_id,
-        }
+        })
         logger.info("  进入回复流程, prompt_msg_id=%d", prompt_msg_id)
 
     elif action == CALLBACK_AI_REPLY and ctx:
@@ -311,18 +308,18 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
             logger.info("  AI 判定无需回复，隐藏")
             return
 
-        _conversations[chat_id] = {
+        save_conversation(chat_id, {
             "state": "ai_reply_selection",
             "summary_msg_id": msg_id,
             "ai_suggestions": suggestions,
             "all_suggestions": list(suggestions),
-        }
+        })
         logger.info("  显示 AI 回复建议 (%d 条)", len(suggestions))
         _show_ai_suggestions(tg_cfg, chat_id, msg_id, suggestions)
 
     elif action == CALLBACK_SELECT_REPLY and ctx:
         idx = int(parts[2])
-        conv = _conversations.get(chat_id)
+        conv = load_conversation(chat_id)
         if not conv or conv.get("summary_msg_id") != msg_id:
             return
         suggestions = conv.get("ai_suggestions", [])
@@ -332,7 +329,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
         logger.info("  选择回复 #%d: %s", idx, sel["text"][:60])
         _do_send_reply(tg_cfg, cfg, chat_id, msg_id, ctx, sel["text"])
         send_text(tg_cfg, str(chat_id), f"✅ 已发送: {sel['text']}")
-        _conversations.pop(chat_id, None)
+        delete_conversation(chat_id)
         can_reply = ctx.get("can_reply", True)
         _tg_req(tg_cfg, "editMessageText", {
             "chat_id": chat_id, "message_id": msg_id,
@@ -341,7 +338,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
         })
 
     elif action == CALLBACK_REGEN and ctx:
-        conv = _conversations.get(chat_id)
+        conv = load_conversation(chat_id)
         if not conv or conv.get("summary_msg_id") != msg_id:
             return
         logger.info("  换一批 (已有 %d 条历史)", len(conv.get("all_suggestions", [])))
@@ -369,14 +366,14 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
             tg_cfg, str(chat_id), "请输入你对回复的提示/要求：",
             reply_markup=_cancel_keyboard(),
         )
-        conv = _conversations.get(chat_id)
+        conv = load_conversation(chat_id)
         if conv and conv.get("summary_msg_id") == msg_id:
             conv["state"] = "awaiting_ai_hint"
             conv["prompt_msg_id"] = prompt_msg_id
             logger.info("  等待用户输入 AI 提示")
 
     elif action == CALLBACK_CANCEL:
-        _conversations.pop(chat_id, None)
+        delete_conversation(chat_id)
         if ctx:
             can_reply = ctx.get("can_reply", True)
             _tg_req(tg_cfg, "editMessageText", {
@@ -424,7 +421,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
 def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
     chat_id = msg["chat"]["id"]
     text = msg["text"].strip()
-    conv = _conversations.get(chat_id)
+    conv = load_conversation(chat_id)
 
     if not conv:
         return
@@ -437,9 +434,9 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
     logger.info("TG 消息: state=%s text=%s", state, text[:60])
 
     if state == "awaiting_reply":
-        ctx = _email_contexts.get(conv["summary_msg_id"])
+        ctx = load_email_context(conv["summary_msg_id"])
         if not ctx:
-            _conversations.pop(chat_id, None)
+            delete_conversation(chat_id)
             delete_message(tg_cfg, chat_id_str, user_msg_id)
             if prompt_msg_id:
                 delete_message(tg_cfg, chat_id_str, prompt_msg_id)
@@ -451,13 +448,13 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
             delete_message(tg_cfg, chat_id_str, prompt_msg_id)
         _do_send_reply(tg_cfg, cfg, chat_id, conv["summary_msg_id"], ctx, text)
         send_text(tg_cfg, chat_id_str, f"✅ 回复已发送: {text}")
-        _conversations.pop(chat_id, None)
+        delete_conversation(chat_id)
         logger.info("  回复发送完成")
 
     elif state == "awaiting_ai_hint":
-        ctx = _email_contexts.get(conv["summary_msg_id"])
+        ctx = load_email_context(conv["summary_msg_id"])
         if not ctx:
-            _conversations.pop(chat_id, None)
+            delete_conversation(chat_id)
             delete_message(tg_cfg, chat_id_str, user_msg_id)
             if prompt_msg_id:
                 delete_message(tg_cfg, chat_id_str, prompt_msg_id)
@@ -473,7 +470,7 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
             )
         except Exception as e:
             send_text(tg_cfg, chat_id_str, f"生成失败: {e}")
-            _conversations.pop(chat_id, None)
+            delete_conversation(chat_id)
             return
 
         if new_suggestions: