Genarrative/rust/crates/agent-host/tests/message_persistence.rs

use std::collections::BTreeSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use agent_host::{AgentHost, HostError};
use agent_provider_fake::{FakeProvider, FakeStep, FakeToolCall};
use agent_runtime_core::{
    ApprovalDecision, ApprovalError, ApprovalPolicy, ContentPart, Message, ProviderErrorKind,
    RuntimeSnapshot, ToolCall, ToolResult, reduce,
};
use agent_runtime_engine::{CompressionRequest, ContextCompressor, EngineError};
use serde_json::json;

fn assert_snapshot_matches_output(host: &AgentHost, result: &agent_host::HostRunOutput) {
    // 结果消息、持久化快照和事件重放必须描述同一条消息历史。
    let snapshot = host
        .load_runtime_snapshot(&result.runtime_id)
        .expect("读取 RuntimeSnapshot")
        .expect("RuntimeSnapshot 应存在");
    let run = snapshot.run(&result.run_id).expect("run 应存在");
    assert_eq!(run.messages(), result.output.messages.as_slice());
    // 消息正确还不够，派生工具索引也必须与当前上下文逐项一致。
    let mut calls = Vec::new();
    let mut results = Vec::new();
    for part in run.messages().iter().flat_map(|message| message.content()) {
        match part {
            ContentPart::ToolCall {
                id,
                name,
                arguments,
            } => {
                calls.push(ToolCall::try_new(id, name, arguments.clone()).unwrap());
            }
            ContentPart::ToolResult {
                tool_call_id,
                output,
                is_error,
            } => {
                results.push(ToolResult::try_new(tool_call_id, output.clone(), *is_error).unwrap());
            }
            _ => {}
        }
    }
    assert_eq!(run.tool_calls(), calls);
    assert_eq!(run.tool_results(), results);

    let events = host
        .list_runtime_events(&result.runtime_id)
        .expect("读取 runtime events");
    let mut replayed = RuntimeSnapshot::try_new(&result.runtime_id).expect("创建空快照");
    for event in events {
        replayed = reduce(&replayed, &event).expect("runtime event 应可重放");
    }
    assert_eq!(replayed, snapshot);
}

fn assert_tool_rows_are_terminal_and_unique(host: &AgentHost, run_id: &str, expected: usize) {
    // 每个工具调用只能有一行最终结果，不能把中间 requested 状态当成完成。
    let rows = host.list_tool_calls(run_id).expect("读取工具调用记录");
    assert_eq!(rows.len(), expected);
    let mut ids = BTreeSet::new();
    for row in rows {
        assert!(ids.insert(row.id), "工具调用 id 不应重复");
        assert_ne!(row.status, "requested", "工具调用不应停留在 requested");
        assert!(row.result.is_some(), "已完成工具调用应有结果");
    }
}

#[test]
fn normal_complete_persists_exactly_one_message_history_and_replay() {
    let host = AgentHost::in_memory()
        .expect("创建 Host")
        .with_provider(Arc::new(FakeProvider::text("普通完成")), "fake");
    let result = host
        .run_with_messages("普通完成", vec![Message::user("普通完成").unwrap()])
        .expect("普通运行应完成");

    assert_eq!(result.output.text, "普通完成");
    assert_snapshot_matches_output(&host, &result);
    assert_tool_rows_are_terminal_and_unique(&host, &result.run_id, 0);
}

#[test]
fn normal_stream_persists_exactly_one_message_history_and_replay() {
    let host = AgentHost::in_memory().expect("创建 Host").with_provider(
        Arc::new(FakeProvider::new([FakeStep::stream_text(["流", "式"])])),
        "fake",
    );
    let result = host
        .run_with_messages_streaming("流式完成", vec![Message::user("流式完成").unwrap()])
        .expect("流式运行应完成");

    assert_eq!(result.output.text, "流式");
    assert!(!result.output.stream_events.is_empty());
    assert_snapshot_matches_output(&host, &result);
    assert_tool_rows_are_terminal_and_unique(&host, &result.run_id, 0);
}

#[test]
fn automatic_tool_batches_and_later_round_match_engine_in_both_modes() {
    for streaming in [false, true] {
        let provider = Arc::new(FakeProvider::new([
            FakeStep::ToolCalls {
                text: "同一模型响应中的说明".to_owned(),
                calls: (1..=3)
                    .map(|index| FakeToolCall {
                        id: format!("automatic-{index}"),
                        name: "echo".to_owned(),
                        arguments: json!({"index": index}),
                    })
                    .collect(),
            },
            FakeStep::tool_call("automatic-next", "echo", json!({"index": 4})),
            FakeStep::text("同一模型响应中的说明"),
        ]));
        let host = AgentHost::in_memory()
            .unwrap()
            .with_provider(provider.clone(), "fake");
        let result = if streaming {
            host.run_with_messages_streaming("自动批次", vec![Message::user("自动批次").unwrap()])
        } else {
            host.run("自动批次")
        }
        .expect("自动放行的连续工具轮次应完成");
        assert_eq!(result.output.messages.len(), 8);
        assert_eq!(provider.requests().snapshot().len(), 3);
        assert_snapshot_matches_output(&host, &result);
        assert_tool_rows_are_terminal_and_unique(&host, &result.run_id, 4);
    }
}

struct AllowThenAsk {
    asked: AtomicBool,
}

impl AllowThenAsk {
    fn new() -> Self {
        Self {
            asked: AtomicBool::new(false),
        }
    }
}

impl ApprovalPolicy for AllowThenAsk {
    fn decide(
        &self,
        request: &agent_runtime_core::ApprovalRequest,
    ) -> Result<ApprovalDecision, ApprovalError> {
        if request.call().id() == "batch-call-3" && !self.asked.swap(true, Ordering::AcqRel) {
            Ok(ApprovalDecision::Ask)
        } else {
            Ok(ApprovalDecision::Allow)
        }
    }
}

#[test]
fn three_tool_batch_and_next_tool_round_keep_one_message_history() {
    let provider = Arc::new(FakeProvider::new([
        FakeStep::tool_calls([
            FakeToolCall {
                id: "batch-call-1".to_owned(),
                name: "echo".to_owned(),
                arguments: json!({"index": 1}),
            },
            FakeToolCall {
                id: "batch-call-2".to_owned(),
                name: "echo".to_owned(),
                arguments: json!({"index": 2}),
            },
            FakeToolCall {
                id: "batch-call-3".to_owned(),
                name: "echo".to_owned(),
                arguments: json!({"index": 3}),
            },
        ]),
        FakeStep::tool_call("next-round-call", "echo", json!({"index": 4})),
        FakeStep::text("三工具批次完成"),
    ]));
    let host = AgentHost::in_memory()
        .expect("创建 Host")
        .with_provider(provider.clone(), "fake")
        .with_approval(Arc::new(AllowThenAsk::new()));
    let handle = host.prepare_run("三工具批次").expect("创建 run");

    let first = host
        .run_existing(&handle.run_id)
        .expect_err("第三个调用应 Ask");
    assert!(matches!(
        first,
        HostError::Engine(EngineError::ApprovalRequired { ref call_id, .. })
            if call_id == "batch-call-3"
    ));
    let checkpoint = host.read_checkpoint(&handle.run_id).unwrap().unwrap();
    let checkpoint_messages: Vec<Message> = serde_json::from_value(checkpoint.messages).unwrap();
    let pending_snapshot = host
        .load_runtime_snapshot(&handle.runtime_id)
        .unwrap()
        .unwrap();
    assert_eq!(
        pending_snapshot.run(&handle.run_id).unwrap().messages(),
        checkpoint_messages
    );
    assert_eq!(checkpoint_messages.len(), 4);
    let approval = host
        .list_approvals(&handle.run_id)
        .expect("读取 approval")
        .into_iter()
        .find(|item| item.tool_call_id.as_deref() == Some("batch-call-3"))
        .expect("第三个调用应有 pending approval");
    host.resolve_approval(&approval.id, ApprovalDecision::Allow)
        .expect("允许第三个调用");
    host.resume_approval(&approval.id).expect("重新排队");

    // 恢复时不能重新请求产生首批工具调用的 Provider；只继续未完成的调用和下一轮。
    let result = host.run_existing(&handle.run_id).expect("恢复后应完成");
    assert_eq!(result.output.text, "三工具批次完成");
    assert_snapshot_matches_output(&host, &result);
    assert_tool_rows_are_terminal_and_unique(&host, &result.run_id, 4);
    assert_eq!(provider.remaining_steps(), 0);
}

#[test]
fn provider_error_keeps_reconciling_snapshot_replay_consistent() {
    let host = AgentHost::in_memory().expect("创建 Host").with_provider(
        Arc::new(FakeProvider::new([
            FakeStep::tool_call("before-error", "echo", json!({"index": 1})),
            FakeStep::Error {
                kind: ProviderErrorKind::Stream,
                message: "provider fixture error".to_owned(),
            },
        ])),
        "fake",
    );
    let handle = host.prepare_run("Provider 错误").expect("创建 run");
    let error = host
        .run_existing(&handle.run_id)
        .expect_err("Provider 错误应返回");
    assert!(matches!(error, HostError::Engine(EngineError::Provider(_))));

    let record = host
        .get_run(&handle.run_id)
        .expect("读取 run")
        .expect("run 应存在");
    assert_eq!(record.status, "reconciling");
    let checkpoint = host
        .read_checkpoint(&handle.run_id)
        .expect("读取 checkpoint")
        .expect("Provider 错误应保留 checkpoint");
    assert_eq!(checkpoint.phase, "provider_in_flight");
    let snapshot = host
        .load_runtime_snapshot(&handle.runtime_id)
        .expect("读取 RuntimeSnapshot")
        .expect("RuntimeSnapshot 应存在");
    let events = host
        .list_runtime_events(&handle.runtime_id)
        .expect("读取 runtime events");
    let mut replayed = RuntimeSnapshot::try_new(&handle.runtime_id).expect("创建空快照");
    for event in events {
        replayed = reduce(&replayed, &event).expect("runtime event 应可重放");
    }
    assert_eq!(replayed, snapshot);
    let checkpoint_messages: Vec<Message> = serde_json::from_value(checkpoint.messages).unwrap();
    assert_eq!(
        snapshot.run(&handle.run_id).unwrap().messages(),
        checkpoint_messages
    );
    assert_eq!(checkpoint_messages.len(), 3);
    assert_tool_rows_are_terminal_and_unique(&host, &handle.run_id, 1);
}

struct ShortCompressor;

impl ContextCompressor for ShortCompressor {
    fn compress(&self, _request: &CompressionRequest) -> Result<Vec<Message>, EngineError> {
        Ok(vec![Message::user("压缩后的历史").expect("构造摘要消息")])
    }
}

#[test]
fn tool_after_context_compression_does_not_restore_old_history() {
    let host = AgentHost::in_memory()
        .expect("创建 Host")
        .with_provider(
            Arc::new(FakeProvider::new([
                FakeStep::ToolCalls {
                    text: "旧历史".repeat(10_000),
                    calls: vec![FakeToolCall {
                        id: "before-compression".to_owned(),
                        name: "echo".to_owned(),
                        arguments: json!({"index": 0}),
                    }],
                },
                FakeStep::ToolCalls {
                    text: "第二轮旧历史".repeat(10_000),
                    calls: vec![FakeToolCall {
                        id: "compressed-call".to_owned(),
                        name: "echo".to_owned(),
                        arguments: json!({"ok": true}),
                    }],
                },
                FakeStep::text("压缩后完成"),
            ])),
            "fake",
        )
        .with_context_compressor(Arc::new(ShortCompressor));
    let result = host
        .run_with_messages("压缩测试", vec![Message::user("开始工具后压缩").unwrap()])
        .expect("压缩后运行应完成");

    assert_eq!(result.output.text, "压缩后完成");
    assert_eq!(
        result
            .output
            .context_observations
            .iter()
            .filter(|observation| observation.compression_attempted)
            .count(),
        2
    );
    assert_snapshot_matches_output(&host, &result);
    assert_tool_rows_are_terminal_and_unique(&host, &result.run_id, 2);
    let runtime = host
        .load_runtime_snapshot(&result.runtime_id)
        .expect("读取 RuntimeSnapshot")
        .expect("RuntimeSnapshot 应存在");
    let run = runtime.run(&result.run_id).expect("run");
    let messages = run.messages();
    assert!(messages.iter().any(|message| {
        message
            .content()
            .iter()
            .any(|part| part.as_text() == Some("压缩后的历史"))
    }));
    assert!(!messages.iter().any(|message| {
        message
            .content()
            .iter()
            .any(|part| part.as_text().is_some_and(|text| text.contains("旧历史")))
    }));
}

struct FailingCompressor;

impl ContextCompressor for FailingCompressor {
    fn compress(&self, _request: &CompressionRequest) -> Result<Vec<Message>, EngineError> {
        Err(EngineError::ContextOverflow("压缩失败 fixture".into()))
    }
}

#[test]
fn failed_compression_after_tool_keeps_full_history_without_repeating_results() {
    let provider = Arc::new(FakeProvider::new([
        FakeStep::ToolCalls {
            text: "待压缩".repeat(10_000),
            calls: vec![FakeToolCall {
                id: "before-failed-compression".into(),
                name: "echo".into(),
                arguments: json!({"ok": true}),
            }],
        },
        FakeStep::text("不能调用这一步"),
    ]));
    let host = AgentHost::in_memory()
        .unwrap()
        .with_provider(provider.clone(), "fake")
        .with_context_compressor(Arc::new(FailingCompressor));
    let handle = host.prepare_run("工具后压缩失败").unwrap();
    assert!(matches!(
        host.run_existing(&handle.run_id),
        Err(HostError::Engine(EngineError::ContextOverflow(_)))
    ));
    assert_eq!(provider.requests().snapshot().len(), 1);
    assert_eq!(
        host.get_run(&handle.run_id).unwrap().unwrap().status,
        "reconciling"
    );
    let checkpoint = host.read_checkpoint(&handle.run_id).unwrap().unwrap();
    assert_eq!(checkpoint.phase, "compacting");
    let messages: Vec<Message> = serde_json::from_value(checkpoint.messages).unwrap();
    assert_eq!(messages.len(), 3);
    let snapshot = host
        .load_runtime_snapshot(&handle.runtime_id)
        .unwrap()
        .unwrap();
    let run = snapshot.run(&handle.run_id).unwrap();
    assert_eq!(run.messages(), messages);
    assert_eq!(run.tool_calls().len(), 1);
    assert_eq!(run.tool_results().len(), 1);
    assert_tool_rows_are_terminal_and_unique(&host, &handle.run_id, 1);
    let replayed = host
        .list_runtime_events(&handle.runtime_id)
        .unwrap()
        .iter()
        .fold(
            RuntimeSnapshot::try_new(&handle.runtime_id).unwrap(),
            |snapshot, event| reduce(&snapshot, event).unwrap(),
        );
    assert_eq!(replayed, snapshot);
}