"""PipelineCheckpoint — 阶段级检查点与断点续跑 (U7) 在 TeamOrchestrator 阶段完成后保存 checkpoint 到 Redis（或内存降级）， 崩溃后可通过 resume(plan_id) 从最后完成阶段恢复。 复用 PipelineStateRedis 的 _safe_redis_call 模式： Redis 失败时降级到内存 dict，不阻断执行。 键命名：agentkit:pipeline:checkpoint:{plan_id}:{phase_id} TTL：7 天（与 PipelineStateRedis._TTL_SECONDS 一致） """ from __future__ import annotations import json import logging import time from dataclasses import asdict, dataclass, field from datetime import datetime, timezone from typing import Any logger = logging.getLogger(__name__) _TTL_SECONDS = 7 * 24 * 3600 # 7 days _KEY_PREFIX = "agentkit:pipeline:checkpoint" @dataclass class CheckpointData: """单个阶段的 checkpoint 数据。""" plan_id: str phase_id: str phase_name: str phase_status: str phase_result: dict[str, Any] | None = None plan_status: str = "" saved_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) def to_dict(self) -> dict[str, Any]: return asdict(self) @classmethod def from_dict(cls, data: dict[str, Any]) -> CheckpointData: return cls( plan_id=data.get("plan_id", ""), phase_id=data.get("phase_id", ""), phase_name=data.get("phase_name", ""), phase_status=data.get("phase_status", ""), phase_result=data.get("phase_result"), plan_status=data.get("plan_status", ""), saved_at=data.get("saved_at", ""), ) class PipelineCheckpoint: """阶段级检查点存储 — Redis 优先，内存降级。 Usage:: checkpoint = PipelineCheckpoint(redis_client=redis) await checkpoint.save(plan.id, phase, plan.status.value) last = await checkpoint.load(plan.id) if last: # resume from last completed phase """ def __init__( self, redis_client: Any = None, prefix: str = _KEY_PREFIX, ttl_seconds: int = _TTL_SECONDS, ) -> None: self._redis = redis_client self._prefix = prefix self._ttl = ttl_seconds # 内存降级存储：plan_id → {phase_id → CheckpointData} # P1 #6: 改用 dict keyed by phase_id，避免重复 append self._memory: dict[str, dict[str, CheckpointData]] = {} # 内存降级存储：plan_id → (plan_dict, saved_timestamp) self._memory_plans: dict[str, tuple[dict[str, Any], float]] = {} def _is_expired(self, saved_at: str) -> bool: """检查 checkpoint 是否已过期（内存模式 TTL）。""" if not saved_at: return False try: saved_time = datetime.fromisoformat(saved_at) age = (datetime.now(timezone.utc) - saved_time).total_seconds() return age > self._ttl except (ValueError, TypeError): return False def _key(self, plan_id: str, phase_id: str) -> str: return f"{self._prefix}:{plan_id}:{phase_id}" def _index_key(self, plan_id: str) -> str: """Redis Sorted Set 索引键，用于列出某 plan 的所有 checkpoint。""" return f"{self._prefix}:index:{plan_id}" def _plan_key(self, plan_id: str) -> str: """完整 plan JSON 的存储键。""" return f"{self._prefix}:plan:{plan_id}" async def save_plan(self, plan: Any) -> None: """保存完整 TeamPlan（用于 resume 重建）。 Args: plan: TeamPlan 对象（需要有 to_dict() 方法） """ plan_id = plan.id plan_dict = plan.to_dict() if hasattr(plan, "to_dict") else {"id": plan_id} # 内存降级（带 TTL 时间戳） self._memory_plans[plan_id] = (plan_dict, time.time()) # 尝试写入 Redis if self._redis is not None: try: await self._redis.set(self._plan_key(plan_id), json.dumps(plan_dict), ex=self._ttl) except Exception as e: logger.warning(f"PipelineCheckpoint.save_plan Redis failed for plan {plan_id}: {e}") async def load_plan(self, plan_id: str) -> dict[str, Any] | None: """加载完整 plan JSON。""" # 优先 Redis if self._redis is not None: try: raw = await self._redis.get(self._plan_key(plan_id)) if raw: return json.loads(raw) except Exception as e: logger.warning(f"PipelineCheckpoint.load_plan Redis failed for plan {plan_id}: {e}") # 内存降级（检查 TTL 过期） entry = self._memory_plans.get(plan_id) if entry is None: return None plan_dict, saved_at = entry if time.time() - saved_at > self._ttl: # 已过期，清除 del self._memory_plans[plan_id] return None return plan_dict async def save(self, plan_id: str, phase: Any, plan_status: str) -> None: """保存阶段 checkpoint。 Args: plan_id: 计划 ID phase: PlanPhase 对象（需要有 id, name, status, result 属性） plan_status: 计划当前状态 """ phase_id = getattr(phase, "id", str(phase)) phase_name = getattr(phase, "name", "") phase_status = getattr(phase, "status", "") if hasattr(phase_status, "value"): phase_status = phase_status.value # 序列化 phase.result（可能是 offloaded dict with _ref_key） phase_result = getattr(phase, "result", None) if phase_result is not None and not isinstance(phase_result, dict): phase_result = {"content": str(phase_result)} data = CheckpointData( plan_id=plan_id, phase_id=phase_id, phase_name=phase_name, phase_status=str(phase_status), phase_result=phase_result, plan_status=plan_status, ) # P1 #6: 内存降级用 dict keyed by phase_id，覆盖重复 checkpoint self._memory.setdefault(plan_id, {})[phase_id] = data # 尝试写入 Redis if self._redis is not None: try: score = time.time() pipe = self._redis.pipeline() pipe.set(self._key(plan_id, phase_id), json.dumps(data.to_dict()), ex=self._ttl) pipe.zadd(self._index_key(plan_id), {phase_id: score}) await pipe.execute() except Exception as e: logger.warning( f"PipelineCheckpoint.save Redis failed for plan {plan_id}, " f"phase {phase_id}: {e} — using memory fallback" ) async def load(self, plan_id: str) -> CheckpointData | None: """加载最后完成的阶段 checkpoint。 Returns: 最后一个 COMPLETED 阶段的 CheckpointData，或 None。 """ checkpoints = await self.list_checkpoints(plan_id) if not checkpoints: return None # 返回最后一个 COMPLETED 阶段 completed = [c for c in checkpoints if c.phase_status == "completed"] if not completed: return None return completed[-1] async def list_checkpoints(self, plan_id: str) -> list[CheckpointData]: """列出某 plan 的所有 checkpoint（按保存时间排序）。""" # 优先从 Redis 读取 if self._redis is not None: try: phase_ids = await self._redis.zrange(self._index_key(plan_id), 0, -1) if not phase_ids: # Redis 无数据，检查内存（过滤过期） return [ c for c in self._memory.get(plan_id, {}).values() if not self._is_expired(c.saved_at) ] # 批量 GET（pipeline 避免 N+1 往返） pipe = self._redis.pipeline() for phase_id in phase_ids: pipe.get(self._key(plan_id, phase_id)) raws = await pipe.execute() results: list[CheckpointData] = [] for raw in raws: if raw: data = json.loads(raw) results.append(CheckpointData.from_dict(data)) return results except Exception as e: logger.warning( f"PipelineCheckpoint.list_checkpoints Redis failed for " f"plan {plan_id}: {e} — using memory fallback" ) # 内存降级（过滤过期 checkpoint） return [ c for c in self._memory.get(plan_id, {}).values() if not self._is_expired(c.saved_at) ] async def clear(self, plan_id: str) -> None: """清除某 plan 的所有 checkpoint。""" # 清除内存 self._memory.pop(plan_id, None) self._memory_plans.pop(plan_id, None) # 清除 Redis if self._redis is not None: try: phase_ids = await self._redis.zrange(self._index_key(plan_id), 0, -1) pipe = self._redis.pipeline() for phase_id in phase_ids: pipe.delete(self._key(plan_id, phase_id)) pipe.delete(self._index_key(plan_id)) pipe.delete(self._plan_key(plan_id)) await pipe.execute() except Exception as e: logger.warning(f"PipelineCheckpoint.clear Redis failed for plan {plan_id}: {e}")