boxmaker/crates/memoryd/src/index.rs

//! The index: SQLite FTS5 over the session logs and `notes.md` (M5a spec, section 6). Disposable:
//! it is rebuilt from those files. What is indexed: the owner's messages, the assistant's
//! replies, and each tool call by name and arguments; never a tool's result.

use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};

use bxcore::{DataClass, LogRecord};
use rusqlite::{Connection, params};

use crate::notes::{self, Kind};

const SCHEMA: &str = "
CREATE VIRTUAL TABLE IF NOT EXISTS hits USING fts5(text, tokenize = 'unicode61');
CREATE TABLE IF NOT EXISTS sources (
  rowid INTEGER PRIMARY KEY, kind TEXT NOT NULL, session TEXT NOT NULL, epoch INTEGER NOT NULL,
  record INTEGER NOT NULL, time TEXT NOT NULL, role TEXT NOT NULL, class TEXT NOT NULL,
  untrusted INTEGER NOT NULL);
CREATE TABLE IF NOT EXISTS progress (log TEXT PRIMARY KEY, records INTEGER NOT NULL);
";

/// One indexed item.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Source {
    /// `session`, `noted`, `moved` or `dropped`.
    pub kind: String,
    pub session: String,
    pub epoch: i64,
    /// The record's line in its log, or the note's place in `notes.md` (from 0).
    pub record: i64,
    pub time: String,
    /// `owner`, `assistant` or `call` for a session record; the file (`memory.md`, `user.md`) for
    /// a moved or dropped note; empty for a noted one.
    pub role: String,
    pub class: DataClass,
    pub untrusted: bool,
}

/// How much a reindex read.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Counts {
    pub logs: u64,
    pub records: u64,
    pub notes: u64,
}

pub struct Index {
    conn: Connection,
    max_class: DataClass,
}

fn db(e: rusqlite::Error) -> String {
    format!("index: {e}")
}

pub fn class_name(class: DataClass) -> &'static str {
    match class {
        DataClass::Public => "public",
        DataClass::Private => "private",
        DataClass::Secret => "secret",
    }
}

pub fn class_of(name: &str) -> DataClass {
    match name {
        "public" => DataClass::Public,
        "secret" => DataClass::Secret,
        _ => DataClass::Private,
    }
}

fn int(n: usize) -> i64 {
    i64::try_from(n).unwrap_or(i64::MAX)
}

impl Index {
    /// Open or create the index at `path` (its directory 0700, the file 0600).
    pub fn open(path: &Path, max_class: DataClass) -> Result<Index, String> {
        if let Some(dir) = path.parent() {
            std::fs::create_dir_all(dir).map_err(|e| format!("{}: {e}", dir.display()))?;
            std::fs::set_permissions(dir, std::fs::Permissions::from_mode(0o700))
                .map_err(|e| format!("{}: {e}", dir.display()))?;
        }
        let conn = Connection::open(path).map_err(|e| format!("{}: {e}", path.display()))?;
        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
            .map_err(|e| format!("{}: {e}", path.display()))?;
        conn.execute_batch(SCHEMA)
            .map_err(|e| format!("{}: {e}", path.display()))?;
        Ok(Index { conn, max_class })
    }

    /// The connection, for `recall`.
    pub fn conn(&self) -> &Connection {
        &self.conn
    }

    fn progress(&self, log: &str) -> Result<usize, String> {
        // `self.conn.query_row("SELECT records FROM progress WHERE log = ?1", params![log], |row|
        // row.get(0))` as an Option<i64>: Err(rusqlite::Error::QueryReturnedNoRows) means None,
        // another error goes through `db`. Return it as a usize (`usize::try_from`, 0 when it does
        // not fit), 0 for None.
        let value: Option<i64> = match self.conn.query_row(
            "SELECT records FROM progress WHERE log = ?1",
            params![log],
            |row| row.get(0),
        ) {
            Ok(v) => Some(v),
            Err(rusqlite::Error::QueryReturnedNoRows) => None,
            Err(e) => return Err(db(e)),
        };
        let n = match value {
            Some(v) => usize::try_from(v).unwrap_or(0),
            None => 0,
        };
        Ok(n)
    }

    fn add(&self, text: &str, source: &Source) -> Result<(), String> {
        // 1. `execute("INSERT INTO hits (text) VALUES (?1)", params![text])`.
        // 2. `let rowid = self.conn.last_insert_rowid();`
        // 3. `execute("INSERT INTO sources (rowid, kind, session, epoch, record, time, role, class,
        //    untrusted) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)", params![rowid, kind, session,
        //    epoch, record, time, role, class_name(class), untrusted])`. Errors through `db`.
        self.conn
            .execute("INSERT INTO hits (text) VALUES (?1)", params![text])
            .map_err(db)?;
        let rowid = self.conn.last_insert_rowid();
        self.conn
            .execute(
                "INSERT INTO sources (rowid, kind, session, epoch, record, time, role, class, untrusted) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
                params![
                    rowid,
                    source.kind.clone(),
                    source.session.clone(),
                    source.epoch,
                    source.record,
                    source.time.clone(),
                    source.role.clone(),
                    class_name(source.class),
                    source.untrusted,
                ],
            )
            .map_err(db)?;
        Ok(())
    }

    fn set_progress(&self, log: &str, records: usize) -> Result<(), String> {
        // `execute("INSERT INTO progress (log, records) VALUES (?1, ?2) ON CONFLICT(log) DO UPDATE
        //    SET records = excluded.records", params![log, int(records)])`, errors through `db`.
        self.conn
            .execute(
                "INSERT INTO progress (log, records) VALUES (?1, ?2) ON CONFLICT(log) DO UPDATE SET records = excluded.records",
                params![log, int(records)],
            )
            .map_err(db)?;
        Ok(())
    }

    /// Index the records of one session log past what was indexed before; how many were read.
    /// The label is walked from the log's first record: the highest class of any tool result
    /// before a record (at least `private`), and whether any was untrusted.
    pub fn index_log(&self, path: &Path, session: &str, epoch: i64) -> Result<usize, String> {
        // 1. `read_to_string(path)` (an error names the path). `key` =
        //    format!("{session}/{epoch}"); `done = self.progress(&key)?`.
        // 2. The label starts at class Private, untrusted false; `read` at 0.
        // 3. For each line with its place n (from 0): parse it as a LogRecord (an error ->
        //    Err(format!("{}: line {}: {e}", path.display(), n + 1))). When n >= done: read += 1,
        //    and `self.index_record(&record, &base)?` where base is a Source with kind "session",
        //    session, epoch, record int(n), time and role empty, and the label as it stands. Then,
        //    for every record (indexed or not): a ToolResult raises the label: class =
        //    class.max(its class), untrusted |= its untrusted.
        // 4. `self.set_progress(&key, done + read)?`; Ok(read).
        let text = std::fs::read_to_string(path).map_err(|e| format!("{}: {e}", path.display()))?;
        let key = format!("{session}/{epoch}");
        let done = self.progress(&key)?;
        let mut class = DataClass::Private;
        let mut untrusted = false;
        let mut read = 0;
        for (n, line) in text.lines().enumerate() {
            let record: LogRecord = serde_json::from_str(line)
                .map_err(|e| format!("{}: line {}: {e}", path.display(), n + 1))?;
            if n >= done {
                read += 1;
                let base = Source {
                    kind: "session".to_string(),
                    session: session.to_string(),
                    epoch,
                    record: int(n),
                    time: String::new(),
                    role: String::new(),
                    class,
                    untrusted,
                };
                self.index_record(&record, &base)?;
            }
            if let LogRecord::ToolResult {
                class: c,
                untrusted: u,
                ..
            } = &record
            {
                class = class.max(*c);
                untrusted |= *u;
            }
        }
        self.set_progress(&key, done + read)?;
        Ok(read)
    }

    /// Index what one log record holds: the owner's message, the assistant's reply, and each tool
    /// call by name and arguments. `base` carries the kind, session, epoch, line and label; each item
    /// gets its own role and time. Nothing when the label's class is above `max_class`, and nothing
    /// for an item with no text.
    fn index_record(&self, record: &LogRecord, base: &Source) -> Result<(), String> {
        // When base.class > self.max_class: Ok(()) and nothing indexed. Otherwise each item is
        // `self.add(text, &Source { time: time.to_rfc3339(), role, ..base.clone() })`, skipped when
        // text.trim() is empty. User -> its content, role "owner". Assistant -> its content when
        // Some, role "assistant"; then each tool call as format!("{} {}", name, arguments), role
        //    "call"; all with the record's time. Every other record: nothing.
        if base.class > self.max_class {
            return Ok(());
        }
        match record {
            LogRecord::User { time, content, .. } if !content.trim().is_empty() => {
                let src = Source {
                    time: time.to_rfc3339(),
                    role: "owner".to_string(),
                    ..base.clone()
                };
                self.add(content, &src)?;
            }
            LogRecord::Assistant {
                time,
                content,
                tool_calls,
                ..
            } => {
                if let Some(c) = content
                    && !c.trim().is_empty()
                {
                    let src = Source {
                        time: time.to_rfc3339(),
                        role: "assistant".to_string(),
                        ..base.clone()
                    };
                    self.add(c, &src)?;
                }
                for call in tool_calls {
                    let text = format!("{} {}", call.name, call.arguments);
                    if !text.trim().is_empty() {
                        let src = Source {
                            time: time.to_rfc3339(),
                            role: "call".to_string(),
                            ..base.clone()
                        };
                        self.add(&text, &src)?;
                    }
                }
            }
            _ => {}
        }
        Ok(())
    }

    /// Index the entries of `notes.md` past what was indexed before; how many were read.
    pub fn index_notes(&self, text: &str) -> Result<usize, String> {
        // `done = self.progress("notes.md")?`. Parse `text` with notes::parse and skip the first
        // `done`. For each remaining note (with its place n): read += 1; its kind, role, class and
        // untrusted: Noted -> ("noted", "", its class, its untrusted); Moved(file) -> ("moved",
        // file.name(), Private, false); Dropped(file) -> ("dropped", file.name(), Private, false).
        // Skip it when class > self.max_class. Otherwise `self.add(&note.text, &Source { kind,
        // session: note.session, epoch: 0, record: int(n), time: note.time, role, class, untrusted
        //    })?`. Then `self.set_progress("notes.md", done + read)?`; Ok(read).
        let done = self.progress("notes.md")?;
        let notes = notes::parse(text);
        let mut read = 0;
        for (n, note) in notes.iter().enumerate() {
            if n < done {
                continue;
            }
            read += 1;
            let (kind, role, class, untrusted) = match &note.kind {
                Kind::Noted { class, untrusted } => {
                    ("noted".to_string(), String::new(), *class, *untrusted)
                }
                Kind::Moved(file) => (
                    "moved".to_string(),
                    file.name().to_string(),
                    DataClass::Private,
                    false,
                ),
                Kind::Dropped(file) => (
                    "dropped".to_string(),
                    file.name().to_string(),
                    DataClass::Private,
                    false,
                ),
            };
            if class > self.max_class {
                continue;
            }
            self.add(
                &note.text,
                &Source {
                    kind,
                    session: note.session.clone(),
                    epoch: 0,
                    record: int(n),
                    time: note.time.clone(),
                    role,
                    class,
                    untrusted,
                },
            )?;
        }
        self.set_progress("notes.md", done + read)?;
        Ok(read)
    }

    /// Index everything new: every `sessions/<id>/<epoch>.jsonl`, then `notes.md`. A log that
    /// cannot be read or parsed is skipped and named in the returned problems; it is `loopd`'s.
    pub fn catch_up(
        &self,
        sessions: &Path,
        notes_md: &Path,
    ) -> Result<(Counts, Vec<String>), String> {
        // For each (session, epoch, path) of `logs(sessions)?`: `self.index_log(&path, &session,
        // epoch)`; Ok(n) -> counts.logs += 1 and counts.records += n (u64::try_from, u64::MAX when
        // it does not fit); Err(e) -> push e onto `problems` and go on. Then read `notes_md`:
        //    NotFound is an empty text, another error -> Err(format!("{}: {e}", notes_md.display()));
        //    counts.notes = `self.index_notes(&text)?` as u64 (u64::try_from). Ok((counts, problems)).
        let mut counts = Counts::default();
        let mut problems: Vec<String> = Vec::new();
        for (session, epoch, path) in logs(sessions)? {
            match self.index_log(&path, &session, epoch) {
                Ok(n) => {
                    counts.logs += 1;
                    counts.records += u64::try_from(n).unwrap_or(u64::MAX);
                }
                Err(e) => problems.push(e),
            }
        }
        let text = match std::fs::read_to_string(notes_md) {
            Ok(t) => t,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => String::new(),
            Err(e) => return Err(format!("{}: {e}", notes_md.display())),
        };
        counts.notes = u64::try_from(self.index_notes(&text)?).unwrap_or(u64::MAX);
        Ok((counts, problems))
    }
}

/// The logs of one session: `(epoch, path)`, sorted. A session with no directory has none.
pub fn session_logs(sessions: &Path, session: &str) -> Result<Vec<(i64, PathBuf)>, String> {
    let dir = sessions.join(session);
    let files = match std::fs::read_dir(&dir) {
        Ok(files) => files,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
        Err(e) => return Err(format!("{}: {e}", dir.display())),
    };
    let mut found = Vec::new();
    for file in files.filter_map(|f| f.ok()) {
        let name = file.file_name().to_string_lossy().into_owned();
        if let Some(Ok(epoch)) = name.strip_suffix(".jsonl").map(str::parse::<i64>) {
            found.push((epoch, file.path()));
        }
    }
    found.sort();
    Ok(found)
}

/// Every session log under `sessions`: `(session, epoch, path)`, sorted.
pub fn logs(sessions: &Path) -> Result<Vec<(String, i64, PathBuf)>, String> {
    // Read the directory `sessions` (NotFound -> Ok(empty); another error -> Err naming it). For
    // each entry, its name is the session; read that directory (skip it when that fails); every
    //    file named `<n>.jsonl` where n parses as an i64 gives (session, n, the file's path). Sort the
    //    list and return it.
    let entries = match std::fs::read_dir(sessions) {
        Ok(e) => e,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
        Err(e) => return Err(format!("{}: {e}", sessions.display())),
    };
    let mut out: Vec<(String, i64, PathBuf)> = Vec::new();
    for entry in entries {
        let entry = match entry {
            Ok(e) => e,
            Err(_) => continue,
        };
        let name = entry.file_name();
        let session = match name.to_str() {
            Some(s) => s,
            None => continue,
        };
        let dir = entry.path();
        let sub = match std::fs::read_dir(&dir) {
            Ok(d) => d,
            Err(_) => continue,
        };
        for sub_entry in sub {
            let sub_entry = match sub_entry {
                Ok(e) => e,
                Err(_) => continue,
            };
            let sub_name = sub_entry.file_name();
            let file_name = match sub_name.to_str() {
                Some(s) => s,
                None => continue,
            };
            if let Some(stem) = file_name.strip_suffix(".jsonl")
                && let Ok(epoch) = stem.parse::<i64>()
            {
                out.push((session.to_string(), epoch, sub_entry.path()));
            }
        }
    }
    out.sort();
    Ok(out)
}

/// Build a fresh index at `path` from `sessions` and `notes.md`: into a new file beside it,
/// renamed over the old one, so a running `recall` keeps its index until the new one is whole.
pub fn rebuild(
    path: &Path,
    max_class: DataClass,
    sessions: &Path,
    notes_md: &Path,
) -> Result<(Counts, Vec<String>), String> {
    // `fresh = path.with_extension("rebuilding")`; remove it if it is there (ignore the result).
    //    `Index::open(&fresh, max_class)?`, then `catch_up(sessions, notes_md)?` on it, and drop that
    //    Index. `std::fs::rename(&fresh, path)` (an error names path). Return the result of catch_up.
    let fresh = path.with_extension("rebuilding");
    let _ = std::fs::remove_file(&fresh);
    let index = Index::open(&fresh, max_class)?;
    let result = index.catch_up(sessions, notes_md)?;
    drop(index);
    std::fs::rename(&fresh, path).map_err(|e| format!("{}: {e}", path.display()))?;
    Ok(result)
}