life-echo/api/app/core/task_tracker.py

"""
任务追踪服务：追踪 Celery 任务状态（从 services 迁入 core）
"""

import json
from datetime import datetime, timezone
from typing import Any, Dict, List

from app.core.logging import get_logger
from app.core.redis import redis_service
from app.core.runtime_constants import redis_defaults

logger = get_logger(__name__)


class TaskTracker:
    """任务追踪器，使用 Redis 存储任务状态"""

    KEY_PREFIX = "task:user:"

    @property
    def task_ttl(self) -> int:
        return int(redis_defaults.task_tracker_ttl_seconds)

    async def _refresh_key_ttl(self, key: str) -> None:
        client = await redis_service.get_client()
        await client.expire(key, self.task_ttl)

    async def add_task(
        self, user_id: str, task_id: str, task_type: str = "memoir"
    ) -> bool:
        try:
            client = await redis_service.get_client()
            key = f"{self.KEY_PREFIX}{user_id}:tasks"
            task_info = {
                "task_id": task_id,
                "task_type": task_type,
                "status": "pending",
                "created_at": datetime.now(timezone.utc).isoformat(),
            }
            await client.hset(key, task_id, json.dumps(task_info))
            await client.expire(key, self.task_ttl)
            logger.debug("任务已记录: user_id={}, task_id={}", user_id, task_id)
            return True
        except Exception as e:
            logger.error("记录任务失败: {}", e)
            return False

    async def update_task_status(
        self, user_id: str, task_id: str, status: str, result: Any = None
    ) -> bool:
        try:
            client = await redis_service.get_client()
            key = f"{self.KEY_PREFIX}{user_id}:tasks"
            data = await client.hget(key, task_id)
            if data:
                task_info = json.loads(data)
            else:
                task_info = {"task_id": task_id}
            task_info["status"] = status
            task_info["updated_at"] = datetime.now(timezone.utc).isoformat()
            if result is not None:
                task_info["result"] = result
            await client.hset(key, task_id, json.dumps(task_info))
            await self._refresh_key_ttl(key)
            return True
        except Exception as e:
            logger.error("更新任务状态失败: {}", e)
            return False

    async def get_user_tasks(self, user_id: str) -> List[Dict]:
        try:
            client = await redis_service.get_client()
            key = f"{self.KEY_PREFIX}{user_id}:tasks"
            tasks_data = await client.hgetall(key)
            return [json.loads(data) for data in tasks_data.values()]
        except Exception as e:
            logger.error("获取用户任务失败: {}", e)
            return []

    async def get_pending_tasks(self, user_id: str) -> List[Dict]:
        tasks = await self.get_user_tasks(user_id)
        return [t for t in tasks if t.get("status") in ("pending", "running")]

    async def check_tasks_status(self, user_id: str) -> Dict:
        tasks = await self.get_user_tasks(user_id)
        status_counts: Dict[str, int] = {
            "total": len(tasks),
            "pending": 0,
            "running": 0,
            "success": 0,
            "failure": 0,
        }
        for task in tasks:
            status = task.get("status", "pending")
            if status in status_counts:
                status_counts[status] += 1
        status_counts["all_completed"] = (
            status_counts["total"] > 0
            and status_counts["pending"] == 0
            and status_counts["running"] == 0
        )
        return status_counts

    async def clear_user_tasks(self, user_id: str) -> bool:
        try:
            client = await redis_service.get_client()
            key = f"{self.KEY_PREFIX}{user_id}:tasks"
            await client.delete(key)
            return True
        except Exception as e:
            logger.error("清除用户任务失败: {}", e)
            return False

    async def remove_task(self, user_id: str, task_id: str) -> bool:
        try:
            client = await redis_service.get_client()
            key = f"{self.KEY_PREFIX}{user_id}:tasks"
            await client.hdel(key, task_id)
            if await client.exists(key):
                await self._refresh_key_ttl(key)
            return True
        except Exception as e:
            logger.error("移除任务失败: {}", e)
            return False


task_tracker = TaskTracker()