boxmaker/crates/loopd/src/memory.rs

//! Memory, as `loopd` sees it (M5a spec, section 9): the block `memoryd` renders for a new
//! session's baseline, the schemas of the two core memory tools, and the notice after a turn.
//! `loopd` keeps no memory itself; `brokerd` decides the tool calls and `memoryd` runs them.

use std::os::unix::net::UnixStream;
use std::path::PathBuf;
use std::time::Duration;

use bxcore::{
    Empty, Envelope, Message, PROTOCOL_VERSION, SessionId, SessionUpdated, read_frame, write_frame,
};

use crate::llama::ToolSchema;

/// Where a new session's baseline points when there is no memory block.
pub const UNAVAILABLE: &str = "see docs/runbook.md#memory-unavailable";

/// How long `loopd` waits for `memoryd`.
pub const MEMORY_WAIT: Duration = Duration::from_secs(10);

/// Where the memory block comes from, and who hears that a session has new records.
pub trait MemorySource: Send + Sync {
    /// The block for a new session. `Err` ends in `UNAVAILABLE`: no session starts without it.
    fn block(&self) -> Result<String, String>;
    /// The session has new records; best effort, a failure is only logged by the caller.
    fn updated(&self, session: &SessionId) -> Result<(), String>;
}

/// No memory: for tests and the self-test, never for `loopd serve`.
pub struct NoMemory;

impl MemorySource for NoMemory {
    fn block(&self) -> Result<String, String> {
        Ok(String::new())
    }
    fn updated(&self, _session: &SessionId) -> Result<(), String> {
        Ok(())
    }
}

/// A client of `memoryd`'s `loopd` socket. One connection per request.
pub struct MemoryPort {
    socket: PathBuf,
    wait: Duration,
}

impl MemoryPort {
    pub fn new(socket: PathBuf, wait: Duration) -> MemoryPort {
        MemoryPort { socket, wait }
    }

    fn ask(&self, msg: Message) -> Result<Message, String> {
        let mut stream = match UnixStream::connect(&self.socket) {
            Ok(stream) => stream,
            Err(e) => return Err(self.line(&e)),
        };
        if let Err(e) = stream.set_write_timeout(Some(self.wait)) {
            return Err(self.line(&e));
        }
        if let Err(e) = stream.set_read_timeout(Some(self.wait)) {
            return Err(self.line(&e));
        }
        let env = Envelope {
            v: PROTOCOL_VERSION,
            id: 1,
            r#final: true,
            msg,
        };
        if let Err(e) = write_frame(&mut stream, &env) {
            return Err(self.line(&e));
        }
        match read_frame(&mut stream) {
            Ok(env) => Ok(env.msg),
            Err(e) => Err(self.line(&e)),
        }
    }

    fn line(&self, e: &dyn std::fmt::Display) -> String {
        format!("memoryd at {}: {e}; {UNAVAILABLE}", self.socket.display())
    }
}

impl MemorySource for MemoryPort {
    fn block(&self) -> Result<String, String> {
        match self.ask(Message::MemoryBlock(Empty {}))? {
            Message::MemoryBlockReply(reply) => Ok(reply.text),
            Message::Error(e) => Err(format!("memoryd: {}; {UNAVAILABLE}", e.detail)),
            _ => Err(format!(
                "memoryd answered with the wrong message; {UNAVAILABLE}"
            )),
        }
    }

    fn updated(&self, session: &SessionId) -> Result<(), String> {
        match self.ask(Message::SessionUpdated(SessionUpdated {
            session: session.clone(),
        }))? {
            Message::Ok(_) => Ok(()),
            Message::Error(e) => Err(format!("memoryd: {}", e.detail)),
            _ => Err("memoryd answered with the wrong message".to_string()),
        }
    }
}

pub const RECALL: &str = "recall";
pub const MEMORY: &str = "memory";

pub fn recall_schema() -> ToolSchema {
    ToolSchema {
        name: RECALL.to_string(),
        description: "Search your memory: past conversations and notes. Each hit says where it \
                      came from: identity, memory, or memory-tainted (from a session that read \
                      untrusted content)."
            .to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "query": { "type": "string", "description": "Words to search for." },
                "k": { "type": "integer", "description": "How many hits, 1 to 10. Default 5." }
            },
            "required": ["query"]
        }),
    }
}

pub fn memory_schema() -> ToolSchema {
    ToolSchema {
        name: MEMORY.to_string(),
        description: "Write memory. note: record something seen or told. add, replace, remove: \
                      edit memory.md (your work) or user.md (your owner); entries are named by \
                      their exact text. consolidate: rewrite a whole file, moving or dropping \
                      entries to notes. Takes effect in the next session."
            .to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "action": { "type": "string", "enum": ["note", "add", "replace", "remove", "consolidate"] },
                "file": { "type": "string", "enum": ["memory", "user"] },
                "text": { "type": "string", "description": "The new entry, or the note." },
                "old": { "type": "string", "description": "The exact text of the entry to replace or remove." },
                "entries": {
                    "type": "array",
                    "description": "consolidate: every current entry once, as {old, keep: true}, {old, text}, {old, move: true} or {old, drop: true}; new entries as {text}.",
                    "items": { "type": "object" }
                }
            },
            "required": ["action"]
        }),
    }
}