brokerd: a wait of its own for ask_memory, 170 s

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

crates/brokerd/src/memory_port.rs

@@ -1,6 +1,7 @@
-//! Where decided `recall` and `memory` calls go: to `memoryd` over its `brokerd` socket, with the
-//! session's state, instead of a container (M5a spec, section 8). `Routed` sends every other call
-//! to the configured runtime.
+//! Where decided `recall`, `memory` and `ask_memory` calls go: to `memoryd` over its `brokerd`
+//! socket, with the session's state, instead of a container (M5a spec, section 8). `Routed` sends
+//! every other call to the configured runtime. `ask_memory` has a wait of its own (M5c spec,
+//! section 9).
 
 use std::os::unix::net::UnixStream;
 use std::path::PathBuf;
@@ -11,6 +12,7 @@ use bxcore::{
     write_frame,
 };
 
+use crate::args::ToolName;
 use crate::policy::Label;
 use crate::runner::{RunError, RunOutput, RunSpec, Runtime};
 
@@ -18,27 +20,52 @@ use crate::runner::{RunError, RunOutput, RunSpec, Runtime};
 pub const UNREACHABLE: &str =
     "memory is unavailable: memoryd did not answer (see docs/runbook.md#memoryd-unreachable)";
 
-/// How long `brokerd` waits for `memoryd` to answer one call.
+/// How long `brokerd` waits for `memoryd` to answer one `recall` or `memory` call.
 pub const MEMORY_WAIT: Duration = Duration::from_secs(30);
+/// How long it waits for one `ask_memory` call: `memoryd` asks Honcho's chat about each
+/// workspace, both at once, with 150 s each (M5c spec, section 9).
+pub const ASK_MEMORY_WAIT: Duration = Duration::from_secs(170);
 
 /// A client of `memoryd`'s `brokerd` socket. One connection per call.
 pub struct MemoryPort {
     socket: PathBuf,
     wait: Duration,
+    ask_wait: Duration,
 }
 
 impl MemoryPort {
+    /// A port that waits `wait` for `recall` and `memory`, and `ASK_MEMORY_WAIT` for `ask_memory`.
     pub fn new(socket: PathBuf, wait: Duration) -> MemoryPort {
-        MemoryPort { socket, wait }
+        MemoryPort {
+            socket,
+            wait,
+            ask_wait: ASK_MEMORY_WAIT,
+        }
+    }
+
+    /// The same port, waiting `wait` for `ask_memory`. The tests use a short one.
+    pub fn with_ask_wait(self, wait: Duration) -> MemoryPort {
+        MemoryPort {
+            ask_wait: wait,
+            ..self
+        }
+    }
+
+    /// How long one call of `tool` waits for `memoryd`, to write the call and to read the answer.
+    pub fn wait_for(&self, tool: ToolName) -> Duration {
+        match tool {
+            ToolName::AskMemory => self.ask_wait,
+            _ => self.wait,
+        }
     }
 
     fn ask(&self, spec: &RunSpec) -> Option<ToolResponse> {
         // `None` at the first thing that fails. 1. A MemoryCall: request with the spec's session,
         // call, `spec.tool().as_str()`, and `spec.arguments().canonical_json()`; taint and
         // untrusted from `spec.state()`. 2. `UnixStream::connect(&self.socket)`, then read and
-        // write timeouts of `self.wait`. 3. `write_frame` an Envelope { v: PROTOCOL_VERSION, id:
-        // spec.call().0, final: true, msg: Message::MemoryCall(call) }. 4. `read_frame`; a
-        // Message::ToolResponse is Some(it), any other message None.
+        // write timeouts of `self.wait_for(spec.tool())`. 3. `write_frame` an Envelope { v:
+        // PROTOCOL_VERSION, id: spec.call().0, final: true, msg: Message::MemoryCall(call) }. 4.
+        // `read_frame`; a Message::ToolResponse is Some(it), any other message None.
         let call = MemoryCall {
             request: ToolRequest {
                 session: spec.session().clone(),
@@ -49,9 +76,10 @@ impl MemoryPort {
             taint: spec.state().taint,
             untrusted: spec.state().untrusted,
         };
+        let wait = self.wait_for(spec.tool());
         let mut stream = UnixStream::connect(&self.socket).ok()?;
-        stream.set_read_timeout(Some(self.wait)).ok()?;
-        stream.set_write_timeout(Some(self.wait)).ok()?;
+        stream.set_read_timeout(Some(wait)).ok()?;
+        stream.set_write_timeout(Some(wait)).ok()?;
         write_frame(
             &mut stream,
             &Envelope {

crates/brokerd/tests/memory_wait.rs

@@ -0,0 +1,131 @@
+//! How long `brokerd` waits for `memoryd`, per tool: 30 s for `recall` and `memory`, 170 s for
+//! `ask_memory`, whose two Honcho chats have 150 s each (M5c spec, sections 9 and 11). The waits
+//! are shortened here with `with_ask_wait` and `new`'s own wait, so no test waits long. Do not
+//! edit.
+
+#[path = "support/build.rs"]
+mod build;
+#[path = "support/tmp.rs"]
+mod tmp;
+
+use std::os::unix::net::UnixListener;
+use std::path::Path;
+use std::time::Duration;
+
+use brokerd::args::ToolName;
+use brokerd::memory_port::{ASK_MEMORY_WAIT, MEMORY_WAIT, MemoryPort, UNREACHABLE};
+use brokerd::policy::decide;
+use brokerd::runner::run;
+use build::{allowed, now, private, request, set};
+use bxcore::{DataClass, Envelope, Message, ToolRequest, ToolResponse, read_frame, write_frame};
+use tmp::TempDir;
+
+/// How long the slow fake `memoryd` takes to answer.
+const SLOW: Duration = Duration::from_millis(1_500);
+/// A wait shorter than `SLOW`, and one far longer.
+const SHORT: Duration = Duration::from_millis(300);
+const LONG: Duration = Duration::from_secs(20);
+
+fn found() -> ToolResponse {
+    ToolResponse::Result {
+        content: "own: memory · The owner prefers UTC.".to_string(),
+        class: DataClass::Private,
+        untrusted: false,
+        truncated: false,
+    }
+}
+
+/// A fake `memoryd` that answers one call after `SLOW`. A failed write is ignored: the port has
+/// stopped waiting, which is what some tests expect.
+fn slow_memoryd(socket: &Path) {
+    let listener = UnixListener::bind(socket).unwrap();
+    std::thread::spawn(move || {
+        let Ok((mut stream, _)) = listener.accept() else {
+            return;
+        };
+        let Ok(env) = read_frame(&mut stream) else {
+            return;
+        };
+        std::thread::sleep(SLOW);
+        let reply = Envelope {
+            v: 1,
+            id: env.id,
+            r#final: true,
+            msg: Message::ToolResponse(found()),
+        };
+        let _ = write_frame(&mut stream, &reply);
+    });
+}
+
+/// One allowed call of `req` through `port`, to a slow fake `memoryd` of its own.
+fn ask(
+    dir: &TempDir,
+    name: &str,
+    port: impl Fn(&Path) -> MemoryPort,
+    req: ToolRequest,
+) -> ToolResponse {
+    let socket = dir.path().join(name);
+    slow_memoryd(&socket);
+    let decision = allowed(decide(req, &set(vec![]), private(), now()));
+    run(decision, &port(&socket))
+}
+
+fn unreachable() -> ToolResponse {
+    ToolResponse::Failed {
+        message: UNREACHABLE.to_string(),
+    }
+}
+
+const QUESTION: &str = r#"{"question":"What does the owner prefer?"}"#;
+const QUERY: &str = r#"{"query":"owner"}"#;
+
+#[test]
+fn the_waits_and_which_tool_gets_which() {
+    assert_eq!(MEMORY_WAIT, Duration::from_secs(30));
+    assert_eq!(ASK_MEMORY_WAIT, Duration::from_secs(170));
+    let port = MemoryPort::new("memory.sock".into(), MEMORY_WAIT);
+    assert_eq!(port.wait_for(ToolName::Recall), MEMORY_WAIT);
+    assert_eq!(port.wait_for(ToolName::Memory), MEMORY_WAIT);
+    assert_eq!(port.wait_for(ToolName::AskMemory), ASK_MEMORY_WAIT);
+    let port = MemoryPort::new("memory.sock".into(), SHORT);
+    assert_eq!(port.wait_for(ToolName::Recall), SHORT);
+    assert_eq!(port.wait_for(ToolName::Memory), SHORT);
+    assert_eq!(port.wait_for(ToolName::AskMemory), ASK_MEMORY_WAIT);
+    let port = MemoryPort::new("memory.sock".into(), MEMORY_WAIT).with_ask_wait(LONG);
+    assert_eq!(port.wait_for(ToolName::Recall), MEMORY_WAIT);
+    assert_eq!(port.wait_for(ToolName::Memory), MEMORY_WAIT);
+    assert_eq!(port.wait_for(ToolName::AskMemory), LONG);
+}
+
+/// A slow answer reaches `ask_memory`, whose wait is long, and not `recall` or `memory`, whose
+/// wait is short: the wait is chosen by the call's tool.
+#[test]
+fn a_slow_answer_reaches_ask_memory_only_when_its_wait_is_longer() {
+    let dir = TempDir::new("wait-ask");
+    let port = |socket: &Path| MemoryPort::new(socket.to_path_buf(), SHORT).with_ask_wait(LONG);
+    assert_eq!(
+        ask(&dir, "a.sock", port, request("ask_memory", QUESTION)),
+        found()
+    );
+    assert_eq!(
+        ask(&dir, "r.sock", port, request("recall", QUERY)),
+        unreachable()
+    );
+    let note = r#"{"action":"note","text":"x"}"#;
+    assert_eq!(
+        ask(&dir, "m.sock", port, request("memory", note)),
+        unreachable()
+    );
+}
+
+/// And the other way round: a short wait for `ask_memory` gives up, while `recall` waits.
+#[test]
+fn a_slow_answer_reaches_recall_only_when_its_wait_is_longer() {
+    let dir = TempDir::new("wait-recall");
+    let port = |socket: &Path| MemoryPort::new(socket.to_path_buf(), LONG).with_ask_wait(SHORT);
+    assert_eq!(ask(&dir, "r.sock", port, request("recall", QUERY)), found());
+    assert_eq!(
+        ask(&dir, "a.sock", port, request("ask_memory", QUESTION)),
+        unreachable()
+    );
+}