notescraft/crates/engine/tests/assistant_control.rs

//! The assistant gate over a live control channel (SPEC 10.3): the tool call
//! arrives with `source: "assistant"`, is gated by the pump and answers
//! `{"queued": true, …}`; wrong sources and wrong ids are refused; reads pass.

#![allow(dead_code)] // shared fixtures: only some helpers are used here

use std::io::Write;
use std::net::TcpStream;

use notescraft_engine::{CommandResult, ControlClient, ControlHandle, ControlOptions, ControlServer, Error, Registry, Session, register_core};
use serde_json::{Value, json};

mod common;
use common::temp_vault;

const ALIVE: std::time::Duration = std::time::Duration::from_secs(30);

/// The `test.echo` body: params back out unchanged (from `tests/control.rs`).
#[expect(clippy::unnecessary_wraps, reason = "the registry signature demands CommandResult")]
fn echo(_session: &mut Session, params: &Value) -> CommandResult {
    Ok(params.clone())
}

/// A vault, its session, its registry and the live control server over it.
fn spun_up(name: &str) -> (common::TempVault, Session, Registry<Session>, ControlHandle) {
    let temp = temp_vault(name);
    common::seed(temp.path(), &[("alpha.md", "# Alpha\n\nbody\n")]);
    let mut session = Session::open(temp.path()).expect("open the vault");
    // The app's stored assistant settings are conservative: no writes.
    // The app's stored assistant settings gate reads only; tests turn the
    // knobs they need on their own session.
    session.set_assistant_settings(notescraft_engine::AssistantSettings {
        allow_write: true,
        review_writes: true,
        ..notescraft_engine::AssistantSettings::default()
    });
    let mut registry = Registry::new();
    register_core(&mut registry).expect("register the core commands");
    registry
        .register(
            notescraft_engine::CommandSpec::new("test.echo", "Echo", "echoes its params", |_: &Session| true, echo)
                .with_param(notescraft_engine::ParamSpec::optional("question", notescraft_engine::ParamKind::Integer, "number echoed back")),
        )
        .expect("register test.echo");
    let options = ControlOptions { idle_timeout: ALIVE, ..ControlOptions::default() };
    let handle = ControlServer::start_with(session.vault(), options).expect("start the control server");
    (temp, session, registry, handle)
}

/// Pumps until `worker` finishes, then joins it (the test thread is the UI thread).
fn drive<T>(handle: &mut ControlHandle, session: &mut Session, registry: &Registry<Session>, worker: std::thread::JoinHandle<T>) -> T {
    let started = std::time::Instant::now();
    loop {
        handle.pump(session, registry);
        if worker.is_finished() {
            return worker.join().expect("worker thread did not panic");
        }
        assert!(started.elapsed() < std::time::Duration::from_secs(60), "the worker never finished");
        std::thread::sleep(std::time::Duration::from_millis(1));
    }
}

/// A write with `source: "assistant"` is queued, not executed; the pending
/// change lands in the session the pump drives.
#[test]
fn an_assistant_write_is_queued_over_the_wire() {
    let (keep, mut session, registry, mut handle) = spun_up("assistant-control-queue");
    let root = keep.path().to_owned();
    let worker = std::thread::spawn(move || {
        let mut client = ControlClient::connect(&root).expect("connect");
        let queued = client.run_as("note.write", &json!({ "path": "alpha.md", "content": "agent text" }), "assistant").expect("queue");
        let status = client.run_as("assistant.status", &Value::Null, "assistant").expect("status");
        (queued, status)
    });
    let (queued, status) = drive(&mut handle, &mut session, &registry, worker);
    assert_eq!(queued.get("queued").and_then(Value::as_bool), Some(true));
    assert_eq!(queued.get("pending_id").and_then(Value::as_u64), Some(1));
    assert_eq!(queued.get("message").and_then(Value::as_str).map(|m| m.starts_with("NOT written yet")), Some(true));
    assert_eq!(status.get("pending").and_then(Value::as_u64), Some(1), "{status}");
    assert_eq!(session.read_note("alpha.md").expect("read"), "# Alpha\n\nbody\n", "nothing was written");
    assert_eq!(session.pending_changes().len(), 1, "the pump stored the pending change");
    let pending = session.pending_changes().first().cloned().expect("a change");
    assert_eq!(pending.old.as_deref(), Some("# Alpha\n\nbody\n"));
    assert_eq!(pending.new.as_deref(), Some("agent text"));
}

/// A write with `source: "cli"` runs like every other user call: ungated.
#[test]
fn a_cli_source_write_runs_ungated() {
    let (keep, mut session, registry, mut handle) = spun_up("assistant-control-cli");
    let root = keep.path().to_owned();
    let worker = std::thread::spawn(move || {
        let mut client = ControlClient::connect(&root).expect("connect");
        let cli = client.run_as("note.write", &json!({ "path": "alpha.md", "content": "cli text" }), "cli").expect("run");
        let ui = client.run_as("note.write", &json!({ "path": "alpha.md", "content": "ui text" }), "ui").expect("run");
        let default = client.run("note.write", &json!({ "path": "alpha.md", "content": "default text" })).expect("run");
        (cli, ui, default)
    });
    let (cli, ui, default) = drive(&mut handle, &mut session, &registry, worker);
    assert_eq!(cli.get("changed").and_then(Value::as_bool), Some(true));
    assert_eq!(ui.get("changed").and_then(Value::as_bool), Some(true));
    assert_eq!(default.get("changed").and_then(Value::as_bool), Some(true));
    assert_eq!(session.pending_changes().len(), 0, "no source means no gate");
}

/// A refused assistant write comes back as the `assistant_denied` command
/// error, and the session gain the user made is honoured.
#[test]
fn a_refused_write_is_an_error_and_a_grant_queues_it() {
    let (keep, mut session, registry, mut handle) = spun_up("assistant-control-deny");
    // Close the write grant; the session grant the user can still open.
    session.set_assistant_settings(notescraft_engine::AssistantSettings::default());
    let root = keep.path().to_owned();
    let worker = std::thread::spawn(move || {
        let mut client = ControlClient::connect(&root).expect("connect");
        // Reads pass even when everything is closed.
        let read = client.run_as("note.read", &json!({ "path": "alpha.md" }), "assistant").expect("read");
        // A write is refused with the gate's own code.
        let denied = client.run_as("note.write", &json!({ "path": "alpha.md", "content": "no" }), "assistant");
        // The user toggles the session grant; now the reviewed write queues.
        client.run("assistant.grant_write", &json!({ "session": true })).expect("grant");
        let queued = client.run_as("note.write", &json!({ "path": "alpha.md", "content": "granted" }), "assistant").expect("queue");
        // The gate also checks deletes separately.
        let delete_denied = client.run_as("note.delete", &json!({ "path": "alpha.md" }), "assistant").err();
        (read, denied, queued, delete_denied)
    });
    let (_read, denied, queued, delete_denied) = drive(&mut handle, &mut session, &registry, worker);
    match denied {
        Err(Error::Command { source, .. }) => assert_eq!(source.code, "assistant_denied", "{}", source.message),
        other => panic!("expected a command error, found {other:?}"),
    }
    assert!(delete_denied.is_some(), "a delete still needs allow_delete even with the grant");
    assert_eq!(queued.get("queued").and_then(Value::as_bool), Some(true));
    let grant_reason = delete_denied.and_then(|error| match error {
        Error::Command { source, .. } => Some(source.code),
        _ => None,
    });
    assert_eq!(grant_reason.as_deref(), Some("assistant_denied"));
}

/// A wrong `source` value is a protocol refusal, read from a raw socket with
/// the token a client would present.
#[test]
fn a_bad_source_is_a_protocol_refusal() {
    let (_keep, _session, _registry, handle) = spun_up("assistant-control-bad-source");
    let mut stream = TcpStream::connect(("127.0.0.1", handle.port())).expect("connect");
    let line =
        format!("{{\"id\":9,\"token\":\"{}\",\"method\":\"run\",\"params\":{{\"id\":\"note.read\",\"json\":{{}},\"source\":\"spooky\"}}}}\n", handle.token());
    stream.write_all(line.as_bytes()).expect("write the request");
    let reply = read_reply(&mut stream).expect("a reply");
    let reply: Value = serde_json::from_slice(&reply).expect("parse the reply");
    assert_eq!(reply["error"]["code"], json!("invalid_params"), "{reply}");
    assert!(reply["error"]["message"].as_str().is_some_and(|message| message.contains("assistant")), "{reply}");
    assert_eq!(reply["id"], json!(9u64), "the id is echoed");
}

/// One reply line (up to the newline) from a raw socket, `None` at EOF.
fn read_reply(stream: &mut TcpStream) -> Option<Vec<u8>> {
    use std::io::Read;
    let mut gathered = Vec::new();
    let mut chunk = [0u8; 512];
    loop {
        if let Some(pos) = gathered.iter().position(|&b| b == b'\n') {
            gathered.truncate(pos);
            return Some(gathered);
        }
        let read = stream.read(&mut chunk).ok()?;
        if read == 0 {
            return gathered.is_empty().then_some(gathered);
        }
        gathered.extend_from_slice(&chunk[..read]);
    }
}