feat(email): implement cancellation handling in email services

Author: cnlimiter

src/core/register.py

@@ -9,7 +9,7 @@
 import time
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional, Dict, Any, List
+from typing import Optional, Dict, Any, List, Callable
 from enum import Enum
 
 from ..config.constants import EmailServiceType, OPENAI_EMAIL_SENDERS, OTP_CODE_PATTERN, OTP_CODE_SEMANTIC_PATTERN
@@ -139,6 +139,10 @@ def __init__(self, message: str = "当前邮件批次未发现 OpenAI 发件人"
         self.error_code = error_code
 
 
+class EmailServiceCancelledError(EmailServiceError):
+    """邮箱服务在轮询过程中收到取消信号。"""
+
+
 class EmailServiceStatus(Enum):
     """邮箱服务状态"""
     HEALTHY = "healthy"
@@ -168,6 +172,7 @@ def __init__(self, service_type: EmailServiceType, name: str = None):
         self._provider_backoff = reset_adaptive_backoff()
         self._used_verification_codes: Dict[str, set] = {}
         self._seen_verification_messages: Dict[str, set] = {}
+        self.check_cancelled: Optional[Callable[[], bool]] = None
 
     _EMAIL_ADDRESS_PATTERN = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")
 
@@ -190,6 +195,35 @@ def apply_provider_backoff_state(self, state: Optional[EmailProviderBackoffState
         """注入外部持久化的邮箱供应商退避状态"""
         self._provider_backoff = state or reset_adaptive_backoff()
 
+    def set_check_cancelled(self, callback: Optional[Callable[[], bool]]) -> None:
+        """注入外部取消检查回调。"""
+        self.check_cancelled = callback if callable(callback) else None
+
+    def _is_cancelled_requested(self) -> bool:
+        """检查邮箱服务是否收到取消请求。"""
+        callback = self.check_cancelled
+        if not callable(callback):
+            return False
+        try:
+            return bool(callback())
+        except Exception as e:
+            logger.warning(f"检查邮箱服务取消状态失败: {e}")
+            return False
+
+    def _raise_if_cancelled(self, message: str = "任务已取消") -> None:
+        """在轮询/等待阶段响应取消请求。"""
+        if self._is_cancelled_requested():
+            raise EmailServiceCancelledError(message)
+
+    def _sleep_with_cancel(self, seconds: float, chunk_seconds: float = 0.2) -> None:
+        """可响应取消的短分片休眠。"""
+        remaining = max(0.0, float(seconds))
+        while remaining > 0:
+            self._raise_if_cancelled()
+            sleep_for = min(chunk_seconds, remaining)
+            time.sleep(sleep_for)
+            remaining -= sleep_for
+
     @abc.abstractmethod
     def create_email(self, config: Dict[str, Any] = None) -> Dict[str, Any]:
         """
@@ -520,6 +554,7 @@ def wait_for_email(
         last_email_id = None
 
         while time.time() - start_time < timeout:
+            self._raise_if_cancelled("等待邮件时任务已取消")
             try:
                 emails = self.list_emails()
                 for email_info in emails:
@@ -562,7 +597,7 @@ def wait_for_email(
             except Exception as e:
                 logger.warning(f"等待邮件时出错: {e}")
 
-            time.sleep(check_interval)
+            self._sleep_with_cancel(check_interval)
 
         return None
 

src/services/base.py

@@ -236,6 +236,7 @@ def get_verification_code(
         seen_mail_ids: set = set()
 
         while time.time() - start_time < timeout:
+            self._raise_if_cancelled("等待 Cloud Mail 验证码时任务已取消")
             try:
                 token = self._get_public_token()
                 mails = self._make_request(
@@ -253,7 +254,7 @@ def get_verification_code(
                     mails = mails["list"]
 
                 if not isinstance(mails, list):
-                    time.sleep(poll_interval)
+                    self._sleep_with_cancel(poll_interval)
                     continue
 
                 if mails:
@@ -302,7 +303,7 @@ def get_verification_code(
                     raise
                 logger.debug(f"检查 Cloud Mail 邮件时出错: {e}")
 
-            time.sleep(poll_interval)
+            self._sleep_with_cancel(poll_interval)
 
         logger.warning(f"等待 Cloud Mail 验证码超时: {email}")
         return None

src/services/cloud_mail.py

@@ -263,6 +263,7 @@ def get_verification_code(
         seen_message_ids = set()
 
         while time.time() - start_time < timeout:
+            self._raise_if_cancelled("等待 DuckMail 验证码时任务已取消")
             try:
                 response = self._make_request(
                     "GET",
@@ -324,7 +325,7 @@ def get_verification_code(
                     raise
                 logger.debug(f"DuckMail 轮询验证码失败: {e}")
 
-            time.sleep(poll_interval)
+            self._sleep_with_cancel(poll_interval)
 
         return None
 

src/services/duck_mail.py

@@ -214,10 +214,11 @@ def get_verification_code(
         seen_mail_ids: set = set()
 
         while time.time() - start_time < timeout:
+            self._raise_if_cancelled("等待 Freemail 验证码时任务已取消")
             try:
                 mails = self._make_request("GET", "/api/emails", params={"mailbox": email, "limit": 20})
                 if not isinstance(mails, list):
-                    time.sleep(poll_interval)
+                    self._sleep_with_cancel(poll_interval)
                     continue
 
                 ordered_mails = self._sort_items_by_message_time(
@@ -299,7 +300,7 @@ def get_verification_code(
                     raise
                 logger.debug(f"检查 Freemail 邮件时出错: {e}")
 
-            time.sleep(poll_interval)
+            self._sleep_with_cancel(poll_interval)
 
         logger.warning(f"等待 Freemail 验证码超时: {email}")
         return None

src/services/freemail.py

@@ -12,7 +12,7 @@
 from email.header import decode_header
 from typing import Any, Dict, Optional
 
-from .base import BaseEmailService, EmailServiceError, OTPNoOpenAISenderEmailServiceError, get_email_code_settings
+from .base import BaseEmailService, OTPNoOpenAISenderEmailServiceError, get_email_code_settings
 from ..config.constants import (
     EmailServiceType,
     OTP_CODE_SEMANTIC_PATTERN,
@@ -124,11 +124,12 @@ def get_verification_code(
             mail.select("INBOX")
 
             while time.time() - start_time < timeout:
+                self._raise_if_cancelled("等待 IMAP 验证码时任务已取消")
                 try:
                     # 搜索所有未读邮件
                     status, data = mail.search(None, "UNSEEN")
                     if status != "OK" or not data or not data[0]:
-                        time.sleep(poll_interval)
+                        self._sleep_with_cancel(poll_interval)
                         continue
 
                     msg_ids = data[0].split()
@@ -181,13 +182,13 @@ def get_verification_code(
                     except Exception:
                         pass
 
-                time.sleep(poll_interval)
+                self._sleep_with_cancel(poll_interval)
 
         except Exception as e:
             if isinstance(e, OTPNoOpenAISenderEmailServiceError):
                 raise
             logger.warning(f"IMAP 连接/轮询失败: {e}")
-            self.update_status(False, str(e))
+            self.update_status(False, e)
         finally:
             if mail:
                 try:

src/services/imap_mail.py

@@ -308,13 +308,14 @@ def get_verification_code(
         seen_message_ids = set()
 
         while time.time() - start_time < timeout:
+            self._raise_if_cancelled("等待自定义域名邮箱验证码时任务已取消")
             try:
                 # 获取邮件列表
                 response = self._make_request("GET", f"/api/emails/{target_email_id}")
 
                 messages = response.get("messages", [])
                 if not isinstance(messages, list):
-                    time.sleep(poll_interval)
+                    self._sleep_with_cancel(poll_interval)
                     continue
 
                 ordered_messages = self._sort_items_by_message_time(
@@ -380,7 +381,7 @@ def get_verification_code(
                 logger.debug(f"检查邮件时出错: {e}")
 
             # 等待一段时间再检查
-            time.sleep(poll_interval)
+            self._sleep_with_cancel(poll_interval)
 
         logger.warning(f"等待验证码超时: {email}")
         return None

src/services/moe_mail.py

@@ -337,6 +337,7 @@ def get_verification_code(
         poll_count = 0
 
         while time.time() - start_time < actual_timeout:
+            self._raise_if_cancelled("等待 Outlook 验证码时任务已取消")
             poll_count += 1
 
             # 渐进式邮件检查：前 3 次只检查未读
@@ -387,7 +388,7 @@ def get_verification_code(
                 logger.warning(f"[{email}] 检查出错: {e}")
 
             # 等待下次轮询
-            time.sleep(poll_interval)
+            self._sleep_with_cancel(poll_interval)
 
         elapsed = int(time.time() - start_time)
         logger.warning(f"[{email}] 验证码超时 ({actual_timeout}s)，共轮询 {poll_count} 次")

src/services/outlook/service.py

@@ -309,6 +309,7 @@ def get_verification_code(
         # jwt = cached.get("jwt")
 
         while time.time() - start_time < timeout:
+            self._raise_if_cancelled("等待 TempMail 验证码时任务已取消")
             try:
                 # if jwt:
                 #     response = self._make_request(
@@ -327,7 +328,7 @@ def get_verification_code(
                 # /user_api/mails 和 /admin/mails 返回格式相同: {"results": [...], "total": N}
                 mails = response.get("results", [])
                 if not isinstance(mails, list):
-                    time.sleep(poll_interval)
+                    self._sleep_with_cancel(poll_interval)
                     continue
 
                 ordered_mails = self._sort_items_by_message_time(
@@ -389,7 +390,7 @@ def get_verification_code(
                     raise
                 logger.debug(f"检查 TempMail 邮件时出错: {e}")
 
-            time.sleep(poll_interval)
+            self._sleep_with_cancel(poll_interval)
 
         logger.warning(f"等待 TempMail 验证码超时: {email}")
         return None

src/services/temp_mail.py

@@ -217,6 +217,7 @@ def get_verification_code(
         seen_ids = set()
 
         while time.time() - start_time < timeout:
+            self._raise_if_cancelled("等待 Tempmail 验证码时任务已取消")
             try:
                 # 获取邮件列表
                 response = self.http_client.get(
@@ -226,7 +227,7 @@ def get_verification_code(
                 )
 
                 if response.status_code != 200:
-                    time.sleep(poll_interval)
+                    self._sleep_with_cancel(poll_interval)
                     continue
 
                 data = response.json()
@@ -239,7 +240,7 @@ def get_verification_code(
                 email_list = data.get("emails", []) if isinstance(data, dict) else []
 
                 if not isinstance(email_list, list):
-                    time.sleep(poll_interval)
+                    self._sleep_with_cancel(poll_interval)
                     continue
 
                 ordered_emails = self._sort_items_by_message_time(
@@ -302,7 +303,7 @@ def get_verification_code(
                 logger.debug(f"检查邮件时出错: {e}")
 
             # 等待一段时间再检查
-            time.sleep(poll_interval)
+            self._sleep_with_cancel(poll_interval)
 
         logger.warning(f"等待验证码超时: {email}")
         return None