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. With
//! `[honcho]`, a reported turn is decided after `loopd` has its answer, under a lock of its own
//! (M5c spec, section 5).

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

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

use crate::config::Config;
use crate::edit::{self, Context};
use crate::files::Store;
use crate::honcho::ledger::{self, LedgerError};
use crate::honcho::mirror::Mirror;
use crate::honcho::report;
use crate::honcho::state;
use crate::honcho::turns::TurnSync;
use crate::honcho::worker::{self, Wake};
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>,
    /// With `[honcho]`: the turns reported and not yet decided, and the ledger. Never taken
    /// while `state` is held, and never held by a request that waits for its answer.
    honcho: Option<Arc<Mutex<TurnSync>>>,
    /// Wakes the Honcho worker when a turn has been decided.
    wake: Arc<Wake>,
    /// With `[honcho]`: decides memory writes into the ledger, under the ledger's own lock.
    mirror: Option<Mirror>,
    log: Log,
}

/// Why `start` failed. `main` prints a start failure with the pointer `memoryd-start-failed`, and
/// a ledger that does not read through with its own.
#[derive(Debug)]
pub enum StartError {
    Failed(String),
    Ledger(LedgerError),
    /// A Honcho key that cannot be read: `bxcore::secrets`' line, with its own pointer.
    Secret(String),
}

impl std::fmt::Display for StartError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StartError::Failed(why) | StartError::Secret(why) => write!(f, "{why}"),
            // `Damaged` names its pointer itself.
            StartError::Ledger(e @ LedgerError::Damaged { .. }) => write!(f, "{e}"),
            StartError::Ledger(e) => write!(f, "{e}; {}", ledger::RUNBOOK),
        }
    }
}

impl std::error::Error for StartError {}

impl From<String> for StartError {
    fn from(why: String) -> StartError {
        StartError::Failed(why)
    }
}

/// 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,
}

/// With `[honcho]`, read the ledger through and the keys first. Then create what is missing
/// under `memory/` (and, the first time, write when sync began), open the index and catch it up,
/// bind the three sockets, and start the Honcho worker. The error names what failed.
pub fn start(config: Config, log: Log) -> Result<Running, StartError> {
    let opened = match config.honcho_config()? {
        Some(honcho) => Some(worker::open(&config, &honcho, &log)?),
        None => None,
    };
    let store = Store::new(&config.memory_dir());
    store.init(&config.memory.soul_starter)?;
    if let Some(opened) = &opened {
        opened.begin()?;
    }
    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),
    ];
    let wake = Arc::new(Wake::default());
    let (honcho, mirror) = match opened {
        Some(opened) => {
            let sync = Arc::clone(&opened.sync);
            let mirror = opened.mirror();
            worker::spawn(opened, &config, Arc::clone(&wake), Arc::clone(&log))?;
            (Some(sync), Some(mirror))
        }
        None => (None, None),
    };
    Ok(Running {
        memoryd: Arc::new(Memoryd {
            config,
            store,
            state: Mutex::new(State { index }),
            honcho,
            wake,
            mirror,
            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 update = match (which, &request.msg) {
            (Which::Loop, Message::SessionUpdated(update)) => Some(update.clone()),
            _ => None,
        };
        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);
        // `loopd` has its answer before `brokerd` is asked: it sends the turn's reply only after
        // this frame, and `brokerd` may take seconds.
        drop(stream);
        if let Some(update) = update {
            self.decide_turns(&update);
        }
    }

    /// With `[honcho]`: decide journalled edit `id` (the latest when `None`) into the ledger. Called
    /// under the index lock, so edits are decided in the journal's order (M5c spec, section 4).
    /// What it decided is sent at once: the worker is woken.
    fn mirror(&self, id: Option<&str>) {
        if let Some(mirror) = &self.mirror
            && mirror.journalled(&self.store, id, Timestamp::now(), &*self.log)
        {
            self.wake.notify();
        }
    }

    /// With `[honcho]`: decide the reported turn, and every other that waits (`TurnSync`).
    fn decide_turns(&self, update: &SessionUpdated) {
        let Some(honcho) = &self.honcho else {
            return;
        };
        let socket = self.config.state_socket();
        let ask = |session: &SessionId| state::ask(&socket, session, state::WAIT);
        let mut sync = honcho.lock().unwrap_or_else(|p| p.into_inner());
        sync.updated(
            &update.session,
            update.turn,
            &ask,
            Timestamp::now(),
            &*self.log,
        );
        drop(sync);
        self.wake.notify();
    }

    /// 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, refusal) in &block.left_out {
                    (self.log)(
                        format!(
                            "memoryd: an entry of {} was left out of the block by the rule {}; see docs/runbook.md#memory-entry-excluded",
                            file.name(),
                            refusal.rule
                        )
                        .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() {
            bxcore::memory::RECALL => self.recall(&state, &call),
            bxcore::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`: refused, 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. A refusal and a flag are logged with the rule's name, never with the text.
    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(),
        };
        if let Some(refusal) = edit::refused(&args) {
            (self.log)(
                format!(
                    "memoryd: a memory write from session {} was refused by the rule {}; see docs/runbook.md#memory-write-refused",
                    ctx.session, refusal.rule
                )
                .as_str(),
            );
            return Err(refusal.sentence.to_string());
        }
        // A write that matches a flag rule becomes notes, and notes are never staged.
        let answer = if self.config.memory.review
            && !matches!(args, MemoryArgs::Note { .. })
            && !call.untrusted
            && edit::flagged(&args).is_none()
        {
            review::submit(&self.store, &args, &call.request.arguments, &ctx)?
        } else {
            let outcome = edit::apply(&self.store, &args, &ctx)?;
            let id = journal::record(&self.store, &ctx, &outcome)?;
            self.mirror(Some(&id));
            if let Some(rule) = outcome.flag {
                (self.log)(
                    format!(
                        "memoryd: a memory write from session {} matches the rule {rule}; it was saved in notes.md as an untrusted note (edit {id}); see docs/runbook.md#memory-write-flagged",
                        ctx.session
                    )
                    .as_str(),
                );
            }
            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),
            MemoryCommand::Sync => Ok(match &self.mirror {
                Some(mirror) => mirror.report(),
                None => report::NOT_CONFIGURED.to_string(),
            }),
        };
        if result.is_ok() && matches!(command, MemoryCommand::Undo | MemoryCommand::Accept) {
            self.mirror(None);
        }
        self.index_notes(&state);
        match result {
            Ok(text) => Message::MemoryAdminReply(MemoryAdminReply { text }),
            Err(e) => error(e),
        }
    }
}