boxmaker/crates/brokerd/src/memory_port.rs

//! 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;
use std::time::Duration;

use bxcore::{
    Envelope, MemoryCall, Message, PROTOCOL_VERSION, ToolRequest, ToolResponse, read_frame,
    write_frame,
};

use crate::args::ToolName;
use crate::policy::Label;
use crate::runner::{RunError, RunOutput, RunSpec, Runtime};

/// The sentence when `memoryd` cannot be reached or answers wrongly: fixed, it reaches the model.
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 `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,
            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_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(),
                call: spec.call(),
                tool: spec.tool().as_str().to_string(),
                arguments: spec.arguments().canonical_json(),
            },
            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(wait)).ok()?;
        stream.set_write_timeout(Some(wait)).ok()?;
        write_frame(
            &mut stream,
            &Envelope {
                v: PROTOCOL_VERSION,
                id: spec.call().0,
                r#final: true,
                msg: Message::MemoryCall(call),
            },
        )
        .ok()?;
        let env = read_frame(&mut stream).ok()?;
        match env.msg {
            Message::ToolResponse(response) => Some(response),
            _ => None,
        }
    }
}

impl Runtime for MemoryPort {
    fn run(&self, spec: &RunSpec) -> Result<RunOutput, RunError> {
        self.run_labelled(spec).map(|(output, _)| output)
    }

    fn run_labelled(&self, spec: &RunSpec) -> Result<(RunOutput, Option<Label>), RunError> {
        // `self.ask(spec)`: Some(Result { content, class, untrusted, truncated }) -> Ok((RunOutput
        // { content, truncated }, Some(Label { class, untrusted }))); Some(Failed { message }) ->
        // Err(RunError::Failed(message)); anything else, None included ->
        // Err(RunError::Unavailable(UNREACHABLE.to_string())).
        match self.ask(spec) {
            Some(ToolResponse::Result {
                content,
                class,
                untrusted,
                truncated,
            }) => Ok((
                RunOutput { content, truncated },
                Some(Label { class, untrusted }),
            )),
            Some(ToolResponse::Failed { message }) => Err(RunError::Failed(message)),
            _ => Err(RunError::Unavailable(UNREACHABLE.to_string())),
        }
    }
}

/// The memory tools to `memoryd`, everything else to `other`.
pub struct Routed {
    pub memory: MemoryPort,
    pub other: Box<dyn Runtime>,
}

impl Runtime for Routed {
    fn run(&self, spec: &RunSpec) -> Result<RunOutput, RunError> {
        self.run_labelled(spec).map(|(output, _)| output)
    }

    fn run_labelled(&self, spec: &RunSpec) -> Result<(RunOutput, Option<Label>), RunError> {
        // `spec.tool().is_memory()` -> `self.memory.run_labelled(spec)`; otherwise
        // `self.other.run_labelled(spec)`.
        if spec.tool().is_memory() {
            self.memory.run_labelled(spec)
        } else {
            self.other.run_labelled(spec)
        }
    }
}