boxmaker/crates/loopd/tests/support/broker.rs

//! A fake `brokerd` for the `BrokerPort` tests. Included with `#[path]`. Do not edit.
//!
//! The fake behaves as the real one will: it reads one request frame, answers on the same
//! connection, never half-closes, and closes after the final frame.

#![allow(dead_code)] // each test file uses a different part of this module

use std::os::unix::net::{UnixListener, UnixStream};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use loopd::broker_port::{BrokerPort, UNAVAILABLE};
use loopd::tools::{Pending, ToolPort};
use proto::{
    CallId, DataClass, Envelope, Message, PROTOCOL_VERSION, SessionId, Timestamp, ToolRequest,
    ToolResponse, read_frame, write_frame,
};

static NEXT: AtomicU32 = AtomicU32::new(0);

/// A fresh socket path, removed when the value is dropped (tip T44).
pub struct SocketPath(PathBuf);

impl std::ops::Deref for SocketPath {
    type Target = Path;
    fn deref(&self) -> &Path {
        &self.0
    }
}

impl AsRef<Path> for SocketPath {
    fn as_ref(&self) -> &Path {
        &self.0
    }
}

impl Drop for SocketPath {
    fn drop(&mut self) {
        let _ = std::fs::remove_file(&self.0);
    }
}

pub fn socket_path() -> SocketPath {
    let n = NEXT.fetch_add(1, Ordering::SeqCst);
    SocketPath(std::env::temp_dir().join(format!("loopd-broker-{}-{n}.sock", std::process::id())))
}

/// Accepts one connection, reads the request frame and hands both to `script`. The connection
/// closes when `script` returns.
pub fn broker<F>(script: F) -> (PathBuf, JoinHandle<Envelope>)
where
    F: FnOnce(&mut UnixStream, &Envelope) + Send + 'static,
{
    let path = socket_path();
    let _ = std::fs::remove_file(&path);
    let listener = UnixListener::bind(&path).unwrap();
    let returned = path.to_path_buf();
    // The thread owns the path, so the socket file goes once the exchange is over.
    let handle = thread::spawn(move || {
        let _path = path;
        let (mut stream, _) = listener.accept().unwrap();
        let request = read_frame(&mut stream).unwrap();
        script(&mut stream, &request);
        request
    });
    (returned, handle)
}

pub fn frame(id: u64, r#final: bool, msg: Message) -> Envelope {
    Envelope {
        v: PROTOCOL_VERSION,
        id,
        r#final,
        msg,
    }
}

pub fn send(stream: &mut UnixStream, id: u64, r#final: bool, response: ToolResponse) {
    // The port may already have given up and gone; the fake does not care.
    let _ = write_frame(stream, &frame(id, r#final, Message::ToolResponse(response)));
}

pub fn result(content: &str) -> ToolResponse {
    ToolResponse::Result {
        content: content.to_string(),
        class: DataClass::Secret,
        untrusted: true,
        truncated: false,
    }
}

pub fn request() -> ToolRequest {
    ToolRequest {
        session: SessionId::new("chat-1").unwrap(),
        call: CallId(7),
        tool: "read_file".to_string(),
        arguments: r#"{"path":"/home/kyle/notes/a.md"}"#.to_string(),
    }
}

pub fn in_ms(ms: u64) -> Timestamp {
    Timestamp::from_unix_millis(Timestamp::now().unix_millis() + ms).unwrap()
}

pub struct Call {
    pub response: ToolResponse,
    pub pending: Vec<Pending>,
    pub lines: Vec<String>,
    pub took: Duration,
}

pub fn call(socket: PathBuf, timeout_ms: u64, request: &ToolRequest) -> Call {
    let lines = Arc::new(Mutex::new(Vec::new()));
    let sink = lines.clone();
    let port = BrokerPort::with_log(
        socket,
        Duration::from_millis(timeout_ms),
        Box::new(move |line| sink.lock().unwrap().push(line.to_string())),
    );
    let mut pending = Vec::new();
    let started = Instant::now();
    let response = port.call(request, &mut |p| pending.push(*p));
    let took = started.elapsed();
    let lines = lines.lock().unwrap().clone();
    Call {
        response,
        pending,
        lines,
        took,
    }
}

/// The answer is the fixed failure, and `loopd` printed exactly one line, with the pointer.
pub fn assert_unavailable(call: &Call, why: &str) {
    assert_eq!(
        call.response,
        ToolResponse::Failed {
            message: UNAVAILABLE.to_string()
        },
        "{why}"
    );
    assert_eq!(
        call.lines.len(),
        1,
        "{why}: one line per failed call: {:?}",
        call.lines
    );
    let line = &call.lines[0];
    assert!(
        line.starts_with("loopd: the tool broker is unavailable: "),
        "{why}: {line}"
    );
    assert!(
        line.ends_with("; see docs/runbook.md#broker-unavailable"),
        "{why}: {line}"
    );
    assert!(call.pending.is_empty() || why.contains("pending"), "{why}");
}