boxmaker/crates/brokerd/tests/consult_allow.rs

//! An allowed `consult`: fake HTTP 200, audit body is the allowlisted JSON, the tool result is
//! the answer text. 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, Mutex};
use std::time::Duration;

use brokerd::consult::{self, Consult, ConsultRuntime, ConsultTransport};
use bxcore::secrets::Secret;
use bxcore::{AuditEvent, DataClass, DecisionRecord, ResultStatus, ToolRequest, ToolResponse};
use client::call;
use rig::{Rig, request};

struct SharedFake {
    reply: Vec<u8>,
    seen: Arc<Mutex<Option<Vec<u8>>>>,
}

impl ConsultTransport for SharedFake {
    fn post(
        &self,
        _url: &str,
        _bearer: &str,
        body: &[u8],
        _timeout: Duration,
    ) -> Result<Vec<u8>, String> {
        *self.seen.lock().unwrap() = Some(body.to_vec());
        Ok(self.reply.clone())
    }
}

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) {
    rig.grant(
        "consult",
        "tool = \"consult\"\nmode = \"auto\"\nmax_taint = \"private\"\n",
    );
}

fn question_req(question: &str) -> ToolRequest {
    request(
        "s1",
        1,
        "consult",
        &format!(
            r#"{{"question":{}}}"#,
            serde_json::to_string(question).unwrap()
        ),
    )
}

fn broker_with_fake(
    rig: &Rig,
    cfg: Consult,
    reply: Vec<u8>,
    seen: Arc<Mutex<Option<Vec<u8>>>>,
) -> Arc<brokerd::broker::Broker> {
    let runtime = ConsultRuntime::new(
        cfg,
        Secret::new("token".to_string()),
        SharedFake { reply, seen },
    );
    Arc::new(brokerd::broker::Broker::new(
        rig.cfg.clone(),
        rig.ledger(),
        Box::new(runtime),
        rig.lines.sink(),
    ))
}

#[test]
fn an_allowed_consult_posts_only_the_question_and_returns_the_answer() {
    let mut rig = Rig::new("consult-allow");
    let cfg = consult_cfg();
    rig.cfg.consult = Some(cfg.clone());
    grant_consult(&rig);
    let tag = concat!("<", "think", ">");
    let raw_answer = format!("blue {tag} sky");
    let reply = serde_json::json!({
        "id": "ignored",
        "choices": [{"index": 0, "message": {"role": "assistant", "content": raw_answer}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": 1}
    })
    .to_string()
    .into_bytes();
    let recorded = Arc::new(Mutex::new(None));
    let broker = broker_with_fake(&rig, cfg.clone(), reply, Arc::clone(&recorded));
    let question = "what colour is the sky?";
    let frames = call(&broker, question_req(question));
    let want_body = consult::wire_body(Some(&cfg), &format!(r#"{{"question":"{question}"}}"#));
    match &frames[0].msg {
        bxcore::Message::ToolResponse(ToolResponse::Result {
            content,
            class,
            untrusted,
            truncated,
        }) => {
            assert!(content.contains("blue"), "{content}");
            assert!(content.contains("sky"), "{content}");
            assert!(!content.contains(tag), "{content}");
            assert!(content.contains("&lt;"), "{content}");
            assert_eq!(*class, DataClass::Private);
            assert!(*untrusted);
            assert!(!*truncated);
        }
        other => panic!("{other:?}"),
    }
    let posted = recorded.lock().unwrap().clone().expect("posted");
    assert_eq!(posted, want_body.as_bytes());
    match rig.events().as_slice() {
        [
            AuditEvent::Decision {
                outcome,
                body: Some(body),
                tool,
                ..
            },
            AuditEvent::Result {
                status: ResultStatus::Result,
                ..
            },
        ] => {
            assert_eq!(*outcome, DecisionRecord::Allowed {});
            assert_eq!(tool, "consult");
            assert_eq!(body, &want_body);
            let v: serde_json::Value = serde_json::from_str(body).unwrap();
            assert_eq!(v["model"], "test-model");
            assert_eq!(v["temperature"], 0.6);
            assert_eq!(v["messages"].as_array().unwrap().len(), 1);
            assert_eq!(v["messages"][0]["role"], "user");
            assert_eq!(v["messages"][0]["content"], question);
            assert!(v.get("tools").is_none());
        }
        other => panic!("{other:?}"),
    }
}

/// The body is only `question`, even when the caller had earlier user text in its own session.
#[test]
fn the_consult_body_is_only_the_question() {
    let mut rig = Rig::new("consult-nohist");
    let cfg = consult_cfg();
    rig.cfg.consult = Some(cfg.clone());
    grant_consult(&rig);
    let reply = serde_json::json!({
        "choices": [{"message": {"content": "ok"}}]
    })
    .to_string()
    .into_bytes();
    let recorded = Arc::new(Mutex::new(None));
    let broker = broker_with_fake(&rig, cfg, reply, Arc::clone(&recorded));
    let arguments = r#"{"question":"only this"}"#;
    let _ = call(&broker, request("s1", 4, "consult", arguments));
    let posted = recorded.lock().unwrap().clone().expect("posted");
    let v: serde_json::Value = serde_json::from_slice(&posted).unwrap();
    assert_eq!(v["messages"].as_array().unwrap().len(), 1);
    assert_eq!(v["messages"][0]["content"], "only this");
    match &rig.events()[0] {
        AuditEvent::Decision {
            body: Some(body), ..
        } => {
            assert_eq!(body.as_bytes(), posted.as_slice());
        }
        other => panic!("{other:?}"),
    }
}