boxmaker/crates/loopd/tests/end_to_end.rs

//! End to end in two processes: `loopd`'s turn loop, with a `BrokerPort`, against the real
//! `brokerd` binary on a temporary home, and the fake llama server. Do not edit.
//!
//! It cannot be one process: `loopd` may not depend on `brokerd`, even for tests. The path of
//! the `brokerd` binary comes from `BOXMAKER_BROKERD`; without it the test is ignored, and
//! `make gate` builds the workspace and runs it with the variable set.

mod support;

use std::os::unix::net::UnixStream;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::time::{Duration, Instant};

use loopd::baseline::Baseline;
use loopd::broker_port::BrokerPort;
use loopd::llama::Client;
use loopd::session::Session;
use loopd::tools::Registry;
use loopd::turn::{Runtime, run_turn};
use proto::{
    AuditEvent, AuditRecord, ChainVerifier, DecisionRecord, DenyReason, SessionId, TurnEvent,
};
use support::{FakeServer, Home, Reply};

const CHAT: &str = "/v1/chat/completions";

/// `brokerd serve`, killed when dropped.
struct Brokerd(Child);

impl Drop for Brokerd {
    fn drop(&mut self) {
        let _ = self.0.kill();
        let _ = self.0.wait();
    }
}

fn start_brokerd(home: &Path) -> (Brokerd, PathBuf) {
    let binary = std::env::var_os("BOXMAKER_BROKERD")
        .expect("set BOXMAKER_BROKERD to the brokerd binary; `make gate` does");
    std::fs::create_dir_all(home.join("grants")).unwrap();
    let config = home.join("brokerd.toml");
    let text = format!(
        "[paths]\nhome = \"{h}\"\ngrants = \"{h}/grants\"\n",
        h = home.display()
    );
    std::fs::write(&config, text).unwrap();
    let child = Command::new(binary)
        .args(["serve", "--config"])
        .arg(&config)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .unwrap();
    let brokerd = Brokerd(child);
    let socket = home.join("run/loop-broker/broker.sock");
    let until = Instant::now() + Duration::from_secs(10);
    while UnixStream::connect(&socket).is_err() {
        assert!(
            Instant::now() < until,
            "brokerd never listened on {}",
            socket.display()
        );
        std::thread::sleep(Duration::from_millis(20));
    }
    (brokerd, socket)
}

/// Every record under `audit/`, after checking that the chain verifies.
fn audit(dir: &Path) -> Vec<AuditRecord> {
    let mut names: Vec<String> = std::fs::read_dir(dir)
        .unwrap()
        .map(|e| e.unwrap().file_name().into_string().unwrap())
        .filter(|n| n.ends_with(".jsonl"))
        .collect();
    names.sort();
    let mut verifier = ChainVerifier::new();
    let mut records = Vec::new();
    for name in &names {
        let bytes = std::fs::read(dir.join(name)).unwrap();
        verifier.feed(name, &bytes);
        for line in String::from_utf8(bytes).unwrap().lines() {
            records.push(serde_json::from_str(line).unwrap());
        }
    }
    let report = verifier.finish();
    assert!(report.failure.is_none(), "{:?}", report.failure);
    assert!(report.torn_tail.is_none());
    assert_eq!(report.records, records.len() as u64);
    records
}

#[test]
#[ignore = "needs the brokerd binary in BOXMAKER_BROKERD; `make gate` runs it"]
fn a_call_no_grant_allows_reaches_the_model_as_a_denial_and_is_audited() {
    let home = Home::new();
    let broker_home = home.dir.join("broker-home");
    let (_brokerd, socket) = start_brokerd(&broker_home);

    let server = FakeServer::start();
    // The recorded model calls `read_file` on /etc/hostname, then answers in plain text.
    server.route(
        CHAT,
        vec![Reply::fixture("tool_call"), Reply::fixture("plain")],
    );
    let cfg = home.config(&server.socket);
    let client = Client::new(cfg.clone());
    let port = BrokerPort::new(socket, Duration::from_secs(10));
    let registry = Registry::m2b();
    let baseline = Baseline::assemble(&cfg, &registry).unwrap();
    let id = SessionId::new("e2e").unwrap();
    let mut session = Session::create(&home.dir, id, baseline, cfg.slots.main).unwrap();
    let runtime = Runtime {
        cfg: &cfg,
        client: &client,
        port: &port,
        registry: &registry,
    };
    let mut events = Vec::new();
    let outcome = run_turn(
        &mut session,
        &runtime,
        "what is this host called?",
        &mut |e| events.push(e.clone()),
    );
    assert!(
        outcome.is_ok(),
        "the turn goes on after a denial: {outcome:?}"
    );
    assert!(
        events.contains(&TurnEvent::ToolDenied {
            name: "read_file".to_string(),
            reason: DenyReason::NoGrant,
        }),
        "{events:?}"
    );
    // The model reads the denial in its next request.
    let second = server.requests_to(CHAT)[1].json();
    assert_eq!(
        second["messages"][3]["content"],
        "Denied: no grant allows this call."
    );

    let records = audit(&broker_home.join("audit"));
    assert_eq!(records.len(), 1, "{records:?}");
    match &records[0].event {
        AuditEvent::Decision {
            session,
            tool,
            arguments,
            outcome,
            grant,
            ..
        } => {
            assert_eq!(session.as_str(), "e2e");
            assert_eq!(tool, "read_file");
            assert_eq!(arguments, r#"{"path":"/etc/hostname"}"#);
            assert_eq!(
                *outcome,
                DecisionRecord::Denied {
                    reason: DenyReason::NoGrant
                }
            );
            assert_eq!(*grant, None);
        }
        other => panic!("{other:?}"),
    }
}