fix: replace std::sync::Mutex with parking_lot::Mutex (#350)
Merges #422
This commit is contained in:
Argenis
GitHub
parent
bff0507132
commit
15e1d50a5d
12 changed files with 1595 additions and 17 deletions
|
|
@ -5,6 +5,8 @@ pub mod hygiene;
|
|||
pub mod lucid;
|
||||
pub mod markdown;
|
||||
pub mod none;
|
||||
pub mod response_cache;
|
||||
pub mod snapshot;
|
||||
pub mod sqlite;
|
||||
pub mod traits;
|
||||
pub mod vector;
|
||||
|
|
@ -17,6 +19,7 @@ pub use backend::{
|
|||
pub use lucid::LucidMemory;
|
||||
pub use markdown::MarkdownMemory;
|
||||
pub use none::NoneMemory;
|
||||
pub use response_cache::ResponseCache;
|
||||
pub use sqlite::SqliteMemory;
|
||||
pub use traits::Memory;
|
||||
#[allow(unused_imports)]
|
||||
|
|
@ -63,6 +66,32 @@ pub fn create_memory(
|
|||
tracing::warn!("memory hygiene skipped: {e}");
|
||||
}
|
||||
|
||||
// If snapshot_on_hygiene is enabled, export core memories during hygiene.
|
||||
if config.snapshot_enabled && config.snapshot_on_hygiene {
|
||||
if let Err(e) = snapshot::export_snapshot(workspace_dir) {
|
||||
tracing::warn!("memory snapshot skipped: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
// Auto-hydration: if brain.db is missing but MEMORY_SNAPSHOT.md exists,
|
||||
// restore the "soul" from the snapshot before creating the backend.
|
||||
if config.auto_hydrate
|
||||
&& matches!(classify_memory_backend(&config.backend), MemoryBackendKind::Sqlite | MemoryBackendKind::Lucid)
|
||||
&& snapshot::should_hydrate(workspace_dir)
|
||||
{
|
||||
tracing::info!("🧬 Cold boot detected — hydrating from MEMORY_SNAPSHOT.md");
|
||||
match snapshot::hydrate_from_snapshot(workspace_dir) {
|
||||
Ok(count) => {
|
||||
if count > 0 {
|
||||
tracing::info!("🧬 Hydrated {count} core memories from snapshot");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("memory hydration failed: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_sqlite_memory(
|
||||
config: &MemoryConfig,
|
||||
workspace_dir: &Path,
|
||||
|
|
@ -113,6 +142,35 @@ pub fn create_memory_for_migration(
|
|||
)
|
||||
}
|
||||
|
||||
/// Factory: create an optional response cache from config.
|
||||
pub fn create_response_cache(
|
||||
config: &MemoryConfig,
|
||||
workspace_dir: &Path,
|
||||
) -> Option<ResponseCache> {
|
||||
if !config.response_cache_enabled {
|
||||
return None;
|
||||
}
|
||||
|
||||
match ResponseCache::new(
|
||||
workspace_dir,
|
||||
config.response_cache_ttl_minutes,
|
||||
config.response_cache_max_entries,
|
||||
) {
|
||||
Ok(cache) => {
|
||||
tracing::info!(
|
||||
"💾 Response cache enabled (TTL: {}min, max: {} entries)",
|
||||
config.response_cache_ttl_minutes,
|
||||
config.response_cache_max_entries
|
||||
);
|
||||
Some(cache)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Response cache disabled due to error: {e}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
|
|||
371
src/memory/response_cache.rs
Normal file
371
src/memory/response_cache.rs
Normal file
|
|
@ -0,0 +1,371 @@
|
|||
//! Response cache — avoid burning tokens on repeated prompts.
|
||||
//!
|
||||
//! Stores LLM responses in a separate SQLite table keyed by a SHA-256 hash of
|
||||
//! `(model, system_prompt_hash, user_prompt)`. Entries expire after a
|
||||
//! configurable TTL (default: 1 hour). The cache is optional and disabled by
|
||||
//! default — users opt in via `[memory] response_cache_enabled = true`.
|
||||
|
||||
use anyhow::Result;
|
||||
use chrono::{Duration, Local};
|
||||
use rusqlite::{params, Connection};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Mutex;
|
||||
|
||||
/// Response cache backed by a dedicated SQLite database.
|
||||
///
|
||||
/// Lives alongside `brain.db` as `response_cache.db` so it can be
|
||||
/// independently wiped without touching memories.
|
||||
pub struct ResponseCache {
|
||||
conn: Mutex<Connection>,
|
||||
#[allow(dead_code)]
|
||||
db_path: PathBuf,
|
||||
ttl_minutes: i64,
|
||||
max_entries: usize,
|
||||
}
|
||||
|
||||
impl ResponseCache {
|
||||
/// Open (or create) the response cache database.
|
||||
pub fn new(workspace_dir: &Path, ttl_minutes: u32, max_entries: usize) -> Result<Self> {
|
||||
let db_dir = workspace_dir.join("memory");
|
||||
std::fs::create_dir_all(&db_dir)?;
|
||||
let db_path = db_dir.join("response_cache.db");
|
||||
|
||||
let conn = Connection::open(&db_path)?;
|
||||
|
||||
conn.execute_batch(
|
||||
"PRAGMA journal_mode = WAL;
|
||||
PRAGMA synchronous = NORMAL;
|
||||
PRAGMA temp_store = MEMORY;",
|
||||
)?;
|
||||
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS response_cache (
|
||||
prompt_hash TEXT PRIMARY KEY,
|
||||
model TEXT NOT NULL,
|
||||
response TEXT NOT NULL,
|
||||
token_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL,
|
||||
accessed_at TEXT NOT NULL,
|
||||
hit_count INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_rc_accessed ON response_cache(accessed_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_rc_created ON response_cache(created_at);",
|
||||
)?;
|
||||
|
||||
Ok(Self {
|
||||
conn: Mutex::new(conn),
|
||||
db_path,
|
||||
ttl_minutes: i64::from(ttl_minutes),
|
||||
max_entries,
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a deterministic cache key from model + system prompt + user prompt.
|
||||
pub fn cache_key(model: &str, system_prompt: Option<&str>, user_prompt: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(model.as_bytes());
|
||||
hasher.update(b"|");
|
||||
if let Some(sys) = system_prompt {
|
||||
hasher.update(sys.as_bytes());
|
||||
}
|
||||
hasher.update(b"|");
|
||||
hasher.update(user_prompt.as_bytes());
|
||||
let hash = hasher.finalize();
|
||||
format!("{:064x}", hash)
|
||||
}
|
||||
|
||||
/// Look up a cached response. Returns `None` on miss or expired entry.
|
||||
pub fn get(&self, key: &str) -> Result<Option<String>> {
|
||||
let conn = self
|
||||
.conn
|
||||
.lock()
|
||||
.map_err(|e| anyhow::anyhow!("Lock error: {e}"))?;
|
||||
|
||||
let now = Local::now();
|
||||
let cutoff = (now - Duration::minutes(self.ttl_minutes)).to_rfc3339();
|
||||
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT response FROM response_cache
|
||||
WHERE prompt_hash = ?1 AND created_at > ?2",
|
||||
)?;
|
||||
|
||||
let result: Option<String> = stmt
|
||||
.query_row(params![key, cutoff], |row| row.get(0))
|
||||
.ok();
|
||||
|
||||
if result.is_some() {
|
||||
// Bump hit count and accessed_at
|
||||
let now_str = now.to_rfc3339();
|
||||
conn.execute(
|
||||
"UPDATE response_cache
|
||||
SET accessed_at = ?1, hit_count = hit_count + 1
|
||||
WHERE prompt_hash = ?2",
|
||||
params![now_str, key],
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Store a response in the cache.
|
||||
pub fn put(
|
||||
&self,
|
||||
key: &str,
|
||||
model: &str,
|
||||
response: &str,
|
||||
token_count: u32,
|
||||
) -> Result<()> {
|
||||
let conn = self
|
||||
.conn
|
||||
.lock()
|
||||
.map_err(|e| anyhow::anyhow!("Lock error: {e}"))?;
|
||||
|
||||
let now = Local::now().to_rfc3339();
|
||||
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO response_cache
|
||||
(prompt_hash, model, response, token_count, created_at, accessed_at, hit_count)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, 0)",
|
||||
params![key, model, response, token_count, now, now],
|
||||
)?;
|
||||
|
||||
// Evict expired entries
|
||||
let cutoff = (Local::now() - Duration::minutes(self.ttl_minutes)).to_rfc3339();
|
||||
conn.execute(
|
||||
"DELETE FROM response_cache WHERE created_at <= ?1",
|
||||
params![cutoff],
|
||||
)?;
|
||||
|
||||
// LRU eviction if over max_entries
|
||||
#[allow(clippy::cast_possible_wrap)]
|
||||
let max = self.max_entries as i64;
|
||||
conn.execute(
|
||||
"DELETE FROM response_cache WHERE prompt_hash IN (
|
||||
SELECT prompt_hash FROM response_cache
|
||||
ORDER BY accessed_at ASC
|
||||
LIMIT MAX(0, (SELECT COUNT(*) FROM response_cache) - ?1)
|
||||
)",
|
||||
params![max],
|
||||
)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Return cache statistics: (total_entries, total_hits, total_tokens_saved).
|
||||
pub fn stats(&self) -> Result<(usize, u64, u64)> {
|
||||
let conn = self
|
||||
.conn
|
||||
.lock()
|
||||
.map_err(|e| anyhow::anyhow!("Lock error: {e}"))?;
|
||||
|
||||
let count: i64 =
|
||||
conn.query_row("SELECT COUNT(*) FROM response_cache", [], |row| row.get(0))?;
|
||||
|
||||
let hits: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COALESCE(SUM(hit_count), 0) FROM response_cache",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
|
||||
let tokens_saved: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COALESCE(SUM(token_count * hit_count), 0) FROM response_cache",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
|
||||
#[allow(clippy::cast_sign_loss)]
|
||||
Ok((count as usize, hits as u64, tokens_saved as u64))
|
||||
}
|
||||
|
||||
/// Wipe the entire cache (useful for `zeroclaw cache clear`).
|
||||
pub fn clear(&self) -> Result<usize> {
|
||||
let conn = self
|
||||
.conn
|
||||
.lock()
|
||||
.map_err(|e| anyhow::anyhow!("Lock error: {e}"))?;
|
||||
|
||||
let affected = conn.execute("DELETE FROM response_cache", [])?;
|
||||
Ok(affected)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn temp_cache(ttl_minutes: u32) -> (TempDir, ResponseCache) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cache = ResponseCache::new(tmp.path(), ttl_minutes, 1000).unwrap();
|
||||
(tmp, cache)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_deterministic() {
|
||||
let k1 = ResponseCache::cache_key("gpt-4", Some("sys"), "hello");
|
||||
let k2 = ResponseCache::cache_key("gpt-4", Some("sys"), "hello");
|
||||
assert_eq!(k1, k2);
|
||||
assert_eq!(k1.len(), 64); // SHA-256 hex
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_varies_by_model() {
|
||||
let k1 = ResponseCache::cache_key("gpt-4", None, "hello");
|
||||
let k2 = ResponseCache::cache_key("claude-3", None, "hello");
|
||||
assert_ne!(k1, k2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_varies_by_system_prompt() {
|
||||
let k1 = ResponseCache::cache_key("gpt-4", Some("You are helpful"), "hello");
|
||||
let k2 = ResponseCache::cache_key("gpt-4", Some("You are rude"), "hello");
|
||||
assert_ne!(k1, k2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_varies_by_prompt() {
|
||||
let k1 = ResponseCache::cache_key("gpt-4", None, "hello");
|
||||
let k2 = ResponseCache::cache_key("gpt-4", None, "goodbye");
|
||||
assert_ne!(k1, k2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn put_and_get() {
|
||||
let (_tmp, cache) = temp_cache(60);
|
||||
let key = ResponseCache::cache_key("gpt-4", None, "What is Rust?");
|
||||
|
||||
cache
|
||||
.put(&key, "gpt-4", "Rust is a systems programming language.", 25)
|
||||
.unwrap();
|
||||
|
||||
let result = cache.get(&key).unwrap();
|
||||
assert_eq!(
|
||||
result.as_deref(),
|
||||
Some("Rust is a systems programming language.")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn miss_returns_none() {
|
||||
let (_tmp, cache) = temp_cache(60);
|
||||
let result = cache.get("nonexistent_key").unwrap();
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_entry_returns_none() {
|
||||
let (_tmp, cache) = temp_cache(0); // 0-minute TTL → everything is instantly expired
|
||||
let key = ResponseCache::cache_key("gpt-4", None, "test");
|
||||
|
||||
cache.put(&key, "gpt-4", "response", 10).unwrap();
|
||||
|
||||
// The entry was created with created_at = now(), but TTL is 0 minutes,
|
||||
// so cutoff = now() - 0 = now(). The entry's created_at is NOT > cutoff.
|
||||
let result = cache.get(&key).unwrap();
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hit_count_incremented() {
|
||||
let (_tmp, cache) = temp_cache(60);
|
||||
let key = ResponseCache::cache_key("gpt-4", None, "hello");
|
||||
|
||||
cache.put(&key, "gpt-4", "Hi!", 5).unwrap();
|
||||
|
||||
// 3 hits
|
||||
for _ in 0..3 {
|
||||
let _ = cache.get(&key).unwrap();
|
||||
}
|
||||
|
||||
let (_, total_hits, _) = cache.stats().unwrap();
|
||||
assert_eq!(total_hits, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tokens_saved_calculated() {
|
||||
let (_tmp, cache) = temp_cache(60);
|
||||
let key = ResponseCache::cache_key("gpt-4", None, "explain rust");
|
||||
|
||||
cache.put(&key, "gpt-4", "Rust is...", 100).unwrap();
|
||||
|
||||
// 5 cache hits × 100 tokens = 500 tokens saved
|
||||
for _ in 0..5 {
|
||||
let _ = cache.get(&key).unwrap();
|
||||
}
|
||||
|
||||
let (_, _, tokens_saved) = cache.stats().unwrap();
|
||||
assert_eq!(tokens_saved, 500);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lru_eviction() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cache = ResponseCache::new(tmp.path(), 60, 3).unwrap(); // max 3 entries
|
||||
|
||||
for i in 0..5 {
|
||||
let key = ResponseCache::cache_key("gpt-4", None, &format!("prompt {i}"));
|
||||
cache
|
||||
.put(&key, "gpt-4", &format!("response {i}"), 10)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let (count, _, _) = cache.stats().unwrap();
|
||||
assert!(count <= 3, "Should have at most 3 entries after eviction");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear_wipes_all() {
|
||||
let (_tmp, cache) = temp_cache(60);
|
||||
|
||||
for i in 0..10 {
|
||||
let key = ResponseCache::cache_key("gpt-4", None, &format!("prompt {i}"));
|
||||
cache
|
||||
.put(&key, "gpt-4", &format!("response {i}"), 10)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let cleared = cache.clear().unwrap();
|
||||
assert_eq!(cleared, 10);
|
||||
|
||||
let (count, _, _) = cache.stats().unwrap();
|
||||
assert_eq!(count, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stats_empty_cache() {
|
||||
let (_tmp, cache) = temp_cache(60);
|
||||
let (count, hits, tokens) = cache.stats().unwrap();
|
||||
assert_eq!(count, 0);
|
||||
assert_eq!(hits, 0);
|
||||
assert_eq!(tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overwrite_same_key() {
|
||||
let (_tmp, cache) = temp_cache(60);
|
||||
let key = ResponseCache::cache_key("gpt-4", None, "question");
|
||||
|
||||
cache.put(&key, "gpt-4", "answer v1", 20).unwrap();
|
||||
cache.put(&key, "gpt-4", "answer v2", 25).unwrap();
|
||||
|
||||
let result = cache.get(&key).unwrap();
|
||||
assert_eq!(result.as_deref(), Some("answer v2"));
|
||||
|
||||
let (count, _, _) = cache.stats().unwrap();
|
||||
assert_eq!(count, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unicode_prompt_handling() {
|
||||
let (_tmp, cache) = temp_cache(60);
|
||||
let key = ResponseCache::cache_key("gpt-4", None, "日本語のテスト 🦀");
|
||||
|
||||
cache.put(&key, "gpt-4", "はい、Rustは素晴らしい", 30).unwrap();
|
||||
|
||||
let result = cache.get(&key).unwrap();
|
||||
assert_eq!(result.as_deref(), Some("はい、Rustは素晴らしい"));
|
||||
}
|
||||
}
|
||||
467
src/memory/snapshot.rs
Normal file
467
src/memory/snapshot.rs
Normal file
|
|
@ -0,0 +1,467 @@
|
|||
//! Memory snapshot — export/import core memories as human-readable Markdown.
|
||||
//!
|
||||
//! **Atomic Soul Export**: dumps `MemoryCategory::Core` from SQLite into
|
||||
//! `MEMORY_SNAPSHOT.md` so the agent's "soul" is always Git-visible.
|
||||
//!
|
||||
//! **Auto-Hydration**: if `brain.db` is missing but `MEMORY_SNAPSHOT.md` exists,
|
||||
//! re-indexes all entries back into a fresh SQLite database.
|
||||
|
||||
use anyhow::Result;
|
||||
use chrono::Local;
|
||||
use rusqlite::{params, Connection};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Filename for the snapshot (lives at workspace root for Git visibility).
|
||||
pub const SNAPSHOT_FILENAME: &str = "MEMORY_SNAPSHOT.md";
|
||||
|
||||
/// Header written at the top of every snapshot file.
|
||||
const SNAPSHOT_HEADER: &str = "# 🧠 ZeroClaw Memory Snapshot\n\n\
|
||||
> Auto-generated by ZeroClaw. Do not edit manually unless you know what you're doing.\n\
|
||||
> This file is the \"soul\" of your agent — if `brain.db` is lost, start the agent\n\
|
||||
> in this workspace and it will auto-hydrate from this file.\n\n";
|
||||
|
||||
/// Export all `Core` memories from SQLite → `MEMORY_SNAPSHOT.md`.
|
||||
///
|
||||
/// Returns the number of entries exported.
|
||||
pub fn export_snapshot(workspace_dir: &Path) -> Result<usize> {
|
||||
let db_path = workspace_dir.join("memory").join("brain.db");
|
||||
if !db_path.exists() {
|
||||
tracing::debug!("snapshot export skipped: brain.db does not exist");
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let conn = Connection::open(&db_path)?;
|
||||
conn.execute_batch("PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL;")?;
|
||||
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT key, content, category, created_at, updated_at
|
||||
FROM memories
|
||||
WHERE category = 'core'
|
||||
ORDER BY updated_at DESC",
|
||||
)?;
|
||||
|
||||
let rows: Vec<(String, String, String, String, String)> = stmt
|
||||
.query_map([], |row| {
|
||||
Ok((
|
||||
row.get(0)?,
|
||||
row.get(1)?,
|
||||
row.get(2)?,
|
||||
row.get(3)?,
|
||||
row.get(4)?,
|
||||
))
|
||||
})?
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
|
||||
if rows.is_empty() {
|
||||
tracing::debug!("snapshot export: no core memories to export");
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let mut output = String::with_capacity(rows.len() * 200);
|
||||
output.push_str(SNAPSHOT_HEADER);
|
||||
|
||||
let now = Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
||||
output.push_str(&format!("**Last exported:** {now}\n\n"));
|
||||
output.push_str(&format!("**Total core memories:** {}\n\n---\n\n", rows.len()));
|
||||
|
||||
for (key, content, _category, created_at, updated_at) in &rows {
|
||||
output.push_str(&format!("### 🔑 `{key}`\n\n"));
|
||||
output.push_str(&format!("{content}\n\n"));
|
||||
output.push_str(&format!(
|
||||
"*Created: {created_at} | Updated: {updated_at}*\n\n---\n\n"
|
||||
));
|
||||
}
|
||||
|
||||
let snapshot_path = snapshot_path(workspace_dir);
|
||||
fs::write(&snapshot_path, output)?;
|
||||
|
||||
tracing::info!(
|
||||
"📸 Memory snapshot exported: {} core memories → {}",
|
||||
rows.len(),
|
||||
snapshot_path.display()
|
||||
);
|
||||
|
||||
Ok(rows.len())
|
||||
}
|
||||
|
||||
/// Import memories from `MEMORY_SNAPSHOT.md` into SQLite.
|
||||
///
|
||||
/// Called during cold-boot when `brain.db` doesn't exist but the snapshot does.
|
||||
/// Returns the number of entries hydrated.
|
||||
pub fn hydrate_from_snapshot(workspace_dir: &Path) -> Result<usize> {
|
||||
let snapshot = snapshot_path(workspace_dir);
|
||||
if !snapshot.exists() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let content = fs::read_to_string(&snapshot)?;
|
||||
let entries = parse_snapshot(&content);
|
||||
|
||||
if entries.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// Ensure the memory directory exists
|
||||
let db_dir = workspace_dir.join("memory");
|
||||
fs::create_dir_all(&db_dir)?;
|
||||
|
||||
let db_path = db_dir.join("brain.db");
|
||||
let conn = Connection::open(&db_path)?;
|
||||
conn.execute_batch("PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL;")?;
|
||||
|
||||
// Initialize schema (same as SqliteMemory::init_schema)
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS memories (
|
||||
id TEXT PRIMARY KEY,
|
||||
key TEXT NOT NULL UNIQUE,
|
||||
content TEXT NOT NULL,
|
||||
category TEXT NOT NULL DEFAULT 'core',
|
||||
embedding BLOB,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_key ON memories(key);
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_cat ON memories(category);
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_updated ON memories(updated_at);
|
||||
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts
|
||||
USING fts5(key, content, content='memories', content_rowid='rowid');
|
||||
|
||||
CREATE TABLE IF NOT EXISTS embedding_cache (
|
||||
content_hash TEXT PRIMARY KEY,
|
||||
embedding BLOB NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);",
|
||||
)?;
|
||||
|
||||
let now = Local::now().to_rfc3339();
|
||||
let mut hydrated = 0;
|
||||
|
||||
for (key, content) in &entries {
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
let result = conn.execute(
|
||||
"INSERT OR IGNORE INTO memories (id, key, content, category, created_at, updated_at)
|
||||
VALUES (?1, ?2, ?3, 'core', ?4, ?5)",
|
||||
params![id, key, content, now, now],
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok(changed) if changed > 0 => {
|
||||
// Populate FTS5
|
||||
let _ = conn.execute(
|
||||
"INSERT INTO memories_fts(key, content) VALUES (?1, ?2)",
|
||||
params![key, content],
|
||||
);
|
||||
hydrated += 1;
|
||||
}
|
||||
Ok(_) => {
|
||||
tracing::debug!("hydrate: key '{key}' already exists, skipping");
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("hydrate: failed to insert key '{key}': {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"🧬 Memory hydration complete: {} entries restored from {}",
|
||||
hydrated,
|
||||
snapshot.display()
|
||||
);
|
||||
|
||||
Ok(hydrated)
|
||||
}
|
||||
|
||||
/// Check if we should auto-hydrate on startup.
|
||||
///
|
||||
/// Returns `true` if:
|
||||
/// 1. `brain.db` does NOT exist (or is empty)
|
||||
/// 2. `MEMORY_SNAPSHOT.md` DOES exist
|
||||
pub fn should_hydrate(workspace_dir: &Path) -> bool {
|
||||
let db_path = workspace_dir.join("memory").join("brain.db");
|
||||
let snapshot = snapshot_path(workspace_dir);
|
||||
|
||||
let db_missing_or_empty = if db_path.exists() {
|
||||
// DB exists but might be empty (freshly created)
|
||||
fs::metadata(&db_path)
|
||||
.map(|m| m.len() < 4096) // SQLite header is ~4096 bytes minimum
|
||||
.unwrap_or(true)
|
||||
} else {
|
||||
true
|
||||
};
|
||||
|
||||
db_missing_or_empty && snapshot.exists()
|
||||
}
|
||||
|
||||
/// Path to the snapshot file.
|
||||
fn snapshot_path(workspace_dir: &Path) -> PathBuf {
|
||||
workspace_dir.join(SNAPSHOT_FILENAME)
|
||||
}
|
||||
|
||||
/// Parse the structured markdown snapshot back into (key, content) pairs.
|
||||
fn parse_snapshot(input: &str) -> Vec<(String, String)> {
|
||||
let mut entries = Vec::new();
|
||||
let mut current_key: Option<String> = None;
|
||||
let mut current_content = String::new();
|
||||
|
||||
for line in input.lines() {
|
||||
let trimmed = line.trim();
|
||||
|
||||
// Match: ### 🔑 `key_name`
|
||||
if trimmed.starts_with("### 🔑 `") && trimmed.ends_with('`') {
|
||||
// Save previous entry
|
||||
if let Some(key) = current_key.take() {
|
||||
let content = current_content.trim().to_string();
|
||||
if !content.is_empty() {
|
||||
entries.push((key, content));
|
||||
}
|
||||
}
|
||||
|
||||
// Extract new key
|
||||
let key = trimmed
|
||||
.strip_prefix("### 🔑 `")
|
||||
.and_then(|s| s.strip_suffix('`'))
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if !key.is_empty() {
|
||||
current_key = Some(key);
|
||||
current_content = String::new();
|
||||
}
|
||||
} else if current_key.is_some() {
|
||||
// Skip metadata lines and separators
|
||||
if trimmed.starts_with("*Created:") || trimmed == "---" {
|
||||
continue;
|
||||
}
|
||||
// Accumulate content
|
||||
if !current_content.is_empty() || !trimmed.is_empty() {
|
||||
if !current_content.is_empty() {
|
||||
current_content.push('\n');
|
||||
}
|
||||
current_content.push_str(line);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Don't forget the last entry
|
||||
if let Some(key) = current_key {
|
||||
let content = current_content.trim().to_string();
|
||||
if !content.is_empty() {
|
||||
entries.push((key, content));
|
||||
}
|
||||
}
|
||||
|
||||
entries
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn parse_snapshot_basic() {
|
||||
let input = r#"# 🧠 ZeroClaw Memory Snapshot
|
||||
|
||||
> Auto-generated by ZeroClaw.
|
||||
|
||||
**Last exported:** 2025-01-15 14:30:00
|
||||
|
||||
**Total core memories:** 2
|
||||
|
||||
---
|
||||
|
||||
### 🔑 `identity`
|
||||
|
||||
I am ZeroClaw, a self-preserving AI agent.
|
||||
|
||||
*Created: 2025-01-15 | Updated: 2025-01-15*
|
||||
|
||||
---
|
||||
|
||||
### 🔑 `preference_lang`
|
||||
|
||||
The user prefers Rust for systems programming.
|
||||
|
||||
*Created: 2025-01-14 | Updated: 2025-01-15*
|
||||
|
||||
---
|
||||
"#;
|
||||
|
||||
let entries = parse_snapshot(input);
|
||||
assert_eq!(entries.len(), 2);
|
||||
assert_eq!(entries[0].0, "identity");
|
||||
assert!(entries[0].1.contains("self-preserving"));
|
||||
assert_eq!(entries[1].0, "preference_lang");
|
||||
assert!(entries[1].1.contains("Rust"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_snapshot_empty() {
|
||||
let input = "# 🧠 ZeroClaw Memory Snapshot\n\n> Nothing here.\n";
|
||||
let entries = parse_snapshot(input);
|
||||
assert!(entries.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_snapshot_multiline_content() {
|
||||
let input = r#"### 🔑 `rules`
|
||||
|
||||
Rule 1: Always be helpful.
|
||||
Rule 2: Never lie.
|
||||
Rule 3: Protect the user.
|
||||
|
||||
*Created: 2025-01-15 | Updated: 2025-01-15*
|
||||
|
||||
---
|
||||
"#;
|
||||
|
||||
let entries = parse_snapshot(input);
|
||||
assert_eq!(entries.len(), 1);
|
||||
assert!(entries[0].1.contains("Rule 1"));
|
||||
assert!(entries[0].1.contains("Rule 3"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn export_no_db_returns_zero() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let count = export_snapshot(tmp.path()).unwrap();
|
||||
assert_eq!(count, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn export_and_hydrate_roundtrip() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let workspace = tmp.path();
|
||||
|
||||
// Create a brain.db manually with some core memories
|
||||
let db_dir = workspace.join("memory");
|
||||
fs::create_dir_all(&db_dir).unwrap();
|
||||
let db_path = db_dir.join("brain.db");
|
||||
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute_batch(
|
||||
"PRAGMA journal_mode = WAL;
|
||||
CREATE TABLE IF NOT EXISTS memories (
|
||||
id TEXT PRIMARY KEY,
|
||||
key TEXT NOT NULL UNIQUE,
|
||||
content TEXT NOT NULL,
|
||||
category TEXT NOT NULL DEFAULT 'core',
|
||||
embedding BLOB,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_key ON memories(key);",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let now = Local::now().to_rfc3339();
|
||||
conn.execute(
|
||||
"INSERT INTO memories (id, key, content, category, created_at, updated_at)
|
||||
VALUES ('id1', 'identity', 'I am a test agent', 'core', ?1, ?2)",
|
||||
params![now, now],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO memories (id, key, content, category, created_at, updated_at)
|
||||
VALUES ('id2', 'preference', 'User likes Rust', 'core', ?1, ?2)",
|
||||
params![now, now],
|
||||
)
|
||||
.unwrap();
|
||||
// Non-core entry (should NOT be exported)
|
||||
conn.execute(
|
||||
"INSERT INTO memories (id, key, content, category, created_at, updated_at)
|
||||
VALUES ('id3', 'conv1', 'Random convo', 'conversation', ?1, ?2)",
|
||||
params![now, now],
|
||||
)
|
||||
.unwrap();
|
||||
drop(conn);
|
||||
|
||||
// Export snapshot
|
||||
let exported = export_snapshot(workspace).unwrap();
|
||||
assert_eq!(exported, 2, "Should export only core memories");
|
||||
|
||||
// Verify the file exists and is readable
|
||||
let snapshot = workspace.join(SNAPSHOT_FILENAME);
|
||||
assert!(snapshot.exists());
|
||||
let content = fs::read_to_string(&snapshot).unwrap();
|
||||
assert!(content.contains("identity"));
|
||||
assert!(content.contains("I am a test agent"));
|
||||
assert!(content.contains("preference"));
|
||||
assert!(!content.contains("Random convo"));
|
||||
|
||||
// Simulate catastrophic failure: delete brain.db
|
||||
fs::remove_file(&db_path).unwrap();
|
||||
assert!(!db_path.exists());
|
||||
|
||||
// Verify should_hydrate detects the scenario
|
||||
assert!(should_hydrate(workspace));
|
||||
|
||||
// Hydrate from snapshot
|
||||
let hydrated = hydrate_from_snapshot(workspace).unwrap();
|
||||
assert_eq!(hydrated, 2, "Should hydrate both core memories");
|
||||
|
||||
// Verify brain.db was recreated
|
||||
assert!(db_path.exists());
|
||||
|
||||
// Verify the data is actually in the new database
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
let count: i64 = conn
|
||||
.query_row("SELECT COUNT(*) FROM memories", [], |row| row.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(count, 2);
|
||||
|
||||
let identity: String = conn
|
||||
.query_row(
|
||||
"SELECT content FROM memories WHERE key = 'identity'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(identity, "I am a test agent");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_hydrate_only_when_needed() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let workspace = tmp.path();
|
||||
|
||||
// No DB, no snapshot → false
|
||||
assert!(!should_hydrate(workspace));
|
||||
|
||||
// Create snapshot but no DB → true
|
||||
let snapshot = workspace.join(SNAPSHOT_FILENAME);
|
||||
fs::write(&snapshot, "### 🔑 `test`\n\nHello\n").unwrap();
|
||||
assert!(should_hydrate(workspace));
|
||||
|
||||
// Create a real DB → false
|
||||
let db_dir = workspace.join("memory");
|
||||
fs::create_dir_all(&db_dir).unwrap();
|
||||
let db_path = db_dir.join("brain.db");
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS memories (
|
||||
id TEXT PRIMARY KEY,
|
||||
key TEXT NOT NULL UNIQUE,
|
||||
content TEXT NOT NULL,
|
||||
category TEXT NOT NULL DEFAULT 'core',
|
||||
embedding BLOB,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
INSERT INTO memories VALUES('x','x','x','core',NULL,'2025-01-01','2025-01-01');",
|
||||
)
|
||||
.unwrap();
|
||||
drop(conn);
|
||||
assert!(!should_hydrate(workspace));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hydrate_no_snapshot_returns_zero() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let count = hydrate_from_snapshot(tmp.path()).unwrap();
|
||||
assert_eq!(count, 0);
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue