boxmaker/crates/memoryd/src/serve.rs

//! Serving: three sockets, one request per connection, each answered by one frame (M5a spec,
//! section 7). One lock covers the files and the index, so writes never interleave.

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

use proto::{
    DataClass, Empty, Envelope, ErrorCode, MemoryAdminReply, MemoryArgs, MemoryBlockReply,
    MemoryCall, MemoryCommand, Message, PROTOCOL_VERSION, RecallArgs, Reindexed, Timestamp,
    ToolResponse, WireError, read_frame, write_frame,
};

use crate::config::Config;
use crate::edit::{self, Context};
use crate::files::Store;
use crate::index::{self, Index};
use crate::recall::Recalled;
use crate::{journal, recall, review};

/// How long a peer may take to send its request.
pub const READ_WAIT: Duration = Duration::from_secs(10);

/// Where log lines go: stderr in `main`, a buffer in tests.
pub type Log = Arc<dyn Fn(&str) + Send + Sync>;

struct State {
    index: Index,
}

pub struct Memoryd {
    config: Config,
    store: Store,
    state: Mutex<State>,
    log: Log,
}

/// A started `memoryd`: its three listeners, bound and waiting.
pub struct Running {
    memoryd: Arc<Memoryd>,
    listeners: Vec<(UnixListener, Which)>,
}

#[derive(Clone, Copy)]
enum Which {
    Loop,
    Broker,
    Admin,
}

/// Create what is missing under `memory/`, open the index and catch it up, then bind the three
/// sockets. The error names what failed; the caller points at `memoryd-start-failed`.
pub fn start(config: Config, log: Log) -> Result<Running, String> {
    let store = Store::new(&config.memory_dir());
    store.init(&config.memory.soul_starter)?;
    let index = Index::open(&config.index_path(), config.index.max_class)?;
    let (_counts, problems) = index.catch_up(&config.sessions_dir(), &store.path("notes.md"))?;
    for problem in problems {
        (log)(format!("memoryd: not indexed: {problem}").as_str());
    }
    let loop_listener = crate::socket::listen(&config.loop_socket())?;
    let broker_listener = crate::socket::listen(&config.broker_socket())?;
    let admin_listener = crate::socket::listen(&config.admin_socket())?;
    let listeners = vec![
        (loop_listener, Which::Loop),
        (broker_listener, Which::Broker),
        (admin_listener, Which::Admin),
    ];
    Ok(Running {
        memoryd: Arc::new(Memoryd {
            config,
            store,
            state: Mutex::new(State { index }),
            log,
        }),
        listeners,
    })
}

impl Running {
    /// Accept on every socket, one thread per connection, until a listener fails.
    pub fn run(self) -> Result<(), String> {
        let (tx, rx) = std::sync::mpsc::channel::<String>();
        for (listener, which) in self.listeners {
            let tx = tx.clone();
            let memoryd = Arc::clone(&self.memoryd);
            std::thread::Builder::new()
                .spawn(move || {
                    for stream in listener.incoming() {
                        match stream {
                            Ok(stream) => {
                                let connection = Arc::clone(&memoryd);
                                match std::thread::Builder::new()
                                    .spawn(move || connection.connection(stream, which))
                                {
                                    Ok(_) => {}
                                    Err(e) => (memoryd.log)(
                                        format!("memoryd: cannot start a thread: {e}").as_str(),
                                    ),
                                }
                            }
                            Err(e) => {
                                let _ = tx.send(format!("a listener failed: {e}"));
                                return;
                            }
                        }
                    }
                })
                .map_err(|e| format!("cannot start a thread: {e}"))?;
        }
        drop(tx);
        match rx.recv() {
            Ok(message) => Err(message),
            Err(_) => Ok(()),
        }
    }
}

fn error(detail: String) -> Message {
    Message::Error(WireError {
        code: ErrorCode::Internal,
        detail,
    })
}

fn bad(detail: &str) -> Message {
    Message::Error(WireError {
        code: ErrorCode::BadMessage,
        detail: detail.to_string(),
    })
}

impl Memoryd {
    fn lock(&self) -> MutexGuard<'_, State> {
        self.state.lock().unwrap_or_else(|p| p.into_inner())
    }

    fn connection(&self, mut stream: UnixStream, which: Which) {
        let _ = stream.set_read_timeout(Some(READ_WAIT));
        let request = match read_frame(&mut stream) {
            Ok(request) => request,
            Err(_) => return,
        };
        let reply = self.handle(request.msg, which);
        let env = Envelope {
            v: PROTOCOL_VERSION,
            id: request.id,
            r#final: true,
            msg: reply,
        };
        let _ = write_frame(&mut stream, &env);
    }

    /// The answer to one request on one socket. A message sent to the wrong socket is refused.
    fn handle(&self, msg: Message, which: Which) -> Message {
        match (which, msg) {
            (Which::Loop, Message::MemoryBlock(_)) => self.block(),
            (Which::Loop, Message::SessionUpdated(update)) => {
                self.session_updated(update.session.as_str())
            }
            (Which::Broker, Message::MemoryCall(call)) => Message::ToolResponse(self.call(call)),
            (Which::Admin, Message::Reindex(_)) => self.reindex(),
            (Which::Admin, Message::MemoryAdmin(admin)) => self.admin(admin.command, &admin.arg),
            _ => bad("memoryd does not take this message on this socket"),
        }
    }

    fn block(&self) -> Message {
        let _state = self.lock();
        match self.store.block() {
            Ok(block) => {
                for (file, refused) in &block.left_out {
                    (self.log)(
                        format!(
                            "memoryd: an entry of {} was left out of the block: {}",
                            file.name(),
                            refused.sentence()
                        )
                        .as_str(),
                    );
                }
                Message::MemoryBlockReply(MemoryBlockReply { text: block.text })
            }
            Err(e) => error(e),
        }
    }

    fn session_updated(&self, session: &str) -> Message {
        let state = self.lock();
        match index::session_logs(&self.config.sessions_dir(), session) {
            Ok(logs) => {
                for (epoch, path) in logs {
                    if let Err(e) = state.index.index_log(&path, session, epoch) {
                        (self.log)(format!("memoryd: not indexed: {e}").as_str());
                    }
                }
            }
            Err(e) => return error(e),
        }
        Message::Ok(Empty {})
    }

    /// A decided call from `brokerd`: `recall` or `memory`, under the one lock. `Failed` carries the
    /// sentence for the model.
    fn call(&self, call: MemoryCall) -> ToolResponse {
        let state = self.lock();
        let result = match call.request.tool.as_str() {
            proto::memory::RECALL => self.recall(&state, &call),
            proto::memory::MEMORY => self.memory(&state, &call),
            _ => {
                return ToolResponse::Failed {
                    message: "memoryd runs only recall and memory".to_string(),
                };
            }
        };
        match result {
            Ok(found) => ToolResponse::Result {
                content: found.text,
                class: found.class,
                untrusted: found.untrusted,
                truncated: false,
            },
            Err(message) => ToolResponse::Failed { message },
        }
    }

    /// `recall`: search, labelled by what it found.
    fn recall(&self, state: &State, call: &MemoryCall) -> Result<Recalled, String> {
        let args: RecallArgs = match serde_json::from_str(&call.request.arguments) {
            Ok(args) => args,
            Err(e) => return Err(format!("recall's arguments: {e}")),
        };
        recall::recall(
            &state.index,
            &self.store,
            &args,
            call.request.session.as_str(),
            Timestamp::now(),
        )
    }

    /// `memory`: staged for review, or applied and journalled; then what it appended to `notes.md`
    /// is indexed. Its answer is labelled private and trusted: it is `memoryd`'s own sentence.
    fn memory(&self, state: &State, call: &MemoryCall) -> Result<Recalled, String> {
        let args: MemoryArgs = match serde_json::from_str(&call.request.arguments) {
            Ok(args) => args,
            Err(e) => return Err(format!("memory's arguments: {e}")),
        };
        let ctx = Context {
            session: call.request.session.as_str().to_string(),
            call: call.request.call.0,
            taint: call.taint,
            untrusted: call.untrusted,
            now: Timestamp::now(),
        };
        let answer = if self.config.memory.review
            && !matches!(args, MemoryArgs::Note { .. })
            && !call.untrusted
        {
            review::submit(&self.store, &args, &call.request.arguments, &ctx)?
        } else {
            let outcome = edit::apply(&self.store, &args, &ctx)?;
            journal::record(&self.store, &ctx, &outcome)?;
            outcome.answer
        };
        self.index_notes(state);
        Ok(Recalled {
            text: answer,
            class: DataClass::Private,
            untrusted: false,
        })
    }

    /// Index what was appended to `notes.md`; a failure is logged, and the next start catches up.
    fn index_notes(&self, state: &State) {
        let text = match std::fs::read_to_string(self.store.path("notes.md")) {
            Ok(text) => text,
            Err(e) => {
                (self.log)(format!("memoryd: notes not indexed: {e}").as_str());
                return;
            }
        };
        if let Err(e) = state.index.index_notes(&text) {
            (self.log)(format!("memoryd: notes not indexed: {e}").as_str());
        }
    }

    fn reindex(&self) -> Message {
        let mut state = self.lock();
        let max_class = self.config.index.max_class;
        match index::rebuild(
            &self.config.index_path(),
            max_class,
            &self.config.sessions_dir(),
            &self.store.path("notes.md"),
        ) {
            Ok((counts, problems)) => {
                for problem in problems {
                    (self.log)(format!("memoryd: not indexed: {problem}").as_str());
                }
                match Index::open(&self.config.index_path(), max_class) {
                    Ok(index) => {
                        state.index = index;
                        Message::Reindexed(Reindexed {
                            logs: counts.logs,
                            records: counts.records,
                            notes: counts.notes,
                        })
                    }
                    Err(e) => error(e),
                }
            }
            Err(e) => error(e),
        }
    }

    fn admin(&self, command: MemoryCommand, arg: &str) -> Message {
        let state = self.lock();
        let now = Timestamp::now();
        let result = match command {
            MemoryCommand::Log => {
                let n = if arg.is_empty() {
                    10
                } else {
                    match arg.parse::<usize>() {
                        Ok(n) => n,
                        Err(_) => return error(format!("log takes a count; {arg:?} is not one")),
                    }
                };
                journal::log(&self.store, n)
            }
            MemoryCommand::Undo => journal::undo(&self.store, arg, now),
            MemoryCommand::Pending => review::list(&self.store),
            MemoryCommand::Accept => review::accept(&self.store, arg, now),
            MemoryCommand::Reject => review::reject(&self.store, arg, now),
        };
        self.index_notes(&state);
        match result {
            Ok(text) => Message::MemoryAdminReply(MemoryAdminReply { text }),
            Err(e) => error(e),
        }
    }
}