brokerd: ask_memory, a built-in memory tool that memoryd runs
Co-Authored-By: ornith-1.5-35b-a3b <agent@wntrmute.dev>
This commit is contained in:
ornith-1.5-35b-a3b
@@ -0,0 +1,340 @@
|
||||
//! `ask_memory` in `brokerd`: its name and arguments, the built-in grant that allows it in every
|
||||
//! session, grant files that name it refused, its calls routed to `memoryd` and never to a
|
||||
//! container, and the labels of `memoryd`'s answer raising the session (M5c spec, section 9).
|
||||
//! Do not edit.
|
||||
|
||||
#[path = "support/build.rs"]
|
||||
mod build;
|
||||
#[path = "support/client.rs"]
|
||||
mod client;
|
||||
#[path = "support/fake_podman.rs"]
|
||||
mod fake_podman;
|
||||
#[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::os::unix::net::UnixListener;
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, mpsc};
|
||||
use std::time::Duration;
|
||||
|
||||
use brokerd::args::{ArgsError, ToolArgs, ToolName, parse};
|
||||
use brokerd::broker::Broker;
|
||||
use brokerd::builtin::{ASK_MEMORY, BUILT_IN};
|
||||
use brokerd::container::{IN_MEMORYD, Podman};
|
||||
use brokerd::grants::load;
|
||||
use brokerd::memory_port::{MemoryPort, Routed};
|
||||
use brokerd::policy::{Label, SessionState, decide};
|
||||
use brokerd::runner::run;
|
||||
use build::{allowed, now, private, secret, set};
|
||||
use bxcore::{
|
||||
AskMemoryArgs, DataClass, Envelope, MemoryCall, Message, SessionId, ToolResponse, read_frame,
|
||||
sha256, write_frame,
|
||||
};
|
||||
use fake_podman::{Fake, Lines, serial};
|
||||
use rig::Rig;
|
||||
use runtime::{Recording, Shared};
|
||||
use tmp::TempDir;
|
||||
|
||||
/// A fake `memoryd`: answers one call with `answer` and hands back what it was sent. A failed
|
||||
/// write is ignored: the test has its answer or has already failed.
|
||||
fn fake_memoryd(socket: &Path, answer: ToolResponse) -> mpsc::Receiver<MemoryCall> {
|
||||
let listener = UnixListener::bind(socket).unwrap();
|
||||
let (tx, rx) = mpsc::channel();
|
||||
std::thread::spawn(move || {
|
||||
let Ok((mut stream, _)) = listener.accept() else {
|
||||
return;
|
||||
};
|
||||
let Ok(env) = read_frame(&mut stream) else {
|
||||
return;
|
||||
};
|
||||
let Message::MemoryCall(call) = env.msg else {
|
||||
return;
|
||||
};
|
||||
let _ = tx.send(call);
|
||||
let reply = Envelope {
|
||||
v: 1,
|
||||
id: env.id,
|
||||
r#final: true,
|
||||
msg: Message::ToolResponse(answer),
|
||||
};
|
||||
let _ = write_frame(&mut stream, &reply);
|
||||
});
|
||||
rx
|
||||
}
|
||||
|
||||
fn answer(class: DataClass, untrusted: bool) -> ToolResponse {
|
||||
ToolResponse::Result {
|
||||
content: "own: memory · The owner prefers UTC.".to_string(),
|
||||
class,
|
||||
untrusted,
|
||||
truncated: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn routed(socket: &Path, other: &Arc<Recording>) -> Routed {
|
||||
Routed {
|
||||
memory: MemoryPort::new(socket.to_path_buf(), Duration::from_secs(5)),
|
||||
other: Box::new(Shared(Arc::clone(other))),
|
||||
}
|
||||
}
|
||||
|
||||
const QUESTION: &str = r#"{"question":"What does the owner prefer?"}"#;
|
||||
|
||||
#[test]
|
||||
fn the_name_and_the_arguments() {
|
||||
assert_eq!(ToolName::parse("ask_memory"), Some(ToolName::AskMemory));
|
||||
assert_eq!(ToolName::AskMemory.as_str(), "ask_memory");
|
||||
assert!(ToolName::AskMemory.is_memory());
|
||||
assert!(
|
||||
!ToolName::ALL.contains(&ToolName::AskMemory),
|
||||
"ALL is the container tools"
|
||||
);
|
||||
let args = parse(
|
||||
ToolName::AskMemory,
|
||||
r#"{ "workspaces": ["fleet", "boxmaker1"], "question": "Who wrote it?" }"#,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(args.tool(), ToolName::AskMemory);
|
||||
assert_eq!(
|
||||
args,
|
||||
ToolArgs::AskMemory(AskMemoryArgs {
|
||||
question: "Who wrote it?".to_string(),
|
||||
workspaces: vec!["fleet".to_string(), "boxmaker1".to_string()],
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
args.canonical_json(),
|
||||
r#"{"question":"Who wrote it?","workspaces":["fleet","boxmaker1"]}"#
|
||||
);
|
||||
let default = parse(ToolName::AskMemory, QUESTION).unwrap();
|
||||
assert_eq!(
|
||||
default.canonical_json(),
|
||||
r#"{"question":"What does the owner prefer?","workspaces":[]}"#,
|
||||
"every field is written, as recall's k is"
|
||||
);
|
||||
}
|
||||
|
||||
/// The refusals are `bxcore`'s sentences, shown to the model after `invalid argument shape: `.
|
||||
#[test]
|
||||
fn bad_arguments_are_refused_with_the_checks_sentences() {
|
||||
let long = format!(r#"{{"question":"{}"}}"#, "é".repeat(2_001));
|
||||
for (json, sentence) in [
|
||||
(
|
||||
r#"{"question":" "}"#,
|
||||
"question is empty; write the question to ask",
|
||||
),
|
||||
(
|
||||
long.as_str(),
|
||||
"question has 2001 characters, over the limit of 2000; ask a shorter question",
|
||||
),
|
||||
(
|
||||
r#"{"question":"q","workspaces":["own","untrusted","fleet"]}"#,
|
||||
"workspaces names 3, over the limit of 2; ask about the others in another call",
|
||||
),
|
||||
(
|
||||
r#"{"question":"q","workspaces":["fleet","fleet"]}"#,
|
||||
"workspaces names fleet twice; name each workspace once",
|
||||
),
|
||||
(
|
||||
r#"{"question":"q","workspaces":[""]}"#,
|
||||
"a workspace name is empty; name a workspace, or leave workspaces out to ask own",
|
||||
),
|
||||
(
|
||||
r#"{"question":"q","workspaces":["a.b"]}"#,
|
||||
"a workspace name has 1 to 512 letters, digits, _ or -; use own, untrusted, fleet or \
|
||||
another agent's name",
|
||||
),
|
||||
] {
|
||||
let error = parse(ToolName::AskMemory, json).unwrap_err();
|
||||
assert_eq!(error, ArgsError::Shape(sentence.to_string()), "{json}");
|
||||
assert_eq!(
|
||||
error.detail(),
|
||||
format!("invalid argument shape: {sentence}")
|
||||
);
|
||||
}
|
||||
for json in [
|
||||
r#"{"question":"q","reasoning":"low"}"#,
|
||||
r#"{"workspaces":["own"]}"#,
|
||||
"not json",
|
||||
] {
|
||||
let error = parse(ToolName::AskMemory, json).unwrap_err();
|
||||
assert!(matches!(error, ArgsError::Shape(_)), "{json}: {error:?}");
|
||||
}
|
||||
}
|
||||
|
||||
/// The built-in grant allows `ask_memory` in every session, `secret` and untrusted ones too: a
|
||||
/// question is not written to Honcho (M5c spec, section 8). Its results are labelled `private`
|
||||
/// and trusted; `memoryd`'s own labels can only make them stricter.
|
||||
#[test]
|
||||
fn the_built_in_grant_allows_it_in_every_session() {
|
||||
assert_eq!(
|
||||
ASK_MEMORY,
|
||||
"tool = \"ask_memory\"\nmode = \"auto\"\nmax_taint = \"secret\"\nresult_class = \
|
||||
\"private\"\nuntrusted = false\n"
|
||||
);
|
||||
let grants = set(vec![]);
|
||||
let untrusted = SessionState {
|
||||
taint: DataClass::Secret,
|
||||
untrusted: true,
|
||||
};
|
||||
for state in [private(), secret(), untrusted] {
|
||||
let d = allowed(decide(
|
||||
build::request("ask_memory", QUESTION),
|
||||
&grants,
|
||||
state,
|
||||
now(),
|
||||
));
|
||||
assert_eq!(d.grant(), "builtin-ask-memory");
|
||||
assert_eq!(d.grant_sha256(), sha256(ASK_MEMORY.as_bytes()).unwrap());
|
||||
assert_eq!(
|
||||
d.label(),
|
||||
Label {
|
||||
class: DataClass::Private,
|
||||
untrusted: false
|
||||
}
|
||||
);
|
||||
}
|
||||
let ids: Vec<&str> = grants.matchable().map(|g| g.id.as_str()).collect();
|
||||
assert_eq!(
|
||||
ids,
|
||||
["builtin-memory", "builtin-recall", "builtin-ask-memory"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_grant_file_cannot_name_it() {
|
||||
assert_eq!(
|
||||
BUILT_IN,
|
||||
"the policy of recall, memory and ask_memory is built in; remove this grant"
|
||||
);
|
||||
let dir = TempDir::new("ask-grants");
|
||||
std::fs::write(
|
||||
dir.path().join("mine.toml"),
|
||||
"tool = \"ask_memory\"\nmode = \"deny\"\nmax_taint = \"secret\"\n",
|
||||
)
|
||||
.unwrap();
|
||||
let problems = load(dir.path()).unwrap_err();
|
||||
let got: Vec<(String, String)> = problems.into_iter().map(|p| (p.file, p.problem)).collect();
|
||||
assert_eq!(got, [("mine.toml".to_string(), BUILT_IN.to_string())]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn it_goes_to_memoryd_with_its_canonical_arguments() {
|
||||
let dir = TempDir::new("ask-route");
|
||||
let socket = dir.path().join("memory.sock");
|
||||
let seen = fake_memoryd(&socket, answer(DataClass::Private, false));
|
||||
let other = Recording::answering("container");
|
||||
let asked = r#"{"workspaces":["own","fleet"],"question":"Who wrote it?"}"#;
|
||||
let got = run(
|
||||
allowed(decide(
|
||||
build::request("ask_memory", asked),
|
||||
&set(vec![]),
|
||||
secret(),
|
||||
now(),
|
||||
)),
|
||||
&routed(&socket, &other),
|
||||
);
|
||||
assert_eq!(got, answer(DataClass::Private, false));
|
||||
let call = seen.recv_timeout(Duration::from_secs(5)).unwrap();
|
||||
assert_eq!(call.request.tool, "ask_memory");
|
||||
assert_eq!(
|
||||
call.request.arguments,
|
||||
r#"{"question":"Who wrote it?","workspaces":["own","fleet"]}"#
|
||||
);
|
||||
assert_eq!((call.taint, call.untrusted), (DataClass::Secret, false));
|
||||
assert_eq!(other.count(), 0, "nothing reached the other runtime");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_container_runtime_never_runs_it() {
|
||||
let _s = serial();
|
||||
let fake = Fake::new("ask", "exit 0");
|
||||
let log = Lines::default();
|
||||
let podman = Podman::new(fake.runner(""), fake.dir.join("egress"), log.sink());
|
||||
let got = run(
|
||||
allowed(decide(
|
||||
build::request("ask_memory", QUESTION),
|
||||
&set(vec![]),
|
||||
private(),
|
||||
now(),
|
||||
)),
|
||||
&podman,
|
||||
);
|
||||
assert_eq!(
|
||||
got,
|
||||
ToolResponse::Failed {
|
||||
message: IN_MEMORYD.to_string()
|
||||
}
|
||||
);
|
||||
assert!(fake.calls().is_empty(), "podman was never called");
|
||||
}
|
||||
|
||||
/// One `ask_memory` call through `broker::handle`, whose runtime sends it to a fake `memoryd`
|
||||
/// answering `reply`; the session's state afterwards.
|
||||
fn state_after(tag: &str, reply: ToolResponse) -> SessionState {
|
||||
let rig = Rig::new(tag);
|
||||
let socket = rig.dir.path().join("memory.sock");
|
||||
let _seen = fake_memoryd(&socket, reply);
|
||||
let other = Recording::answering("unused");
|
||||
let broker = Arc::new(Broker::new(
|
||||
rig.cfg.clone(),
|
||||
rig.ledger(),
|
||||
Box::new(routed(&socket, &other)),
|
||||
rig.lines.sink(),
|
||||
));
|
||||
let frames = client::call(&broker, rig::request("s1", 1, "ask_memory", QUESTION));
|
||||
assert!(
|
||||
matches!(
|
||||
frames.as_slice(),
|
||||
[Envelope {
|
||||
msg: Message::ToolResponse(ToolResponse::Result { .. }),
|
||||
..
|
||||
}]
|
||||
),
|
||||
"{frames:?}"
|
||||
);
|
||||
rig.state().read(&SessionId::new("s1").unwrap()).unwrap()
|
||||
}
|
||||
|
||||
/// T58: an answer from a workspace that is not trusted makes a clean session untrusted, as a
|
||||
/// `memory-tainted` recall hit does.
|
||||
#[test]
|
||||
fn an_untrusted_answer_makes_the_session_untrusted() {
|
||||
assert_eq!(
|
||||
state_after("ask-untrusted", answer(DataClass::Private, true)),
|
||||
SessionState {
|
||||
taint: DataClass::Private,
|
||||
untrusted: true
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/// T58: an answer labelled `secret` raises the session to `secret`.
|
||||
#[test]
|
||||
fn a_secret_answer_raises_the_session_to_secret() {
|
||||
assert_eq!(
|
||||
state_after("ask-secret", answer(DataClass::Secret, false)),
|
||||
SessionState {
|
||||
taint: DataClass::Secret,
|
||||
untrusted: false
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/// The grant's label is the floor: a `public` answer still leaves the session `private`.
|
||||
#[test]
|
||||
fn a_public_answer_leaves_the_session_private() {
|
||||
assert_eq!(
|
||||
state_after("ask-public", answer(DataClass::Public, false)),
|
||||
SessionState {
|
||||
taint: DataClass::Private,
|
||||
untrusted: false
|
||||
}
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user