LQAgentPlatform/core/base/progress_manager.py

"""
任务进度管理器
负责任务进度的存储、更新和查询
"""

import json
import asyncio
from typing import Dict, Any, Optional
from datetime import datetime

from foundation.logger.loggering import server_logger as logger
from foundation.base.config import config_handler

class SSECallbackManager:
    """SSE回调管理器 - 单例模式管理全局SSE回调"""
    _instance = None
    _callbacks = {}  # {callback_task_id: callback_function}

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def register_callback(self, callback_task_id: str, callback_func):
        """注册SSE回调函数"""
        self._callbacks[callback_task_id] = callback_func
        logger.info(f"SSE回调注册, 当前注册数: {len(self._callbacks)}")

    def unregister_callback(self, callback_task_id: str):
        """注销SSE回调函数"""
        if callback_task_id in self._callbacks:
            del self._callbacks[callback_task_id]
            logger.info(f"SSE回调注销, 剩余注册数: {len(self._callbacks)}")

    async def trigger_callback(self, callback_task_id: str, current_data: dict):
        """触发SSE回调"""
        if callback_task_id in self._callbacks:
            try:
                # 直接异步执行回调，保持trace上下文
                await self._callbacks[callback_task_id](callback_task_id, current_data)
                logger.debug(f"SSE回调执行成功: {callback_task_id}")

                logger.debug(f"SSE回调已触发: {callback_task_id}, 当前注册回调数: {len(self._callbacks)}")
                return True

            except Exception as e:
                logger.error(f"SSE回调执行失败: {callback_task_id}, {e}")
                return False
        else:
            logger.debug(f"未找到SSE回调: {callback_task_id}, 当前注册回调数: {len(self._callbacks)}, 已注册ID: {list(self._callbacks.keys())}")
            return False

    def get_callbacks_count(self):
        """获取当前回调数量"""
        return len(self._callbacks)

    def clear_all_callbacks(self):
        """清空所有回调"""
        self._callbacks.clear()
        logger.info("已清空所有SSE回调")

# 全局SSE回调管理器实例
sse_callback_manager = SSECallbackManager()

class ProgressManager:
    """任务进度管理器 - 增长型进度管理版本"""

    def __init__(self):
        self.redis_client = None
        self.redis_connected = False
        self._init_redis()

    def _init_redis(self):
        """初始化Redis连接"""
        try:
            import redis

            redis_host = config_handler.get('redis', 'REDIS_HOST', 'localhost')
            redis_port = config_handler.get('redis', 'REDIS_PORT', '6379')
            redis_password = config_handler.get('redis', 'REDIS_PASSWORD', '')
            redis_db = config_handler.get('redis', 'REDIS_DB', '0')

            # 构建Redis连接URL
            if redis_password:
                redis_url = f"redis://:{redis_password}@{redis_host}:{redis_port}/{redis_db}"
            else:
                redis_url = f"redis://{redis_host}:{redis_port}/{redis_db}"

            logger.debug(f"ProgressManager连接Redis: {redis_url}")

            # 连接Redis
            self.redis_client = redis.from_url(redis_url, decode_responses=True)

            # 测试连接
            self.redis_client.ping()
            self.redis_connected = True
            logger.debug(f"ProgressManager Redis连接成功: {redis_host}:{redis_port}")

        except Exception as e:
            logger.error(f"ProgressManager Redis连接失败: {e}")
            self.redis_connected = False
            logger.warning("ProgressManager将使用内存存储作为备选方案")
            self.current_data = {}  # 备选内存存储

    async def _get_redis_key(self, callback_task_id: str) -> str:
        """获取Redis键名"""
        return f"current:{callback_task_id}"

    async def initialize_progress(self, callback_task_id: str, user_id: str, stages: list):
        """初始化进度记录"""
        try:

            # 设置总量为100（百分比模式）
            stage_name = stages[0]["stage_name"] if stages else ""
            message = "任务开始"

            current_data = {
                "user_id": user_id,
                "current": 0,
                "stage_name": "",
                "status": "准备开始",
                "message": "任务开始",
                "updated_at": datetime.now().isoformat(),
                "overall_task_status": "pending"
            }

            if self.redis_connected:
                # 使用同步Redis操作避免异步任务销毁问题
                try:
                    redis_key = await self._get_redis_key(callback_task_id)
                    self.redis_client.setex(
                        redis_key,
                        3600,  # 1小时过期
                        json.dumps(current_data)
                    )
                    logger.info(f"初始化任务进度列表")
                except Exception as redis_e:
                    logger.warning(f"初始化进度到Redis失败: {callback_task_id}, {redis_e}")
                    # 降级到内存存储
                    if not hasattr(self, 'current_data'):
                        self.current_data = {}
                    self.current_data[callback_task_id] = current_data
                    logger.info(f"降级使用内存存储: {callback_task_id}")
            else:
                # 使用内存存储
                if not hasattr(self, 'current_data'):
                    self.current_data = {}
                self.current_data[callback_task_id] = current_data
                logger.info(f"初始化任务进度到内存: {callback_task_id}")

        except Exception as e:
            logger.error(f"初始化进度失败: {str(e)}")
            raise

    async def update_stage_progress(self, callback_task_id: str, stage_name: str, current: int, status: str, message: str = ""):
        """更新阶段进度"""
        try:
            task_progress = None

            if self.redis_connected:
                # 从Redis读取
                redis_key = await self._get_redis_key(callback_task_id)
                progress_json = self.redis_client.get(redis_key)
                if progress_json:
                    task_progress = json.loads(progress_json)
                else:
                    logger.warning(f"Redis中未找到任务进度: {callback_task_id}")
                    return
            else:
                # 从内存读取
                if callback_task_id in self.current_data:
                    task_progress = self.current_data[callback_task_id]
                else:
                    logger.warning(f"内存中未找到任务进度: {callback_task_id}")
                    return

            # 更新进度数据
            task_progress["current"] = current
            task_progress["stage_name"] = stage_name
            task_progress["status"] = status
            task_progress["message"] = message
            task_progress["updated_at"] = datetime.now().isoformat()

            # 保留overall_task_status字段，不要被普通进度更新覆盖
            if "overall_task_status" not in task_progress:
                task_progress["overall_task_status"] = "processing"

            try:
                if self.redis_connected:
                    try:
                        self.redis_client.setex(
                            redis_key,
                            3600,  # 1小时过期
                            json.dumps(task_progress)
                        )
                        logger.debug(f"更新进度到Redis: {callback_task_id}, 进度: {current}%")
                    except Exception as sync_e:
                        logger.warning(f"同步Redis操作失败: {callback_task_id}, {sync_e}")
                        # 同步操作也失败时，降级到内存存储
                        if not hasattr(self, 'current_data'):
                            self.current_data = {}
                        self.current_data[callback_task_id] = task_progress
                        logger.debug(f"降级使用内存存储: {callback_task_id}")
                else:
                    if not hasattr(self, 'current_data'):
                        self.current_data = {}
                    self.current_data[callback_task_id] = task_progress
                    logger.debug(f"更新进度到内存: {callback_task_id}, 进度: {current}%")
            except Exception as e:
                logger.error(f"保存进度数据异常: {callback_task_id}, {e}")
                if not hasattr(self, 'current_data'):
                    self.current_data = {}
                self.current_data[callback_task_id] = task_progress

            # 触发SSE推送 - 使用全局回调管理器
            logger.debug(f"触发SSE推送: {callback_task_id}")
            updated_progress = await self.get_progress(callback_task_id)
            if updated_progress:
                await sse_callback_manager.trigger_callback(callback_task_id, updated_progress)

        except Exception as e:
            logger.error(f"更新阶段进度失败: {str(e)}")
            raise

    async def get_progress(self, callback_task_id: str) -> Optional[Dict[str, Any]]:
        """获取任务进度"""
        try:
            logger.debug(f"开始获取进度: {callback_task_id}, Redis连接状态: {self.redis_connected}")
            task_progress = None

            if self.redis_connected:
                # 从Redis读取
                redis_key = await self._get_redis_key(callback_task_id)
                logger.debug(f"Redis键: {redis_key}")
                progress_json = self.redis_client.get(redis_key)
                logger.debug(f"从Redis读取数据: {progress_json is not None}")
                if progress_json:
                    task_progress = json.loads(progress_json)
                else:
                    logger.debug(f"Redis中未找到任务进度: {callback_task_id}")
                    return None
            else:
                # 从内存读取
                if hasattr(self, 'current_data') and callback_task_id in self.current_data:
                    task_progress = self.current_data[callback_task_id]
                else:
                    logger.debug(f"内存中未找到任务进度: {callback_task_id}")
                    return None

            # 获取overall_task_status，默认为"pending"
            overall_task_status = task_progress.get("overall_task_status", "pending")

            # 转换时间戳
            updated_at = task_progress["updated_at"]
            if isinstance(updated_at, str):
                updated_at_timestamp = int(datetime.fromisoformat(updated_at).timestamp())
            else:
                updated_at_timestamp = int(updated_at.timestamp())

            return {
                "callback_task_id": callback_task_id,
                "user_id": task_progress["user_id"],
                "current": task_progress["current"],
                "stage_name": task_progress["stage_name"],
                "status": task_progress["status"],
                "message": task_progress["message"],
                "overall_task_status": overall_task_status,
                "updated_at": updated_at_timestamp
            }

        except Exception as e:
            logger.error(f"获取进度失败: {str(e)}")
            return None

    async def complete_task(self, callback_task_id: str):
        """标记任务完成"""
        try:
            task_progress = None
            logger.info(f"通知sse连接关闭: {callback_task_id}")
            if self.redis_connected:
                redis_key = await self._get_redis_key(callback_task_id)
                progress_json = self.redis_client.get(redis_key)
                if progress_json:
                    task_progress = json.loads(progress_json)
                else:
                    logger.warning(f"Redis中未找到任务进度: {callback_task_id}")
                    return
            else:
                # 从内存读取
                if hasattr(self, 'current_data') and callback_task_id in self.current_data:
                    task_progress = self.current_data[callback_task_id]
                else:
                    logger.warning(f"内存中未找到任务进度: {callback_task_id}")
                    return

            task_progress["status"] = "completed"
            task_progress["overall_task_status"] = "completed"
            task_progress["message"] = "任务已全部完成"
            task_progress["updated_at"] = datetime.now().isoformat()


            # 保存更新后的数据
            if self.redis_connected:
                self.redis_client.setex(
                    redis_key,
                    3600,
                    json.dumps(task_progress)
                )
            else:
                if hasattr(self, 'current_data'):
                    self.current_data[callback_task_id] = task_progress

            # 触发SSE进度更新推送
            completed_progress = await self.get_progress(callback_task_id)
            if completed_progress:
                await sse_callback_manager.trigger_callback(callback_task_id, completed_progress)
                logger.debug(f"SSE完成进度已推送: {callback_task_id}")
            else:
                logger.warning(f"无法获取完成进度数据: {callback_task_id}")
        except Exception as e:
            logger.error(f"标记任务完成失败: {str(e)}")
            raise