boxmaker/crates/loopd/tests/memory.rs

//! `loopd` and memory: the memory block in a new session's baseline and not in a running one's,
//! no session without it, `memoryd` told after a finished turn, and the two core memory tools
//! (M5a spec, section 9). Do not edit.

mod support;

use std::os::unix::net::{UnixListener, UnixStream};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use loopd::channel::{Context, serve};
use loopd::llama::Client;
use loopd::memory::{MemorySource, UNAVAILABLE};
use loopd::tools::Registry;
use proto::{
    Envelope, ErrorCode, Message, PROTOCOL_VERSION, SessionId, Turn, read_frame, write_frame,
};
use support::{FakeServer, Home, Reply, ScriptedPort};

const CHAT: &str = "/v1/chat/completions";

/// A `memoryd` stand-in: the block it serves, and the sessions it was told about.
struct Recorder {
    block: Mutex<Result<String, String>>,
    updated: Mutex<Vec<String>>,
}

/// Lets the test keep its `Arc<Recorder>` while the context owns a `Box<dyn MemorySource>`.
struct Shared(Arc<Recorder>);

impl MemorySource for Shared {
    fn block(&self) -> Result<String, String> {
        self.0.block.lock().unwrap().clone()
    }
    fn updated(&self, session: &SessionId) -> Result<(), String> {
        self.0
            .updated
            .lock()
            .unwrap()
            .push(session.as_str().to_string());
        Ok(())
    }
}

struct Loopd {
    server: FakeServer,
    socket: PathBuf,
    memory: Arc<Recorder>,
    // Declared last so it is dropped last: fields drop in order, and the home must outlive
    // everything that may still write into it (tip T44).
    _home: Home,
}

fn start(block: Result<&str, &str>) -> Loopd {
    let home = Home::new();
    let server = FakeServer::start();
    let cfg = home.config(&server.socket);
    let socket = cfg.channel_socket();
    let listener = UnixListener::bind(&socket).unwrap();
    let memory = Arc::new(Recorder {
        block: Mutex::new(block.map(str::to_string).map_err(str::to_string)),
        updated: Mutex::new(Vec::new()),
    });
    let ctx = Arc::new(Context::new(
        cfg.clone(),
        Client::new(cfg),
        Box::new(ScriptedPort::new(vec![])),
        Registry::m2b(),
        Box::new(Shared(Arc::clone(&memory))),
    ));
    std::thread::spawn(move || serve(listener, ctx));
    Loopd {
        _home: home,
        server,
        socket,
        memory,
    }
}

/// One turn; the last frame.
fn turn(l: &Loopd, session: &str, resume: bool) -> Message {
    let mut stream = UnixStream::connect(&l.socket).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(10)))
        .unwrap();
    let msg = Message::Turn(Turn {
        session: SessionId::new(session).unwrap(),
        content: "hello".to_string(),
        resume,
    });
    let env = Envelope {
        v: PROTOCOL_VERSION,
        id: 7,
        r#final: true,
        msg,
    };
    write_frame(&mut stream, &env).unwrap();
    loop {
        let frame = read_frame(&mut stream).unwrap();
        if frame.r#final {
            return frame.msg;
        }
    }
}

fn system_prompt(l: &Loopd, n: usize) -> String {
    let requests = l.server.requests_to(CHAT);
    requests[n].json()["messages"][0]["content"]
        .as_str()
        .unwrap()
        .to_string()
}

#[test]
fn a_new_session_starts_with_the_block_and_keeps_it() {
    let l = start(Ok("# soul.md\nYou are kind."));
    l.server
        .route(CHAT, vec![Reply::fixture("plain"), Reply::fixture("plain")]);
    assert!(matches!(turn(&l, "a", false), Message::TurnDone(_)));
    assert_eq!(
        system_prompt(&l, 0),
        "You are Boxmaker, a test agent.\n\n# soul.md\nYou are kind."
    );
    *l.memory.block.lock().unwrap() = Ok("# soul.md\nYou changed.".to_string());
    assert!(matches!(turn(&l, "a", true), Message::TurnDone(_)));
    assert_eq!(
        system_prompt(&l, 1),
        "You are Boxmaker, a test agent.\n\n# soul.md\nYou are kind.",
        "a running session keeps the bytes it started with"
    );
}

#[test]
fn no_session_starts_without_the_block() {
    let l = start(Err(
        "memoryd at /x: refused; see docs/runbook.md#memory-unavailable",
    ));
    match turn(&l, "b", false) {
        Message::Error(e) => {
            assert_eq!(e.code, ErrorCode::Internal);
            assert!(e.detail.ends_with(UNAVAILABLE), "{}", e.detail);
        }
        other => panic!("started without memory: {other:?}"),
    }
    assert!(l.server.requests_to(CHAT).is_empty(), "no request was made");
}

#[test]
fn memoryd_is_told_after_a_finished_turn_only() {
    let l = start(Ok(""));
    l.server.route(
        CHAT,
        vec![Reply::fixture("plain"), Reply::fixture("bad_request")],
    );
    turn(&l, "c", false);
    assert!(
        matches!(turn(&l, "d", false), Message::Error(_)),
        "a turn that ran and failed"
    );
    assert!(matches!(
        turn(&l, "nobody-made-this", true),
        Message::Error(_)
    ));
    assert_eq!(*l.memory.updated.lock().unwrap(), ["c"]);
}

#[test]
fn recall_and_memory_are_core_tools_of_the_m5a_registry() {
    let names: Vec<String> = Registry::m5a()
        .core_schemas()
        .into_iter()
        .map(|s| s.name)
        .collect();
    assert_eq!(
        names,
        ["clock", "recall", "memory", "find_tool", "call_tool"]
    );
    let m3a: Vec<String> = Registry::m3a()
        .core_schemas()
        .into_iter()
        .map(|s| s.name)
        .collect();
    assert_eq!(m3a, ["clock", "find_tool", "call_tool"], "m3a is unchanged");
    assert!(Registry::m5a().get("read_file").is_some());
}