""" WorkflowStateManager - 步骤级状态持久化管理器 提供工作流状态的原子写入、备份轮转和崩溃恢复功能。 Purpose: 为工作流提供可靠的状态管理,支持断点续传和崩溃恢复 """ from __future__ import annotations import json import os import tempfile from pathlib import Path from typing import Optional from kb_workflow.pipeline.schemas import WorkflowState class WorkflowStateManager: """工作流状态管理器 管理工作流状态的持久化,包括: - 原子写入(使用临时文件 + os.replace) - 备份轮转(workflow_state.json -> workflow_state.json.bak) - 主文件损坏时从备份恢复 Attributes: STATE_FILENAME: 状态文件名 content_dir: 内容目录路径 state_path: 状态文件完整路径 backup_path: 备份文件完整路径 """ STATE_FILENAME = "workflow_state.json" def __init__(self, content_dir: Path): """初始化状态管理器 Args: content_dir: 内容目录路径,状态文件将保存在此目录下 """ self.content_dir = Path(content_dir) self.state_path = self.content_dir / self.STATE_FILENAME self.backup_path = self.content_dir / f"{self.STATE_FILENAME}.bak" def load(self) -> Optional[WorkflowState]: """加载工作流状态 加载策略: 1. 尝试加载主文件 (workflow_state.json) 2. 如果主文件损坏或不存在,尝试加载备份文件 3. 如果都不存在,返回 None Returns: WorkflowState 实例,如果文件不存在则返回 None Raises: ValueError: 如果文件格式无效且备份也无法恢复 """ def _try_load(path: Path) -> Optional[WorkflowState]: """尝试从指定路径加载状态""" if not path.exists(): return None try: with open(path, 'r', encoding='utf-8') as f: data = json.load(f) return WorkflowState.model_validate(data) except (json.JSONDecodeError, ValueError): return None # 先尝试主文件 state = _try_load(self.state_path) if state is not None: return state # 主文件失败,尝试备份 state = _try_load(self.backup_path) return state def save(self, state: WorkflowState) -> None: """原子保存工作流状态 保存策略(原子写入 + 备份轮转): 1. 创建父目录(如果不存在) 2. 如果状态文件存在,将其重命名为备份文件 3. 写入临时文件 4. 原子替换(os.replace) 原子写入确保: - 读取者永远不会看到部分写入的文件 - 崩溃时要么是旧文件,要么是新文件 Args: state: 要保存的工作流状态 Raises: OSError: 文件操作失败(权限、磁盘空间等) """ # 1. 创建父目录 self.content_dir.mkdir(parents=True, exist_ok=True) # 2. 备份轮转:如果状态文件存在,移动到备份 if self.state_path.exists(): # 删除旧备份 if self.backup_path.exists(): self.backup_path.unlink() # 将当前状态移动到备份 self.state_path.rename(self.backup_path) # 3. 写入临时文件(在相同目录下,确保原子替换有效) fd, temp_path = tempfile.mkstemp(dir=self.content_dir, suffix='.tmp') try: with os.fdopen(fd, 'w', encoding='utf-8') as f: f.write(state.model_dump_json(indent=2)) # 4. 原子替换 os.replace(temp_path, self.state_path) except Exception: # 清理临时文件 if os.path.exists(temp_path): os.unlink(temp_path) raise def get_or_create( self, content_id: str, platform: str, source_url: str ) -> WorkflowState: """获取现有状态或创建新状态 幂等操作:如果状态已存在,返回现有状态;否则创建新状态。 Args: content_id: 内容唯一标识 platform: 平台名称 source_url: 原始链接 Returns: WorkflowState 实例 """ existing = self.load() if existing is not None: return existing # 创建新状态 return WorkflowState( content_id=content_id, platform=platform, source_url=source_url )