Files
group_fqcd_jr/app/worker/offsite_mail_worker.py
T

656 lines
26 KiB
Python

"""场外基金收件箱独立 Worker。"""
from __future__ import annotations
import asyncio
import logging
from collections.abc import Awaitable, Callable
from contextlib import suppress
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Protocol
from uuid import uuid4
from sqlalchemy import func, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.config import Settings
from app.core.contracts import RequestContext
from app.core.offsite_fund_contracts import (
ReceiveRecognizedMailRequest,
RecognizedAttachment,
)
from app.infrastructure.db import SessionFactory
from app.model.audit import InteractionAudit
from app.model.offsite_fund import (
OffsiteFundMail,
OffsiteMailCursor,
OffsiteNotification,
OffsiteRecognitionAttempt,
)
from app.service.offsite_document_recognition_adapter import (
REQUIRED_FIELDS,
OffsiteDocumentRecognitionAdapter,
RecognitionSourceFile,
StructuredRecognitionResult,
)
from app.service.offsite_fund_service import OffsiteFundService
from app.service.offsite_mail_adapter import (
OffsiteImapReceiver,
OffsiteMailStorage,
RawMailMessage,
SavedMailAttachment,
)
logger = logging.getLogger(__name__)
class MailReceiver(Protocol):
last_scanned_uid: str | None
def health_check(self) -> dict[str, object]: ...
def fetch_since(self, last_uid: str | None, *, limit: int) -> tuple[RawMailMessage, ...]: ...
def wait_for_new_mail(self, timeout_seconds: float) -> bool: ...
def close(self) -> None: ...
class MailRecognizer(Protocol):
async def recognize(self, source: RecognitionSourceFile) -> StructuredRecognitionResult: ...
class MailService(Protocol):
async def receive_recognized_mail(
self, payload: ReceiveRecognizedMailRequest, context: RequestContext
) -> dict[str, object]: ...
class CursorLease:
def __init__(self, lease_id: str, last_uid: str) -> None:
self.lease_id = lease_id
self.last_uid = last_uid
class OffsiteMailWorker:
"""按 UID 顺序处理邮件,失败时停在当前 UID 并等待补偿。"""
def __init__(
self,
settings: Settings,
*,
receiver: MailReceiver | None = None,
storage: OffsiteMailStorage | None = None,
recognizer: MailRecognizer | None = None,
identity_resolver: Callable[[RequestContext], Awaitable[RequestContext]] | None = None,
session_factory: Callable[[], AsyncSession] = SessionFactory,
service_factory: Callable[[AsyncSession], MailService] | None = None,
) -> None:
self.settings = settings
self.receiver = receiver or OffsiteImapReceiver(settings)
self.storage = storage or OffsiteMailStorage(settings.offsite_mail_storage_dir)
self.recognizer = recognizer or OffsiteDocumentRecognitionAdapter(settings)
self.identity_resolver = identity_resolver
self.session_factory = session_factory
self.service_factory = service_factory or OffsiteFundService
async def run_once(self) -> bool:
if not self.settings.offsite_mail_worker_enabled:
return False
recovered = await self.recover_stale_notifications()
if not self.settings.offsite_imap_enabled:
return recovered
if not self._real_recognition_ready():
logger.error("场外 Worker 拒绝处理:真实 IMAP 不得使用 Mock 识别结果")
return recovered
if not self.settings.offsite_worker_user_id:
logger.error("场外 Worker 未配置合法操作用户 ID,拒绝自动写入业务数据")
return recovered
try:
context = await self._resolve_worker_context()
lease = await self._claim_cursor()
if lease is None:
return recovered
return await self._process_batch(lease, context) or recovered
except asyncio.CancelledError:
raise
except Exception:
logger.exception("场外邮件 Worker 执行失败")
return True
async def close(self) -> None:
self.receiver.close()
close = getattr(self.recognizer, "close", None)
if callable(close):
result = close()
if asyncio.iscoroutine(result):
await result
def _real_recognition_ready(self) -> bool:
if not isinstance(self.recognizer, OffsiteDocumentRecognitionAdapter):
return True
if not self.settings.offsite_ocr_enabled or not self.settings.offsite_deepseek_enabled:
return False
health = self.recognizer.health_check()
ocr = health.get("ocr")
deepseek = health.get("deepseek")
return (
isinstance(ocr, dict)
and ocr.get("status") == "ok"
and isinstance(deepseek, dict)
and deepseek.get("status") == "ok"
)
async def recover_stale_notifications(self) -> bool:
now = datetime.now(UTC).replace(tzinfo=None)
cutoff = now - timedelta(
seconds=self.settings.offsite_notification_sending_timeout_seconds
)
async with self.session_factory() as session, session.begin():
activity = func.coalesce(
OffsiteNotification.updated_at, OffsiteNotification.created_at
)
result = await session.execute(
update(OffsiteNotification)
.where(
OffsiteNotification.status == "发送中",
activity < cutoff,
)
.values(
status="发送失败",
failure_reason="发送状态不确定,需人工核验外部邮箱后再决定是否重试",
updated_at=now,
)
)
count = int(getattr(result, "rowcount", 0) or 0)
if count:
session.add(
InteractionAudit(
actor_type="system",
actor_id=None,
target_customer_id=None,
session_id=None,
portal="worker",
action_type="offsite.notification_timeout_recovered",
detail={"count": count},
created_at=now,
)
)
return count > 0
async def _resolve_worker_context(self) -> RequestContext:
identity = RequestContext(
user_id=self.settings.offsite_worker_user_id,
trace_id=f"offsite-worker-{uuid4()}",
portal="worker",
)
resolver = self.identity_resolver
if resolver is None:
from app.service.identity_service import IdentityService
resolver = IdentityService().resolve
context = await resolver(identity)
if not context.roles or not context.permissions:
raise RuntimeError("场外 Worker 操作用户未通过角色和权限校验")
return context
async def _claim_cursor(self) -> CursorLease | None:
for attempt in range(2):
lease_id = str(uuid4())
now = datetime.now(UTC).replace(tzinfo=None)
try:
async with self.session_factory() as session, session.begin():
cursor = await session.scalar(
select(OffsiteMailCursor)
.where(
OffsiteMailCursor.mailbox == self.settings.offsite_mailbox,
OffsiteMailCursor.folder == OffsiteImapReceiver.inbox_name,
)
.with_for_update()
)
if cursor is None:
cursor = OffsiteMailCursor(
mailbox=self.settings.offsite_mailbox,
folder=OffsiteImapReceiver.inbox_name,
last_uid="0",
status="idle",
retry_count=0,
created_at=now,
updated_at=now,
)
session.add(cursor)
await session.flush()
if cursor.status == "blocked":
return None
if cursor.lease_until is not None and cursor.lease_until > now:
return None
if cursor.next_retry_at is not None and cursor.next_retry_at > now:
return None
cursor.status = "processing"
cursor.lease_id = lease_id
cursor.lease_until = now + timedelta(
seconds=self.settings.worker_lease_seconds
)
cursor.updated_at = now
return CursorLease(lease_id, cursor.last_uid)
except IntegrityError:
if attempt == 1:
raise
return None
async def _process_batch(self, lease: CursorLease, context: RequestContext) -> bool:
task = asyncio.current_task()
if task is None:
raise RuntimeError("场外 Worker 无法建立当前任务租约")
heartbeat = asyncio.create_task(self._cursor_heartbeat(lease.lease_id, task))
try:
return await self._process_batch_work(lease, context)
finally:
heartbeat.cancel()
with suppress(asyncio.CancelledError):
await heartbeat
async def _process_batch_work(self, lease: CursorLease, context: RequestContext) -> bool:
try:
health = await asyncio.to_thread(self.receiver.health_check)
if health.get("status") != "ok":
message = str(health.get("message") or health.get("status") or "IMAP不可用")
raise RuntimeError(f"场外 IMAP 健康检查失败:{message}")
messages = await self._fetch_with_idle_compensation(lease.last_uid)
current_uid = lease.last_uid
for mail in sorted(messages, key=self._uid_sort_key):
if not self._uid_after(mail.imap_uid, current_uid):
continue
try:
await self._process_message(mail, context)
except asyncio.CancelledError:
raise
except Exception as exc:
await self._record_failure(lease, mail, exc)
return True
current_uid = mail.imap_uid
await self._advance_cursor(lease.lease_id, current_uid, release=False)
scanned_uid = self.receiver.last_scanned_uid
if scanned_uid is not None and self._uid_after(scanned_uid, current_uid):
current_uid = scanned_uid
await self._advance_cursor(lease.lease_id, current_uid, release=True)
return bool(messages) or scanned_uid is not None
except asyncio.CancelledError:
raise
except Exception as exc:
await self._record_cursor_failure(lease, exc)
self.receiver.close()
return True
async def _fetch_with_idle_compensation(
self, last_uid: str
) -> tuple[RawMailMessage, ...]:
messages = await asyncio.to_thread(
self.receiver.fetch_since,
last_uid,
limit=self.settings.offsite_mail_worker_batch_size,
)
if (
messages
or self.receiver.last_scanned_uid is not None
or not self.settings.offsite_imap_idle_enabled
):
return messages
wait_for_new_mail = getattr(self.receiver, "wait_for_new_mail", None)
if not callable(wait_for_new_mail):
raise RuntimeError("已启用 IMAP IDLE,但收件适配器未提供 IDLE 等待能力")
has_new_mail = await asyncio.to_thread(
wait_for_new_mail,
self.settings.offsite_imap_idle_timeout_seconds,
)
if not has_new_mail:
return ()
return await asyncio.to_thread(
self.receiver.fetch_since,
last_uid,
limit=self.settings.offsite_mail_worker_batch_size,
)
async def _cursor_heartbeat(self, lease_id: str, task: asyncio.Task[bool]) -> None:
try:
while True:
await asyncio.sleep(self.settings.worker_lease_seconds / 3)
now = datetime.now(UTC).replace(tzinfo=None)
async with self.session_factory() as session, session.begin():
result = await session.execute(
update(OffsiteMailCursor)
.where(OffsiteMailCursor.lease_id == lease_id)
.values(
lease_until=now + timedelta(
seconds=self.settings.worker_lease_seconds
),
updated_at=now,
)
)
if int(getattr(result, "rowcount", 0) or 0) != 1:
task.cancel()
return
except asyncio.CancelledError:
raise
except Exception:
logger.exception("场外 Worker 游标租约续期失败")
task.cancel()
async def _process_message(self, mail: RawMailMessage, context: RequestContext) -> None:
saved = await asyncio.to_thread(self.storage.save, mail)
attachments: list[RecognizedAttachment] = []
for saved_attachment in saved.attachments:
recognition = await self._recognize_with_retry(
mail, saved_attachment
)
attachments.append(
self._recognized_attachment(saved_attachment, recognition)
)
request = ReceiveRecognizedMailRequest(
imap_uid=saved.imap_uid,
message_id=saved.message_id,
sender=saved.sender,
return_path=saved.return_path,
auth_result=saved.auth_result,
eml_path=saved.eml_path,
attachments=tuple(attachments),
)
async with self.session_factory() as session:
service = self.service_factory(session)
result = await service.receive_recognized_mail(request, context)
if result.get("code") != 0:
raise RuntimeError(f"场外业务入库失败:{result.get('message', '未知错误')}")
async def _recognize_with_retry(
self, mail: RawMailMessage, saved_attachment: SavedMailAttachment
) -> StructuredRecognitionResult:
payload = await asyncio.to_thread(Path(saved_attachment.original_file_path).read_bytes)
source = RecognitionSourceFile(
filename=saved_attachment.filename,
media_type=saved_attachment.media_type,
payload=payload,
file_hash=saved_attachment.file_hash,
original_file_path=saved_attachment.original_file_path,
)
result: StructuredRecognitionResult | None = None
for attempt_no in range(1, 3):
started_at = datetime.now(UTC).replace(tzinfo=None)
error_message: str | None = None
try:
result = await self.recognizer.recognize(source)
except asyncio.CancelledError:
raise
except Exception as exc:
error_message = self._error_message(exc)
result = None
finished_at = datetime.now(UTC).replace(tzinfo=None)
await self._record_recognition_attempt(
mail=mail,
saved_attachment=saved_attachment,
attempt_no=attempt_no,
result=result,
source="automatic",
operator_id=None,
started_at=started_at,
finished_at=finished_at,
error_message=error_message,
)
if result is not None and not self._recognition_requires_retry(result):
return result
if attempt_no == 2:
break
if result is None or self._recognition_failed(result):
raise RuntimeError("附件识别失败,已保留原始文件并等待补偿")
return result
async def _record_recognition_attempt(
self,
*,
mail: RawMailMessage,
saved_attachment: SavedMailAttachment,
attempt_no: int,
result: StructuredRecognitionResult | None,
source: str,
operator_id: str | None,
started_at: datetime,
finished_at: datetime,
error_message: str | None,
task_id: str | None = None,
mail_id: str | None = None,
attachment_id: str | None = None,
) -> None:
fields = result.extracted_fields if result is not None else {}
confidence = (
{key: str(value) for key, value in result.field_confidence.items()}
if result is not None
else {}
)
missing_fields = list(result.missing_fields) if result is not None else []
low_confidence_fields = (
list(result.low_confidence_fields) if result is not None else []
)
if result is None:
status = "error"
document_type = "other"
ocr_status = "error"
llm_status = "error"
page_evidence: dict[str, object] = {}
error_message = error_message or "识别器未返回结果"
else:
status = self._recognition_attempt_status(result)
document_type = result.document_type
ocr_status = result.ocr_status
llm_status = result.llm_status
page_evidence = result.page_evidence
error_message = error_message or result.error_message
async with self.session_factory() as session, session.begin():
session.add(
OffsiteRecognitionAttempt(
task_id=task_id,
mail_id=mail_id,
attachment_id=attachment_id,
imap_uid=mail.imap_uid,
message_id=mail.message_id,
filename=saved_attachment.filename,
file_hash=saved_attachment.file_hash,
original_file_path=saved_attachment.original_file_path,
attempt_no=attempt_no,
source=source,
operator_id=operator_id,
document_type=document_type,
ocr_status=ocr_status,
llm_status=llm_status,
extracted_fields=fields,
field_confidence=confidence,
missing_fields=missing_fields,
low_confidence_fields=low_confidence_fields,
page_evidence=page_evidence,
status=status,
error_message=error_message,
started_at=started_at,
finished_at=finished_at,
created_at=finished_at,
)
)
async def _advance_cursor(
self, lease_id: str, last_uid: str, *, release: bool
) -> None:
now = datetime.now(UTC).replace(tzinfo=None)
async with self.session_factory() as session, session.begin():
cursor = await session.scalar(
select(OffsiteMailCursor)
.where(OffsiteMailCursor.lease_id == lease_id)
.with_for_update()
)
if cursor is None:
return
cursor.last_uid = last_uid
cursor.status = "idle" if release else "processing"
if release:
cursor.retry_count = 0
cursor.last_error = None
cursor.blocked_uid = None
cursor.blocked_message_id = None
cursor.next_retry_at = None
cursor.lease_id = None
cursor.lease_until = None
cursor.updated_at = now
async def _record_failure(
self, lease: CursorLease, mail: RawMailMessage, exc: Exception
) -> None:
now = datetime.now(UTC).replace(tzinfo=None)
message = self._error_message(exc)
async with self.session_factory() as session, session.begin():
cursor = await session.scalar(
select(OffsiteMailCursor)
.where(OffsiteMailCursor.lease_id == lease.lease_id)
.with_for_update()
)
if cursor is None:
return
cursor.retry_count += 1
cursor.status = (
"blocked"
if cursor.retry_count >= self.settings.offsite_max_retry_count
else "failed"
)
cursor.blocked_uid = mail.imap_uid
cursor.blocked_message_id = mail.message_id
cursor.last_error = message
cursor.next_retry_at = (
None
if cursor.status == "blocked"
else now + timedelta(seconds=min(300, 2**cursor.retry_count))
)
cursor.lease_id = None
cursor.lease_until = None
cursor.updated_at = now
if cursor.status == "blocked":
# 游标锁死必须留痕:前端"收件箱异常"提示读的就是这条记录。
# 锁死后 _claim_cursor() 会直接返回,因此这里每次锁死只会写一条。
session.add(
InteractionAudit(
actor_type="system",
actor_id=None,
target_customer_id=None,
session_id=None,
portal="worker",
action_type="offsite.mail_cursor_blocked",
detail={
"mailbox": self.settings.offsite_mailbox,
"imap_uid": mail.imap_uid,
"message_id": mail.message_id,
"retry_count": cursor.retry_count,
"error": message,
},
created_at=now,
)
)
await session.execute(
update(OffsiteFundMail)
.where(
OffsiteFundMail.imap_uid == mail.imap_uid,
OffsiteFundMail.message_id == mail.message_id,
)
.values(
retry_count=OffsiteFundMail.retry_count + 1,
last_error=message,
next_retry_at=cursor.next_retry_at,
last_attempt_at=now,
updated_at=now,
)
)
logger.error(
"场外邮件处理失败 uid=%s message_id=%s error=%s",
mail.imap_uid,
mail.message_id,
message,
)
async def _record_cursor_failure(self, lease: CursorLease, exc: Exception) -> None:
message = self._error_message(exc)
now = datetime.now(UTC).replace(tzinfo=None)
async with self.session_factory() as session, session.begin():
cursor = await session.scalar(
select(OffsiteMailCursor)
.where(OffsiteMailCursor.lease_id == lease.lease_id)
.with_for_update()
)
if cursor is None:
return
cursor.retry_count += 1
cursor.status = "failed"
cursor.last_error = message
cursor.next_retry_at = now + timedelta(
seconds=min(300, 2**cursor.retry_count)
)
cursor.lease_id = None
cursor.lease_until = None
cursor.updated_at = now
logger.error("场外邮件批处理失败 error=%s", message)
@staticmethod
def _uid_sort_key(mail: RawMailMessage) -> tuple[int, str]:
return (int(mail.imap_uid) if mail.imap_uid.isdigit() else 2**63 - 1, mail.imap_uid)
@staticmethod
def _uid_after(candidate: str, current: str) -> bool:
if candidate.isdigit() and current.isdigit():
return int(candidate) > int(current)
return candidate > current
@staticmethod
def _recognition_failed(result: object) -> bool:
error = getattr(result, "error_message", None)
ocr_status = getattr(result, "ocr_status", None)
llm_status = getattr(result, "llm_status", None)
failed = {"error", "misconfigured"}
return bool(error and (ocr_status in failed or llm_status in failed))
@staticmethod
def _recognition_requires_retry(result: StructuredRecognitionResult) -> bool:
if OffsiteMailWorker._recognition_failed(result):
return True
if result.document_type not in REQUIRED_FIELDS:
return False
fields = result.extracted_fields
required = REQUIRED_FIELDS[result.document_type]
if any(not str(fields.get(name) or "").strip() for name in required):
return True
if not (fields.get("投资者名称") or fields.get("客户标识")):
return True
return bool(result.missing_fields or result.low_confidence_fields)
@staticmethod
def _recognition_attempt_status(result: StructuredRecognitionResult) -> str:
if OffsiteMailWorker._recognition_failed(result):
return "error"
if OffsiteMailWorker._recognition_requires_retry(result):
return "recognition_exception"
return "success"
@staticmethod
def _recognized_attachment(
saved: SavedMailAttachment, result: StructuredRecognitionResult
) -> RecognizedAttachment:
return RecognizedAttachment(
filename=saved.filename,
file_hash=saved.file_hash,
original_file_path=saved.original_file_path,
media_type=saved.media_type,
size_bytes=saved.size_bytes,
document_type=result.document_type,
extracted_fields=result.extracted_fields,
field_confidence=result.field_confidence,
ocr_text=result.ocr_text,
page_evidence=result.page_evidence,
)
@staticmethod
def _error_message(exc: Exception) -> str:
return f"{type(exc).__name__}: {str(exc)[:450]}"