LQConstPlanWriterAgent/core/base/progress_manager.py

import json
import time
import asyncio
from typing import Dict, Any, Optional, List
from datetime import datetime
from urllib.parse import quote

from foundation.observability.logger.loggering import write_logger as logger
from foundation.infrastructure.config import config_handler


class ProgressManager:
    """任务进度管理器"""

    def __init__(self):
        self.redis_client = None
        self.redis_connected = False
        self.current_data = {}
        self.stream_events = {}
        self._missing_progress_log_state = {}
        self._init_redis()

    def _log_missing_progress(self, callback_task_id: str, storage_name: str, log_func=logger.debug):
        """Avoid flooding logs while SSE polling waits for worker progress."""
        now = time.time()
        state = self._missing_progress_log_state.get(callback_task_id, {"count": 0, "last_log_at": 0})
        state["count"] += 1

        should_log = state["count"] == 1 or now - state["last_log_at"] >= 30
        if should_log:
            suffix = ""
            if state["count"] > 1:
                suffix = f" (已连续 {state['count']} 次，后续每30秒提示一次)"
            log_func(f"{storage_name}中未找到任务进度: {callback_task_id}{suffix}")
            state["last_log_at"] = now

        self._missing_progress_log_state[callback_task_id] = state

    def _clear_missing_progress_state(self, callback_task_id: str):
        self._missing_progress_log_state.pop(callback_task_id, None)

    def _init_redis(self):
        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')

            if redis_password:
                redis_url = f"redis://:{quote(redis_password, safe='')}@{redis_host}:{redis_port}/{redis_db}"
            else:
                redis_url = f"redis://{redis_host}:{redis_port}/{redis_db}"

            logger.debug(f"ProgressManager connecting to Redis: {redis_host}:{redis_port}/{redis_db}")

            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:
        return f"current:{callback_task_id}"

    async def _get_stream_event_key(self, callback_task_id: str) -> str:
        return f"stream_events:{callback_task_id}"

    async def initialize_progress(self, callback_task_id: str, user_id: str, stages: list):
        try:
            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:
                try:
                    redis_key = await self._get_redis_key(callback_task_id)
                    existing_progress_json = self.redis_client.get(redis_key)
                    if existing_progress_json:
                        try:
                            existing_progress = json.loads(existing_progress_json)
                            current_data.update(existing_progress)
                        except Exception as parse_e:
                            logger.warning(f"解析已有任务进度失败，将重新初始化: {callback_task_id}, {parse_e}")
                    self.redis_client.setex(
                        redis_key,
                        3600,
                        json.dumps(current_data)
                    )
                    stream_event_key = await self._get_stream_event_key(callback_task_id)
                    self.redis_client.delete(stream_event_key)
                    logger.info(f"初始化任务进度列表")
                except Exception as redis_e:
                    logger.warning(f"初始化进度到Redis失败: {callback_task_id}, {redis_e}")
                    self.current_data[callback_task_id] = current_data
                    self.stream_events[callback_task_id] = []
                    logger.info(f"降级使用内存存储: {callback_task_id}")
            else:
                if callback_task_id in self.current_data:
                    current_data.update(self.current_data[callback_task_id])
                self.current_data[callback_task_id] = current_data
                self.stream_events[callback_task_id] = []
                logger.info(f"初始化任务进度到内存: {callback_task_id}")

            self._clear_missing_progress_state(callback_task_id)

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

    async def append_stream_event(self, callback_task_id: str, event_data: Dict[str, Any]):
        """追加流式事件，供 SSE 轮询端按事件顺序吐给前端。"""
        try:
            if self.redis_connected:
                stream_event_key = await self._get_stream_event_key(callback_task_id)
                self.redis_client.rpush(
                    stream_event_key,
                    json.dumps(event_data, ensure_ascii=False)
                )
                self.redis_client.expire(stream_event_key, 3600)
            else:
                if not hasattr(self, 'stream_events'):
                    self.stream_events = {}
                self.stream_events.setdefault(callback_task_id, []).append(event_data)
        except Exception as e:
            logger.error(f"追加流式事件失败: {callback_task_id}, {str(e)}")
            raise

    async def pop_stream_events(self, callback_task_id: str, max_events: int = 100) -> List[Dict[str, Any]]:
        """按 FIFO 顺序取出待发送的流式事件。"""
        events: List[Dict[str, Any]] = []
        try:
            if self.redis_connected:
                stream_event_key = await self._get_stream_event_key(callback_task_id)
                for _ in range(max_events):
                    event_json = self.redis_client.lpop(stream_event_key)
                    if not event_json:
                        break
                    try:
                        events.append(json.loads(event_json))
                    except Exception as parse_e:
                        logger.warning(f"解析流式事件失败: {callback_task_id}, {parse_e}")
            else:
                if not hasattr(self, 'stream_events'):
                    self.stream_events = {}
                queued_events = self.stream_events.get(callback_task_id, [])
                events = queued_events[:max_events]
                self.stream_events[callback_task_id] = queued_events[max_events:]
        except Exception as e:
            logger.error(f"获取流式事件失败: {callback_task_id}, {str(e)}")

        return events




    async def update_stage_progress(self, callback_task_id: str, stage_name: str = None, current: int = None, status: str = None, message: str = None, issues=None, user_id: str = None, overall_task_status: str = None, event_type: str = "processing"):
        """更新阶段进度 - 除callback_task_id外，其他参数都可选

        Args:
            callback_task_id: 回调任务ID（必需）
            stage_name: 阶段名称（可选）
            current: 当前进度（可选）
            status: 状态（可选）
            message: 消息（可选）
            issues: 问题列表（可选）
            user_id: 用户ID（可选）
            overall_task_status: 整体任务状态（可选）
            event_type: SSE事件类型（可选，默认为"processing"）
        """
        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)
                    self._clear_missing_progress_state(callback_task_id)
                else:
                    self._log_missing_progress(callback_task_id, "Redis", logger.warning)
                    return
            else:
                # 从内存读取
                if callback_task_id in self.current_data:
                    task_progress = self.current_data[callback_task_id]
                    self._clear_missing_progress_state(callback_task_id)
                else:
                    self._log_missing_progress(callback_task_id, "内存", logger.warning)
                    return

            # 更新进度数据 - 只有非空参数才更新
            if current is not None:
                task_progress["current"] = current
            if stage_name is not None:
                task_progress["stage_name"] = stage_name
            if status is not None:
                task_progress["status"] = status
            if message is not None:
                task_progress["message"] = message
            task_progress["updated_at"] = datetime.now().isoformat()
            if issues is not None:
                task_progress["issues"] = issues
            else:
                task_progress["issues"] = []
            if user_id is not None:
                task_progress["user_id"] = user_id
            if overall_task_status is not None:
                task_progress["overall_task_status"] = overall_task_status
            elif "overall_task_status" not in task_progress:
                task_progress["overall_task_status"] = "processing"
            if event_type is not None:
                task_progress["event_type"] = event_type
                logger.debug(f"设置event_type: {event_type} for {callback_task_id}")
            else:
                logger.debug(f"event_type为None，不设置 for {callback_task_id}")

            try:
                if self.redis_connected:
                    try:
                        self.redis_client.setex(
                            redis_key,
                            3600,  # 1小时过期
                            json.dumps(task_progress)
                        )
                        actual_current = task_progress.get("current")
                        if current is not None:
                            logger.debug(f"更新进度到Redis: {callback_task_id}, 进度: {actual_current}%")
                        else:
                            logger.debug(f"更新进度到Redis: {callback_task_id}, 进度保持: {actual_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
                    actual_current = task_progress.get("current")
                    if current is not None:
                        logger.debug(f"更新进度到内存: {callback_task_id}, 进度: {actual_current}%")
                    else:
                        logger.debug(f"更新进度到内存: {callback_task_id}, 进度保持: {actual_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

            # 进度已保存到 Redis，SSE 由主进程通过轮询获取
            actual_current = task_progress.get("current")
            if current is not None:
                logger.debug(f"进度已更新到Redis: {callback_task_id}, current={actual_current}%")
            else:
                logger.debug(f"进度已更新到Redis: {callback_task_id}, current={actual_current}% (保持)")

        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)
                    self._clear_missing_progress_state(callback_task_id)
                else:
                    self._log_missing_progress(callback_task_id, "Redis", logger.debug)
                    return None
            else:
                # 从内存读取
                if hasattr(self, 'current_data') and callback_task_id in self.current_data:
                    task_progress = self.current_data[callback_task_id]
                    self._clear_missing_progress_state(callback_task_id)
                else:
                    self._log_missing_progress(callback_task_id, "内存", logger.debug)
                    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())
            elif isinstance(updated_at, (int, float)):
                updated_at_timestamp = int(updated_at)
            else:
                updated_at_timestamp = int(updated_at.timestamp())

            # 构建返回数据
            result = {
                "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
            }

            # 添加可选字段
            if "issues" in task_progress:
                result["issues"] = task_progress["issues"]
            if "event_type" in task_progress:
                result["event_type"] = task_progress["event_type"]

            return result

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

    async def complete_task(self, callback_task_id: str, user_id: str = None, current_data: dict = None):
        """标记任务完成 - 使用单一同步强制关闭逻辑

        Args:
            callback_task_id: 回调任务ID
            user_id: 用户ID
            current_data: 包含 overall_task_status 等数据的字典
        """
        try:
            # 先更新 Redis 中的状态（让 SSE 轮询能检测到）
            if current_data and current_data.get("overall_task_status"):
                progress_data = await self.get_progress(callback_task_id)
                if progress_data:
                    # 更新状态字段
                    progress_data["overall_task_status"] = current_data.get("overall_task_status")
                    progress_data["status"] = current_data.get("status", current_data.get("overall_task_status"))
                    progress_data["message"] = current_data.get("message", "任务已完成")
                    progress_data["updated_at"] = datetime.now().isoformat()
                    if current_data.get("error"):
                        progress_data["error"] = current_data.get("error")

                    # 保存回 Redis
                    if self.redis_connected:
                        redis_key = await self._get_redis_key(callback_task_id)
                        self.redis_client.setex(
                            redis_key,
                            3600,
                            json.dumps(progress_data)
                        )
                        logger.info(f"任务状态已更新到Redis: {callback_task_id}, status={current_data.get('overall_task_status')}")
                    else:
                        if not hasattr(self, 'current_data'):
                            self.current_data = {}
                        self.current_data[callback_task_id] = progress_data

            # SSE 连接由主进程管理（通过 Redis 轮询），Worker 无需关闭
            logger.info(f"任务状态已更新，等待主进程轮询检测: {callback_task_id}")
        except Exception as e:
            logger.error(f"标记任务完成失败: {str(e)}")
            raise