notescraft/crates/engine/tests/assistant_checkpoint.rs

//! Checkpoints (assistant-ai-elements 4.3): the journal of assistant writes,
//! the reverse restore, conflicts, caps, and the commands staying user-only.

use std::collections::BTreeMap;
use std::path::Path;

use notescraft_engine::{AssistantSettings, Error, Registry, Session, register_core, run_as_assistant};
use serde_json::{Value, json};

mod common;

const CHAT: &str = "chat-1";
const ALPHA: &str = "# Alpha\n\nfirst\n";
const BETA: &str = "# Beta\n\nsecond\n";
const DELTA: &str = "# Delta\n\nfourth\n";

fn spun_up(name: &str, review: bool, extra: &[(&str, &str)]) -> (common::TempVault, Session, Registry<Session>) {
    let vault = common::temp_vault(name);
    let mut files = vec![("alpha.md", ALPHA), ("beta.md", BETA), ("delta.md", DELTA)];
    files.extend_from_slice(extra);
    common::seed(vault.path(), &files);
    let mut session = Session::open(vault.path()).expect("open the vault");
    session.set_assistant_settings(AssistantSettings { allow_write: true, allow_delete: true, review_writes: review, ..AssistantSettings::default() });
    let mut registry = Registry::new();
    register_core(&mut registry).expect("register the core commands");
    (vault, session, registry)
}

/// Every visible note with its text, to compare a vault before and after.
fn tree(root: &Path) -> BTreeMap<String, String> {
    fn walk(dir: &Path, prefix: &str, out: &mut BTreeMap<String, String>) {
        for entry in std::fs::read_dir(dir).expect("list").flatten() {
            let name = entry.file_name().to_string_lossy().into_owned();
            if name.starts_with('.') {
                continue;
            }
            let path = if prefix.is_empty() { name.clone() } else { format!("{prefix}/{name}") };
            if entry.path().is_dir() {
                walk(&entry.path(), &path, out);
            } else {
                out.insert(path, std::fs::read_to_string(entry.path()).expect("read"));
            }
        }
    }
    let mut out = BTreeMap::new();
    walk(root, "", &mut out);
    out
}

fn turn(session: &mut Session, registry: &Registry<Session>, number: u64) {
    registry.run("assistant.checkpoint_turn", session, &json!({ "id": CHAT, "turn": number })).expect("set the turn");
}

fn agent(session: &mut Session, registry: &Registry<Session>, id: &str, params: &Value) {
    run_as_assistant(session, registry, id, params).unwrap_or_else(|error| panic!("{id}: {error}"));
}

fn restore(session: &mut Session, registry: &Registry<Session>, number: u64) -> Value {
    registry.run("assistant.checkpoint_restore", session, &json!({ "id": CHAT, "turn": number })).expect("restore")
}

fn paths(reply: &Value, key: &str) -> Vec<String> {
    reply[key].as_array().expect("array").iter().map(|row| row["path"].as_str().expect("path").to_owned()).collect()
}

#[test]
fn a_turn_of_create_modify_delete_rename_restores_the_vault() {
    let (vault, mut session, registry) = spun_up("cp-all", false, &[]);
    let before = tree(vault.path());
    turn(&mut session, &registry, 1);
    agent(&mut session, &registry, "note.create", &json!({ "path": "fresh.md", "content": "new" }));
    agent(&mut session, &registry, "note.write", &json!({ "path": "alpha.md", "content": "rewritten" }));
    agent(&mut session, &registry, "note.delete", &json!({ "path": "delta.md" }));
    agent(&mut session, &registry, "note.rename", &json!({ "from": "beta.md", "to": "gamma.md" }));
    assert_ne!(tree(vault.path()), before);

    let listed = registry.run("assistant.checkpoint_list", &mut session, &json!({ "id": CHAT })).expect("list");
    assert_eq!(listed["turns"][0]["count"], 4);

    let reply = restore(&mut session, &registry, 1);
    assert_eq!(paths(&reply, "conflicts"), Vec::<String>::new(), "{reply}");
    assert_eq!(reply["restored"].as_array().map(Vec::len), Some(4));
    assert_eq!(tree(vault.path()), before);
    let again = registry.run("assistant.checkpoint_list", &mut session, &json!({ "id": CHAT })).expect("list");
    assert_eq!(again["turns"], json!([]), "restored entries leave the journal");
}

#[test]
fn accepted_pending_changes_are_journaled_and_restore_to_the_old_text() {
    let (vault, mut session, registry) = spun_up("cp-accept", true, &[]);
    turn(&mut session, &registry, 1);
    agent(&mut session, &registry, "note.write", &json!({ "path": "alpha.md", "content": "proposed" }));
    let listed = registry.run("assistant.checkpoint_list", &mut session, &json!({ "id": CHAT })).expect("list");
    assert_eq!(listed["turns"], json!([]), "a queued change is not a write yet");
    registry.run("assistant.accept", &mut session, &json!({ "all": true })).expect("accept");
    assert_eq!(std::fs::read_to_string(vault.path().join("alpha.md")).expect("read"), "proposed");
    restore(&mut session, &registry, 1);
    assert_eq!(std::fs::read_to_string(vault.path().join("alpha.md")).expect("read"), ALPHA);
}

#[test]
fn a_note_edited_since_is_a_conflict_and_left_alone() {
    let (vault, mut session, registry) = spun_up("cp-conflict", false, &[]);
    turn(&mut session, &registry, 1);
    agent(&mut session, &registry, "note.write", &json!({ "path": "alpha.md", "content": "agent text" }));
    agent(&mut session, &registry, "note.write", &json!({ "path": "beta.md", "content": "agent beta" }));
    session.write_note("alpha.md", "the user typed this").expect("user edit");

    let reply = restore(&mut session, &registry, 1);
    assert_eq!(paths(&reply, "conflicts"), ["alpha.md"]);
    assert_eq!(paths(&reply, "restored"), ["beta.md"]);
    assert_eq!(std::fs::read_to_string(vault.path().join("alpha.md")).expect("read"), "the user typed this");
    assert_eq!(std::fs::read_to_string(vault.path().join("beta.md")).expect("read"), BETA);
    let left = registry.run("assistant.checkpoint_list", &mut session, &json!({ "id": CHAT })).expect("list");
    assert_eq!(left["turns"][0]["count"], 1, "the conflict stays journaled");
}

#[test]
fn restoring_goes_back_through_every_later_turn() {
    let (vault, mut session, registry) = spun_up("cp-turns", false, &[]);
    turn(&mut session, &registry, 1);
    agent(&mut session, &registry, "note.write", &json!({ "path": "alpha.md", "content": "one" }));
    turn(&mut session, &registry, 2);
    agent(&mut session, &registry, "note.write", &json!({ "path": "alpha.md", "content": "two" }));
    restore(&mut session, &registry, 2);
    assert_eq!(std::fs::read_to_string(vault.path().join("alpha.md")).expect("read"), "one");
    turn(&mut session, &registry, 3);
    agent(&mut session, &registry, "note.write", &json!({ "path": "alpha.md", "content": "three" }));
    let reply = restore(&mut session, &registry, 1);
    assert_eq!(reply["restored"].as_array().map(Vec::len), Some(2));
    assert_eq!(std::fs::read_to_string(vault.path().join("alpha.md")).expect("read"), ALPHA);
}

#[test]
fn an_entry_over_the_cap_is_not_restorable_and_the_note_stays() {
    let big = "x".repeat(notescraft_engine::MAX_PENDING_TEXT + 1);
    let (vault, mut session, registry) = spun_up("cp-cap", false, &[("big.md", &big)]);
    turn(&mut session, &registry, 1);
    agent(&mut session, &registry, "note.write", &json!({ "path": "big.md", "content": "small" }));
    let listed = registry.run("assistant.checkpoint_list", &mut session, &json!({ "id": CHAT })).expect("list");
    assert_eq!(listed["turns"][0]["files"][0]["restorable"], false);
    let reply = restore(&mut session, &registry, 1);
    assert_eq!(paths(&reply, "skipped"), ["big.md"]);
    assert_eq!(std::fs::read_to_string(vault.path().join("big.md")).expect("read"), "small");
}

#[test]
fn restore_snapshots_the_text_it_replaces() {
    let (_vault, mut session, registry) = spun_up("cp-history", false, &[]);
    turn(&mut session, &registry, 1);
    agent(&mut session, &registry, "note.write", &json!({ "path": "alpha.md", "content": "agent text" }));
    restore(&mut session, &registry, 1);
    let versions = session.note_snapshots("alpha.md");
    assert_eq!(versions.len(), 1, "version history holds the agent's text");
    assert_eq!(session.read_snapshot("alpha.md", versions[0].id).expect("snapshot"), "agent text");
}

#[test]
fn writes_outside_a_turn_are_not_journaled() {
    let (_vault, mut session, registry) = spun_up("cp-noturn", false, &[]);
    agent(&mut session, &registry, "note.write", &json!({ "path": "alpha.md", "content": "x" }));
    let listed = registry.run("assistant.checkpoint_list", &mut session, &json!({ "id": CHAT })).expect("list");
    assert_eq!(listed["turns"], json!([]));
    let missing = registry.run("assistant.checkpoint_restore", &mut session, &json!({ "id": CHAT, "turn": 1 }));
    assert!(matches!(missing, Err(Error::Command { .. })), "{missing:?}");
}

#[test]
fn the_assistant_cannot_call_the_checkpoint_commands() {
    for review in [false, true] {
        let (_vault, mut session, registry) = spun_up("cp-denied", review, &[]);
        for id in ["assistant.checkpoint_turn", "assistant.checkpoint_list", "assistant.checkpoint_restore"] {
            let outcome = run_as_assistant(&mut session, &registry, id, &json!({ "id": CHAT, "turn": 1 }));
            match outcome {
                Err(Error::Command { source, .. }) => assert_eq!(source.code, "assistant_denied", "{id}"),
                other => panic!("{id} must be denied, found {other:?}"),
            }
        }
        assert_eq!(session.pending_changes().len(), 0, "a denied call is not queued");
    }
}