From 624cd765388a97504e5b811742a80f58d5ce0e72 Mon Sep 17 00:00:00 2001
From: starmen1265 <starmen1265@163.com>
Date: Fri, 13 Mar 2026 17:16:30 +0800
Subject: [PATCH] feat(01-01): implement WorkflowStateManager with atomic write

GREEN phase: Implement state management with:
- Atomic write pattern (temp file + os.replace)
- Backup rotation (workflow_state.json -> .bak)
- Backup fallback on corruption
- get_or_create idempotent creation

Fix test to properly create backup before corruption test.

All 20 tests pass.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 kb_pipeline/state.py              | 152 ++++++++++++++++++++++++++++++
 tests/test_pipeline/test_state.py |  15 ++-
 2 files changed, 165 insertions(+), 2 deletions(-)
 create mode 100644 kb_pipeline/state.py

diff --git a/kb_pipeline/state.py b/kb_pipeline/state.py
new file mode 100644
index 0000000..598c6a5
--- /dev/null
+++ b/kb_pipeline/state.py
@@ -0,0 +1,152 @@
+"""
+WorkflowStateManager - 步骤级状态持久化管理器
+
+提供工作流状态的原子写入、备份轮转和崩溃恢复功能。
+
+Purpose: 为工作流提供可靠的状态管理，支持断点续传和崩溃恢复
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Optional
+
+from kb_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
+        )
diff --git a/tests/test_pipeline/test_state.py b/tests/test_pipeline/test_state.py
index 38f6c8a..e427547 100644
--- a/tests/test_pipeline/test_state.py
+++ b/tests/test_pipeline/test_state.py
@@ -281,7 +281,7 @@ class TestBackupFallback:
         state_file = tmp_path / "workflow_state.json"
         backup_file = tmp_path / "workflow_state.json.bak"
 
-        # Create valid backup
+        # Create valid backup by saving twice (backup is created on second save)
         valid_state = WorkflowState(
             content_id="backup-valid-001",
             platform="bilibili",
@@ -289,11 +289,22 @@ class TestBackupFallback:
         )
         manager.save(valid_state)
 
+        # Save again to create backup
+        valid_state2 = WorkflowState(
+            content_id="backup-valid-002",
+            platform="bilibili",
+            source_url="https://www.bilibili.com/video/BV888"
+        )
+        manager.save(valid_state2)
+
+        # Now backup exists with "backup-valid-001", main has "backup-valid-002"
+        assert backup_file.exists()
+
         # Corrupt main file
         with open(state_file, 'w', encoding='utf-8') as f:
             f.write("{ invalid json content }")
 
-        # Should fall back to backup
+        # Should fall back to backup (which has backup-valid-001)
         loaded = manager.load()
 
         assert loaded is not None