taiji-AI-PAD/services/mcp-server/app/events.py

"""
NATS事件采集模块
计费事件消息队列处理
"""

import json
import uuid
import asyncio
from datetime import datetime
from typing import Dict, List, Optional, Any, Callable
from sqlalchemy import select, update, and_
from sqlalchemy.ext.asyncio import AsyncSession
import structlog

from models import BillingEvent, User, Agent, Execution

logger = structlog.get_logger(__name__)


# 事件类型定义
EVENT_TYPES = {
    "execution.start": "执行开始",
    "execution.end": "执行结束",
    "execution.error": "执行错误",
    "quota.check": "配额检查",
    "quota.exceeded": "配额超限",
    "balance.deduct": "余额扣款",
    "balance.recharge": "余额充值",
    "rate_limit.hit": "速率限制触发",
}


class EventPublisher:
    """事件发布者"""
    
    def __init__(self, nats_client=None):
        self.nats = nats_client
        self._handlers: Dict[str, List[Callable]] = {}
    
    def register_handler(self, event_type: str, handler: Callable):
        """注册事件处理器"""
        if event_type not in self._handlers:
            self._handlers[event_type] = []
        self._handlers[event_type].append(handler)
    
    async def publish(
        self,
        event_type: str,
        user_id: str,
        payload: Dict,
        agent_id: Optional[str] = None,
        execution_id: Optional[str] = None,
        db: Optional[AsyncSession] = None
    ) -> str:
        """
        发布事件
        
        Args:
            event_type: 事件类型
            user_id: 用户ID
            payload: 事件数据
            agent_id: Agent ID
            execution_id: 执行ID
            db: 数据库会话
            
        Returns:
            事件ID
        """
        event_id = str(uuid.uuid4())
        
        event_data = {
            "event_id": event_id,
            "event_type": event_type,
            "user_id": user_id,
            "agent_id": agent_id,
            "execution_id": execution_id,
            "payload": payload,
            "timestamp": datetime.utcnow().isoformat(),
        }
        
        # 保存到数据库
        if db:
            billing_event = BillingEvent(
                event_type=event_type,
                event_id=event_id,
                user_id=user_id,
                agent_id=agent_id,
                execution_id=execution_id,
                payload=payload,
                status="pending",
            )
            db.add(billing_event)
            await db.commit()
        
        # 发布到NATS（如果配置了）
        if self.nats:
            try:
                subject = f"billing.{event_type.replace('.', '_')}"
                await self.nats.publish(subject, json.dumps(event_data).encode())
                logger.info("事件发布成功", event_id=event_id, event_type=event_type)
            except Exception as e:
                logger.error("NATS发布失败", error=str(e), event_id=event_id)
        
        # 触发本地处理器
        handlers = self._handlers.get(event_type, [])
        for handler in handlers:
            try:
                await handler(event_data)
            except Exception as e:
                logger.error("事件处理器错误", error=str(e), event_id=event_id)
        
        return event_id
    
    async def publish_execution_start(
        self,
        user_id: str,
        agent_id: str,
        execution_id: str,
        method: str,
        params: Dict,
        db: Optional[AsyncSession] = None
    ) -> str:
        """发布执行开始事件"""
        return await self.publish(
            event_type="execution.start",
            user_id=user_id,
            agent_id=agent_id,
            execution_id=execution_id,
            payload={
                "method": method,
                "params_keys": list(params.keys()) if params else [],
                "started_at": datetime.utcnow().isoformat(),
            },
            db=db,
        )
    
    async def publish_execution_end(
        self,
        user_id: str,
        agent_id: str,
        execution_id: str,
        success: bool,
        duration_ms: int,
        eu_consumed: float,
        db: Optional[AsyncSession] = None
    ) -> str:
        """发布执行结束事件"""
        return await self.publish(
            event_type="execution.end",
            user_id=user_id,
            agent_id=agent_id,
            execution_id=execution_id,
            payload={
                "success": success,
                "duration_ms": duration_ms,
                "eu_consumed": eu_consumed,
                "ended_at": datetime.utcnow().isoformat(),
            },
            db=db,
        )
    
    async def publish_balance_event(
        self,
        user_id: str,
        event_subtype: str,
        amount: float,
        balance_before: float,
        balance_after: float,
        description: str,
        db: Optional[AsyncSession] = None
    ) -> str:
        """发布余额变动事件"""
        return await self.publish(
            event_type=f"balance.{event_subtype}",
            user_id=user_id,
            payload={
                "amount": amount,
                "balance_before": balance_before,
                "balance_after": balance_after,
                "description": description,
            },
            db=db,
        )
    
    async def publish_quota_event(
        self,
        user_id: str,
        quota_type: str,
        current_value: float,
        limit_value: float,
        db: Optional[AsyncSession] = None
    ) -> str:
        """发布配额事件"""
        exceeded = current_value >= limit_value
        return await self.publish(
            event_type="quota.exceeded" if exceeded else "quota.check",
            user_id=user_id,
            payload={
                "quota_type": quota_type,
                "current_value": current_value,
                "limit_value": limit_value,
                "usage_percent": round(current_value / limit_value * 100, 2) if limit_value > 0 else 0,
            },
            db=db,
        )


class EventConsumer:
    """事件消费者"""
    
    def __init__(self, nats_client=None, db_session_factory=None):
        self.nats = nats_client
        self.db_factory = db_session_factory
        self._running = False
        self._subscriptions = []
    
    async def start(self):
        """启动事件消费"""
        if not self.nats:
            logger.warning("NATS未配置，跳过事件消费")
            return
        
        self._running = True
        
        # 订阅计费相关事件
        subjects = [
            "billing.execution_start",
            "billing.execution_end",
            "billing.balance_*",
            "billing.quota_*",
        ]
        
        for subject in subjects:
            try:
                sub = await self.nats.subscribe(subject, cb=self._handle_message)
                self._subscriptions.append(sub)
                logger.info("订阅成功", subject=subject)
            except Exception as e:
                logger.error("订阅失败", subject=subject, error=str(e))
    
    async def stop(self):
        """停止事件消费"""
        self._running = False
        for sub in self._subscriptions:
            await sub.unsubscribe()
        self._subscriptions.clear()
    
    async def _handle_message(self, msg):
        """处理消息"""
        try:
            data = json.loads(msg.data.decode())
            event_type = data.get("event_type")
            event_id = data.get("event_id")
            
            logger.info("收到事件", event_type=event_type, event_id=event_id)
            
            # 更新事件状态
            if self.db_factory:
                async with self.db_factory() as db:
                    await self._mark_event_processed(event_id, db)
            
        except Exception as e:
            logger.error("消息处理失败", error=str(e))
    
    async def _mark_event_processed(self, event_id: str, db: AsyncSession):
        """标记事件已处理"""
        await db.execute(
            update(BillingEvent)
            .where(BillingEvent.event_id == event_id)
            .values(status="completed", processed_at=datetime.utcnow())
        )
        await db.commit()


async def get_pending_events(
    limit: int,
    db: AsyncSession
) -> List[Dict]:
    """
    获取待处理的事件
    
    Args:
        limit: 数量限制
        db: 数据库会话
        
    Returns:
        待处理事件列表
    """
    result = await db.execute(
        select(BillingEvent)
        .where(BillingEvent.status == "pending")
        .order_by(BillingEvent.created_at)
        .limit(limit)
    )
    
    events = result.scalars().all()
    
    return [
        {
            "id": str(event.id),
            "eventId": event.event_id,
            "eventType": event.event_type,
            "userId": str(event.user_id),
            "agentId": str(event.agent_id) if event.agent_id else None,
            "payload": event.payload,
            "status": event.status,
            "createdAt": event.created_at.isoformat(),
        }
        for event in events
    ]


async def retry_failed_events(
    max_retries: int,
    db: AsyncSession
) -> int:
    """
    重试失败的事件
    
    Args:
        max_retries: 最大重试次数
        db: 数据库会话
        
    Returns:
        重试的事件数
    """
    result = await db.execute(
        select(BillingEvent)
        .where(
            and_(
                BillingEvent.status == "failed",
                BillingEvent.retry_count < max_retries,
            )
        )
        .limit(100)
    )
    
    events = result.scalars().all()
    count = 0
    
    for event in events:
        event.status = "pending"
        event.retry_count += 1
        count += 1
    
    await db.commit()
    return count


async def get_event_stats(
    start_date: datetime,
    end_date: datetime,
    db: AsyncSession
) -> Dict:
    """
    获取事件统计
    
    Args:
        start_date: 开始日期
        end_date: 结束日期
        db: 数据库会话
        
    Returns:
        事件统计
    """
    from sqlalchemy import func
    
    # 按状态统计
    status_stats = await db.execute(
        select(
            BillingEvent.status,
            func.count(BillingEvent.id).label("count"),
        )
        .where(
            and_(
                BillingEvent.created_at >= start_date,
                BillingEvent.created_at <= end_date,
            )
        )
        .group_by(BillingEvent.status)
    )
    
    status_breakdown = {row.status: row.count for row in status_stats.all()}
    
    # 按事件类型统计
    type_stats = await db.execute(
        select(
            BillingEvent.event_type,
            func.count(BillingEvent.id).label("count"),
        )
        .where(
            and_(
                BillingEvent.created_at >= start_date,
                BillingEvent.created_at <= end_date,
            )
        )
        .group_by(BillingEvent.event_type)
    )
    
    type_breakdown = {row.event_type: row.count for row in type_stats.all()}
    
    return {
        "startDate": start_date.isoformat(),
        "endDate": end_date.isoformat(),
        "byStatus": status_breakdown,
        "byType": type_breakdown,
        "total": sum(status_breakdown.values()),
    }


# 全局事件发布者实例
_event_publisher: Optional[EventPublisher] = None


def get_event_publisher() -> EventPublisher:
    """获取事件发布者实例"""
    global _event_publisher
    if _event_publisher is None:
        _event_publisher = EventPublisher()
    return _event_publisher


def set_event_publisher(publisher: EventPublisher):
    """设置事件发布者实例"""
    global _event_publisher
    _event_publisher = publisher


async def setup_nats_handlers():
    """
    设置NATS事件处理器
    
    在服务启动时调用，初始化事件发布和消费
    """
    try:
        from .state import get_state
        
        state = get_state()
        nats_client = state.nats_client
        
        if nats_client:
            # 初始化带NATS的事件发布者
            publisher = EventPublisher(nats_client=nats_client)
            set_event_publisher(publisher)
            logger.info("NATS事件发布者初始化成功")
        else:
            # 无NATS时使用本地发布者
            publisher = EventPublisher()
            set_event_publisher(publisher)
            logger.warning("NATS未连接，使用本地事件发布者")
        
    except Exception as e:
        logger.error("NATS事件处理器初始化失败", error=str(e))
        # 使用默认本地发布者
        publisher = EventPublisher()
        set_event_publisher(publisher)