taiji-AI-PAD/services/mcp-server/fix_fake_quota.py

#!/usr/bin/env python3
"""
修复脚本：清理历史假用量数据

此脚本用于清理因旧接口（/api/user/tools/generate）产生的假用量数据。
旧接口会在数据库中创建Agent记录并扣除配额，但不实际部署到K8s，导致quota表中记录了用量但实际没有Pod运行。

脚本逻辑：
1. 备份当前quota数据到JSON文件
2. 遍历所有租户的TenantCustomAgentQuota记录
3. 对每个租户，重新统计agents表中type='custom'且status='active'的真实运行Agent
4. 更新quota表的cpu_used、memory_used、agent_count为真实值
5. 生成修复报告

使用方法：
在mcp-server容器内运行：
python fix_fake_quota.py

注意：此脚本会修改数据库，请务必先备份数据库！
"""

import asyncio
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any

from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession

from database import get_db
from models import Agent, TenantCustomAgentQuota

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class QuotaFixer:
    """配额修复器"""

    def __init__(self):
        self.backup_file = Path("/app/logs/quota_backup.json")
        self.report_file = Path("/app/logs/quota_fix_report.json")

    async def backup_current_quota(self, db: AsyncSession) -> Dict[str, Any]:
        """备份当前quota数据"""
        logger.info("开始备份当前quota数据...")

        result = await db.execute(
            select(TenantCustomAgentQuota)
        )
        quotas = result.scalars().all()

        backup_data = {
            "backup_time": datetime.utcnow().isoformat(),
            "quotas": []
        }

        for quota in quotas:
            backup_data["quotas"].append({
                "tenant_id": str(quota.tenant_id),
                "cpu_quota": float(quota.cpu_quota or 0),
                "memory_quota": float(quota.memory_quota or 0),
                "cpu_used": float(quota.cpu_used or 0),
                "memory_used": float(quota.memory_used or 0),
                "agent_count": quota.agent_count or 0
            })

        # 确保logs目录存在
        self.backup_file.parent.mkdir(exist_ok=True)

        # 保存备份
        with open(self.backup_file, 'w', encoding='utf-8') as f:
            json.dump(backup_data, f, indent=2, ensure_ascii=False)

        logger.info(f"备份完成，已保存到: {self.backup_file}")
        return backup_data

    async def calculate_real_usage(self, db: AsyncSession, tenant_id: str) -> Dict[str, float]:
        """计算租户的真实资源使用量

        只统计type='custom'且status='active'的Agent
        """
        # 查询该租户所有运行中的自定义Agent
        result = await db.execute(
            select(
                func.count(Agent.id).label('agent_count'),
                func.sum(Agent.cpu).label('total_cpu'),
                func.sum(Agent.memory).label('total_memory')
            )
            .where(Agent.owner_id == tenant_id)
            .where(Agent.type == 'custom')
            .where(Agent.status == 'active')
        )

        row = result.first()

        real_usage = {
            'agent_count': row.agent_count or 0,
            'cpu_used': float(row.total_cpu or 0),
            'memory_used': float(row.total_memory or 0)
        }

        logger.info(f"租户 {tenant_id} 真实用量: {real_usage}")
        return real_usage

    async def fix_tenant_quota(self, db: AsyncSession, quota: TenantCustomAgentQuota) -> Dict[str, Any]:
        """修复单个租户的quota"""
        tenant_id = str(quota.tenant_id)

        # 获取当前记录的用量
        current_usage = {
            'cpu_used': float(quota.cpu_used or 0),
            'memory_used': float(quota.memory_used or 0),
            'agent_count': quota.agent_count or 0
        }

        # 计算真实用量
        real_usage = await self.calculate_real_usage(db, tenant_id)

        # 计算差异
        diff = {
            'cpu_diff': real_usage['cpu_used'] - current_usage['cpu_used'],
            'memory_diff': real_usage['memory_used'] - current_usage['memory_used'],
            'count_diff': real_usage['agent_count'] - current_usage['agent_count']
        }

        # 更新quota
        quota.cpu_used = real_usage['cpu_used']
        quota.memory_used = real_usage['memory_used']
        quota.agent_count = real_usage['agent_count']

        await db.commit()
        await db.refresh(quota)

        logger.info(f"租户 {tenant_id} quota已更新: {current_usage} -> {real_usage}")

        return {
            'tenant_id': tenant_id,
            'before': current_usage,
            'after': real_usage,
            'diff': diff
        }

    async def fix_all_quotas(self, db: AsyncSession) -> Dict[str, Any]:
        """修复所有租户的quota"""
        logger.info("开始修复所有租户的quota...")

        # 获取所有quota记录
        result = await db.execute(
            select(TenantCustomAgentQuota)
        )
        quotas = result.scalars().all()

        logger.info(f"找到 {len(quotas)} 个租户的quota记录")

        # 备份数据
        backup_data = await self.backup_current_quota(db)

        # 修复报告
        fix_report = {
            "fix_time": datetime.utcnow().isoformat(),
            "total_tenants": len(quotas),
            "fixed_tenants": [],
            "summary": {
                "total_fake_cpu": 0.0,
                "total_fake_memory": 0.0,
                "total_fake_agents": 0
            }
        }

        for quota in quotas:
            try:
                fix_result = await self.fix_tenant_quota(db, quota)
                fix_report["fixed_tenants"].append(fix_result)

                # 累加假用量
                fix_report["summary"]["total_fake_cpu"] += fix_result["diff"]["cpu_diff"]
                fix_report["summary"]["total_fake_memory"] += fix_result["diff"]["memory_diff"]
                fix_report["summary"]["total_fake_agents"] += fix_result["diff"]["count_diff"]

            except Exception as e:
                logger.error(f"修复租户 {quota.tenant_id} 失败: {str(e)}")
                fix_report["fixed_tenants"].append({
                    "tenant_id": str(quota.tenant_id),
                    "error": str(e)
                })

        # 保存修复报告
        with open(self.report_file, 'w', encoding='utf-8') as f:
            json.dump(fix_report, f, indent=2, ensure_ascii=False)

        logger.info(f"修复完成！报告已保存到: {self.report_file}")
        logger.info(f"总共清理假用量: CPU {fix_report['summary']['total_fake_cpu']:.2f}核, "
                   f"内存 {fix_report['summary']['total_fake_memory']:.2f}GB, "
                   f"Agent {fix_report['summary']['total_fake_agents']}个")

        return fix_report

    async def run(self):
        """运行修复任务"""
        logger.info("开始执行配额修复任务...")

        async for db in get_db():
            try:
                report = await self.fix_all_quotas(db)
                logger.info("配额修复任务完成！")
                return report
            except Exception as e:
                logger.error(f"修复任务失败: {str(e)}")
                raise
            finally:
                await db.close()


async def main():
    """主函数"""
    fixer = QuotaFixer()
    await fixer.run()


if __name__ == "__main__":
    asyncio.run(main())