brokerd: consult posts the allowlisted body

Co-Authored-By: ornith-1.5-35b-a3b <agent@wntrmute.dev>

crates/brokerd/tests/consult_allow.rs

@@ -0,0 +1,185 @@
+//! 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:?}"),
+    }
+}