code/ai_finance_agent/backend/app/application/worker.py

"""持久 outbox 至少一次消费；租约用于恢复崩溃任务，业务副作用另做幂等。"""
import asyncio
import logging
from datetime import timedelta
from sqlalchemy import select, update, or_
from app.domain.models import ApprovalRecord, ApprovalRound, AuditRun, KnowledgeChunk, KnowledgeDocument, OutboxEvent, Reimbursement
from app.application.services import alert, get, log
from app.infrastructure.ai import ai
from app.infrastructure.database import Session, transaction
from app.shared.common import BusinessError, utcnow
from app.shared.config import settings

logger = logging.getLogger("audit.worker")


async def index_document(document_id: int, version: int):
    """向量先在事务外生成，版本未变化才整批写入；任何失败不写随机向量。"""
    async with Session() as session:
        document = await get(session, KnowledgeDocument, document_id)
        if document.version != version:
            return
        blocks = (await session.scalars(select(KnowledgeChunk).where(KnowledgeChunk.document_id == document_id, KnowledgeChunk.document_version == version).order_by(KnowledgeChunk.position))).all()
        identifiers, texts = [item.id for item in blocks], [item.text for item in blocks]
    vectors = []
    for offset in range(0, len(texts), 32):
        vectors.extend(await ai.embeddings(texts[offset:offset + 32]))
    async with transaction() as session:
        document = await get(session, KnowledgeDocument, document_id)
        if document.version != version:
            return
        for identifier, vector in zip(identifiers, vectors):
            block = await get(session, KnowledgeChunk, identifier)
            block.vector_json, block.dimension = vector, settings.embedding_dimension
        document.index_status, document.embedding_model = "indexed", settings.embedding_model
        log(session, None, "knowledge.indexed", "knowledge", document_id, after={"version": version, "chunks": len(vectors)})


class Worker:
    def __init__(self):
        self.task = None
        self.stop_event = asyncio.Event()
        self.last_scan = None
        self.current = None

    async def start(self):
        """重启时本部署只有一个消费者，旧 processing 租约可安全回收。"""
        self.stop_event.clear()
        async with transaction() as session:
            await session.execute(update(OutboxEvent).where(OutboxEvent.status == "processing")
                                  .values(status="pending", lease_until=None, next_attempt_at=utcnow()))
        if settings.worker_enabled:
            self.task = asyncio.create_task(self.loop())

    async def close(self):
        """给当前任务短暂收尾机会，超时取消保留租约，下一次启动继续处理。"""
        self.stop_event.set()
        if self.task:
            try:
                await asyncio.wait_for(asyncio.shield(self.task), 10)
            except asyncio.TimeoutError:
                self.task.cancel()
                await asyncio.gather(self.task, return_exceptions=True)

    async def claim(self):
        async with transaction() as session:
            now = utcnow()
            row = await session.scalar(select(OutboxEvent).where(
                or_(OutboxEvent.status == "pending", (OutboxEvent.status == "processing") & (OutboxEvent.lease_until < now)),
                OutboxEvent.next_attempt_at <= now).order_by(OutboxEvent.id).limit(1))
            if not row:
                return None
            row.status, row.lease_until = "processing", now + timedelta(minutes=30)
            row.attempts += 1
            return {"id": row.id, "type": row.event_type, "payload": row.payload_json, "attempts": row.attempts}

    async def process_one(self):
        """返回是否领取事件，供实际消费者和集成测试共用同一个入口。"""
        event = await self.claim()
        if not event:
            return False
        self.current = event["id"]
        try:
            if event["type"] == "audit":
                from app.application.pipeline import run_audit
                await run_audit(event["payload"]["audit_id"])
            elif event["type"] == "approval":
                from app.application.approvals import start_approval
                await start_approval(event["payload"]["audit_id"])
            elif event["type"] == "index":
                await index_document(event["payload"]["document_id"], event["payload"]["version"])
            else:
                raise BusinessError("未知事件类型", code="EVENT_INVALID")
            async with transaction() as session:
                row = await get(session, OutboxEvent, event["id"])
                row.status, row.lease_until, row.error_code = "done", None, None
        except Exception as exc:
            code = exc.code if isinstance(exc, BusinessError) else type(exc).__name__
            logger.error("event_failed id=%s code=%s", event["id"], code)
            exhausted = event["attempts"] >= settings.outbox_attempts
            async with transaction() as session:
                row = await get(session, OutboxEvent, event["id"])
                row.status, row.lease_until, row.error_code = "failed" if exhausted else "pending", None, code
                row.next_attempt_at = utcnow() + timedelta(seconds=2 ** event["attempts"])
                audit_id = event["payload"].get("audit_id")
                audit = await session.get(AuditRun, audit_id) if audit_id else None
                if exhausted and audit and event["type"] == "audit":
                    reimbursement = await get(session, Reimbursement, audit.reimbursement_id)
                    # 撤回或新轮次已使任务失效时，失败处理也不能覆盖取消的终态。
                    if reimbursement.revision == audit.revision and reimbursement.status == "pending_audit" and audit.status not in {"completed", "cancelled"}:
                        audit.status, audit.error_code, audit.finished_at = "failed", code, utcnow()
                if exhausted and event["type"] == "index":
                    document = await get(session, KnowledgeDocument, event["payload"]["document_id"])
                    if document.version == event["payload"]["version"]:
                        document.index_status = "failed"
                if exhausted:
                    await alert(session, f"event:{event['id']}", "task_failed", "任务多次失败，错误码：" + code,
                        "error", audit_id, audit.reimbursement_id if audit else None)
        finally:
            self.current = None
        return True

    async def scan_timeouts(self):
        """只扫描当前 active 轮次未决记录，退回/驳回后的旧记录不会继续催办。"""
        cutoff = utcnow() - timedelta(hours=settings.approval_timeout_hours)
        async with transaction() as session:
            records = (await session.scalars(select(ApprovalRecord).join(ApprovalRound, ApprovalRecord.round_id == ApprovalRound.id)
                .where(ApprovalRound.status == "active", ApprovalRecord.result.is_(None), ApprovalRecord.created_at < cutoff))).all()
            for record in records:
                round_item = await get(session, ApprovalRound, record.round_id)
                await alert(session, f"timeout:{record.id}", "approval_timeout", "审批记录超过配置时限，请指定审批人处理", "warning", round_item.audit_id, round_item.reimbursement_id)
        self.last_scan = utcnow()

    async def loop(self):
        while not self.stop_event.is_set():
            try:
                if self.last_scan is None or (utcnow() - self.last_scan).total_seconds() >= settings.timeout_scan_seconds:
                    await self.scan_timeouts()
                if await self.process_one():
                    continue
            except Exception:
                logger.exception("worker_loop_failed")
            try:
                await asyncio.wait_for(self.stop_event.wait(), timeout=settings.worker_interval)
            except asyncio.TimeoutError:
                pass


worker = Worker()