use std::collections::HashMap;
use chrono::{DateTime, Utc};
use rusqlite::Connection;
use crate::Memory;
const CMD_NS: i64 = 1_000_000_000;
const TURN_NS: i64 = 2_000_000_000;
const NOTE_NS: i64 = 3_000_000_000;
const SESSION_NS: i64 = 4_000_000_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Source {
Document,
Command,
Conversation,
Note,
Session,
}
impl Source {
pub fn tag(self) -> &'static str {
match self {
Source::Document => "document",
Source::Command => "command",
Source::Conversation => "conversation",
Source::Note => "note",
Source::Session => "session",
}
}
}
pub fn decode_id(id: i64) -> (Source, i64) {
if id >= SESSION_NS {
(Source::Session, id - SESSION_NS)
} else if id >= NOTE_NS {
(Source::Note, id - NOTE_NS)
} else if id >= TURN_NS {
(Source::Conversation, id - TURN_NS)
} else if id >= CMD_NS {
(Source::Command, id - CMD_NS)
} else {
(Source::Document, id)
}
}
fn parse_ts(ts: &str) -> DateTime<Utc> {
DateTime::parse_from_rfc3339(ts)
.map(|d| d.with_timezone(&Utc))
.unwrap_or_else(|_| Utc::now())
}
pub fn load_memories(conn: &Connection) -> Vec<Memory> {
let mut mems = Vec::new();
mems.extend(documents(conn));
mems.extend(command_runs(conn));
mems.extend(agent_turns(conn));
mems.extend(bookmarks(conn));
mems.extend(sessions(conn));
mems
}
fn documents(conn: &Connection) -> Vec<Memory> {
let Ok(mut stmt) =
conn.prepare("SELECT id, title, content, source_type, created_at FROM documents")
else {
return Vec::new();
};
let rows = stmt.query_map([], |r| {
Ok((
r.get::<_, i64>(0)?,
r.get::<_, String>(1)?,
r.get::<_, String>(2)?,
r.get::<_, String>(3)?,
r.get::<_, String>(4)?,
))
});
let Ok(rows) = rows else { return Vec::new() };
rows.flatten()
.map(|(id, title, content, source_type, created)| {
let when = parse_ts(&created);
Memory {
id,
text: format!("{title}. {content}"),
created_at: when,
last_used: when,
mentions: 1,
importance: 0.5,
tags: vec!["document".into(), source_type],
embedding: None,
}
})
.collect()
}
fn command_runs(conn: &Connection) -> Vec<Memory> {
let Ok(mut stmt) = conn.prepare(
"SELECT id, command, argv_json, notes, stderr, exit_code, created_at FROM command_runs",
) else {
return Vec::new();
};
let rows = stmt.query_map([], |r| {
Ok((
r.get::<_, i64>(0)?,
r.get::<_, String>(1)?,
r.get::<_, String>(2)?,
r.get::<_, String>(3)?,
r.get::<_, String>(4)?,
r.get::<_, Option<i64>>(5)?,
r.get::<_, String>(6)?,
))
});
let Ok(rows) = rows else { return Vec::new() };
let runs: Vec<_> = rows.flatten().collect();
let mut counts: HashMap<String, u32> = HashMap::new();
for (_, cmd, ..) in &runs {
*counts.entry(cmd.to_lowercase()).or_insert(0) += 1;
}
runs.into_iter()
.map(|(id, command, argv_json, notes, stderr, exit_code, created)| {
let when = parse_ts(&created);
let ok = exit_code == Some(0);
Memory {
id: CMD_NS + id,
text: format!("{command} {argv_json} {notes} {stderr}"),
created_at: when,
last_used: when,
mentions: *counts.get(&command.to_lowercase()).unwrap_or(&1),
importance: if exit_code.unwrap_or(0) != 0 { 0.65 } else { 0.4 },
tags: vec!["command".into(), if ok { "ok".into() } else { "error".into() }],
embedding: None,
}
})
.collect()
}
fn agent_turns(conn: &Connection) -> Vec<Memory> {
let Ok(mut stmt) = conn.prepare("SELECT id, ts, direction, text FROM agent_turns") else {
return Vec::new();
};
let rows = stmt.query_map([], |r| {
Ok((
r.get::<_, i64>(0)?,
r.get::<_, String>(1)?,
r.get::<_, String>(2)?,
r.get::<_, String>(3)?,
))
});
let Ok(rows) = rows else { return Vec::new() };
let turns: Vec<_> = rows.flatten().filter(|(_, _, _, t)| !t.trim().is_empty()).collect();
let mut counts: HashMap<String, u32> = HashMap::new();
for (_, _, _, t) in &turns {
*counts.entry(t.trim().to_lowercase()).or_insert(0) += 1;
}
turns
.into_iter()
.map(|(id, ts, direction, text)| {
let when = parse_ts(&ts);
let importance = match direction.as_str() {
"user" => 0.6_f32,
"agent" => 0.45,
_ => 0.3,
};
let mentions = *counts.get(&text.trim().to_lowercase()).unwrap_or(&1);
Memory {
id: TURN_NS + id,
text,
created_at: when,
last_used: when,
mentions,
importance,
tags: vec!["conversation".into(), direction],
embedding: None,
}
})
.collect()
}
fn bookmarks(conn: &Connection) -> Vec<Memory> {
let mut stmt = match conn
.prepare("SELECT id, label, created_at FROM bookmarks WHERE approved = 1")
{
Ok(stmt) => stmt,
Err(_) => match conn.prepare("SELECT id, label, created_at FROM bookmarks") {
Ok(stmt) => stmt,
Err(_) => return Vec::new(),
},
};
let rows = stmt.query_map([], |r| {
Ok((
r.get::<_, i64>(0)?,
r.get::<_, String>(1)?,
r.get::<_, String>(2)?,
))
});
let Ok(rows) = rows else { return Vec::new() };
rows.flatten()
.filter(|(_, label, _)| !label.trim().is_empty())
.map(|(id, label, created)| {
let when = parse_ts(&created);
Memory {
id: NOTE_NS + id,
text: label,
created_at: when,
last_used: when,
mentions: 1,
importance: 0.55,
tags: vec!["note".into()],
embedding: None,
}
})
.collect()
}
fn sessions(conn: &Connection) -> Vec<Memory> {
let Ok(mut stmt) = conn.prepare("SELECT id, notes, transcript, started_at FROM sessions") else {
return Vec::new();
};
let rows = stmt.query_map([], |r| {
Ok((
r.get::<_, i64>(0)?,
r.get::<_, String>(1)?,
r.get::<_, String>(2)?,
r.get::<_, String>(3)?,
))
});
let Ok(rows) = rows else { return Vec::new() };
rows.flatten()
.filter_map(|(id, notes, transcript, started)| {
let text = format!("{notes}\n{transcript}").trim().to_string();
if text.is_empty() {
return None;
}
let when = parse_ts(&started);
Some(Memory {
id: SESSION_NS + id,
text,
created_at: when,
last_used: when,
mentions: 1,
importance: 0.5,
tags: vec!["session".into()],
embedding: None,
})
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn fixture() -> Connection {
let conn = Connection::open_in_memory().unwrap();
conn.execute_batch(
"CREATE TABLE command_runs (id INTEGER PRIMARY KEY, command TEXT, argv_json TEXT,
notes TEXT, stderr TEXT, exit_code INTEGER, created_at TEXT);
CREATE TABLE bookmarks (id INTEGER PRIMARY KEY, label TEXT, created_at TEXT);
INSERT INTO command_runs VALUES
(1, 'cargo', '[\"cargo\",\"build\"]', '', 'error: linker failed', 1, '2026-06-20T12:00:00+00:00'),
(2, 'cargo', '[\"cargo\",\"build\"]', '', '', 0, '2026-06-26T12:00:00+00:00');
INSERT INTO bookmarks VALUES
(1, 'the linker failure was a missing -lstdc++', '2026-06-25T12:00:00+00:00');",
)
.unwrap();
conn
}
#[test]
fn tolerates_missing_tables_and_namespaces_ids() {
let mems = load_memories(&fixture());
assert_eq!(mems.len(), 3);
assert!(mems.iter().any(|m| decode_id(m.id) == (Source::Command, 1)));
assert!(mems.iter().any(|m| decode_id(m.id) == (Source::Note, 1)));
let cmd = mems.iter().find(|m| decode_id(m.id) == (Source::Command, 1)).unwrap();
assert_eq!(cmd.mentions, 2);
}
#[test]
fn decode_roundtrips_each_source() {
assert_eq!(decode_id(5), (Source::Document, 5));
assert_eq!(decode_id(CMD_NS + 5), (Source::Command, 5));
assert_eq!(decode_id(TURN_NS + 5), (Source::Conversation, 5));
assert_eq!(decode_id(NOTE_NS + 5), (Source::Note, 5));
assert_eq!(decode_id(SESSION_NS + 5), (Source::Session, 5));
}
#[test]
fn cli_and_desktop_rank_identically_for_fixed_now() {
use crate::engine::{BuiltinEngine, MemoryEngine};
use crate::Query;
use chrono::{TimeZone, Utc};
let conn = fixture();
let now = Utc.with_ymd_and_hms(2026, 6, 27, 12, 0, 0).unwrap();
let cli = BuiltinEngine::new(load_memories(&conn));
let desktop = BuiltinEngine::new(load_memories(&conn));
let q = Query::new("linker failure", now);
let cli_rank: Vec<(i64, u32)> =
cli.retrieve(&q, 20).iter().map(|s| (s.memory.id, s.percent())).collect();
let desktop_rank: Vec<(i64, u32)> =
desktop.retrieve(&q, 20).iter().map(|s| (s.memory.id, s.percent())).collect();
assert_eq!(cli_rank, desktop_rank, "CLI and desktop must rank identically");
assert_eq!(decode_id(cli_rank[0].0).0, Source::Note);
}
}