//! 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) { 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 { fn walk(dir: &Path, prefix: &str, out: &mut BTreeMap) { 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, number: u64) { registry.run("assistant.checkpoint_turn", session, &json!({ "id": CHAT, "turn": number })).expect("set the turn"); } fn agent(session: &mut Session, registry: &Registry, 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, number: u64) -> Value { registry.run("assistant.checkpoint_restore", session, &json!({ "id": CHAT, "turn": number })).expect("restore") } fn paths(reply: &Value, key: &str) -> Vec { 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, ®istry, 1); agent(&mut session, ®istry, "note.create", &json!({ "path": "fresh.md", "content": "new" })); agent(&mut session, ®istry, "note.write", &json!({ "path": "alpha.md", "content": "rewritten" })); agent(&mut session, ®istry, "note.delete", &json!({ "path": "delta.md" })); agent(&mut session, ®istry, "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, ®istry, 1); assert_eq!(paths(&reply, "conflicts"), Vec::::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, ®istry, 1); agent(&mut session, ®istry, "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, ®istry, 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, ®istry, 1); agent(&mut session, ®istry, "note.write", &json!({ "path": "alpha.md", "content": "agent text" })); agent(&mut session, ®istry, "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, ®istry, 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, ®istry, 1); agent(&mut session, ®istry, "note.write", &json!({ "path": "alpha.md", "content": "one" })); turn(&mut session, ®istry, 2); agent(&mut session, ®istry, "note.write", &json!({ "path": "alpha.md", "content": "two" })); restore(&mut session, ®istry, 2); assert_eq!(std::fs::read_to_string(vault.path().join("alpha.md")).expect("read"), "one"); turn(&mut session, ®istry, 3); agent(&mut session, ®istry, "note.write", &json!({ "path": "alpha.md", "content": "three" })); let reply = restore(&mut session, ®istry, 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, ®istry, 1); agent(&mut session, ®istry, "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, ®istry, 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, ®istry, 1); agent(&mut session, ®istry, "note.write", &json!({ "path": "alpha.md", "content": "agent text" })); restore(&mut session, ®istry, 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, ®istry, "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, ®istry, 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"); } }