徐涛
5baa170508
- README 添加 feature 组合表 + 模块级 features 清单 + 升级指南 - 18 个 example 顶部添加 Required features 注释 - roadmap.md 和 roadmap-v0.3.2.md 同步 Phase 26-27 完成状态 - cargo fmt 全量格式化（修复预存格式问题，CI format job 可通过）
260 lines
8.8 KiB
Rust
260 lines
8.8 KiB
Rust
//! engine_demo —— SessionManager + Checkpointer 端到端示例。
|
||
//! Required features: cargo run --example engine_demo --features "engine"
|
||
//!
|
||
//! 演示：
|
||
//! 1. SessionManager::create 创建 session
|
||
//! 2. SessionManager::submit_turn（auto_checkpoint=true 自动写 checkpoint）
|
||
//! 3. SessionManager::create_child 创建子 session
|
||
//! 4. children() / parent() 树形查询
|
||
//! 5. Checkpointer::list_checkpoints 列出所有 checkpoint
|
||
//! 6. SessionManager::recover 从 checkpoint 恢复（模拟进程重启）
|
||
//! 7. AgentSession::to_snapshot + SessionManager::replace 演示 rollback 流程
|
||
//! 8. SessionManager::destroy 清理
|
||
//!
|
||
//! 运行：`cargo run --example engine_demo`
|
||
|
||
use std::sync::Arc;
|
||
|
||
use agcore::agent::{Agent, AgentBuilder, AgentSession};
|
||
use agcore::engine::SessionManager;
|
||
use agcore::llm::hooks::HookExecutor;
|
||
use agcore::llm::mock::MockProvider;
|
||
use agcore::llm::types::Usage;
|
||
use agcore::llm::types::message::{ContentBlock, Message};
|
||
use agcore::llm::types::response_v2::{MessageResponse, StopReason};
|
||
use agcore::memory::store::InMemoryStore;
|
||
use agcore::tools::ToolRegistry;
|
||
|
||
struct DemoAgent;
|
||
|
||
impl Agent for DemoAgent {
|
||
fn name(&self) -> &str {
|
||
"demo"
|
||
}
|
||
fn system_prompt(&self) -> Option<&str> {
|
||
Some("你是 demo agent，每轮回复一句话。")
|
||
}
|
||
}
|
||
|
||
fn assistant_text(text: &str) -> MessageResponse {
|
||
MessageResponse {
|
||
id: String::new(),
|
||
model: String::new(),
|
||
message: Message::Assistant {
|
||
content: vec![ContentBlock::Text { text: text.into() }],
|
||
},
|
||
usage: Usage::from_input_output(8, 4),
|
||
stop_reason: StopReason::Stop,
|
||
extra: Default::default(),
|
||
}
|
||
}
|
||
|
||
#[tokio::main]
|
||
async fn main() {
|
||
// 1. 准备底层组件
|
||
let store: Arc<dyn agcore::memory::store::MemoryStore> = Arc::new(InMemoryStore::new());
|
||
let provider = Arc::new(MockProvider::new(vec![
|
||
assistant_text("turn 1 response"),
|
||
assistant_text("turn 2 response"),
|
||
assistant_text("turn 3 response"),
|
||
assistant_text("child turn 1 response"),
|
||
assistant_text("recovered turn response"),
|
||
]));
|
||
|
||
let bundle = Arc::new(
|
||
AgentBuilder::new()
|
||
.provider(provider.clone())
|
||
.tool_registry(Arc::new(ToolRegistry::new()))
|
||
.hook_executor(Arc::new(HookExecutor::new()))
|
||
.session_memory_backend(store.clone())
|
||
.build()
|
||
.expect("RuntimeBundle 装配失败"),
|
||
);
|
||
|
||
let agent: Arc<dyn Agent> = Arc::new(DemoAgent);
|
||
|
||
// 2. 构造 SessionManager（auto_checkpoint 默认 true）
|
||
let sm = SessionManager::new(store.clone());
|
||
println!("=== SessionManager 创建 ===");
|
||
|
||
// 3. create + submit_turn（auto_checkpoint 触发）
|
||
println!("\n=== 创建根 session + 跑 3 轮 ===");
|
||
let parent_id = sm
|
||
.create(agent.clone(), bundle.clone())
|
||
.await
|
||
.expect("create 失败");
|
||
println!("parent_id = {parent_id}");
|
||
|
||
for i in 1..=3 {
|
||
let _resp = sm
|
||
.submit_turn(&parent_id, format!("turn {i}"))
|
||
.await
|
||
.expect("submit_turn 失败");
|
||
}
|
||
|
||
// 写入自定义 session memory 数据（演示持久层往返）
|
||
sm.get(&parent_id)
|
||
.await
|
||
.unwrap()
|
||
.lock()
|
||
.await
|
||
.set_session_data("design", "PostgreSQL")
|
||
.await
|
||
.unwrap();
|
||
|
||
// 4. 显式 checkpoint（覆盖 auto_checkpoint 的 turn-level，写入额外快照）
|
||
println!("\n=== Checkpointer 显式 checkpoint ===");
|
||
let ckpt_id = sm
|
||
.checkpointer()
|
||
.checkpoint(&*sm.get(&parent_id).await.unwrap().lock().await)
|
||
.await
|
||
.expect("checkpoint 失败");
|
||
println!("explicit ckpt_id = {ckpt_id}");
|
||
|
||
// 5. list_checkpoints
|
||
let metas = sm
|
||
.checkpointer()
|
||
.list_checkpoints(&parent_id)
|
||
.await
|
||
.expect("list_checkpoints 失败");
|
||
println!("parent session 有 {} 个 checkpoint", metas.len());
|
||
for m in &metas {
|
||
println!(
|
||
" - ckpt_id={}, turn_index={}, created_at={}",
|
||
m.ckpt_id, m.turn_index, m.created_at
|
||
);
|
||
}
|
||
|
||
// 6. create_child
|
||
println!("\n=== 创建子 session ===");
|
||
let child_id = sm
|
||
.create_child(&parent_id, agent.clone())
|
||
.await
|
||
.expect("create_child 失败");
|
||
println!("child_id = {child_id}");
|
||
|
||
let children = sm.children(&parent_id).await.expect("children 失败");
|
||
assert_eq!(children, vec![child_id.clone()]);
|
||
println!("children(parent) = {children:?}");
|
||
|
||
let parent_of_child = sm.parent(&child_id).await.expect("parent 失败");
|
||
assert_eq!(parent_of_child, Some(parent_id.clone()));
|
||
println!("parent({child_id}) = {parent_of_child:?}");
|
||
|
||
// 7. recover（模拟"进程重启"——新建 SessionManager 实例，但 store 复用）
|
||
println!("\n=== 从存储恢复 session（模拟进程重启）===");
|
||
let sm2 = SessionManager::new(store.clone());
|
||
let recovered = sm2
|
||
.recover(&parent_id, agent.clone(), bundle.clone())
|
||
.await
|
||
.expect("recover 失败");
|
||
let recovered_session = recovered.lock().await;
|
||
let v = recovered_session
|
||
.get_session_data("design")
|
||
.await
|
||
.expect("get_session_data 失败");
|
||
println!("recovered session_memory['design'] = {v:?}");
|
||
assert_eq!(v, Some("PostgreSQL".into()));
|
||
|
||
// 8. 演示 rollback 流程：先记录当前 turn_index，再 rollback 到一个早期 checkpoint，
|
||
// 验证 session_memory 和 turn_index 已恢复到 checkpoint 时刻
|
||
println!("\n=== rollback 流程 ===");
|
||
let (before_turn, before_cost) = {
|
||
let s = sm.get(&parent_id).await.unwrap();
|
||
let g = s.lock().await;
|
||
(g.turn_index(), g.usage().total().total_tokens)
|
||
};
|
||
println!(
|
||
"rollback 前 turn_index={}, total_tokens={}",
|
||
before_turn, before_cost
|
||
);
|
||
|
||
let metas = sm
|
||
.checkpointer()
|
||
.list_checkpoints(&parent_id)
|
||
.await
|
||
.expect("list_checkpoints 失败");
|
||
assert!(metas.len() >= 2, "至少 2 个 checkpoint 才能演示 rollback");
|
||
// 取第二个 checkpoint（不是最新的）作为 rollback 目标
|
||
let rollback_ckpt = &metas[metas.len() - 2];
|
||
println!("rollback 到 ckpt_id={}", rollback_ckpt.ckpt_id);
|
||
|
||
let snapshot = sm
|
||
.checkpointer()
|
||
.rollback_load(&parent_id, &rollback_ckpt.ckpt_id)
|
||
.await
|
||
.expect("rollback_load 失败");
|
||
let snapshot_turn = snapshot.turn_index;
|
||
let snapshot_data_count = snapshot.session_memory_data.len();
|
||
println!(
|
||
"checkpoint 时刻 turn_index={}, session_memory 条目数={}",
|
||
snapshot_turn, snapshot_data_count
|
||
);
|
||
|
||
let mut rolled_back = AgentSession::from_snapshot(snapshot, agent.clone(), bundle.clone())
|
||
.expect("from_snapshot");
|
||
rolled_back
|
||
.restore_memory()
|
||
.await
|
||
.expect("restore_memory 失败");
|
||
sm.replace(&parent_id, rolled_back)
|
||
.await
|
||
.expect("replace 失败");
|
||
|
||
// 验证 rollback 后状态与 checkpoint 一致
|
||
let (after_turn, after_cost) = {
|
||
let s = sm.get(&parent_id).await.unwrap();
|
||
let g = s.lock().await;
|
||
(g.turn_index(), g.usage().total().total_tokens)
|
||
};
|
||
println!(
|
||
"rollback 后 turn_index={}, total_tokens={}",
|
||
after_turn, after_cost
|
||
);
|
||
assert!(
|
||
after_turn <= before_turn,
|
||
"rollback 后 turn_index({}) 应 ≤ rollback 前({})",
|
||
after_turn,
|
||
before_turn
|
||
);
|
||
assert_eq!(
|
||
after_turn, snapshot_turn,
|
||
"rollback 后 turn_index 应等于 checkpoint 时刻值"
|
||
);
|
||
assert!(
|
||
after_cost <= before_cost,
|
||
"rollback 后 cost 应 ≤ rollback 前"
|
||
);
|
||
println!("✓ rollback + replace 一致性验证通过");
|
||
|
||
// 9. destroy 父子 session
|
||
println!("\n=== 销毁 session ===");
|
||
sm.destroy(&child_id).await.expect("destroy child 失败");
|
||
sm.destroy(&parent_id).await.expect("destroy parent 失败");
|
||
|
||
// 验证清理
|
||
assert!(sm.get(&parent_id).await.is_err());
|
||
assert!(
|
||
sm.checkpointer()
|
||
.list_checkpoints(&parent_id)
|
||
.await
|
||
.unwrap()
|
||
.is_empty()
|
||
);
|
||
println!("✓ parent 已彻底清理（内存 + meta + checkpoints）");
|
||
|
||
// 验证孤儿语义：父被销毁后子仍存在但 parent() 返回 None
|
||
println!("\n=== 孤儿策略演示（先创建父子，再仅销毁父）===");
|
||
let p_id = sm.create(agent.clone(), bundle.clone()).await.unwrap();
|
||
let c_id = sm.create_child(&p_id, agent.clone()).await.unwrap();
|
||
sm.destroy(&p_id).await.unwrap();
|
||
let p_of_c = sm.parent(&c_id).await.expect("parent 失败");
|
||
assert_eq!(p_of_c, None, "父被销毁后 child.parent() 应为 None");
|
||
println!("✓ child({c_id}) 仍是孤儿 session，parent() = None");
|
||
|
||
// 清理孤儿
|
||
sm.destroy(&c_id).await.unwrap();
|
||
|
||
println!("\n✓ engine_demo 完成");
|
||
}
|