notescraft/crates/ui-egui/tests/assistant_flow.rs

//! Offline ACP-to-egui conversation, permission and cancellation integration.

use egui::{Context, Event as InputEvent, Key, Modifiers, RawInput, Rect};
use notescraft_acp::{AgentCommand, Event, McpServerConfig, PermissionOption, PromptBlock, RequestId, SessionUpdate, ToolCallInfo};
use notescraft_ui_egui::{
    app::NotesApp,
    assistant_panel::{
        AssistantPanel,
        worker::agent::{Agent, Spawner},
    },
};
use std::{
    path::Path,
    sync::{Arc, Mutex, mpsc},
    time::{Duration, Instant},
};

/// One test at a time: `agent_gets_the_notescraft_mcp_server...` plants a stand-in `notescraft-cli` beside the
/// test binary, which would otherwise change what the prompts of the other tests say.
static SERIAL: Mutex<()> = Mutex::new(());

fn serial() -> std::sync::MutexGuard<'static, ()> {
    SERIAL.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
}

#[derive(Default)]
struct Script {
    events: Option<mpsc::Sender<Event>>,
    prompts: Vec<Vec<PromptBlock>>,
    answers: Vec<(RequestId, Option<String>)>,
    cancelled: bool,
    spawns: usize,
    servers: Vec<Vec<McpServerConfig>>,
}

fn script_lock(script: &Mutex<Script>) -> std::sync::MutexGuard<'_, Script> {
    script.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
}

struct FakeSpawner(Arc<Mutex<Script>>);
struct FakeAgent(Arc<Mutex<Script>>, mpsc::Receiver<Event>);

impl Spawner for FakeSpawner {
    fn spawn(&mut self, _command: AgentCommand) -> Result<Box<dyn Agent>, String> {
        let (sender, receiver) = mpsc::channel();
        let mut script = script_lock(&self.0);
        script.events = Some(sender);
        script.spawns += 1;
        Ok(Box::new(FakeAgent(self.0.clone(), receiver)))
    }
}

impl Agent for FakeAgent {
    fn initialize(&self, _cancelled: &dyn Fn() -> bool) -> Result<(), String> {
        Ok(())
    }
    fn new_session(&self, _vault: &Path, servers: &[McpServerConfig], _cancelled: &dyn Fn() -> bool) -> Result<String, String> {
        script_lock(&self.0).servers.push(servers.to_vec());
        Ok("offline".into())
    }
    fn events(&self) -> &mpsc::Receiver<Event> {
        &self.1
    }
    fn stderr_tail(&self) -> String {
        "offline agent stderr diagnostic".into()
    }
    fn prompt(&self, _session: &str, blocks: &[PromptBlock]) -> Result<(), String> {
        script_lock(&self.0).prompts.push(blocks.to_vec());
        Ok(())
    }
    fn cancel(&self, _session: &str) -> Result<(), String> {
        script_lock(&self.0).cancelled = true;
        Ok(())
    }
    fn answer(&self, id: &RequestId, option: Option<&str>) -> Result<(), String> {
        script_lock(&self.0).answers.push((id.clone(), option.map(str::to_owned)));
        Ok(())
    }
}

struct Harness {
    app: NotesApp,
    ctx: Context,
    script: Arc<Mutex<Script>>,
    registry: notescraft_engine::Registry<notescraft_engine::Session>,
}

impl Harness {
    fn new(read_only: bool) -> Option<Self> {
        use std::sync::atomic::{AtomicU64, Ordering};
        static COUNT: AtomicU64 = AtomicU64::new(0);
        let root = std::env::temp_dir().join(format!("notescraft-assistant-flow-{}-{}", std::process::id(), COUNT.fetch_add(1, Ordering::Relaxed)));
        assert!(std::fs::create_dir_all(&root).is_ok());
        let mut session = notescraft_engine::Session::open(&root).ok()?;
        assert!(session.create_note("note.md", "# Context\n").is_ok());
        if read_only {
            drop(session);
            session = notescraft_engine::Session::open_read_only(&root).ok()?;
        }
        session.use_settings_dir(root.join(".notescraft/test-config"));
        let mut app = NotesApp::new(session, notescraft_themes::builtin("catppuccin-mocha")?);
        let mut early = notescraft_engine::Registry::new();
        assert!(notescraft_engine::register_core(&mut early).is_ok());
        let set = early.run(
            "settings.set",
            &mut app.session,
            &serde_json::json!({"key":"assistant.connections","value":[{"id":"offline","name":"Offline","kind":"acp","command":"offline-agent","env":{},"icon":"custom","base_url":"","model":"","api_key_env":"","system_prompt":""}]}),
        );
        assert!(set.is_ok(), "{set:?}");
        let set = early.run("settings.set", &mut app.session, &serde_json::json!({"key":"assistant.active","value":"offline"}));
        assert!(set.is_ok(), "{set:?}");
        assert_eq!(app.session.assistant_settings().connections.first().map(|c| c.command.as_str()), Some("offline-agent"));
        app.state.animations = false;
        app.state.right_width = 560.0;
        let script = Arc::new(Mutex::new(Script::default()));
        app.assistant = AssistantPanel::with_spawner(Box::new(FakeSpawner(script.clone())));
        app.open_note("note.md");
        app.run_command("ui.toggle_assistant", &serde_json::Value::Null);
        let mut registry = notescraft_engine::Registry::new();
        assert!(notescraft_engine::register_core(&mut registry).is_ok());
        let mut harness = Self { app, ctx: Context::default(), script, registry };
        for _ in 0..3 {
            harness.frame(Vec::new());
        }
        Some(harness)
    }

    fn frame(&mut self, events: Vec<InputEvent>) -> Vec<(String, egui::Pos2)> {
        let input = RawInput { screen_rect: Some(Rect::from_min_size(egui::Pos2::ZERO, egui::vec2(1440.0, 900.0))), events, ..Default::default() };
        let mut output = self.ctx.run_ui(input, |ui| self.app.ui(ui.ctx()));
        let mut labels: Vec<(String, egui::Pos2)> = output
            .shapes
            .iter()
            .filter_map(|shape| match &shape.shape {
                egui::epaint::Shape::Text(text) => Some((text.galley.text().to_owned(), text.pos)),
                _ => None,
            })
            .collect();
        // The send and stop buttons are icons without painted text: name them by their ids.
        for (id, name) in [("assistant-send", "Send"), ("assistant-stop", "Stop")] {
            if let Some(button) = self.ctx.read_response(egui::Id::new(id)) {
                labels.push((name.to_owned(), button.rect.min));
            }
        }
        output.textures_delta.clear();
        labels
    }

    fn wait(&mut self, mut condition: impl FnMut(&mut Self, &[(String, egui::Pos2)]) -> bool) {
        let deadline = Instant::now() + Duration::from_secs(3);
        loop {
            let labels = self.frame(Vec::new());
            if condition(self, &labels) {
                break;
            }
            assert!(Instant::now() < deadline, "assistant flow deadline; labels: {labels:?}");
            std::thread::sleep(Duration::from_millis(5));
        }
    }

    fn emit(&self, event: Event) {
        assert!(script_lock(&self.script).events.as_ref().is_some_and(|sender| sender.send(event).is_ok()));
    }

    fn click_label(&mut self, target: &str) {
        let labels = self.frame(Vec::new());
        let position = labels.iter().find(|(text, _)| text == target).map(|(_, position)| *position + egui::vec2(5.0, 5.0));
        assert!(position.is_some(), "missing {target}: {labels:?}");
        let Some(pos) = position else { return };
        for pressed in [true, false] {
            self.frame(vec![
                InputEvent::PointerMoved(pos),
                InputEvent::PointerButton { pos, button: egui::PointerButton::Primary, pressed, modifiers: Modifiers::NONE },
            ]);
        }
    }

    fn key(&mut self, key: Key, modifiers: Modifiers) {
        for pressed in [true, false] {
            self.frame(vec![InputEvent::Key { key, physical_key: None, pressed, repeat: false, modifiers }]);
        }
    }

    fn send(&mut self, prompt: &str) {
        let previous = script_lock(&self.script).prompts.len();
        assert!(self.app.assistant.set_composer(prompt).is_ok());
        self.ctx.memory_mut(|memory| memory.request_focus(egui::Id::new("assistant-composer")));
        self.key(Key::Enter, Modifiers::NONE);
        self.wait(|harness, _| script_lock(&harness.script).prompts.len() > previous);
    }

    fn propose(&mut self, content: &str) {
        let queued = notescraft_engine::assistant::run_as_assistant(
            &mut self.app.session,
            &self.registry,
            "note.write",
            &serde_json::json!({"path":"note.md","content":content}),
        );
        assert!(queued.is_ok(), "queue write: {queued:?}");
        assert_eq!(queued.ok().and_then(|value| value.get("queued").and_then(serde_json::Value::as_bool)), Some(true));
        self.wait(|_, labels| labels.iter().any(|(label, _)| label == "Accept"));
    }
}

#[test]
fn send_stream_tools_permission_and_cancel_cross_the_real_panel_seam() {
    let _serial = serial();
    let harness = Harness::new(false);
    assert!(harness.is_some());
    let Some(mut harness) = harness else { return };
    assert!(harness.app.assistant.set_composer("Summarise context").is_ok());
    harness.ctx.memory_mut(|memory| memory.request_focus(egui::Id::new("assistant-composer")));
    harness.frame(vec![InputEvent::Key { key: Key::Enter, physical_key: None, pressed: true, repeat: false, modifiers: Modifiers::NONE }]);
    harness.wait(|harness, labels| !script_lock(&harness.script).prompts.is_empty() && labels.iter().any(|(text, _)| text == "Summarise context"));
    assert!(
        script_lock(&harness.script)
            .prompts
            .first()
            .is_some_and(|blocks| blocks.iter().any(|block| matches!(block, PromptBlock::ResourceLink { uri, .. } if uri == "note://note.md")))
    );
    harness.emit(Event::Update(SessionUpdate::AgentMessageChunk("Here is ".into())));
    harness.emit(Event::Update(SessionUpdate::AgentMessageChunk("your summary.".into())));
    harness.emit(Event::Update(SessionUpdate::ToolCall(ToolCallInfo {
        id: "read-1".into(),
        title: "Read note".into(),
        kind: "read".into(),
        status: "in_progress".into(),
        locations: vec!["note.md".into()],
        raw_input: None,
        parent_id: None,
        tool_name: None,
    })));
    harness.wait(|_, labels| {
        labels.iter().any(|(text, _)| text.starts_with("Reading note.md")) && labels.iter().any(|(text, _)| text.contains("Here is your summary."))
    });
    harness.emit(Event::Update(SessionUpdate::ToolCallUpdate {
        id: "read-1".into(),
        status: Some("completed".into()),
        title: None,
        content_summary: Some("Read 2 lines".into()),
    }));
    harness.emit(Event::PermissionRequest {
        request_id: RequestId::Number(42),
        tool_call: ToolCallInfo { title: "Continue reading".into(), ..Default::default() },
        options: vec![PermissionOption { id: "once-42".into(), name: "Allow this read".into(), kind: "allow_once".into() }],
    });
    harness.wait(|_, labels| labels.iter().any(|(text, _)| text == "Allow this read"));
    harness.click_label("Allow this read");
    harness.wait(|harness, _| !script_lock(&harness.script).answers.is_empty());
    assert_eq!(script_lock(&harness.script).answers.first(), Some(&(RequestId::Number(42), Some("once-42".into()))));
    harness.click_label("Stop");
    harness.wait(|harness, _| script_lock(&harness.script).cancelled);
}

#[test]
fn process_death_retry_and_history_preserve_stream_order() {
    let _serial = serial();
    let harness = Harness::new(false);
    assert!(harness.is_some());
    let Some(mut harness) = harness else { return };
    harness.send("Original prompt");
    for event in [
        Event::Update(SessionUpdate::AgentMessageChunk("First ".into())),
        Event::Update(SessionUpdate::AgentMessageChunk("reply".into())),
        Event::Update(SessionUpdate::ToolCall(ToolCallInfo {
            id: "read".into(),
            title: "Read note".into(),
            status: "in_progress".into(),
            locations: vec!["note.md".into()],
            ..Default::default()
        })),
        Event::Update(SessionUpdate::ToolCallUpdate {
            id: "read".into(),
            status: Some("completed".into()),
            title: None,
            content_summary: Some("Read context".into()),
        }),
        Event::Exited { code: Some(1) },
    ] {
        harness.emit(event);
    }
    harness.wait(|_, labels| {
        labels.iter().any(|(label, _)| label.contains("offline agent stderr diagnostic")) && labels.iter().any(|(label, _)| label == "Retry")
    });
    harness.click_label("Retry");
    harness.wait(|harness, _| script_lock(&harness.script).spawns == 2 && script_lock(&harness.script).prompts.len() == 2);
    {
        let script = script_lock(&harness.script);
        assert_eq!(script.prompts.first(), script.prompts.last());
    }
    harness.emit(Event::Update(SessionUpdate::AgentMessageChunk("Recovered".into())));
    harness.emit(Event::PromptDone { stop_reason: "end_turn".into() });
    harness.wait(|_, labels| labels.iter().any(|(label, _)| label == "Recovered") && labels.iter().any(|(label, _)| label == "Send"));
    let sessions = harness.app.session.chat_list();
    assert!(sessions.as_ref().is_ok_and(|sessions| sessions.len() == 1), "chat list: {sessions:?}");
    let Some(id) = sessions.ok().and_then(|sessions| sessions.first().map(|entry| entry.id.clone())) else { return };
    let lines = harness.app.session.chat_read(&id);
    assert!(lines.is_ok());
    let Ok(lines) = lines else { return };
    assert_eq!(lines.iter().map(|line| line.role.name()).collect::<Vec<_>>(), ["user", "agent", "agent", "tool", "tool", "user", "agent"]);
    assert_eq!(lines.get(4).and_then(|line| line.tool.as_ref()).map(|tool| (tool.status.as_str(), tool.path.as_deref())), Some(("DONE", Some("note.md"))));
    // The header history control is an icon button without a text label.
    harness.app.assistant.handle(notescraft_ui_egui::assistant_panel::Action::ToggleHistory, &mut harness.app.session, &harness.registry, None);
    harness.click_label("Original prompt");
    harness.wait(|_, labels| labels.iter().any(|(label, _)| label == "Viewing a past chat"));
    assert!(
        harness
            .app
            .assistant
            .messages()
            .iter()
            .any(|message| matches!(message, notescraft_ui_egui::assistant_panel::Message::Agent(text) if text == "First reply"))
    );
    assert!(harness.app.assistant.set_composer("History stays read only").is_ok());
    harness.click_label("Send");
    assert_eq!(script_lock(&harness.script).prompts.len(), 2);
    assert_eq!(harness.app.assistant.composer(), "History stays read only");
}

#[test]
fn write_grant_accept_reject_and_conflict_use_real_review_buttons() {
    let _serial = serial();
    let harness = Harness::new(false);
    assert!(harness.is_some());
    let Some(mut harness) = harness else { return };
    harness.send("Change context");
    harness.emit(Event::PermissionRequest {
        request_id: RequestId::Number(8),
        tool_call: ToolCallInfo { title: "Edit note".into(), kind: "edit".into(), ..Default::default() },
        options: vec![],
    });
    harness.wait(|_, labels| labels.iter().any(|(label, _)| label == "This agent is asking to change your vault."));
    harness.click_label("Grant write access for this session");
    assert!(harness.app.session.session_write_granted());
    harness.emit(Event::PromptDone { stop_reason: "end_turn".into() });
    harness.wait(|_, labels| labels.iter().any(|(label, _)| label == "Write access granted"));
    harness.propose("# Accepted\n");
    assert_eq!(harness.app.session.read_note("note.md").ok().as_deref(), Some("# Context\n"));
    harness.click_label("Accept");
    assert_eq!(harness.app.session.pending_changes(), []);
    assert_eq!(harness.app.session.read_note("note.md").ok().as_deref(), Some("# Accepted\n"));
    harness.propose("# Rejected\n");
    harness.click_label("Reject");
    assert_eq!(harness.app.session.pending_changes(), []);
    assert_eq!(harness.app.session.read_note("note.md").ok().as_deref(), Some("# Accepted\n"));
    harness.propose("# Conflicting\n");
    harness.app.run_command("note.write", &serde_json::json!({"path":"note.md","content":"# User edit\n"}));
    harness.click_label("Accept");
    harness.wait(|_, labels| labels.iter().any(|(label, _)| label.contains("changed since it was queued")));
    assert_eq!(harness.app.session.pending_changes().len(), 1);
    assert_eq!(harness.app.session.read_note("note.md").ok().as_deref(), Some("# User edit\n"));
    harness.click_label("Reject");
    assert_eq!(harness.app.session.pending_changes(), []);
}

#[test]
fn read_only_vault_sends_recalls_and_shows_engine_refusals_without_writing() {
    let _serial = serial();
    let harness = Harness::new(true);
    assert!(harness.is_some());
    let Some(mut harness) = harness else { return };
    harness.send("Read-only question");
    harness.emit(Event::PromptDone { stop_reason: "end_turn".into() });
    harness.wait(|_, labels| labels.iter().any(|(label, _)| label == "Send"));
    harness.ctx.memory_mut(|memory| memory.request_focus(egui::Id::new("assistant-composer")));
    harness.frame(Vec::new());
    harness.key(Key::ArrowUp, Modifiers::NONE);
    assert_eq!(harness.app.assistant.composer(), "Read-only question");
    harness.key(Key::Enter, Modifiers::SHIFT);
    assert!(
        harness.app.assistant.composer().contains('\n'),
        "Shift+Enter composer: {:?}; focus: {:?}",
        harness.app.assistant.composer(),
        harness.ctx.memory(egui::Memory::focused)
    );
    assert_eq!(script_lock(&harness.script).prompts.len(), 1);
    harness.key(Key::Escape, Modifiers::NONE);
    assert_eq!(harness.app.assistant.composer(), "\nRead-only question", "Escape leaves the field and keeps the draft");
    let refused = harness.registry.run("note.write", &mut harness.app.session, &serde_json::json!({"path":"note.md","content":"forbidden"}));
    assert!(refused.is_err());
    let Err(error) = refused else { return };
    let notice = error.to_string();
    harness.emit(Event::Error(notice.clone()));
    harness.wait(|_, labels| labels.iter().any(|(label, _)| label == &notice));
    let labels = harness.frame(Vec::new());
    assert!(labels.iter().any(|(label, _)| label == "Read only"));
    assert!(!labels.iter().any(|(label, _)| label == "Grant write"));
    assert_eq!(harness.app.session.read_note("note.md").ok().as_deref(), Some("# Context\n"));
    assert!(harness.app.session.chat_list().is_ok_and(|sessions| sessions.is_empty()));
}

fn first_prompt_text(harness: &Harness) -> String {
    let script = script_lock(&harness.script);
    let blocks = script.prompts.first().cloned().unwrap_or_default();
    blocks.iter().filter_map(|block| if let PromptBlock::Text(text) = block { Some(text.as_str()) } else { None }).collect::<Vec<_>>().join("\n")
}

#[test]
fn agent_gets_the_notescraft_mcp_server_and_a_prompt_naming_the_vault_tools_and_grant() {
    let _serial = serial();
    // The panel looks for `notescraft-cli` beside the running binary; stand one in so the test does not depend on a prior CLI build.
    let beside =
        std::env::current_exe().ok().and_then(|exe| exe.parent().map(|dir| dir.join(if cfg!(windows) { "notescraft-cli.exe" } else { "notescraft-cli" })));
    let planted = beside.filter(|path| !path.exists() && std::fs::write(path, b"").is_ok());
    let harness = Harness::new(false);
    assert!(harness.is_some());
    let Some(mut harness) = harness else { return };
    harness.send("Create a note called ideas");
    let denied = first_prompt_text(&harness);
    assert!(denied.contains("NotesCraft") && denied.contains("not Obsidian") && denied.contains("create_note"), "{denied}");
    assert!(denied.contains("Write access is OFF"), "{denied}");
    let servers = script_lock(&harness.script).servers.last().cloned().unwrap_or_default();
    assert_eq!(servers.len(), 1, "exactly the NotesCraft MCP server is handed over");
    assert_eq!(servers[0].name, "notescraft");
    assert!(servers[0].args.iter().any(|arg| arg == "--via-control") && servers[0].args.iter().any(|arg| arg == "mcp"));
    assert!(!servers[0].args.iter().any(|arg| arg == "--write"), "no write flag before the grant: {:?}", servers[0].args);
    harness.emit(Event::PromptDone { stop_reason: "end_turn".into() });
    harness.wait(|_, labels| labels.iter().any(|(label, _)| label == "Send"));
    harness.app.run_command("assistant.grant_write", &serde_json::json!({"session": true}));
    assert!(harness.app.session.session_write_granted());
    harness.send("Create a note called ideas");
    harness.wait(|harness, _| script_lock(&harness.script).prompts.len() >= 2);
    let script = script_lock(&harness.script);
    let text = script.prompts.last().map(|blocks| format!("{blocks:?}")).unwrap_or_default();
    assert!(text.contains("Write access is granted"), "{text}");
    let last = script.servers.last().cloned().unwrap_or_default();
    assert!(last.iter().any(|server| server.args.iter().any(|arg| arg == "--write")), "{last:?}");
    if let Some(path) = planted {
        let _ = std::fs::remove_file(path);
    }
}

#[test]
fn without_the_cli_the_agent_is_told_it_has_no_vault_tools() {
    let _serial = serial();
    let harness = Harness::new(false);
    assert!(harness.is_some());
    let Some(mut harness) = harness else { return };
    harness.send("Create a note called ideas");
    let text = first_prompt_text(&harness);
    assert!(text.contains("NotesCraft"), "{text}");
    assert!(text.contains("create_note") != text.contains("could not be started"), "exactly one of tools listed or honest failure: {text}");
}