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 五个状态和六种转移,transition 的 match 分支不到 10 个。加了"暂停/恢复"和"人工审批"两个需求后,状态从 5 个涨到 9 个,可能的转移路径从 6 条涨到 23 条。
23 条匹配臂的 transition 函数已经无法一眼看完。我的解法是把复杂的 transition 拆成多个小函数:handle_planning_event、handle_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 状态设计有什么想法?欢迎在评论区交流。

385

被折叠的 条评论
为什么被折叠?



