boxmaker/docs/plans/M6b/files/crates/brokerd/tests/consult_deny.rs

//! Consult denials: not configured, taint too high, unknown keys. Each records the body. Do not edit.

#[path = "support/client.rs"]
mod client;
#[path = "support/rig.rs"]
mod rig;
#[path = "support/runtime.rs"]
mod runtime;
#[path = "support/sink.rs"]
mod sink;
#[path = "support/tmp.rs"]
mod tmp;

use std::sync::Arc;

use brokerd::consult::{self, Consult};
use brokerd::policy::{Label, SessionState};
use bxcore::{
    AuditEvent, DataClass, DecisionRecord, DenyReason, SessionId, ToolRequest, ToolResponse,
};
use client::{Serve, call};
use rig::{Rig, request};
use runtime::Recording;

fn consult_req(session: &str, call: u64, arguments: &str) -> ToolRequest {
    request(session, call, "consult", arguments)
}

fn consult_cfg() -> Consult {
    Consult {
        url: "http://127.0.0.1:9/v1/chat/completions".to_string(),
        model: "test-model".to_string(),
        temperature_milli: 600,
        timeout_ms: 1_000,
        secret: "consult".to_string(),
    }
}

fn grant_consult(rig: &Rig, max_taint: &str) {
    rig.grant(
        "consult",
        &format!("tool = \"consult\"\nmode = \"auto\"\nmax_taint = \"{max_taint}\"\n"),
    );
}

fn decision_body(events: &[AuditEvent]) -> Option<&str> {
    match events.first() {
        Some(AuditEvent::Decision { body, .. }) => body.as_deref(),
        _ => None,
    }
}

#[test]
fn no_consult_table_is_denied_with_the_runbook_and_records_the_body() {
    let rig = Rig::new("consult-noconfig");
    let rt = Recording::answering("no");
    let broker = rig.broker(&rt);
    let question = r#"{"question":"what colour is the sky?"}"#;
    let frames = call(&broker, consult_req("s1", 1, question));
    assert_eq!(rt.count(), 0);
    match &frames[0].msg {
        bxcore::Message::ToolResponse(ToolResponse::Denied { reason, detail }) => {
            assert_eq!(*reason, DenyReason::InvalidArguments);
            let detail = detail.as_deref().unwrap_or("");
            assert!(detail.contains("consult-not-configured"), "{detail}");
            assert!(detail.contains("docs/runbook.md"), "{detail}");
        }
        other => panic!("{other:?}"),
    }
    let events = rig.events();
    let body = decision_body(&events).expect("body");
    let want = consult::wire_body(None, question);
    assert_eq!(body, want);
    assert!(body.contains("what colour is the sky?"), "{body}");
    assert!(!body.contains("tools"), "{body}");
}

#[test]
fn a_secret_session_is_taint_too_high_and_records_the_body() {
    let mut rig = Rig::new("consult-taint");
    rig.cfg.consult = Some(consult_cfg());
    grant_consult(&rig, "private");
    let sid = SessionId::new("s1").unwrap();
    rig.state()
        .raise(
            &sid,
            SessionState {
                taint: DataClass::Secret,
                untrusted: false,
            },
            Label {
                class: DataClass::Secret,
                untrusted: false,
            },
        )
        .unwrap();
    let rt = Recording::answering("no");
    let broker = Arc::new(brokerd::broker::Broker::new(
        rig.cfg.clone(),
        rig.ledger(),
        Box::new(runtime::Shared(Arc::clone(&rt))),
        rig.lines.sink(),
    ));
    let question = r#"{"question":"summarise the notes"}"#;
    let frames = call(&broker, consult_req("s1", 2, question));
    assert_eq!(rt.count(), 0);
    match &frames[0].msg {
        bxcore::Message::ToolResponse(ToolResponse::Denied { reason, detail }) => {
            assert_eq!(*reason, DenyReason::TaintTooHigh);
            assert_eq!(*detail, None);
        }
        other => panic!("{other:?}"),
    }
    match rig.events().as_slice() {
        [
            AuditEvent::Decision {
                outcome,
                body: Some(body),
                taint,
                ..
            },
        ] => {
            assert_eq!(
                *outcome,
                DecisionRecord::Denied {
                    reason: DenyReason::TaintTooHigh
                }
            );
            assert_eq!(*taint, DataClass::Secret);
            let want = consult::wire_body(rig.cfg.consult.as_ref(), question);
            assert_eq!(body, &want);
            assert!(body.contains("\"model\":\"test-model\""), "{body}");
            assert_eq!(
                serde_json::from_str::<serde_json::Value>(body).unwrap()["messages"]
                    .as_array()
                    .unwrap()
                    .len(),
                1
            );
        }
        other => panic!("{other:?}"),
    }
}

#[test]
fn an_unknown_field_is_invalid_arguments() {
    let mut rig = Rig::new("consult-unknown");
    rig.cfg.consult = Some(consult_cfg());
    grant_consult(&rig, "secret");
    let rt = Recording::answering("no");
    let broker = Arc::new(brokerd::broker::Broker::new(
        rig.cfg.clone(),
        rig.ledger(),
        Box::new(runtime::Shared(Arc::clone(&rt))),
        rig.lines.sink(),
    ));
    let frames = call(
        &broker,
        consult_req("s1", 3, r#"{"question":"hi","history":["old"]}"#),
    );
    assert_eq!(rt.count(), 0);
    match &frames[0].msg {
        bxcore::Message::ToolResponse(ToolResponse::Denied { reason, .. }) => {
            assert_eq!(*reason, DenyReason::InvalidArguments);
        }
        other => panic!("{other:?}"),
    }
}