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>
This commit is contained in:
2026-03-13 17:16:30 +08:00
parent e789d81230
commit 624cd76538
2 changed files with 165 additions and 2 deletions
+152
View File
@@ -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
)
+13 -2
View File
@@ -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