AI Agent 的状态机设计:用 Rust enum 把 agent 状态建模成有限状态机

AI Agent 的状态机设计:用 Rust enum 把 agent 状态建模成有限状态机

一、AI Agent 开发的核心痛点:状态散落一地

写原型时,Agent 的状态靠几个 bool 变量来跟踪:

// ❌ 原型阶段的状态管理——靠布尔值堆砌
struct Agent {
    is_planning: bool,
    is_executing: bool,
    is_reflecting: bool,
    is_finished: bool,
    tool_call_count: usize,
    max_retries: usize,
    // ... 10 个以上的布尔/计数字段
}

这种写法的问题在第三天的开发中就暴露了:我在 is_executing = true 的同时不小心也保留了 is_planning = true,Agent 进入了一个非法状态,它一边"执行工具调用"一边"重新规划",结果无限循环了。

排查了半天才找到根因:布尔值之间没有互斥约束,编译器不会阻止我同时把它们设为 true

解决方案是:用 Rust 的 enum 把 Agent 状态建模成有限状态机。 Rust 的 enum 天然适合表达"同一时刻只能处于一个状态"的语义。

stateDiagram-v2
    [*] --> Idle: Agent 创建
    Idle --> Planning: 收到用户任务
    Planning --> Executing: 任务计划生成完毕
    Planning --> Idle: 计划失败
    Executing --> Reflecting: 工具调用返回结果
    Executing --> Executing: 继续执行下一工具
    Reflecting --> Executing: 需要重试/继续
    Reflecting --> Finished: 任务完成
    Reflecting --> Idle: 任务失败(不可重试)
    Finished --> Idle: 清理上下文
    Idle --> [*]: Agent 销毁

二、用 Rust enum 实现 Agent 状态机

use std::collections::HashMap;

/// Agent 状态枚举——每个变体携带该状态独有的数据
#[derive(Debug, Clone)]
pub enum AgentState {
    /// 空闲状态,等待用户输入
    Idle,

    /// 正在规划任务分解
    Planning {
        user_task: String,             // 用户的原始任务
        attempts: usize,               // 已尝试规划次数
        max_attempts: usize,           // 最多尝试几次
    },

    /// 正在执行工具调用
    Executing {
        plan: TaskPlan,                // 当前计划(携带子任务列表)
        current_step: usize,           // 执行到第几步
        tool_results: Vec<String>,     // 已收集的工具结果
    },

    /// 正在反思和评估结果
    Reflecting {
        plan: TaskPlan,
        completed_steps: Vec<StepResult>,  // 已完成步骤的结果
        failed_steps: Vec<StepResult>,     // 失败的步骤
        reflection_round: usize,           // 第几轮反思
    },

    /// 任务完成
    Finished {
        summary: String,
        steps_taken: usize,
    },

    /// 出错了(但 Agent 仍存活)
    Error {
        error: String,
        recoverable: bool,  // 是否可恢复
        previous_state: Box<AgentState>, // 保存出错前的状态
    },
}

/// 任务计划,包含子任务列表
#[derive(Debug, Clone)]
pub struct TaskPlan {
    pub steps: Vec<TaskStep>,
}

#[derive(Debug, Clone)]
pub struct TaskStep {
    pub tool_name: String,
    pub parameters: HashMap<String, String>,
    pub description: String,
}

#[derive(Debug, Clone)]
pub struct StepResult {
    pub step_index: usize,
    pub success: bool,
    pub output: String,
}

enum 相较于布尔值的关键优势

flowchart LR
    subgraph bad["❌ 布尔值方案"]
        B1["is_planning: bool"]
        B2["is_executing: bool"]
        B3["is_reflecting: bool"]
        B1 -.-> NB["无法互斥:可能同时为 true"]
    end

    subgraph good["✅ enum 方案"]
        G1["Planning {...}"]
        G2["Executing {...}"]
        G3["Reflecting {...}"]
        G1 ==> GE["编译期保证互斥\n状态独有数据仅在该状态可见"]
    end

三、状态转移的实现

状态机最关键的是"状态转移函数"——确保从一个状态到另一个状态的转换是合法且安全的

impl AgentState {
    /// 核心方法:状态转移
    /// 返回新的 AgentState(consuming self)
    pub fn transition(self, event: AgentEvent) -> Result<Self, TransitionError> {
        match (self, event) {
            // 空闲 → 收到任务 → 开始规划
            (AgentState::Idle, AgentEvent::ReceiveTask { task }) => {
                Ok(AgentState::Planning {
                    user_task: task,
                    attempts: 0,
                    max_attempts: 3,
                })
            }

            // 规划成功 → 开始执行
            (
                AgentState::Planning { user_task, .. },
                AgentEvent::PlanReady { plan },
            ) => {
                println!("任务规划完成: {}", user_task);
                Ok(AgentState::Executing {
                    plan,
                    current_step: 0,
                    tool_results: Vec::new(),
                })
            }

            // 执行中 → 收到工具结果 → 决定下一步
            (
                AgentState::Executing { plan, current_step, mut tool_results },
                AgentEvent::ToolResult { result },
            ) => {
                tool_results.push(result);
                let next_step = current_step + 1;

                if next_step >= plan.steps.len() {
                    // 所有步骤完成 → 进入反思
                    Ok(AgentState::Reflecting {
                        plan,
                        completed_steps: Vec::new(),
                        failed_steps: Vec::new(),
                        reflection_round: 0,
                    })
                } else {
                    // 还有步骤未完成 → 继续执行
                    Ok(AgentState::Executing {
                        plan,
                        current_step: next_step,
                        tool_results,
                    })
                }
            }

            // 反思 → 任务完成
            (
                AgentState::Reflecting { plan, .. },
                AgentEvent::ReflectionDone { summary, success },
            ) => {
                if success {
                    Ok(AgentState::Finished {
                        summary,
                        steps_taken: plan.steps.len(),
                    })
                } else {
                    // 即使失败也可能完成任务(部分成功)
                    Ok(AgentState::Finished {
                        summary: format!("部分完成: {}", summary),
                        steps_taken: plan.steps.len(),
                    })
                }
            }

            // 任何状态 → 发生错误
            (current_state, AgentEvent::ErrorOccurred { error, recoverable }) => {
                Ok(AgentState::Error {
                    error,
                    recoverable,
                    previous_state: Box::new(current_state),
                })
            }

            // 错误 → 恢复
            (
                AgentState::Error { recoverable: true, previous_state, .. },
                AgentEvent::Recover,
            ) => {
                Ok(*previous_state)
            }

            // 非法转移——编译期捕获不到,运行时返回错误
            (current_state, event) => {
                Err(TransitionError::InvalidTransition {
                    from: format!("{:?}", current_state),
                    event: format!("{:?}", event),
                })
            }
        }
    }
}

/// 触发状态转移的事件
#[derive(Debug, Clone)]
pub enum AgentEvent {
    ReceiveTask { task: String },
    PlanReady { plan: TaskPlan },
    ToolResult { result: String },
    ReflectionDone { summary: String, success: bool },
    ErrorOccurred { error: String, recoverable: bool },
    Recover,
}

#[derive(Debug)]
pub enum TransitionError {
    InvalidTransition { from: String, event: String },
}

实战踩坑:状态机膨胀和无效转移的爆炸问题

enum 状态机刚开始让人很兴奋,但加第三个功能时我就遇到了问题。原本只有 Idle→Planning→Executing→Reflecting→Finished 五个状态和六种转移,transitionmatch 分支不到 10 个。加了"暂停/恢复"和"人工审批"两个需求后,状态从 5 个涨到 9 个,可能的转移路径从 6 条涨到 23 条。

23 条匹配臂的 transition 函数已经无法一眼看完。我的解法是把复杂的 transition 拆成多个小函数:handle_planning_eventhandle_executing_event 等,每个只处理特定状态下的转移。Rust 的 enum 保证了嵌套 match 的穷尽性,拆函数不影响安全性。

另一个坑:Error 状态用 Box<AgentState> 保存上一个状态时,如果连续出错(Error→恢复→再Error),会形成一个"Error(恢复(Error(原状态)))"的嵌套链。深度超过 3 层后调试打印会溢出终端,而且反序列化(如果用 serde)会失败。解法:限制 Error 状态不嵌套另一个 Error 状态,如果 Error 状态下再出错就覆盖而不是嵌套。

四、这个设计带来的实际收益

flowchart TD
    subgraph gains["实际收益"]
        G1["类型安全\n编译器检查状态结构"]
        G2["业务逻辑清晰\n一个 match 看清所有转移"]
        G3["测试友好\n给定状态+事件→断言新状态"]
        G4["文档自动生成\n状态机图即设计文档"]
    end

重构前后的对比:用布尔值方案时,每周末跑回归测试都会发现 1-2 个"状态不一致"的 bug——Agent 既不 planning 也不 executing,停在了一个 undefined 的状态里。重构为 enum 状态机后,这种 bug 在过去两个月中出现了 0 次。不是变少了,是直接消失。

测试覆盖方面,状态机让单元测试变得极其容易写。"给定状态 A + 事件 B → 期望状态 C"这种三段式测试,每个转移一条就够。我写了 18 个测试覆盖所有合法转移和 5 个非法转移,总测试代码量不到 80 行。

测试示例

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_idle_to_planning_transition() {
        let state = AgentState::Idle;
        let result = state.transition(AgentEvent::ReceiveTask {
            task: "帮我查询今天的天气".into(),
        });
        assert!(result.is_ok());
        match result.unwrap() {
            AgentState::Planning { user_task, .. } => {
                assert_eq!(user_task, "帮我查询今天的天气");
            }
            _ => panic!("状态转移结果不对:期望 Planning"),
        }
    }

    #[test]
    fn test_invalid_transition_returns_error() {
        // Planning 状态下直接收到 ReflectionDone 是违法的
        let state = AgentState::Planning {
            user_task: "test".into(),
            attempts: 0,
            max_attempts: 3,
        };
        let result = state.transition(AgentEvent::ReflectionDone {
            summary: "ok".into(),
            success: true,
        });
        assert!(result.is_err());
    }
}

我在项目里就踩过这个坑:忘写 wildcard 匹配分支,一个不该出现的事件没有处理,代码直接 panic 了。后来给所有 match 都加了 catch-all 分支,非法状态转移变成可打印的错误信息。

五、总结

用 Rust enum 实现 Agent 状态机,我最大的体会是:

当你的代码中出现 3 个以上互斥的布尔值时,就是引入状态 enum 的信号。 不要等到出了 bug 再改。

另外有几点建议:

  • 状态数据只放在对应变体里Planning 状态的 attempts 不会出现在 Executing 状态里——不可能在 Executing 时"误读" planning 的数据
  • Box<AgentState> 用于自引用Error 状态保存上一个状态时用它,避免 enum 自身无限递归
  • 给所有"从任何状态都可能的转移"显式处理:比如"出错"事件可以从任何状态触发学设计模式可能头大,但状态机这个概念其实在编程中非常自然——你玩过游戏就知道,角色永远是"站立→走路→跑步→跳跃"这样的状态链条。把这个直觉搬到 Agent 设计里就对了。

你对 Agent 状态设计有什么想法?欢迎在评论区交流。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值