""" 任务追踪服务:追踪 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()