agcore/examples/engine_demo.rs

//! engine_demo —— SessionManager + Checkpointer 端到端示例。
//!
//! 演示：
//! 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 完成");
}