AI Agent 工作流持久化:从状态快照到故障恢复的工程实践

一、Agent 工作流的"脆弱性":一次 OOM 杀掉 30 分钟的推理链

AI Agent 在执行多步骤工作流时,状态全部驻留在内存中。某自动化运维 Agent 执行一个 12 步的故障排查流程,在第 9 步时因内存溢出被 OOM Killer 终止,前 8 步的推理结果、工具调用记录和中间状态全部丢失,只能从头开始。更严重的是,某些工具调用(如发送通知、创建工单)是不可逆的,重试会导致重复操作。

Agent 工作流持久化的核心目标是:在任意步骤失败后,能从最近的检查点恢复执行,而非从头开始。这不仅是可靠性问题,更是成本问题——LLM 推理的 Token 消耗在重试中成倍增加。

二、Agent 工作流持久化的架构设计

flowchart TB
    subgraph 执行层["Agent 执行引擎"]
        direction TB
        E1[步骤 1: 意图解析]
        E2[步骤 2: 工具选择]
        E3[步骤 3: 工具执行]
        E4[步骤 4: 结果评估]
        E5[步骤 5: 响应生成]
    end

    subgraph 持久层["状态持久化层"]
        direction TB
        P1[检查点存储<br/>每步完成后写入]
        P2[事件日志<br/>WAL 模式追加写入]
        P3[快照存储<br/>关键节点全量快照]
    end

    subgraph 恢复层["故障恢复层"]
        direction TB
        R1[故障检测<br/>心跳 + 超时]
        R2[最近检查点定位<br/>二分查找有效点]
        R3[状态恢复<br/>反序列化 + 重放]
        R4[幂等性校验<br/>跳过已执行步骤]
    end

    E1 --> P1
    E2 --> P1
    E3 --> P2
    E4 --> P3
    E5 --> P3
    P1 --> R2
    P2 --> R2
    P3 --> R2
    R2 --> R3 --> R4

    style 执行层 fill:#eef,stroke:#333
    style 持久层 fill:#fee,stroke:#333
    style 恢复层 fill:#efe,stroke:#333

三、Agent 工作流持久化的代码实现

from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any, Callable
from enum import Enum
from datetime import datetime
import json
import hashlib


class StepStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"       # 幂等恢复时跳过


class CheckpointType(Enum):
    STEP = "step"             # 步骤级检查点
    SNAPSHOT = "snapshot"     # 全量快照
    EVENT = "event"           # 事件日志


@dataclass
class WorkflowStep:
    """工作流步骤"""
    step_id: str
    name: str
    status: StepStatus = StepStatus.PENDING
    input_data: Dict = field(default_factory=dict)
    output_data: Dict = field(default_factory=dict)
    tool_calls: List[Dict] = field(default_factory=list)
    error: Optional[str] = None
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    idempotency_key: Optional[str] = None  # 幂等键


@dataclass
class Checkpoint:
    """检查点"""
    checkpoint_id: str
    workflow_id: str
    checkpoint_type: CheckpointType
    step_id: str
    state: Dict                # 序列化的工作流状态
    timestamp: datetime
    checksum: str              # 状态校验和


@dataclass
class WorkflowState:
    """工作流全局状态"""
    workflow_id: str
    steps: List[WorkflowStep]
    context: Dict              # 全局上下文(对话历史、中间变量等)
    current_step_index: int = 0
    created_at: datetime = field(default_factory=datetime.now)
    updated_at: datetime = field(default_factory=datetime.now)


class WorkflowPersistenceEngine:
    """
    工作流持久化引擎
    核心机制:检查点 + 事件日志 + 幂等恢复
    """

    def __init__(self, storage_backend=None):
        self._checkpoints: Dict[str, List[Checkpoint]] = {}
        self._event_log: Dict[str, List[Dict]] = {}
        self._workflows: Dict[str, WorkflowState] = {}
        self._storage = storage_backend

    # ============ 检查点管理 ============

    def save_checkpoint(self, workflow_id: str, step_id: str,
                        state: Dict,
                        checkpoint_type: CheckpointType = CheckpointType.STEP) -> str:
        """保存检查点"""
        checkpoint_id = f"cp-{workflow_id}-{step_id}-{datetime.now().strftime('%H%M%S')}"
        state_json = json.dumps(state, sort_keys=True, default=str)
        checksum = hashlib.md5(state_json.encode()).hexdigest()

        checkpoint = Checkpoint(
            checkpoint_id=checkpoint_id,
            workflow_id=workflow_id,
            checkpoint_type=checkpoint_type,
            step_id=step_id,
            state=state,
            timestamp=datetime.now(),
            checksum=checksum,
        )

        if workflow_id not in self._checkpoints:
            self._checkpoints[workflow_id] = []
        self._checkpoints[workflow_id].append(checkpoint)

        return checkpoint_id

    def save_event(self, workflow_id: str, event_type: str,
                   data: Dict):
        """追加事件日志(WAL 模式)"""
        event = {
            "event_id": f"evt-{len(self._event_log.get(workflow_id, [])) + 1}",
            "workflow_id": workflow_id,
            "event_type": event_type,
            "data": data,
            "timestamp": datetime.now().isoformat(),
        }
        if workflow_id not in self._event_log:
            self._event_log[workflow_id] = []
        self._event_log[workflow_id].append(event)

    # ============ 故障恢复 ============

    def find_latest_valid_checkpoint(self, workflow_id: str) -> Optional[Checkpoint]:
        """找到最近的有效检查点"""
        checkpoints = self._checkpoints.get(workflow_id, [])
        if not checkpoints:
            return None

        # 从最新到最旧查找
        for cp in reversed(checkpoints):
            # 验证检查点完整性
            state_json = json.dumps(cp.state, sort_keys=True, default=str)
            expected_checksum = hashlib.md5(state_json.encode()).hexdigest()
            if cp.checksum == expected_checksum:
                return cp

        return None

    def recover_workflow(self, workflow_id: str) -> Optional[WorkflowState]:
        """从检查点恢复工作流状态"""
        checkpoint = self.find_latest_valid_checkpoint(workflow_id)
        if not checkpoint:
            return None

        # 反序列化状态
        state = checkpoint.state
        workflow_state = WorkflowState(
            workflow_id=state.get("workflow_id", workflow_id),
            steps=[WorkflowStep(**s) for s in state.get("steps", [])],
            context=state.get("context", {}),
            current_step_index=state.get("current_step_index", 0),
        )

        # 幂等性校验:标记已完成的步骤为 SKIPPED
        for step in workflow_state.steps:
            if step.status == StepStatus.COMPLETED and step.idempotency_key:
                # 检查该步骤的副作用是否已存在
                if self._check_side_effect(step.idempotency_key):
                    step.status = StepStatus.SKIPPED

        return workflow_state

    def _check_side_effect(self, idempotency_key: str) -> bool:
        """检查副作用是否已存在(幂等性保障)"""
        # 实际实现中查询外部系统
        return False

    # ============ 工作流执行 ============

    def execute_workflow(self, workflow_id: str,
                         steps: List[Dict],
                         step_executor: Callable) -> Dict:
        """执行工作流,支持自动检查点与故障恢复"""
        # 尝试从检查点恢复
        existing = self.recover_workflow(workflow_id)

        if existing:
            workflow_state = existing
            start_index = workflow_state.current_step_index
        else:
            workflow_state = WorkflowState(
                workflow_id=workflow_id,
                steps=[WorkflowStep(
                    step_id=s["step_id"],
                    name=s["name"],
                    idempotency_key=s.get("idempotency_key"),
                ) for s in steps],
                context={},
            )
            start_index = 0

        # 从断点继续执行
        for i in range(start_index, len(workflow_state.steps)):
            step = workflow_state.steps[i]

            if step.status == StepStatus.SKIPPED:
                continue

            if step.status == StepStatus.COMPLETED:
                continue

            step.status = StepStatus.RUNNING
            step.started_at = datetime.now()
            workflow_state.current_step_index = i

            try:
                # 执行步骤
                result = step_executor(step.name, workflow_state.context)

                step.output_data = result.get("output", {})
                step.tool_calls = result.get("tool_calls", [])
                step.status = StepStatus.COMPLETED
                step.completed_at = datetime.now()

                # 更新上下文
                workflow_state.context.update(result.get("context_update", {}))

                # 保存步骤级检查点
                self.save_checkpoint(
                    workflow_id, step.step_id,
                    self._serialize_state(workflow_state),
                    CheckpointType.STEP
                )

                # 记录事件日志
                self.save_event(workflow_id, "step_completed", {
                    "step_id": step.step_id,
                    "output_keys": list(step.output_data.keys()),
                })

            except Exception as e:
                step.status = StepStatus.FAILED
                step.error = str(e)

                # 保存故障点检查点
                self.save_checkpoint(
                    workflow_id, step.step_id,
                    self._serialize_state(workflow_state),
                    CheckpointType.SNAPSHOT
                )

                self.save_event(workflow_id, "step_failed", {
                    "step_id": step.step_id,
                    "error": str(e),
                })

                return {
                    "status": "failed",
                    "failed_step": step.step_id,
                    "error": str(e),
                    "completed_steps": sum(
                        1 for s in workflow_state.steps
                        if s.status == StepStatus.COMPLETED
                    ),
                }

        workflow_state.updated_at = datetime.now()
        return {
            "status": "completed",
            "total_steps": len(workflow_state.steps),
            "context": workflow_state.context,
        }

    def _serialize_state(self, state: WorkflowState) -> Dict:
        """序列化工作流状态"""
        return {
            "workflow_id": state.workflow_id,
            "steps": [
                {
                    "step_id": s.step_id,
                    "name": s.name,
                    "status": s.status.value,
                    "input_data": s.input_data,
                    "output_data": s.output_data,
                    "tool_calls": s.tool_calls,
                    "error": s.error,
                    "idempotency_key": s.idempotency_key,
                }
                for s in state.steps
            ],
            "context": state.context,
            "current_step_index": state.current_step_index,
        }


# ============ 云原生部署适配 ============

class KubernetesWorkflowAdapter:
    """
    Kubernetes 适配器:将工作流状态持久化到 PVC
    支持 Pod 重启后的状态恢复
    """

    def __init__(self, pvc_mount_path: str = "/data/workflows"):
        self._mount_path = pvc_mount_path

    def save_to_pvc(self, workflow_id: str, state: Dict):
        """将状态写入 PVC 挂载的持久卷"""
        import os
        file_path = os.path.join(self._mount_path, f"{workflow_id}.json")
        with open(file_path, "w") as f:
            json.dump(state, f, indent=2, default=str)

    def load_from_pvc(self, workflow_id: str) -> Optional[Dict]:
        """从 PVC 读取状态"""
        import os
        file_path = os.path.join(self._mount_path, f"{workflow_id}.json")
        if os.path.exists(file_path):
            with open(file_path, "r") as f:
                return json.load(f)
        return None

    def get_liveness_config(self, workflow_id: str) -> Dict:
        """
        生成 K8s Liveness Probe 配置
        检测工作流是否卡死(长时间无检查点更新)
        """
        return {
            "exec": {
                "command": [
                    "sh", "-c",
                    f"find {self._mount_path} -name '{workflow_id}.json' "
                    f"-mmin -300 | grep -q ."
                ]
            },
            "initialDelaySeconds": 30,
            "periodSeconds": 60,
            "failureThreshold": 3,
        }

四、Agent 工作流持久化的 Trade-offs

检查点频率与性能开销。每步保存检查点增加了 I/O 开销,尤其当工作流状态较大时(包含长对话历史)。某 Agent 的上下文达 50KB,每步序列化写入耗时 20ms,12 步累计 240ms。解决方案是区分"轻量检查点"(仅保存步骤状态)和"全量快照"(包含完整上下文),前者每步执行,后者仅在关键节点执行。

幂等性保障的实现复杂度。工具调用的幂等性需要外部系统配合。发送通知的幂等需要消息去重,创建工单的幂等需要唯一约束,数据库写入的幂等需要 upsert 语义。每个工具都需要单独设计幂等方案,增加了开发成本。

状态恢复的一致性窗口。检查点保存和实际执行之间存在微小的时间窗口,如果在这个窗口内崩溃,检查点可能不反映最新状态。WAL 模式可以缩小这个窗口,但增加了存储和恢复的复杂度。

云原生环境中的存储依赖。PVC 挂载的持久卷在 Pod 调度到不同节点时可能不可用。StatefulSet 可以保证 Pod 与 PVC 的绑定关系,但限制了调度的灵活性。对象存储(如 S3)是更通用的方案,但增加了网络延迟。

五、总结

Agent 工作流持久化通过检查点、事件日志和幂等恢复三个机制,确保多步骤工作流在任意故障点后可恢复执行。检查点提供状态快照,事件日志提供操作审计,幂等键防止重复执行副作用。云原生部署中,PVC 或对象存储作为持久化后端,Liveness Probe 检测工作流卡死。关键权衡在于检查点频率与性能开销、幂等性保障的实现复杂度、状态恢复的一致性窗口,以及云原生存储的调度约束。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐