kb-skills/kb_workflow/pipeline/state.py

"""
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
        )