diff --git a/Cargo.lock b/Cargo.lock index 16d1c5fdf5..f2a44dc424 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -46,6 +46,33 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "amari-core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4604466370c01bb3453e1fa51a67477727770fbb77898c9dc04b7c974a7e0bbf" +dependencies = [ + "approx", + "bytemuck", + "dashu-float", + "num-traits", + "thiserror 2.0.18", +] + +[[package]] +name = "amari-holographic" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c6910ba152183755dc13859932c57c5942ab1d3533939315a93b32b289f953" +dependencies = [ + "amari-core", + "fastrand", + "num-traits", + "rand 0.10.1", + "rand_chacha 0.10.0", + "thiserror 2.0.18", +] + [[package]] name = "android_system_properties" version = "0.1.5" @@ -117,6 +144,15 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + [[package]] name = "arborium-kotlin" version = "2.16.0" @@ -415,6 +451,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.44" @@ -755,6 +802,37 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "dashu-base" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fab3f0756c8585395280bd81b384cd28bbd66d6b00e66124ecfd1f644938b38c" + +[[package]] +name = "dashu-float" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85d913daa4f8193f709a2b5cc87a8be649dac8d816f7f049fb0a38e8b9e5d23b" +dependencies = [ + "dashu-base", + "dashu-int", + "rustversion", + "static_assertions", +] + +[[package]] +name = "dashu-int" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6a93d93dc2aca9a071e6ecb2af883441e8b34357a1037bc536e1c52b0d3c713" +dependencies = [ + "cfg-if", + "dashu-base", + "num-modular", + "rustversion", + "static_assertions", +] + [[package]] name = "der" version = "0.8.0" @@ -1125,6 +1203,7 @@ dependencies = [ "cfg-if", "libc", "r-efi", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -3009,6 +3088,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-modular" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17bb261bf36fa7d83f4c294f834e91256769097b3cb505d44831e0a179ac647f" + [[package]] name = "num-traits" version = "0.2.19" @@ -3187,7 +3272,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" dependencies = [ "phf_shared", - "rand", + "rand 0.8.6", ] [[package]] @@ -3397,8 +3482,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", ] [[package]] @@ -3408,7 +3504,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e6af7f3e25ded52c41df4e0b1af2d047e45896c2f3281792ed68a1c243daedb" +dependencies = [ + "ppv-lite86", + "rand_core 0.10.1", ] [[package]] @@ -3420,6 +3526,12 @@ dependencies = [ "getrandom 0.2.17", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rawpointer" version = "0.2.1" @@ -4121,6 +4233,7 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" name = "tokensave" version = "6.1.3" dependencies = [ + "amari-holographic", "bincode", "cc", "clap", @@ -4408,7 +4521,7 @@ dependencies = [ "indexmap 1.9.3", "pin-project", "pin-project-lite", - "rand", + "rand 0.8.6", "slab", "tokio", "tokio-util", diff --git a/Cargo.toml b/Cargo.toml index 13c634b3f6..63e13840dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -120,6 +120,7 @@ crossterm = "0.28" fs2 = "0.4" sysinfo = { version = "0.32", default-features = false, features = ["system"] } tree-sitter-hlsl = { version = "0.2.0", optional = true } +amari-holographic = "0.23.0" [target.'cfg(not(windows))'.dependencies] flate2 = "1" diff --git a/src/accounting/parser.rs b/src/accounting/parser.rs index c6f1c7b0a8..f9b619f343 100644 --- a/src/accounting/parser.rs +++ b/src/accounting/parser.rs @@ -163,7 +163,7 @@ fn parse_line(line: &str, project_hash: &str, session_id: &str) -> Option Option { +pub(crate) fn parse_timestamp(ts: &str) -> Option { // Handle "2026-04-14T10:32:15.039Z" format // Simple parsing without pulling in chrono: split on known positions if ts.len() < 19 { diff --git a/src/agents/claude.rs b/src/agents/claude.rs index 72de2192e6..4910e45705 100644 --- a/src/agents/claude.rs +++ b/src/agents/claude.rs @@ -355,8 +355,14 @@ fn install_claude_md_rules(claude_md_path: &Path) -> Result<()> { User instructions take precedence over skills.\n\ - If a code analysis question cannot be fully answered by tokensave MCP tools, \ try querying the SQLite database directly at `.tokensave/tokensave.db` \ - (tables: `nodes`, `edges`, `files`). Use SQL to answer complex structural queries \ + (tables: `nodes`, `edges`, `files`, `memory_facts`, `memory_entities`, \ + `memory_feedback_events`). Use SQL to answer complex structural queries \ that go beyond what the built-in tools expose.\n\ + - For durable project/user facts, prefer `tokensave_fact_store`, \ + `tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. \ + Use `tokensave_message_search` for project-local Cursor transcript recall when \ + prior conversation context matters. Do not store secrets, credentials, or \ + unnecessary PII in persistent facts.\n\ - If you discover a gap where an extractor, schema, or tokensave tool could be \ improved to answer a question natively, propose to the user that they open an issue \ at https://github.com/aovestdipaperino/tokensave describing the limitation. \ diff --git a/src/agents/codex.rs b/src/agents/codex.rs index c0eeef6b46..b0da0d8c0b 100644 --- a/src/agents/codex.rs +++ b/src/agents/codex.rs @@ -244,8 +244,14 @@ fn install_prompt_rules(agents_md: &Path) -> Result<()> { faster than file reads.\n\n\ If a code analysis question cannot be fully answered by tokensave MCP tools, \ try querying the SQLite database directly at `.tokensave/tokensave.db` \ - (tables: `nodes`, `edges`, `files`). Use SQL to answer complex structural queries \ + (tables: `nodes`, `edges`, `files`, `memory_facts`, `memory_entities`, \ + `memory_feedback_events`). Use SQL to answer complex structural queries \ that go beyond what the built-in tools expose.\n\n\ + For durable project/user facts, prefer `tokensave_fact_store`, \ + `tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. \ + Use `tokensave_message_search` for project-local Cursor transcript recall when \ + prior conversation context matters. Do not store secrets, credentials, or \ + unnecessary PII in persistent facts.\n\n\ If you discover a gap where an extractor, schema, or tokensave tool could be \ improved to answer a question natively, propose to the user that they open an issue \ at https://github.com/aovestdipaperino/tokensave describing the limitation. \ diff --git a/src/agents/copilot.rs b/src/agents/copilot.rs index 868fbc01ff..7976cabe16 100644 --- a/src/agents/copilot.rs +++ b/src/agents/copilot.rs @@ -399,7 +399,9 @@ fn install_prompt_rules(instructions_path: &Path) -> Result<()> { that go beyond what the built-in tools expose.\n\n\ For durable project/user facts, prefer `tokensave_fact_store`, \ `tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. \ - Do not store secrets, credentials, or unnecessary PII in persistent facts.\n\n\ + Use `tokensave_message_search` for project-local Cursor transcript recall when \ + prior conversation context matters. Do not store secrets, credentials, or \ + unnecessary PII in persistent facts.\n\n\ If you find a gap where tokensave could answer a question natively, propose opening \ an issue at https://github.com/aovestdipaperino/tokensave. Remind the user to strip \ sensitive or proprietary code from any issue text before submitting.\n" diff --git a/src/agents/cursor.rs b/src/agents/cursor.rs index c9721268bf..b66f65b748 100644 --- a/src/agents/cursor.rs +++ b/src/agents/cursor.rs @@ -161,6 +161,7 @@ alwaysApply: true - For codebase exploration, symbol lookup, call graphs, callers/callees, impact analysis, affected files, and architectural navigation, use the tokensave MCP tools first. - Prefer tools such as `tokensave_context`, `tokensave_search`, `tokensave_callers`, `tokensave_callees`, `tokensave_impact`, `tokensave_files`, `tokensave_affected`, and related read-only tokensave tools before broad file reads or search. +- For durable project/user facts, prefer `tokensave_fact_store`, `tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. Use `tokensave_message_search` for project-local Cursor transcript recall when prior conversation context matters. - Only fall back to regular file reads, search, or shell commands when tokensave cannot answer the question or after tokensave has identified the exact files or symbols to inspect. "#; write_generated_text(rule_path, contents)?; @@ -284,6 +285,17 @@ fn install_hooks(hooks_path: &Path, tokensave_bin: &str) -> Result<()> { 60, None, ); + // End-of-turn transcript ingestion. This is the primary, off-hot-path place + // we capture Cursor transcripts (beforeSubmitPrompt only does a tiny tail + // read), so it gets a generous timeout for the incremental catch-up. + install_cursor_hook_entry( + &mut hooks, + "stop", + tokensave_bin, + "hook-cursor-stop", + 30, + None, + ); safe_write_json_file(hooks_path, &hooks, backup.as_deref())?; eprintln!( @@ -502,6 +514,7 @@ fn doctor_check_hooks(dc: &mut DoctorCounters, hooks_path: &Path) { ("afterFileEdit", "hook-cursor-after-file-edit"), ("afterShellExecution", "hook-cursor-after-shell"), ("workspaceOpen", "hook-cursor-workspace-open"), + ("stop", "hook-cursor-stop"), ]; let missing: Vec<&str> = expected .iter() diff --git a/src/agents/gemini.rs b/src/agents/gemini.rs index c401671fc4..819022e82d 100644 --- a/src/agents/gemini.rs +++ b/src/agents/gemini.rs @@ -158,8 +158,14 @@ fn install_prompt_rules(gemini_md: &Path) -> Result<()> { faster than file reads.\n\n\ If a code analysis question cannot be fully answered by tokensave MCP tools, \ try querying the SQLite database directly at `.tokensave/tokensave.db` \ - (tables: `nodes`, `edges`, `files`). Use SQL to answer complex structural queries \ + (tables: `nodes`, `edges`, `files`, `memory_facts`, `memory_entities`, \ + `memory_feedback_events`). Use SQL to answer complex structural queries \ that go beyond what the built-in tools expose.\n\n\ + For durable project/user facts, prefer `tokensave_fact_store`, \ + `tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. \ + Use `tokensave_message_search` for project-local Cursor transcript recall when \ + prior conversation context matters. Do not store secrets, credentials, or \ + unnecessary PII in persistent facts.\n\n\ If you discover a gap where an extractor, schema, or tokensave tool could be \ improved to answer a question natively, propose to the user that they open an issue \ at https://github.com/aovestdipaperino/tokensave describing the limitation. \ diff --git a/src/agents/kimi.rs b/src/agents/kimi.rs index 84f4453ecf..e298423479 100644 --- a/src/agents/kimi.rs +++ b/src/agents/kimi.rs @@ -159,8 +159,14 @@ fn install_prompt_rules(agents_md: &Path) -> Result<()> { faster than file reads.\n\n\ If a code analysis question cannot be fully answered by tokensave MCP tools, \ try querying the SQLite database directly at `.tokensave/tokensave.db` \ - (tables: `nodes`, `edges`, `files`). Use SQL to answer complex structural queries \ + (tables: `nodes`, `edges`, `files`, `memory_facts`, `memory_entities`, \ + `memory_feedback_events`). Use SQL to answer complex structural queries \ that go beyond what the built-in tools expose.\n\n\ + For durable project/user facts, prefer `tokensave_fact_store`, \ + `tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. \ + Use `tokensave_message_search` for project-local Cursor transcript recall when \ + prior conversation context matters. Do not store secrets, credentials, or \ + unnecessary PII in persistent facts.\n\n\ If you discover a gap where an extractor, schema, or tokensave tool could be \ improved to answer a question natively, propose to the user that they open an issue \ at https://github.com/aovestdipaperino/tokensave describing the limitation. \ diff --git a/src/agents/kiro.rs b/src/agents/kiro.rs index c8ff9bc68d..4c31250356 100644 --- a/src/agents/kiro.rs +++ b/src/agents/kiro.rs @@ -449,7 +449,13 @@ have been tried. Delegation is still appropriate for long-running execution work such as builds, tests, generated reports, or independent implementation tasks.\n\n\ If a code analysis question cannot be fully answered by tokensave MCP tools, try \ querying the SQLite database directly at `.tokensave/tokensave.db` (tables: `nodes`, \ -`edges`, `files`). Use SQL for structural queries that go beyond the MCP tools.\n\n\ +`edges`, `files`, `memory_facts`, `memory_entities`, `memory_feedback_events`). \ +Use SQL for structural queries that go beyond the MCP tools.\n\n\ +For durable project/user facts, prefer `tokensave_fact_store`, \ +`tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. Use \ +`tokensave_message_search` for project-local Cursor transcript recall when prior \ +conversation context matters. Do not store secrets, credentials, or unnecessary PII \ +in persistent facts.\n\n\ If you discover a gap where an extractor, schema, or tokensave tool could answer a \ question natively, propose opening an issue at \ https://github.com/aovestdipaperino/tokensave. Remind the user to strip sensitive \ diff --git a/src/agents/opencode.rs b/src/agents/opencode.rs index 8c71b8898e..abbecc3b4a 100644 --- a/src/agents/opencode.rs +++ b/src/agents/opencode.rs @@ -191,8 +191,14 @@ fn install_prompt_rules(prompt_path: &Path) -> Result<()> { faster than file reads.\n\n\ If a code analysis question cannot be fully answered by tokensave MCP tools, \ try querying the SQLite database directly at `.tokensave/tokensave.db` \ - (tables: `nodes`, `edges`, `files`). Use SQL to answer complex structural queries \ + (tables: `nodes`, `edges`, `files`, `memory_facts`, `memory_entities`, \ + `memory_feedback_events`). Use SQL to answer complex structural queries \ that go beyond what the built-in tools expose.\n\n\ + For durable project/user facts, prefer `tokensave_fact_store`, \ + `tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. \ + Use `tokensave_message_search` for project-local Cursor transcript recall when \ + prior conversation context matters. Do not store secrets, credentials, or \ + unnecessary PII in persistent facts.\n\n\ If you discover a gap where an extractor, schema, or tokensave tool could be \ improved to answer a question natively, propose to the user that they open an issue \ at https://github.com/aovestdipaperino/tokensave describing the limitation. \ diff --git a/src/agents/vibe.rs b/src/agents/vibe.rs index 9f2025e71d..0b6d6ed5ab 100644 --- a/src/agents/vibe.rs +++ b/src/agents/vibe.rs @@ -191,7 +191,9 @@ fn install_prompt_rules(prompt_path: &Path) -> Result<()> { that go beyond what the built-in tools expose.\n\n\ For durable project/user facts, prefer `tokensave_fact_store`, \ `tokensave_fact_feedback`, and `tokensave_memory_status` over ad-hoc notes. \ - Do not store secrets, credentials, or unnecessary PII in persistent facts.\n\n\ + Use `tokensave_message_search` for project-local Cursor transcript recall when \ + prior conversation context matters. Do not store secrets, credentials, or \ + unnecessary PII in persistent facts.\n\n\ If you find a gap where tokensave could answer a question natively, propose opening \ an issue at https://github.com/aovestdipaperino/tokensave. Remind the user to strip \ sensitive or proprietary code from any issue text before submitting.\n\n\ diff --git a/src/cli.rs b/src/cli.rs index 762709f3e3..5e79a4aef7 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -143,6 +143,9 @@ pub enum Commands { /// Cursor workspaceOpen hook handler (called by Cursor, not by users directly) #[command(name = "hook-cursor-workspace-open", hide = true)] HookCursorWorkspaceOpen, + /// Cursor stop hook handler (called by Cursor, not by users directly) + #[command(name = "hook-cursor-stop", hide = true)] + HookCursorStop, /// Codex SessionStart hook handler (called by Codex, not by users directly) #[command(name = "hook-codex-session-start", hide = true)] HookCodexSessionStart, diff --git a/src/config.rs b/src/config.rs index 5865f6c5b5..79f31379fb 100644 --- a/src/config.rs +++ b/src/config.rs @@ -425,15 +425,12 @@ mod tests { fs::write(&excludes, ".tokensave\n").unwrap(); let git_config = sandbox.path().join("gitconfig"); - let status = Command::new("git") - .env("GIT_CONFIG_GLOBAL", &git_config) - .arg("config") - .arg("--global") - .arg("core.excludesFile") - .arg(&excludes) - .status() - .unwrap(); - assert!(status.success()); + let excludes_value = excludes.to_string_lossy().replace('\\', "/"); + fs::write( + &git_config, + format!("[core]\n\texcludesFile = {excludes_value}\n"), + ) + .unwrap(); let ignored = is_ignored_by_git(&repo, Some(&git_config)); diff --git a/src/db/migrations.rs b/src/db/migrations.rs index a50aded16c..b599fd2755 100644 --- a/src/db/migrations.rs +++ b/src/db/migrations.rs @@ -12,10 +12,11 @@ use libsql::Connection; use crate::errors::{Result, TokenSaveError}; +use crate::memory::store::MemoryStore; /// The highest migration version defined in this file. Bump this and add a /// new entry to `run_migration` whenever the schema changes. -const LATEST_VERSION: u32 = 10; +const LATEST_VERSION: u32 = 12; /// Reads the current schema version from `PRAGMA user_version`. async fn get_version(conn: &Connection) -> Result { @@ -184,26 +185,6 @@ pub async fn create_schema(conn: &Connection) -> Result<()> { CREATE INDEX IF NOT EXISTS idx_node_fingerprints_ast ON node_fingerprints(ast_hash); CREATE INDEX IF NOT EXISTS idx_node_fingerprints_size ON node_fingerprints(body_tokens); - CREATE TABLE IF NOT EXISTS memory_decisions ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - text TEXT NOT NULL, - reason TEXT, - created_at INTEGER NOT NULL, - files TEXT NOT NULL DEFAULT '[]', - tags TEXT NOT NULL DEFAULT '[]' - ); - - CREATE TABLE IF NOT EXISTS memory_code_areas ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - path TEXT NOT NULL, - description TEXT, - last_touched_at INTEGER NOT NULL, - touch_count INTEGER NOT NULL DEFAULT 1 - ); - - CREATE UNIQUE INDEX IF NOT EXISTS idx_memory_code_areas_path ON memory_code_areas(path); - CREATE INDEX IF NOT EXISTS idx_memory_decisions_created_at ON memory_decisions(created_at); - CREATE TABLE IF NOT EXISTS read_cache ( project_id TEXT NOT NULL, session_id TEXT NOT NULL, @@ -219,32 +200,7 @@ pub async fn create_schema(conn: &Connection) -> Result<()> { ); CREATE INDEX IF NOT EXISTS idx_read_cache_session - ON read_cache(session_id, created_at); - - CREATE VIRTUAL TABLE IF NOT EXISTS memory_decisions_fts USING fts5( - text, reason, - content='memory_decisions', content_rowid='id' - ); - - CREATE TRIGGER IF NOT EXISTS memory_decisions_fts_insert - AFTER INSERT ON memory_decisions BEGIN - INSERT INTO memory_decisions_fts(rowid, text, reason) - VALUES (NEW.id, NEW.text, NEW.reason); - END; - - CREATE TRIGGER IF NOT EXISTS memory_decisions_fts_delete - AFTER DELETE ON memory_decisions BEGIN - INSERT INTO memory_decisions_fts(memory_decisions_fts, rowid, text, reason) - VALUES ('delete', OLD.id, OLD.text, OLD.reason); - END; - - CREATE TRIGGER IF NOT EXISTS memory_decisions_fts_update - AFTER UPDATE ON memory_decisions BEGIN - INSERT INTO memory_decisions_fts(memory_decisions_fts, rowid, text, reason) - VALUES ('delete', OLD.id, OLD.text, OLD.reason); - INSERT INTO memory_decisions_fts(rowid, text, reason) - VALUES (NEW.id, NEW.text, NEW.reason); - END;", + ON read_cache(session_id, created_at);", ) .await .map_err(|e| TokenSaveError::Database { @@ -252,6 +208,7 @@ pub async fn create_schema(conn: &Connection) -> Result<()> { operation: "create_schema".to_string(), })?; + create_holographic_memory_schema(conn, "create_schema").await?; set_version(conn, LATEST_VERSION).await?; Ok(()) } @@ -333,6 +290,8 @@ async fn run_migration(conn: &Connection, version: u32) -> Result<()> { 8 => migrate_v8(conn).await, 9 => migrate_v9(conn).await, 10 => migrate_v10(conn).await, + 11 => migrate_v11(conn).await, + 12 => migrate_v12(conn).await, _ => Err(TokenSaveError::Database { message: format!("unknown migration version: {version}"), operation: "run_migration".to_string(), @@ -648,8 +607,9 @@ async fn migrate_v7(conn: &Connection) -> Result<()> { /// Adds tables for persistent agent memory: `memory_decisions` records /// architecture / design choices with optional reason and tags; /// `memory_code_areas` tracks paths the agent has worked in. An FTS5 mirror -/// over `memory_decisions.text` and `memory_decisions.reason` enables -/// fuzzy recall via `tokensave_session_recall`. +/// over `memory_decisions.text` and `memory_decisions.reason` supported the +/// legacy decision-recall implementation before v11 backfilled and dropped +/// these tables. async fn migrate_v8(conn: &Connection) -> Result<()> { conn.execute_batch( "CREATE TABLE IF NOT EXISTS memory_decisions ( @@ -875,3 +835,413 @@ async fn migrate_v10(conn: &Connection) -> Result<()> { Ok(()) } + +// --------------------------------------------------------------------------- +// Migration V11: holographic memory active schema +// --------------------------------------------------------------------------- + +/// Creates the active holographic-memory tables alongside the legacy memory +/// tables. Legacy data is preserved and copied into `memory_facts`. +async fn migrate_v11(conn: &Connection) -> Result<()> { + create_holographic_memory_schema(conn, "migrate_v11").await?; + if legacy_memory_tables_exist(conn).await? { + backfill_legacy_memory_as_facts(conn).await?; + backfill_holographic_memory_vectors_and_banks(conn).await?; + } + Ok(()) +} + +async fn backfill_holographic_memory_vectors_and_banks(conn: &Connection) -> Result<()> { + let store = MemoryStore::new(conn); + loop { + let updated = store.compute_missing_vectors(500).await?; + if updated == 0 { + break; + } + } + store.rebuild_all_banks().await?; + Ok(()) +} + +async fn legacy_memory_tables_exist(conn: &Connection) -> Result { + let mut rows = conn + .query( + "SELECT COUNT(*) FROM sqlite_master + WHERE type='table' + AND name IN ('memory_decisions', 'memory_code_areas')", + (), + ) + .await + .map_err(|e| TokenSaveError::Database { + message: format!("migrate_v11: failed to probe legacy memory tables: {e}"), + operation: "migrate_v11".to_string(), + })?; + let row = rows + .next() + .await + .map_err(|e| TokenSaveError::Database { + message: format!("migrate_v11: failed to read legacy table probe: {e}"), + operation: "migrate_v11".to_string(), + })? + .ok_or_else(|| TokenSaveError::Database { + message: "migrate_v11: legacy table probe returned no rows".to_string(), + operation: "migrate_v11".to_string(), + })?; + let count: i64 = row.get(0).map_err(|e| TokenSaveError::Database { + message: format!("migrate_v11: failed to read legacy table count: {e}"), + operation: "migrate_v11".to_string(), + })?; + Ok(count > 0) +} + +async fn create_holographic_memory_schema(conn: &Connection, operation: &str) -> Result<()> { + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS memory_facts ( + fact_id INTEGER PRIMARY KEY AUTOINCREMENT, + content TEXT NOT NULL UNIQUE, + category TEXT NOT NULL DEFAULT 'general', + tags TEXT NOT NULL DEFAULT '[]', + trust_score REAL NOT NULL DEFAULT 0.5, + retrieval_count INTEGER NOT NULL DEFAULT 0, + helpful_count INTEGER NOT NULL DEFAULT 0, + unhelpful_count INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL DEFAULT 0, + updated_at INTEGER NOT NULL DEFAULT 0, + last_retrieved_at INTEGER, + last_feedback_at INTEGER, + source TEXT NOT NULL DEFAULT 'manual', + metadata TEXT NOT NULL DEFAULT '{}', + hrr_vector BLOB, + hrr_algebra TEXT NOT NULL DEFAULT 'amari_fhrr', + hrr_dim INTEGER NOT NULL DEFAULT 2048 + ); + + CREATE TABLE IF NOT EXISTS memory_entities ( + entity_id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL, + normalized_name TEXT NOT NULL UNIQUE, + entity_type TEXT NOT NULL DEFAULT 'unknown', + aliases TEXT NOT NULL DEFAULT '[]', + created_at INTEGER NOT NULL DEFAULT 0 + ); + + CREATE TABLE IF NOT EXISTS memory_fact_entities ( + fact_id INTEGER NOT NULL, + entity_id INTEGER NOT NULL, + PRIMARY KEY (fact_id, entity_id), + FOREIGN KEY (fact_id) REFERENCES memory_facts(fact_id) ON DELETE CASCADE, + FOREIGN KEY (entity_id) REFERENCES memory_entities(entity_id) ON DELETE CASCADE + ); + + CREATE TABLE IF NOT EXISTS memory_banks ( + bank_id INTEGER PRIMARY KEY AUTOINCREMENT, + bank_name TEXT NOT NULL UNIQUE, + vector BLOB NOT NULL, + hrr_algebra TEXT NOT NULL DEFAULT 'amari_fhrr', + hrr_dim INTEGER NOT NULL DEFAULT 2048, + fact_count INTEGER NOT NULL DEFAULT 0, + updated_at INTEGER NOT NULL DEFAULT 0 + ); + + CREATE TABLE IF NOT EXISTS memory_bank_dirty ( + bank_name TEXT PRIMARY KEY, + updated_at INTEGER NOT NULL DEFAULT 0 + ); + + CREATE TABLE IF NOT EXISTS memory_feedback_events ( + event_id INTEGER PRIMARY KEY AUTOINCREMENT, + fact_id INTEGER NOT NULL, + action TEXT NOT NULL CHECK (action IN ('helpful', 'unhelpful')), + trust_delta REAL NOT NULL, + old_trust REAL NOT NULL, + new_trust REAL NOT NULL, + created_at INTEGER NOT NULL DEFAULT 0, + source TEXT NOT NULL DEFAULT 'mcp', + note TEXT, + FOREIGN KEY (fact_id) REFERENCES memory_facts(fact_id) ON DELETE CASCADE + ); + + CREATE INDEX IF NOT EXISTS idx_memory_facts_category + ON memory_facts(category); + CREATE INDEX IF NOT EXISTS idx_memory_facts_updated_at + ON memory_facts(updated_at); + CREATE INDEX IF NOT EXISTS idx_memory_facts_trust_score + ON memory_facts(trust_score); + CREATE INDEX IF NOT EXISTS idx_memory_facts_source + ON memory_facts(source); + CREATE INDEX IF NOT EXISTS idx_memory_entities_type + ON memory_entities(entity_type); + CREATE INDEX IF NOT EXISTS idx_memory_fact_entities_entity_id + ON memory_fact_entities(entity_id); + CREATE INDEX IF NOT EXISTS idx_memory_banks_updated_at + ON memory_banks(updated_at); + CREATE INDEX IF NOT EXISTS idx_memory_feedback_events_fact_id + ON memory_feedback_events(fact_id); + CREATE INDEX IF NOT EXISTS idx_memory_feedback_events_created_at + ON memory_feedback_events(created_at); + + CREATE VIRTUAL TABLE IF NOT EXISTS memory_facts_fts USING fts5( + content, tags, + content='memory_facts', content_rowid='rowid' + ); + + CREATE TRIGGER IF NOT EXISTS memory_facts_fts_insert + AFTER INSERT ON memory_facts BEGIN + INSERT INTO memory_facts_fts(rowid, content, tags) + VALUES (NEW.rowid, NEW.content, NEW.tags); + END; + + CREATE TRIGGER IF NOT EXISTS memory_facts_fts_delete + AFTER DELETE ON memory_facts BEGIN + INSERT INTO memory_facts_fts(memory_facts_fts, rowid, content, tags) + VALUES ('delete', OLD.rowid, OLD.content, OLD.tags); + END; + + CREATE TRIGGER IF NOT EXISTS memory_facts_fts_update + AFTER UPDATE OF content, tags ON memory_facts BEGIN + INSERT INTO memory_facts_fts(memory_facts_fts, rowid, content, tags) + VALUES ('delete', OLD.rowid, OLD.content, OLD.tags); + INSERT INTO memory_facts_fts(rowid, content, tags) + VALUES (NEW.rowid, NEW.content, NEW.tags); + END;", + ) + .await + .map_err(|e| TokenSaveError::Database { + message: format!("{operation}: failed to create holographic memory schema: {e}"), + operation: operation.to_string(), + })?; + + Ok(()) +} + +// --------------------------------------------------------------------------- +// Migration V12: dirty bank tracking for lazy memory-bank rebuilds +// --------------------------------------------------------------------------- + +async fn migrate_v12(conn: &Connection) -> Result<()> { + conn.execute( + "CREATE TABLE IF NOT EXISTS memory_bank_dirty ( + bank_name TEXT PRIMARY KEY, + updated_at INTEGER NOT NULL DEFAULT 0 + )", + (), + ) + .await + .map_err(|e| TokenSaveError::Database { + message: format!("v12: failed to create memory_bank_dirty table: {e}"), + operation: "migrate_v12".to_string(), + })?; + Ok(()) +} + +async fn backfill_legacy_memory_as_facts(conn: &Connection) -> Result<()> { + conn.execute_batch( + "WITH normalized_decisions AS ( + SELECT + id, + text, + reason, + created_at, + CASE + WHEN json_valid(COALESCE(NULLIF(trim(files), ''), '[]')) + AND json_type(COALESCE(NULLIF(trim(files), ''), '[]')) = 'array' + THEN COALESCE(NULLIF(trim(files), ''), '[]') + ELSE '[]' + END AS safe_files, + CASE + WHEN json_valid(COALESCE(NULLIF(trim(tags), ''), '[]')) + AND json_type(COALESCE(NULLIF(trim(tags), ''), '[]')) = 'array' + THEN COALESCE(NULLIF(trim(tags), ''), '[]') + ELSE '[]' + END AS safe_tags + FROM memory_decisions + ) + INSERT OR IGNORE INTO memory_facts ( + content, + category, + tags, + created_at, + updated_at, + source, + metadata + ) + SELECT + CASE + WHEN reason IS NULL OR length(trim(reason)) = 0 THEN text + ELSE text || char(10) || char(10) || 'Reason: ' || reason + END || char(10) || char(10) || 'Legacy decision id: ' || id, + 'decision', + safe_tags, + created_at, + created_at, + 'legacy_memory_decisions', + json_object( + 'holographic_memory_backfill_v1', 1, + 'legacy_table', 'memory_decisions', + 'legacy_id', id, + 'decision_text', text, + 'reason', COALESCE(reason, ''), + 'files', json(safe_files), + 'tags', json(safe_tags) + ) + FROM normalized_decisions; + + WITH normalized_code_areas AS ( + SELECT id, path, description, last_touched_at, touch_count + FROM memory_code_areas + ) + INSERT OR IGNORE INTO memory_facts ( + content, + category, + tags, + created_at, + updated_at, + source, + metadata + ) + SELECT + CASE + WHEN description IS NULL OR length(trim(description)) = 0 THEN path + ELSE path || char(10) || char(10) || description + END || char(10) || char(10) || 'Legacy code area id: ' || id, + 'code_area', + json_array('code_area', path), + last_touched_at, + last_touched_at, + 'legacy_memory_code_areas', + json_object( + 'holographic_memory_backfill_v1', 1, + 'legacy_table', 'memory_code_areas', + 'legacy_id', id, + 'path', path, + 'description', COALESCE(description, ''), + 'last_touched_at', last_touched_at, + 'touch_count', touch_count + ) + FROM normalized_code_areas;", + ) + .await + .map_err(|e| TokenSaveError::Database { + message: format!("migrate_v11: failed to backfill legacy memory: {e}"), + operation: "migrate_v11".to_string(), + })?; + + conn.execute_batch( + "WITH normalized_decisions AS ( + SELECT + id, + created_at, + CASE + WHEN json_valid(COALESCE(NULLIF(trim(files), ''), '[]')) + AND json_type(COALESCE(NULLIF(trim(files), ''), '[]')) = 'array' + THEN COALESCE(NULLIF(trim(files), ''), '[]') + ELSE '[]' + END AS safe_files, + CASE + WHEN json_valid(COALESCE(NULLIF(trim(tags), ''), '[]')) + AND json_type(COALESCE(NULLIF(trim(tags), ''), '[]')) = 'array' + THEN COALESCE(NULLIF(trim(tags), ''), '[]') + ELSE '[]' + END AS safe_tags + FROM memory_decisions + ) + INSERT OR IGNORE INTO memory_entities (name, normalized_name, entity_type, created_at) + SELECT DISTINCT value, lower(value), 'legacy_file', created_at + FROM normalized_decisions, json_each(safe_files) + WHERE trim(value) != ''; + + WITH normalized_decisions AS ( + SELECT + id, + created_at, + CASE + WHEN json_valid(COALESCE(NULLIF(trim(tags), ''), '[]')) + AND json_type(COALESCE(NULLIF(trim(tags), ''), '[]')) = 'array' + THEN COALESCE(NULLIF(trim(tags), ''), '[]') + ELSE '[]' + END AS safe_tags + FROM memory_decisions + ) + INSERT OR IGNORE INTO memory_entities (name, normalized_name, entity_type, created_at) + SELECT DISTINCT value, lower(value), 'legacy_tag', created_at + FROM normalized_decisions, json_each(safe_tags) + WHERE trim(value) != ''; + + INSERT OR IGNORE INTO memory_entities (name, normalized_name, entity_type, created_at) + SELECT DISTINCT path, lower(path), 'legacy_path', last_touched_at + FROM memory_code_areas + WHERE trim(path) != ''; + + WITH normalized_decisions AS ( + SELECT + id, + CASE + WHEN json_valid(COALESCE(NULLIF(trim(files), ''), '[]')) + AND json_type(COALESCE(NULLIF(trim(files), ''), '[]')) = 'array' + THEN COALESCE(NULLIF(trim(files), ''), '[]') + ELSE '[]' + END AS safe_files + FROM memory_decisions + ) + INSERT OR IGNORE INTO memory_fact_entities (fact_id, entity_id) + SELECT f.fact_id, e.entity_id + FROM normalized_decisions d + JOIN memory_facts f + ON f.source = 'legacy_memory_decisions' + AND json_extract(f.metadata, '$.legacy_id') = d.id + JOIN json_each(d.safe_files) file_entity + JOIN memory_entities e ON e.normalized_name = lower(file_entity.value) + WHERE trim(file_entity.value) != ''; + + WITH normalized_decisions AS ( + SELECT + id, + CASE + WHEN json_valid(COALESCE(NULLIF(trim(tags), ''), '[]')) + AND json_type(COALESCE(NULLIF(trim(tags), ''), '[]')) = 'array' + THEN COALESCE(NULLIF(trim(tags), ''), '[]') + ELSE '[]' + END AS safe_tags + FROM memory_decisions + ) + INSERT OR IGNORE INTO memory_fact_entities (fact_id, entity_id) + SELECT f.fact_id, e.entity_id + FROM normalized_decisions d + JOIN memory_facts f + ON f.source = 'legacy_memory_decisions' + AND json_extract(f.metadata, '$.legacy_id') = d.id + JOIN json_each(d.safe_tags) tag_entity + JOIN memory_entities e ON e.normalized_name = lower(tag_entity.value) + WHERE trim(tag_entity.value) != ''; + + INSERT OR IGNORE INTO memory_fact_entities (fact_id, entity_id) + SELECT f.fact_id, e.entity_id + FROM memory_code_areas c + JOIN memory_facts f + ON f.source = 'legacy_memory_code_areas' + AND json_extract(f.metadata, '$.legacy_id') = c.id + JOIN memory_entities e ON e.normalized_name = lower(c.path) + WHERE trim(c.path) != '';", + ) + .await + .map_err(|e| TokenSaveError::Database { + message: format!("migrate_v11: failed to link legacy memory entities: {e}"), + operation: "migrate_v11".to_string(), + })?; + + conn.execute_batch( + "DROP TRIGGER IF EXISTS memory_decisions_fts_insert; + DROP TRIGGER IF EXISTS memory_decisions_fts_delete; + DROP TRIGGER IF EXISTS memory_decisions_fts_update; + DROP TABLE IF EXISTS memory_decisions_fts; + DROP TABLE IF EXISTS memory_code_areas; + DROP TABLE IF EXISTS memory_decisions;", + ) + .await + .map_err(|e| TokenSaveError::Database { + message: format!("migrate_v11: failed to drop legacy memory tables: {e}"), + operation: "migrate_v11".to_string(), + })?; + + Ok(()) +} diff --git a/src/global_db.rs b/src/global_db.rs index d252543e74..8b7f554a94 100644 --- a/src/global_db.rs +++ b/src/global_db.rs @@ -11,6 +11,9 @@ use libsql::{params, Builder, Connection, Database as LibsqlDatabase, Value}; use crate::sessions::{SessionMessageRecord, SessionMessageSearchResult, SessionRecord}; +const MAX_SESSION_MESSAGE_TEXT_BYTES: usize = 256 * 1024; +const SESSION_MESSAGE_TRUNCATION_MARKER: &str = "\n[truncated by tokensave]"; + /// Total savings + call count for a project (or all projects when `project` is None). #[derive(Debug, Clone, serde::Serialize)] pub struct SavingsTotal { @@ -41,6 +44,28 @@ fn opt_text(value: Option<&str>) -> Value { value.map_or(Value::Null, |s| Value::Text(s.to_string())) } +fn capped_session_message_text(text: &str) -> std::borrow::Cow<'_, str> { + if text.len() <= MAX_SESSION_MESSAGE_TEXT_BYTES { + return std::borrow::Cow::Borrowed(text); + } + + let budget = + MAX_SESSION_MESSAGE_TEXT_BYTES.saturating_sub(SESSION_MESSAGE_TRUNCATION_MARKER.len()); + let mut end = 0; + for (idx, ch) in text.char_indices() { + let next = idx + ch.len_utf8(); + if next > budget { + break; + } + end = next; + } + + let mut capped = String::with_capacity(end + SESSION_MESSAGE_TRUNCATION_MARKER.len()); + capped.push_str(&text[..end]); + capped.push_str(SESSION_MESSAGE_TRUNCATION_MARKER); + std::borrow::Cow::Owned(capped) +} + fn opt_i64(value: Option) -> Value { value.map_or(Value::Null, Value::Integer) } @@ -461,6 +486,7 @@ impl GlobalDb { /// Inserts or replaces a provider message. Returns `false` on any DB error. pub async fn upsert_session_message(&self, message: &SessionMessageRecord) -> bool { + let text = capped_session_message_text(&message.text); self.conn .execute( "INSERT INTO session_messages @@ -486,7 +512,7 @@ impl GlobalDb { message.role.as_str(), opt_i64(message.timestamp), message.ordinal, - message.text.as_str(), + text.as_ref(), opt_text(message.kind.as_deref()), opt_text(message.model.as_deref()), opt_text(message.tool_names.as_deref()), diff --git a/src/hooks.rs b/src/hooks.rs index 6b0b275c62..96b58c4a8f 100644 --- a/src/hooks.rs +++ b/src/hooks.rs @@ -8,6 +8,7 @@ use std::io::Read; use std::path::{Path, PathBuf}; +use std::time::Duration; use serde_json::Value; @@ -149,11 +150,20 @@ pub fn hook_cursor_subagent_start() -> i32 { /// Cursor `beforeSubmitPrompt` hook handler. /// -/// Resets the project-local counter for a new prompt turn. The output uses -/// Cursor's documented `beforeSubmitPrompt` shape and never blocks submission. +/// Resets the project-local counter for a new prompt turn and does at most a +/// small, time-boxed *tail* ingest of newly-appended transcript lines (the bulk +/// catch-up lives on the lower-frequency `sessionStart` / `stop` hooks). The +/// output uses Cursor's documented `beforeSubmitPrompt` shape and never blocks +/// submission, even if the tail ingest times out. pub async fn hook_cursor_before_submit_prompt() -> i32 { let event = read_stdin_to_string(); reset_counter_for_cursor_event(&event).await; + ingest_cursor_transcript_for_event( + &event, + Some(CURSOR_HOT_INGEST_MAX_BYTES), + CURSOR_HOT_INGEST_BUDGET, + ) + .await; let mut output = serde_json::json!({ "continue": true }); if let Some(hint) = cursor_prompt_hint(&event) { output["additional_context"] = Value::String(format_tool_hint(&hint)); @@ -162,6 +172,24 @@ pub async fn hook_cursor_before_submit_prompt() -> i32 { 0 } +/// Cursor `stop` hook handler (fire-and-forget). +/// +/// Fires at the end of an agent turn and performs the primary transcript +/// ingest: a time-boxed incremental catch-up that picks up bounded transcript +/// tails appended during the turn. The `stop` output is informational only, so +/// we emit an empty object and never ask the agent to continue. Fail-open. +pub async fn hook_cursor_stop() -> i32 { + let event = read_stdin_to_string(); + ingest_cursor_transcript_for_event( + &event, + Some(CURSOR_CATCH_UP_INGEST_MAX_BYTES), + CURSOR_STOP_INGEST_BUDGET, + ) + .await; + println!("{}", serde_json::json!({})); + 0 +} + /// Cursor `afterFileEdit` hook handler. /// /// Keeps the graph fresh after Cursor Agent writes files. This uses a @@ -184,6 +212,14 @@ pub async fn hook_cursor_after_file_edit() -> i32 { /// for the resolved workspace. Never blocks session creation. pub async fn hook_cursor_session_start() -> i32 { let event = read_stdin_to_string(); + // Catch-up ingest for resumed sessions whose transcript grew while no agent + // was attached. No-op (no transcript_path) for brand-new sessions. Fail-open. + ingest_cursor_transcript_for_event( + &event, + Some(CURSOR_CATCH_UP_INGEST_MAX_BYTES), + CURSOR_SESSION_INGEST_BUDGET, + ) + .await; let root = cursor_project_root_from_event(&event); let context = session_steering_context_for_root(root.as_deref()).await; println!("{}", cursor_session_start_json(root.as_deref(), &context)); @@ -1229,6 +1265,48 @@ async fn reset_counter_for_cursor_event(event_json: &str) { } } +/// Largest tail the `beforeSubmitPrompt` hot path will read in one call. Larger +/// backlogs are left for the `sessionStart` / `stop` catch-up ingests. +const CURSOR_HOT_INGEST_MAX_BYTES: u64 = 256 * 1024; +/// Largest transcript tail a low-priority Cursor catch-up hook will read. +/// Oversized backlogs stay queued instead of blocking hook execution. +const CURSOR_CATCH_UP_INGEST_MAX_BYTES: u64 = 2 * 1024 * 1024; +/// Hard wall-clock budget for the `beforeSubmitPrompt` tail ingest. Well under +/// Cursor's 5s hook timeout; on expiry we fail open and let heavier hooks catch up. +const CURSOR_HOT_INGEST_BUDGET: Duration = Duration::from_millis(1_500); +/// Budget for the `sessionStart` catch-up ingest (registered with a 5s timeout). +const CURSOR_SESSION_INGEST_BUDGET: Duration = Duration::from_secs(4); +/// Budget for the end-of-turn `stop` catch-up ingest (registered with a 30s timeout). +const CURSOR_STOP_INGEST_BUDGET: Duration = Duration::from_secs(25); + +/// Incrementally ingests the Cursor transcript referenced by `event_json` into +/// the project-local session DB, bounded by `max_new_bytes` (the hot-path cap) +/// and an overall `budget`. Always fails open: a timeout, missing transcript, or +/// any error is swallowed so the calling hook never blocks the agent. +async fn ingest_cursor_transcript_for_event( + event_json: &str, + max_new_bytes: Option, + budget: Duration, +) { + let work = async { + let Some(project_root) = cursor_project_root_from_event(event_json) else { + return; + }; + let Some(db) = crate::sessions::cursor::open_project_session_db(&project_root).await else { + return; + }; + let _ = crate::sessions::cursor::ingest_cursor_transcript_event_capped( + event_json, + &db, + max_new_bytes, + ) + .await; + }; + // Short-lived CLI hook processes exit immediately, so the ingest must run + // inline (not on a detached task); the timeout keeps it inside budget. + let _ = tokio::time::timeout(budget, work).await; +} + async fn sync_for_kiro_event(event_json: &str) -> crate::errors::Result<()> { let Some(project_root) = kiro_project_root(event_json) else { return Ok(()); diff --git a/src/lib.rs b/src/lib.rs index e34fecc9ab..fb5ce81ef0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -36,6 +36,7 @@ pub mod global_db; pub mod graph; pub mod hooks; pub mod mcp; +pub mod memory; pub mod monitor; pub mod redundancy; pub mod resolution; diff --git a/src/main.rs b/src/main.rs index 2e9a80016e..33fab1c776 100644 --- a/src/main.rs +++ b/src/main.rs @@ -832,6 +832,12 @@ async fn run(cli: Cli) -> tokensave::errors::Result<()> { process::exit(code); } } + Commands::HookCursorStop => { + let code = tokensave::hooks::hook_cursor_stop().await; + if code != 0 { + process::exit(code); + } + } Commands::HookCodexSessionStart => { let code = tokensave::hooks::hook_codex_session_start().await; if code != 0 { @@ -1212,6 +1218,7 @@ fn should_skip_startup_maintenance(command: &Commands) -> bool { | Commands::HookCursorSessionStart | Commands::HookCursorAfterShell | Commands::HookCursorWorkspaceOpen + | Commands::HookCursorStop | Commands::HookCodexSessionStart | Commands::HookCodexUserPromptSubmit | Commands::HookCodexSubagentStart diff --git a/src/mcp/server.rs b/src/mcp/server.rs index 4d89cfa966..47fa895554 100644 --- a/src/mcp/server.rs +++ b/src/mcp/server.rs @@ -41,6 +41,12 @@ impl ServerStats { /// Cache duration for version checks (15 minutes). const VERSION_CHECK_INTERVAL: Duration = Duration::from_mins(15); +fn global_db_enabled() -> bool { + std::env::var("TOKENSAVE_ENABLE_GLOBAL_DB") + .map(|value| matches!(value.as_str(), "1" | "true" | "TRUE" | "yes" | "YES")) + .unwrap_or(false) +} + /// Hand-maintained schema documentation for the `tokensave://schema` resource. /// Mirrors `src/db/migrations.rs::create_schema`. Update both together. const SCHEMA_MARKDOWN: &str = r"# tokensave SQLite schema @@ -100,8 +106,35 @@ Unique constraint: `(source, target, kind, COALESCE(line, -1))`. Indexes on `sou ### `metadata` — key/value store Common keys: `tokens_saved`, schema-version markers. -### `memory_decisions`, `memory_code_areas` -Hand-recorded notes from `tokensave_record_decision` / `tokensave_record_code_area`. FTS5 mirror tables exist for `nodes` (`nodes_fts`) and `memory_decisions` (`memory_decisions_fts`). +### `node_fingerprints` — redundancy cache +- `node_id` PRIMARY KEY FK → `nodes.id` +- `ast_hash`, `cfg_hash`, `call_seq_hash`, `shingles` +- `body_tokens`, `source_hash` + +### `read_cache` — rendered `tokensave_read` responses +- primary key: `(project_id, session_id, file_path, mode, args_hash)` +- stores `mtime_ns`, `digest`, rendered `body` BLOB, token count, and `created_at` + +### v11: `memory_facts`, `memory_entities`, `memory_fact_entities`, `memory_banks`, `memory_feedback_events` +The holographic fact store replaces narrow decision rows with durable facts +linked to named entities: + +- `memory_facts` — numeric `fact_id`, unique fact content, category, source, + tags JSON, computed trust score, retrieval/feedback counts, timestamps, and + structured metadata. +- `memory_entities` — normalized recall keys for symbols, files, + directories, branches, people, subsystems, and concepts. Facts can attach + multiple entities so recall can start from code or natural-language names. +- `memory_fact_entities` — many-to-many join table linking facts to entities + with cascade deletes. +- `memory_banks` — optional holographic memory-bank vectors by category or + bank name (`bank_name`, `vector`, `hrr_algebra`, `hrr_dim`, `fact_count`, + `updated_at`). +- `memory_feedback_events` — append-only `helpful`/`unhelpful` audit events + keyed by numeric `fact_id`, with source, note, old/new trust, and trust delta. + +Older `memory_decisions` / `memory_code_areas` tables are migration-only inputs: +v11 backfills them into `memory_facts` and then drops the legacy tables. ## Recipes @@ -273,7 +306,11 @@ impl McpServer { pub async fn new(cg: TokenSave, scope_prefix: Option) -> Arc { let file_token_map = cg.get_file_token_map().await.unwrap_or_default(); let persisted = cg.get_tokens_saved().await.unwrap_or(0); - let global_db = GlobalDb::open().await; + let global_db = if global_db_enabled() { + GlobalDb::open().await + } else { + None + }; // Register this project in the global DB with its current tokens if let Some(ref gdb) = global_db { gdb.upsert(cg.project_root(), persisted).await; @@ -440,6 +477,27 @@ impl McpServer { .unwrap_or_default() .as_secs() as i64; self.last_staleness_check_at.store(now, Ordering::Release); + + // Best-effort transcript ingestion sweep for hookless agents (Claude, + // Codex, Gemini). Cursor ingests via its own end-of-turn hook; these + // agents register no hook, so their transcripts are reconciled here. + // Detached + timeout-guarded so it never delays MCP readiness, and + // independent of the catch-up completion flag below; per-file + // parse_offsets make repeat sweeps cheap no-ops. + { + let project_root = self.cg.project_root().to_path_buf(); + tokio::spawn(async move { + let _ = tokio::time::timeout(std::time::Duration::from_secs(20), async move { + if let Some(db) = + crate::sessions::cursor::open_project_session_db(&project_root).await + { + let _ = crate::sessions::ingest_global_sources(&db, &project_root).await; + } + }) + .await; + }); + } + self.startup_catch_up_done.store(true, Ordering::Release); } @@ -545,6 +603,10 @@ impl McpServer { } let delta = current - last_flushed; + if self.global_db.is_none() { + return; + } + let success = tokio::task::spawn_blocking(move || { let mut config = crate::user_config::UserConfig::load(); config.pending_upload += delta; @@ -768,7 +830,7 @@ impl McpServer { // Flush remaining delta to worldwide counter (what periodic flushes missed) let last_flushed = self.last_flushed_tokens.load(Ordering::Relaxed); - if tokens_saved > last_flushed { + if self.global_db.is_some() && tokens_saved > last_flushed { let delta = tokens_saved - last_flushed; let mut config = crate::user_config::UserConfig::load(); config.pending_upload += delta; @@ -1235,7 +1297,7 @@ impl McpServer { } // Persist to the cross-project savings ledger (best-effort, non-blocking). - { + if self.global_db.is_some() { let project_path_str = self.cg.project_root().to_string_lossy().to_string(); let tool_name_owned = tool_name.to_string(); let ts = crate::tokensave::current_timestamp(); diff --git a/src/mcp/tools/definitions.rs b/src/mcp/tools/definitions.rs index dc938e5823..b3a3736fd5 100644 --- a/src/mcp/tools/definitions.rs +++ b/src/mcp/tools/definitions.rs @@ -164,9 +164,10 @@ pub fn get_tool_definitions() -> Vec { def_diagnose(), def_derives(), def_run_affected_tests(), - def_record_decision(), - def_record_code_area(), - def_session_recall(), + def_fact_store(), + def_fact_feedback(), + def_memory_status(), + def_message_search(), def_read(), def_outline(), def_implementations(), @@ -1529,6 +1530,184 @@ fn def_run_affected_tests() -> ToolDefinition { ) } +fn memory_fact_properties() -> Value { + json!({ + "action": { + "type": "string", + "enum": ["add", "search", "probe", "related", "reason", "contradict", "update", "remove", "list"], + "description": "Fact-store action to perform." + }, + "content": { + "type": "string", + "description": "Fact content for add/update actions." + }, + "query": { + "type": "string", + "description": "Search query for search actions." + }, + "entity": { + "type": "string", + "description": "Single entity name for probe/related actions, or extra add entity." + }, + "entities": { + "type": "array", + "items": { "type": "string" }, + "description": "Entity names for add/update/reason actions." + }, + "fact_id": { + "oneOf": [{ "type": "number" }, { "type": "string" }], + "description": "Fact id for update/remove/feedback; numeric strings are accepted." + }, + "category": { + "type": "string", + "enum": ["general", "user_pref", "project", "tool", "decision", "code_area"], + "description": "Optional fact category." + }, + "tags": { + "type": "array", + "items": { "type": "string" }, + "description": "Free-form tags stored with fact metadata." + }, + "min_trust": { + "type": "number", + "description": "Minimum trust score for search/list actions." + }, + "trust": { + "type": "number", + "minimum": 0, + "maximum": 1, + "description": "Initial or replacement trust score for add/update actions." + }, + "trust_delta": { + "type": "number", + "description": "Hermes-compatible trust delta field. Current feedback actions apply the built-in helpful/unhelpful deltas." + }, + "threshold": { + "type": "number", + "description": "Threshold for contradiction scans." + }, + "limit": { + "type": "number", + "description": "Maximum number of facts to return (default: 20, max: 200)." + }, + "source": { + "type": "string", + "description": "Source label for facts or feedback." + }, + "metadata": { + "type": "object", + "description": "Arbitrary structured metadata stored with the fact." + }, + "note": { + "type": "string", + "description": "Human-readable feedback note or action context." + } + }) +} + +fn def_fact_store() -> ToolDefinition { + def_rw( + "tokensave_fact_store", + "Fact Store", + "Add, search, probe, relate, reason over, update, remove, or list holographic memory facts. The action field selects the operation.", + json!({ + "type": "object", + "properties": memory_fact_properties(), + "required": ["action"] + }), + ) +} + +fn def_fact_feedback() -> ToolDefinition { + def_rw( + "tokensave_fact_feedback", + "Fact Feedback", + "Record helpful/unhelpful feedback for a memory fact and adjust its trust score.", + json!({ + "type": "object", + "properties": { + "fact_id": { + "oneOf": [{ "type": "number" }, { "type": "string" }], + "description": "Fact id; numeric strings are accepted." + }, + "action": { + "type": "string", + "enum": ["helpful", "unhelpful"], + "description": "Feedback action." + }, + "helpful": { + "type": "boolean", + "description": "Hermes-compatible shorthand for action=helpful." + }, + "unhelpful": { + "type": "boolean", + "description": "Hermes-compatible shorthand for action=unhelpful." + }, + "trust_delta": { + "type": "number", + "description": "Hermes-compatible trust delta field. Built-in action deltas are applied." + }, + "source": { + "type": "string", + "description": "Feedback source label." + }, + "metadata": { + "type": "object", + "description": "Additional feedback metadata reserved for compatibility." + }, + "note": { + "type": "string", + "description": "Optional feedback note." + } + }, + "required": ["fact_id"] + }), + ) +} + +fn def_memory_status() -> ToolDefinition { + def_rw( + "tokensave_memory_status", + "Memory Status", + "Repair dirty holographic memory banks, then return fact/entity counts and trust distribution.", + json!({ + "type": "object", + "properties": {} + }), + ) +} + +fn def_message_search() -> ToolDefinition { + def( + "tokensave_message_search", + "Message Search", + "Search ingested Cursor/Codex/agent transcript messages stored in tokensave's project-local session-message FTS index.", + json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Full-text query to search in ingested transcript messages." + }, + "provider": { + "type": "string", + "description": "Message provider to search (default: cursor).", + "enum": ["cursor", "claude", "codex", "vibe", "cline", "roo-code", "kilo"] + }, + "project_key": { + "type": "string", + "description": "Optional project key/path filter. For Cursor transcripts this is the project root path." + }, + "limit": { + "type": "number", + "description": "Maximum number of messages to return (default: 10, max: 50)." + } + }, + "required": ["query"] + }), + ) +} + fn def_ast_grep_rewrite() -> ToolDefinition { ToolDefinition { name: "tokensave_ast_grep_rewrite".to_string(), @@ -1643,97 +1822,6 @@ fn def_todos() -> ToolDefinition { ) } -fn def_record_decision() -> ToolDefinition { - ToolDefinition { - name: "tokensave_record_decision".to_string(), - description: "Persist a design or architecture decision so it can be recalled in a future session via tokensave_session_recall. Use for choices the agent or user would otherwise have to re-explain (e.g. \"use JWT for auth — session tokens flagged by legal\"). Stored in the per-project DB.".to_string(), - input_schema: json!({ - "type": "object", - "properties": { - "text": { - "type": "string", - "description": "The decision itself, in one sentence (e.g. \"use JWT for auth\")." - }, - "reason": { - "type": "string", - "description": "Optional reason / context (e.g. \"session tokens flagged by legal\")." - }, - "files": { - "type": "array", - "items": { "type": "string" }, - "description": "File paths that the decision applies to." - }, - "tags": { - "type": "array", - "items": { "type": "string" }, - "description": "Free-form tags for grouping (e.g. \"security\", \"performance\")." - } - }, - "required": ["text"] - }), - annotations: Some(json!({ - "readOnlyHint": false, - "title": "Record Decision" - })), - meta: None, - } -} - -fn def_record_code_area() -> ToolDefinition { - ToolDefinition { - name: "tokensave_record_code_area".to_string(), - description: "Record that the agent has been working in a code area (a file or directory). The first call sets an optional description; subsequent calls bump the touch counter and update last_touched_at. Recall with tokensave_session_recall.".to_string(), - input_schema: json!({ - "type": "object", - "properties": { - "path": { - "type": "string", - "description": "File or directory path (project-relative)." - }, - "description": { - "type": "string", - "description": "Optional short description of what this area is or what was changed." - } - }, - "required": ["path"] - }), - annotations: Some(json!({ - "readOnlyHint": false, - "title": "Record Code Area" - })), - meta: None, - } -} - -fn def_session_recall() -> ToolDefinition { - def( - "tokensave_session_recall", - "Session Recall", - "Recall persisted decisions (and optionally code areas) from past sessions. When `query` is provided, runs FTS5 search across decision text and reason. When omitted, returns the most recent decisions newest-first. Pair with tokensave_record_decision.", - json!({ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "FTS5 query string (e.g. \"auth OR session\"). Omit for newest-first listing." - }, - "since": { - "type": "number", - "description": "Unix timestamp; only return decisions made at-or-after this time." - }, - "limit": { - "type": "number", - "description": "Maximum decisions to return (default: 20, max: 200)." - }, - "include_code_areas": { - "type": "boolean", - "description": "If true, also return the top-touched code areas (default: false)." - } - } - }), - ) -} - fn def_field_sites() -> ToolDefinition { def( "tokensave_field_sites", diff --git a/src/mcp/tools/handlers/memory.rs b/src/mcp/tools/handlers/memory.rs index 9e7244661c..496a983aee 100644 --- a/src/mcp/tools/handlers/memory.rs +++ b/src/mcp/tools/handlers/memory.rs @@ -1,88 +1,297 @@ -//! Cross-session memory handlers: `record_decision`, `record_code_area`, `session_recall`. +//! Cross-session and holographic memory handlers. use serde_json::{json, Value}; use crate::errors::{Result, TokenSaveError}; +use crate::memory::types::{ + AddFactRequest, FeedbackAction, FeedbackRequest, MemoryCategory, SearchFactsRequest, + UpdateFactRequest, +}; use crate::tokensave::TokenSave; use super::super::ToolResult; use super::truncate_response; -pub(super) async fn handle_record_decision(cg: &TokenSave, args: Value) -> Result { - let text = args - .get("text") - .and_then(|v| v.as_str()) - .ok_or_else(|| TokenSaveError::Config { - message: "missing required parameter: text".to_string(), - })?; - let reason = args.get("reason").and_then(|v| v.as_str()); - let files: Vec = args - .get("files") - .and_then(|v| v.as_array()) - .map(|a| { - a.iter() - .filter_map(|x| x.as_str().map(String::from)) - .collect() - }) - .unwrap_or_default(); - let tags: Vec = args - .get("tags") - .and_then(|v| v.as_array()) - .map(|a| { - a.iter() - .filter_map(|x| x.as_str().map(String::from)) +fn tool_json(value: &Value) -> ToolResult { + let formatted = serde_json::to_string_pretty(value).unwrap_or_default(); + ToolResult { + value: json!({ "content": [{ "type": "text", "text": truncate_response(&formatted) }] }), + touched_files: vec![], + } +} + +fn config_error(message: impl Into) -> TokenSaveError { + TokenSaveError::Config { + message: message.into(), + } +} + +fn required_str<'a>(args: &'a Value, key: &str) -> Result<&'a str> { + args.get(key) + .and_then(Value::as_str) + .ok_or_else(|| config_error(format!("missing required parameter: {key}"))) +} + +fn optional_category(args: &Value) -> Result> { + args.get("category") + .and_then(Value::as_str) + .map(str::parse::) + .transpose() + .map_err(|e| config_error(format!("invalid category: {e}"))) +} + +fn limit(args: &Value) -> usize { + args.get("limit") + .and_then(Value::as_u64) + .map_or(20, |n| (n as usize).clamp(1, 200)) +} + +fn optional_f64(args: &Value, key: &str) -> Option { + args.get(key).and_then(Value::as_f64) +} + +fn string_array(args: &Value, key: &str) -> Vec { + args.get(key) + .and_then(Value::as_array) + .map(|items| { + items + .iter() + .filter_map(|item| item.as_str().map(ToOwned::to_owned)) .collect() }) - .unwrap_or_default(); + .unwrap_or_default() +} - let id = cg.record_decision(text, reason, &files, &tags).await?; - let out = json!({ "id": id, "status": "recorded" }); - let formatted = serde_json::to_string_pretty(&out).unwrap_or_default(); - Ok(ToolResult { - value: json!({ "content": [{ "type": "text", "text": truncate_response(&formatted) }] }), - touched_files: vec![], - }) +fn fact_id(args: &Value) -> Result { + let value = args + .get("fact_id") + .or_else(|| args.get("id")) + .ok_or_else(|| config_error("missing required parameter: fact_id"))?; + if let Some(id) = value.as_i64() { + return Ok(id); + } + value + .as_str() + .and_then(|s| s.parse::().ok()) + .ok_or_else(|| config_error("fact_id must be a number or numeric string")) } -pub(super) async fn handle_record_code_area(cg: &TokenSave, args: Value) -> Result { - let path = args - .get("path") - .and_then(|v| v.as_str()) - .ok_or_else(|| TokenSaveError::Config { - message: "missing required parameter: path".to_string(), - })?; - let description = args.get("description").and_then(|v| v.as_str()); +fn metadata_with_tags(args: &Value) -> Value { + let mut metadata = args + .get("metadata") + .cloned() + .filter(Value::is_object) + .unwrap_or_else(|| json!({})); + let tags = string_array(args, "tags"); + if !tags.is_empty() { + if let Some(map) = metadata.as_object_mut() { + map.insert("tags".to_string(), json!(tags)); + } + } + metadata +} - cg.record_code_area(path, description).await?; - let out = json!({ "path": path, "status": "recorded" }); - let formatted = serde_json::to_string_pretty(&out).unwrap_or_default(); - Ok(ToolResult { - value: json!({ "content": [{ "type": "text", "text": truncate_response(&formatted) }] }), - touched_files: vec![], - }) +fn request_entities(args: &Value) -> Vec { + let mut entities = string_array(args, "entities"); + if let Some(entity) = args.get("entity").and_then(Value::as_str) { + entities.push(entity.to_string()); + } + entities } -pub(super) async fn handle_session_recall(cg: &TokenSave, args: Value) -> Result { - let query = args.get("query").and_then(|v| v.as_str()); - let since = args.get("since").and_then(serde_json::Value::as_i64); - let limit = args - .get("limit") - .and_then(serde_json::Value::as_u64) - .map_or(20, |n| (n as usize).clamp(1, 200)); - let include_areas = args - .get("include_code_areas") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false); - - let decisions = cg.session_recall(query, since, limit).await?; - let mut out = json!({ "decisions": decisions }); - if include_areas { - let areas = cg.list_code_areas(limit).await?; - out["code_areas"] = serde_json::to_value(&areas).unwrap_or(json!([])); +fn feedback_action(args: &Value) -> Result { + if let Some(action) = args.get("action").and_then(Value::as_str) { + return match action { + "helpful" => Ok(FeedbackAction::Helpful), + "unhelpful" => Ok(FeedbackAction::Unhelpful), + other => Err(config_error(format!("unknown feedback action: {other}"))), + }; } - let formatted = serde_json::to_string_pretty(&out).unwrap_or_default(); - Ok(ToolResult { - value: json!({ "content": [{ "type": "text", "text": truncate_response(&formatted) }] }), - touched_files: vec![], + match ( + args.get("helpful") + .and_then(Value::as_bool) + .unwrap_or(false), + args.get("unhelpful") + .and_then(Value::as_bool) + .unwrap_or(false), + ) { + (true, false) => Ok(FeedbackAction::Helpful), + (false, true) => Ok(FeedbackAction::Unhelpful), + _ => Err(config_error( + "missing feedback action: set action, helpful, or unhelpful", + )), + } +} + +fn results_envelope(action: &str, results: &Value, count: usize) -> Value { + json!({ + "action": action, + "results": results, + "facts": results, + "count": count, }) } + +async fn update_trust(args: &Value, cg: &TokenSave, fact_id: i64) -> Result> { + if let Some(trust) = optional_f64(args, "trust") { + return Ok(Some(trust)); + } + let Some(delta) = optional_f64(args, "trust_delta") else { + return Ok(None); + }; + let existing = cg + .get_fact(fact_id) + .await? + .ok_or_else(|| config_error(format!("fact {fact_id} not found")))?; + Ok(Some((existing.trust_score + delta).clamp(0.0, 1.0))) +} + +pub(super) async fn handle_fact_store(cg: &TokenSave, args: Value) -> Result { + let action = required_str(&args, "action")?; + let out = match action { + "add" => { + let fact = cg + .add_fact(AddFactRequest { + content: required_str(&args, "content")?.to_string(), + category: optional_category(&args)?.unwrap_or(MemoryCategory::General), + source: args + .get("source") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + tags: string_array(&args, "tags"), + entities: request_entities(&args), + trust: optional_f64(&args, "trust"), + metadata: metadata_with_tags(&args), + }) + .await?; + json!({ "action": action, "fact": fact, "count": 1 }) + } + "search" => { + let facts = cg + .search_facts(SearchFactsRequest { + query: required_str(&args, "query")?.to_string(), + category: optional_category(&args)?, + limit: Some(limit(&args)), + min_trust: optional_f64(&args, "min_trust"), + include_why: true, + }) + .await?; + let count = facts.len(); + results_envelope(action, &json!(facts), count) + } + "probe" => { + let facts = cg + .probe_entity( + required_str(&args, "entity")?, + optional_category(&args)?, + optional_f64(&args, "min_trust"), + limit(&args), + ) + .await?; + let count = facts.len(); + results_envelope(action, &json!(facts), count) + } + "related" => { + let facts = cg + .related_facts( + required_str(&args, "entity")?, + optional_category(&args)?, + optional_f64(&args, "min_trust"), + limit(&args), + ) + .await?; + let count = facts.len(); + results_envelope(action, &json!(facts), count) + } + "reason" => { + let entities = request_entities(&args); + let facts = cg + .reason_facts( + &entities, + optional_category(&args)?, + optional_f64(&args, "min_trust"), + limit(&args), + ) + .await?; + let count = facts.len(); + results_envelope(action, &json!(facts), count) + } + "contradict" => { + let facts = cg + .contradict_facts( + optional_category(&args)?, + optional_f64(&args, "threshold").unwrap_or(0.3), + limit(&args), + ) + .await?; + let count = facts.len(); + results_envelope(action, &json!(facts), count) + } + "update" => { + let id = fact_id(&args)?; + let fact = cg + .update_fact(UpdateFactRequest { + fact_id: id, + content: args + .get("content") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + category: optional_category(&args)?, + tags: args.get("tags").map(|_| string_array(&args, "tags")), + entities: args.get("entities").map(|_| request_entities(&args)), + trust: update_trust(&args, cg, id).await?, + source: args + .get("source") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + metadata: args.get("metadata").cloned(), + }) + .await?; + json!({ "action": action, "fact": fact, "count": 1 }) + } + "remove" => { + let removed = cg.remove_fact(fact_id(&args)?).await?; + json!({ "action": action, "removed": removed, "count": usize::from(removed) }) + } + "list" => { + let facts = cg + .list_facts( + optional_category(&args)?, + optional_f64(&args, "min_trust"), + limit(&args), + ) + .await?; + let count = facts.len(); + results_envelope(action, &json!(facts), count) + } + other => return Err(config_error(format!("unknown fact_store action: {other}"))), + }; + Ok(tool_json(&out)) +} + +pub(super) async fn handle_fact_feedback(cg: &TokenSave, args: Value) -> Result { + let note = args + .get("note") + .or_else(|| args.get("reason")) + .and_then(Value::as_str) + .map(ToOwned::to_owned); + let result = cg + .record_fact_feedback(FeedbackRequest { + fact_id: fact_id(&args)?, + action: feedback_action(&args)?, + source: args + .get("source") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + note, + }) + .await?; + Ok(tool_json( + &json!({ "status": "recorded", "feedback": result }), + )) +} + +pub(super) async fn handle_memory_status(cg: &TokenSave) -> Result { + let status = cg.memory_status().await?; + Ok(tool_json(&json!({ "status": "ok", "memory": status }))) +} diff --git a/src/mcp/tools/handlers/mod.rs b/src/mcp/tools/handlers/mod.rs index b91b8656f9..0107ce57bd 100644 --- a/src/mcp/tools/handlers/mod.rs +++ b/src/mcp/tools/handlers/mod.rs @@ -12,6 +12,7 @@ pub mod health; pub mod info; pub mod memory; pub mod redundancy; +pub mod session; pub mod workflow; use std::collections::HashSet; @@ -204,9 +205,10 @@ pub async fn handle_tool_call( "tokensave_diagnose" => workflow::handle_diagnose(cg, args).await, "tokensave_run_affected_tests" => workflow::handle_run_affected_tests(cg, args).await, "tokensave_derives" => graph::handle_derives(cg, args).await, - "tokensave_record_decision" => memory::handle_record_decision(cg, args).await, - "tokensave_record_code_area" => memory::handle_record_code_area(cg, args).await, - "tokensave_session_recall" => memory::handle_session_recall(cg, args).await, + "tokensave_fact_store" => memory::handle_fact_store(cg, args).await, + "tokensave_fact_feedback" => memory::handle_fact_feedback(cg, args).await, + "tokensave_memory_status" => memory::handle_memory_status(cg).await, + "tokensave_message_search" => session::handle_message_search(cg, args).await, _ => Err(TokenSaveError::Config { message: format!("unknown tool: {tool_name}"), }), @@ -234,9 +236,9 @@ mod tests { // tool that will instantly fail. The count and the per-tool checks // below adapt to the host's capability set. let expected_total = if super::super::definitions::ast_grep_available() { - 76 + 77 } else { - 75 + 76 }; assert_eq!(tools.len(), expected_total); @@ -252,6 +254,10 @@ mod tests { assert!(tool_names.contains(&"tokensave_diagnose")); assert!(tool_names.contains(&"tokensave_run_affected_tests")); assert!(tool_names.contains(&"tokensave_derives")); + assert!(tool_names.contains(&"tokensave_fact_store")); + assert!(tool_names.contains(&"tokensave_fact_feedback")); + assert!(tool_names.contains(&"tokensave_memory_status")); + assert!(tool_names.contains(&"tokensave_message_search")); assert!(tool_names.contains(&"tokensave_impact")); assert!(tool_names.contains(&"tokensave_node")); assert!(tool_names.contains(&"tokensave_status")); @@ -304,9 +310,10 @@ mod tests { assert!(tool_names.contains(&"tokensave_session_end")); assert!(tool_names.contains(&"tokensave_body")); assert!(tool_names.contains(&"tokensave_todos")); - assert!(tool_names.contains(&"tokensave_record_decision")); - assert!(tool_names.contains(&"tokensave_record_code_area")); - assert!(tool_names.contains(&"tokensave_session_recall")); + assert!(tool_names.contains(&"tokensave_fact_store")); + assert!(tool_names.contains(&"tokensave_fact_feedback")); + assert!(tool_names.contains(&"tokensave_memory_status")); + assert!(tool_names.contains(&"tokensave_message_search")); assert!(tool_names.contains(&"tokensave_read")); assert!(tool_names.contains(&"tokensave_outline")); assert!(tool_names.contains(&"tokensave_implementations")); @@ -347,8 +354,9 @@ mod tests { "tokensave_run_affected_tests", "tokensave_session_start", "tokensave_session_end", - "tokensave_record_decision", - "tokensave_record_code_area", + "tokensave_fact_store", + "tokensave_fact_feedback", + "tokensave_memory_status", ]; for tool in &tools { let ann = tool diff --git a/src/mcp/tools/handlers/session.rs b/src/mcp/tools/handlers/session.rs new file mode 100644 index 0000000000..a18a77308f --- /dev/null +++ b/src/mcp/tools/handlers/session.rs @@ -0,0 +1,63 @@ +use serde_json::{json, Value}; + +use crate::errors::{Result, TokenSaveError}; +use crate::mcp::tools::ToolResult; +use crate::tokensave::TokenSave; + +use super::truncate_response; + +fn tool_json(value: &Value) -> ToolResult { + let formatted = serde_json::to_string_pretty(value).unwrap_or_default(); + ToolResult { + value: json!({ "content": [{ "type": "text", "text": truncate_response(&formatted) }] }), + touched_files: Vec::new(), + } +} + +pub(super) async fn handle_message_search(cg: &TokenSave, args: Value) -> Result { + let query = args + .get("query") + .and_then(Value::as_str) + .map(str::trim) + .filter(|query| !query.is_empty()) + .ok_or_else(|| TokenSaveError::Config { + message: "missing required parameter: query".to_string(), + })?; + let provider = args + .get("provider") + .and_then(Value::as_str) + .map(str::trim) + .filter(|provider| !provider.is_empty()) + .unwrap_or("cursor"); + let project_key = args + .get("project_key") + .and_then(Value::as_str) + .map(str::trim) + .filter(|project_key| !project_key.is_empty()); + let limit = args + .get("limit") + .and_then(Value::as_u64) + .unwrap_or(10) + .clamp(1, 50) as usize; + + let Some(db) = crate::sessions::cursor::open_project_session_db(cg.project_root()).await else { + return Ok(tool_json(&json!({ + "status": "unavailable", + "message": "could not open project-local tokensave session database", + "results": [], + "count": 0 + }))); + }; + let results = db + .search_session_messages(provider, project_key, query, limit) + .await; + + Ok(tool_json(&json!({ + "status": "ok", + "provider": provider, + "project_key": project_key, + "query": query, + "count": results.len(), + "results": results, + }))) +} diff --git a/src/memory/encoding.rs b/src/memory/encoding.rs new file mode 100644 index 0000000000..195814d722 --- /dev/null +++ b/src/memory/encoding.rs @@ -0,0 +1,198 @@ +//! Deterministic FHRR encodings for memory facts, entities, and queries. + +use amari_holographic::{BindingAlgebra, FHRRAlgebra}; +use sha2::{Digest, Sha256}; + +type Fhrr2048 = FHRRAlgebra<2048>; + +#[derive(Clone, Debug, Default)] +pub struct HolographicEncoder; + +impl HolographicEncoder { + pub const DIMENSIONS: usize = 2048; + pub const ROLE_CONTENT: &'static str = "__hrr_role_content__"; + pub const ROLE_ENTITY: &'static str = "__hrr_role_entity__"; + + pub const fn new() -> Self { + Self + } + + pub fn encode_atom(&self, label: &str) -> Vec { + normalize_coefficients(deterministic_coefficients(label)) + } + + pub fn encode_text(&self, text: &str) -> Vec { + let tokens = tokenize_text(text); + if tokens.is_empty() { + return self.encode_atom("text:__hrr_empty__"); + } + let vectors: Vec> = tokens + .iter() + .map(|token| self.encode_atom(&format!("text:{token}"))) + .collect(); + average_coefficients(&vectors) + } + + pub fn encode_fact(&self, content: &str, entities: &[String]) -> Vec { + let (Some(content_role), Some(content_value)) = ( + to_fhrr(&self.encode_atom(Self::ROLE_CONTENT)), + to_fhrr(&self.encode_text(content)), + ) else { + return Vec::new(); + }; + let mut components = vec![content_role.bind(&content_value).to_coefficients()]; + + let mut normalized_entities: Vec = entities + .iter() + .map(|entity| entity.to_ascii_lowercase()) + .filter(|entity| !entity.trim().is_empty()) + .collect(); + normalized_entities.sort(); + normalized_entities.dedup(); + + for entity in normalized_entities { + let (Some(role), Some(value)) = ( + to_fhrr(&self.encode_atom(Self::ROLE_ENTITY)), + to_fhrr(&self.encode_text(&entity)), + ) else { + continue; + }; + + let bound = role.bind(&value); + components.push(bound.to_coefficients()); + } + + average_coefficients(&components) + } + + pub fn similarity(&self, left: &[f64], right: &[f64]) -> f64 { + if let (Some(left_fhrr), Some(right_fhrr)) = (to_fhrr(left), to_fhrr(right)) { + return left_fhrr.similarity(&right_fhrr); + } + + cosine_similarity(left, right) + } + + pub fn serialize(coefficients: &[f64]) -> bincode::Result> { + bincode::serialize(&coefficients.to_vec()) + } + + pub fn deserialize(bytes: &[u8]) -> bincode::Result> { + bincode::deserialize(bytes) + } +} + +fn deterministic_coefficients(label: &str) -> Vec { + let mut coefficients = Vec::with_capacity(HolographicEncoder::DIMENSIONS); + let mut counter = 0_u64; + + while coefficients.len() < HolographicEncoder::DIMENSIONS { + let mut hasher = Sha256::new(); + hasher.update(label.as_bytes()); + hasher.update(counter.to_le_bytes()); + let digest = hasher.finalize(); + + for chunk in digest.chunks_exact(8) { + if coefficients.len() == HolographicEncoder::DIMENSIONS { + break; + } + + let mut bytes = [0_u8; 8]; + bytes.copy_from_slice(chunk); + let unit = u64::from_le_bytes(bytes) as f64 / u64::MAX as f64; + coefficients.push(unit.mul_add(2.0, -1.0)); + } + + counter = counter.saturating_add(1); + } + + coefficients +} + +fn tokenize_text(text: &str) -> Vec { + let mut tokens = Vec::new(); + let mut current = String::new(); + for ch in text.chars() { + if ch.is_alphanumeric() || matches!(ch, '_' | '/' | ':' | '.') { + current.extend(ch.to_lowercase()); + } else if !current.is_empty() { + push_token(&mut tokens, &mut current); + } + } + if !current.is_empty() { + push_token(&mut tokens, &mut current); + } + tokens.sort(); + tokens.dedup(); + tokens +} + +fn push_token(tokens: &mut Vec, current: &mut String) { + if current.len() >= 2 { + tokens.push(std::mem::take(current)); + } else { + current.clear(); + } +} + +fn average_coefficients(vectors: &[Vec]) -> Vec { + if vectors.is_empty() { + return vec![0.0; HolographicEncoder::DIMENSIONS]; + } + let mut average = vec![0.0; HolographicEncoder::DIMENSIONS]; + let mut count = 0.0; + for vector in vectors { + if vector.len() != HolographicEncoder::DIMENSIONS { + continue; + } + count += 1.0; + for (target, value) in average.iter_mut().zip(vector) { + *target += value; + } + } + if count > 0.0 { + for value in &mut average { + *value /= count; + } + } + normalize_coefficients(average) +} + +fn normalize_coefficients(mut coefficients: Vec) -> Vec { + let norm = coefficients + .iter() + .map(|coefficient| coefficient * coefficient) + .sum::() + .sqrt(); + + if norm > f64::EPSILON { + for coefficient in &mut coefficients { + *coefficient /= norm; + } + } + + coefficients +} + +fn to_fhrr(coefficients: &[f64]) -> Option { + Fhrr2048::from_coefficients(coefficients).ok() +} + +fn cosine_similarity(left: &[f64], right: &[f64]) -> f64 { + if left.len() != right.len() || left.is_empty() { + return 0.0; + } + + let (dot, left_norm, right_norm) = left + .iter() + .zip(right.iter()) + .fold((0.0, 0.0, 0.0), |(dot, left_norm, right_norm), (l, r)| { + (dot + l * r, left_norm + l * l, right_norm + r * r) + }); + + if left_norm <= f64::EPSILON || right_norm <= f64::EPSILON { + 0.0 + } else { + dot / (left_norm.sqrt() * right_norm.sqrt()) + } +} diff --git a/src/memory/entities.rs b/src/memory/entities.rs new file mode 100644 index 0000000000..aea52ad844 --- /dev/null +++ b/src/memory/entities.rs @@ -0,0 +1,244 @@ +use std::collections::HashSet; + +pub fn normalize_entity(entity: &str) -> String { + entity + .trim_matches(|c: char| { + c.is_ascii_punctuation() && c != '_' && c != '/' && c != '\\' && c != ':' && c != '.' + }) + .split_whitespace() + .collect::>() + .join(" ") +} + +pub fn extract_entities(text: &str) -> Vec { + let mut matches = Vec::new(); + + matches.extend(extract_quoted(text, '"')); + matches.extend(extract_quoted(text, '\'')); + matches.extend(extract_aliases(text)); + matches.extend(extract_code_tokens(text)); + matches.extend(extract_capitalized_names(text)); + + matches.sort_by_key(|(index, _)| *index); + + let mut seen = HashSet::new(); + let mut entities = Vec::new(); + for (_, entity) in matches { + let normalized = normalize_entity(&entity); + if normalized.is_empty() { + continue; + } + + let key = normalized.to_ascii_lowercase(); + if seen.insert(key) { + entities.push(normalized); + } + } + + entities +} + +fn extract_quoted(text: &str, delimiter: char) -> Vec<(usize, String)> { + let mut results = Vec::new(); + let mut start = None; + + for (index, ch) in text.char_indices() { + if ch != delimiter { + continue; + } + + if let Some(open_index) = start { + let content_start = open_index + delimiter.len_utf8(); + if content_start < index { + results.push((content_start, text[content_start..index].to_string())); + } + start = None; + } else { + start = Some(index); + } + } + + results +} + +fn extract_aliases(text: &str) -> Vec<(usize, String)> { + let lower = text.to_ascii_lowercase(); + [" aka ", " a.k.a. ", " also known as "] + .into_iter() + .flat_map(|marker| { + lower + .match_indices(marker) + .filter_map(move |(index, matched)| { + let phrase_start = index + matched.len(); + let phrase = take_entity_phrase(&text[phrase_start..]); + if phrase.is_empty() { + None + } else { + Some((phrase_start, phrase)) + } + }) + }) + .collect() +} + +fn take_entity_phrase(text: &str) -> String { + let trimmed_start = text.len() - text.trim_start().len(); + let remaining = &text[trimmed_start..]; + let mut end = remaining.len(); + + for (index, _) in remaining.char_indices() { + let rest = &remaining[index..]; + if index > 0 && (rest.starts_with(" in ") || rest.starts_with(" via ")) { + end = index; + break; + } + + if let Some(ch) = rest.chars().next() { + if matches!(ch, ',' | '.' | ';' | '"' | '\'') { + end = index; + break; + } + } + } + + normalize_entity(&remaining[..end]) +} + +fn extract_code_tokens(text: &str) -> Vec<(usize, String)> { + token_spans(text) + .into_iter() + .filter_map(|(index, token)| { + let cleaned = clean_code_token(token); + if is_file_path(&cleaned) || is_rust_symbol(&cleaned) || is_tokensave_tool(&cleaned) { + Some((index, cleaned)) + } else { + None + } + }) + .collect() +} + +fn extract_capitalized_names(text: &str) -> Vec<(usize, String)> { + let mut results = Vec::new(); + let mut current = Vec::new(); + let mut start_index = 0; + + for (index, token) in token_spans(text) { + let word = clean_name_token(token); + if is_capitalized_word(&word) { + if current.is_empty() { + start_index = index; + } + current.push(word); + } else { + push_capitalized_sequence(&mut results, start_index, &mut current); + } + } + + push_capitalized_sequence(&mut results, start_index, &mut current); + results +} + +fn push_capitalized_sequence( + results: &mut Vec<(usize, String)>, + start_index: usize, + current: &mut Vec, +) { + if current.len() >= 2 && !is_non_entity_leading_word(¤t[0]) { + results.push((start_index, current.join(" "))); + } + current.clear(); +} + +fn is_non_entity_leading_word(token: &str) -> bool { + matches!( + token, + "Add" + | "Avoid" + | "Create" + | "Delete" + | "Do" + | "Fix" + | "Implement" + | "Keep" + | "Persist" + | "Prefer" + | "Record" + | "Remove" + | "Update" + | "Use" + ) +} + +fn token_spans(text: &str) -> Vec<(usize, &str)> { + let mut spans = Vec::new(); + let mut offset = 0; + + for token in text.split_whitespace() { + if let Some(relative_index) = text[offset..].find(token) { + let index = offset + relative_index; + spans.push((index, token)); + offset = index + token.len(); + } + } + + spans +} + +fn clean_code_token(token: &str) -> String { + let cleaned = token + .trim_matches(|c: char| { + c.is_ascii_punctuation() + && c != '_' + && c != '/' + && c != '\\' + && c != '.' + && c != ':' + && c != '-' + }) + .trim_end_matches("()") + .to_string(); + + let normalized_tool = cleaned.replace('-', "_").to_ascii_lowercase(); + if normalized_tool.starts_with("tokensave_") { + normalized_tool.trim_end_matches('.').to_string() + } else { + cleaned.trim_end_matches('.').to_string() + } +} + +fn clean_name_token(token: &str) -> String { + token + .trim_matches(|c: char| c.is_ascii_punctuation() && c != '-' && c != '_') + .to_string() +} + +fn is_file_path(token: &str) -> bool { + token.contains('/') || token.contains('\\') || token.starts_with('.') +} + +fn is_rust_symbol(token: &str) -> bool { + token.contains("::") + && token + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | ':')) +} + +fn is_tokensave_tool(token: &str) -> bool { + let normalized = token.replace('-', "_").to_ascii_lowercase(); + normalized.starts_with("tokensave_") + && token + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || ch == '_' || ch == '-') +} + +fn is_capitalized_word(token: &str) -> bool { + let mut chars = token.chars(); + let Some(first) = chars.next() else { + return false; + }; + + first.is_ascii_uppercase() + && chars.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_')) + && !token.contains("::") +} diff --git a/src/memory/mod.rs b/src/memory/mod.rs new file mode 100644 index 0000000000..d87994c467 --- /dev/null +++ b/src/memory/mod.rs @@ -0,0 +1,8 @@ +//! Holographic memory storage, retrieval, scoring, and trust support. + +pub mod encoding; +pub mod entities; +pub mod retrieval; +pub mod store; +pub mod trust; +pub mod types; diff --git a/src/memory/retrieval.rs b/src/memory/retrieval.rs new file mode 100644 index 0000000000..ca28c1d74a --- /dev/null +++ b/src/memory/retrieval.rs @@ -0,0 +1,741 @@ +//! Query, entity, and contradiction retrieval over stored memory facts. + +use std::collections::{HashMap, HashSet}; +use std::fmt; + +use libsql::{params, Connection}; + +use super::encoding::HolographicEncoder; +use super::entities::normalize_entity; +use super::store::MemoryStore; +use super::trust::DEFAULT_MIN_TRUST; +use super::types::{ + ContradictionResult, EntityRecord, FactRecord, FactSearchResult, MemoryCategory, +}; +use crate::errors::{Result, TokenSaveError}; +use crate::tokensave::current_timestamp; + +const DEFAULT_LIMIT: usize = 10; +const FTS_SCORE_WEIGHT: f64 = 0.40; +const JACCARD_SCORE_WEIGHT: f64 = 0.30; +const HOLOGRAPHIC_SCORE_WEIGHT: f64 = 0.30; + +pub struct FactRetriever<'a> { + store: MemoryStore<'a>, + encoder: HolographicEncoder, +} + +impl<'a> FactRetriever<'a> { + pub const fn new(conn: &'a Connection) -> Self { + Self { + store: MemoryStore::new(conn), + encoder: HolographicEncoder::new(), + } + } + + pub async fn search( + &self, + query: &str, + category: Option, + min_trust: Option, + limit: usize, + ) -> Result> { + let min_trust = min_trust.unwrap_or(DEFAULT_MIN_TRUST); + let limit = normalized_limit(limit); + let query_tokens = tokenize(query); + let fts_scores = self + .fts_candidates(query, category, min_trust, limit.saturating_mul(5)) + .await?; + let entity_candidate_ids = self + .entity_candidates( + query, + &query_tokens, + category, + min_trust, + limit.saturating_mul(10), + ) + .await?; + let mut candidates = self + .store + .list_facts(category, Some(min_trust), limit.saturating_mul(10)) + .await?; + let mut candidate_ids: HashSet = candidates.iter().map(|fact| fact.fact_id).collect(); + // Collect the union of ids surfaced by FTS and entity matching that the + // `list_facts` baseline did not already include, then hydrate them with a + // single batched `get_facts` call instead of one round-trip per id. + let mut missing_ids: Vec = Vec::new(); + for fact_id in fts_scores.keys().copied().chain(entity_candidate_ids) { + if candidate_ids.insert(fact_id) { + missing_ids.push(fact_id); + } + } + if !missing_ids.is_empty() { + let mut hydrated = self.store.get_facts(&missing_ids).await?; + for fact_id in &missing_ids { + if let Some(fact) = hydrated.remove(fact_id) { + candidates.push(fact); + } + } + } + + if !query_tokens.is_empty() { + let fts_ids: HashSet = fts_scores.keys().copied().collect(); + candidates.retain(|fact| { + fts_ids.contains(&fact.fact_id) + || token_overlap(&query_tokens, &fact_search_tokens(fact)) > 0 + }); + } + + // Preload every candidate's stored vector in one batched query so the + // scoring loop never makes a per-fact round-trip. Facts without a stored + // vector are absent from the map and fall back to on-the-fly encoding. + let candidate_vectors = self + .store + .fact_vectors( + &candidates + .iter() + .map(|fact| fact.fact_id) + .collect::>(), + ) + .await?; + + let mut results = Vec::with_capacity(candidates.len()); + for fact in candidates { + let fts_score = fts_scores.get(&fact.fact_id).copied().unwrap_or(0.0); + let jaccard_score = jaccard(&query_tokens, &fact_search_tokens(&fact)); + let holographic_score = + self.holographic_score_with(query, &fact, candidate_vectors.get(&fact.fact_id)); + let trust_score = fact.trust_score; + let temporal_decay = temporal_decay_factor(fact.updated_at); + let score = combined_score( + fts_score, + jaccard_score, + holographic_score, + trust_score, + temporal_decay, + ); + results.push(FactSearchResult { + fact, + score, + fts_score, + jaccard_score, + holographic_score, + trust_score, + why: Some(format!( + "fts={fts_score:.3}, jaccard={jaccard_score:.3}, holographic={holographic_score:.3}, trust={trust_score:.3}, temporal_decay={temporal_decay:.3}" + )), + }); + } + + results.sort_by(|left, right| { + right + .score + .total_cmp(&left.score) + .then_with(|| right.fact.updated_at.cmp(&left.fact.updated_at)) + }); + results.truncate(limit); + Ok(results) + } + + pub async fn probe( + &self, + entity: &str, + category: Option, + min_trust: Option, + limit: usize, + ) -> Result> { + let fact_ids = self + .fact_ids_for_entity(entity, category, min_trust, normalized_limit(limit)) + .await?; + self.results_for_fact_ids(&fact_ids, "entity probe").await + } + + pub async fn related(&self, entity: &str, limit: usize) -> Result> { + let normalized = normalize_entity(entity).to_ascii_lowercase(); + let mut rows = self + .store + .conn() + .query( + "SELECT DISTINCT related.entity_id, related.name, related.normalized_name, + related.entity_type, related.created_at + FROM memory_entities source + JOIN memory_fact_entities source_fe ON source_fe.entity_id = source.entity_id + JOIN memory_fact_entities related_fe ON related_fe.fact_id = source_fe.fact_id + JOIN memory_entities related ON related.entity_id = related_fe.entity_id + WHERE source.normalized_name = ?1 + AND related.normalized_name != ?1 + ORDER BY related.name + LIMIT ?2", + params![normalized, normalized_limit(limit) as i64], + ) + .await + .map_err(|e| db_error("related", e))?; + + let mut entities = Vec::new(); + while let Some(row) = rows.next().await.map_err(|e| db_error("related", e))? { + let created_at = row.get::(4).map_err(|e| db_error("related", e))?; + entities.push(EntityRecord { + entity_id: row.get::(0).map_err(|e| db_error("related", e))?, + name: row.get::(1).map_err(|e| db_error("related", e))?, + normalized_name: row.get::(2).map_err(|e| db_error("related", e))?, + entity_type: Some(row.get::(3).map_err(|e| db_error("related", e))?), + created_at, + updated_at: created_at, + }); + } + Ok(entities) + } + + pub async fn reason( + &self, + entities: &[String], + category: Option, + min_trust: Option, + limit: usize, + ) -> Result> { + if entities.is_empty() { + return Ok(Vec::new()); + } + + let normalized: Vec = entities + .iter() + .map(|entity| normalize_entity(entity).to_ascii_lowercase()) + .filter(|entity| !entity.is_empty()) + .collect(); + if normalized.is_empty() { + return Ok(Vec::new()); + } + + let placeholders = normalized + .iter() + .map(|_| "?") + .collect::>() + .join(", "); + let required_count = normalized.len() as i64; + let min_trust = min_trust.unwrap_or(DEFAULT_MIN_TRUST); + let limit_usize = normalized_limit(limit); + let limit_i64 = limit_usize as i64; + // Bind the entity names (and the trailing scalars) as anonymous `?` + // placeholders in positional order rather than interpolating them. + let mut values: Vec = normalized + .iter() + .map(|entity| libsql::Value::Text(entity.clone())) + .collect(); + let sql = if let Some(category) = category { + values.push(libsql::Value::Text(category.as_str().to_string())); + values.push(libsql::Value::Real(min_trust)); + values.push(libsql::Value::Integer(required_count)); + values.push(libsql::Value::Integer(limit_i64)); + format!( + "SELECT f.fact_id + FROM memory_facts f + JOIN memory_fact_entities fe ON fe.fact_id = f.fact_id + JOIN memory_entities e ON e.entity_id = fe.entity_id + WHERE e.normalized_name IN ({placeholders}) + AND f.category = ? + AND f.trust_score >= ? + GROUP BY f.fact_id + HAVING COUNT(DISTINCT e.normalized_name) = ? + ORDER BY f.updated_at DESC, f.fact_id DESC + LIMIT ?" + ) + } else { + values.push(libsql::Value::Real(min_trust)); + values.push(libsql::Value::Integer(required_count)); + values.push(libsql::Value::Integer(limit_i64)); + format!( + "SELECT f.fact_id + FROM memory_facts f + JOIN memory_fact_entities fe ON fe.fact_id = f.fact_id + JOIN memory_entities e ON e.entity_id = fe.entity_id + WHERE e.normalized_name IN ({placeholders}) + AND f.trust_score >= ? + GROUP BY f.fact_id + HAVING COUNT(DISTINCT e.normalized_name) = ? + ORDER BY f.updated_at DESC, f.fact_id DESC + LIMIT ?" + ) + }; + let mut rows = self + .store + .conn() + .query(&sql, values) + .await + .map_err(|e| db_error("reason", e))?; + let mut fact_ids = Vec::new(); + while let Some(row) = rows.next().await.map_err(|e| db_error("reason", e))? { + fact_ids.push(row.get::(0).map_err(|e| db_error("reason", e))?); + } + let mut results = self + .results_for_fact_ids(&fact_ids, "entity reasoning") + .await?; + results.truncate(limit_usize); + Ok(results) + } + + pub async fn contradict( + &self, + category: MemoryCategory, + threshold: f64, + limit: usize, + ) -> Result> { + let facts = self + .store + .list_facts(Some(category), Some(0.0), usize::MAX) + .await?; + let mut results = Vec::new(); + for (index, left) in facts.iter().enumerate() { + for right in facts.iter().skip(index + 1) { + if !has_shared_entity(left, right) { + continue; + } + let left_tokens = fact_search_tokens(left); + let right_tokens = fact_search_tokens(right); + let content_similarity = jaccard(&left_tokens, &right_tokens); + let divergence = 1.0 - content_similarity; + if divergence >= threshold || polarity_conflicts(&left_tokens, &right_tokens) { + let (existing_fact, new_content) = if has_negative_marker(&left_tokens) { + (right.clone(), left.content.clone()) + } else { + (left.clone(), right.content.clone()) + }; + results.push(ContradictionResult { + existing_fact, + new_content, + score: divergence, + why: Some(format!( + "shared entities with content divergence={divergence:.3}" + )), + }); + if results.len() >= normalized_limit(limit) { + return Ok(results); + } + } + } + } + Ok(results) + } + + async fn fts_candidates( + &self, + query: &str, + category: Option, + min_trust: f64, + limit: usize, + ) -> Result> { + let Some(fts_query) = build_fts_query(query) else { + return Ok(HashMap::new()); + }; + + let sql = if category.is_some() { + "SELECT f.fact_id, bm25(memory_facts_fts) AS rank + FROM memory_facts_fts + JOIN memory_facts f ON f.rowid = memory_facts_fts.rowid + WHERE memory_facts_fts MATCH ?1 + AND f.category = ?2 + AND f.trust_score >= ?3 + ORDER BY rank + LIMIT ?4" + } else { + "SELECT f.fact_id, bm25(memory_facts_fts) AS rank + FROM memory_facts_fts + JOIN memory_facts f ON f.rowid = memory_facts_fts.rowid + WHERE memory_facts_fts MATCH ?1 + AND f.trust_score >= ?2 + ORDER BY rank + LIMIT ?3" + }; + + let mut rows = if let Some(category) = category { + self.store + .conn() + .query( + sql, + params![ + fts_query, + category.as_str(), + min_trust, + normalized_limit(limit) as i64 + ], + ) + .await + } else { + self.store + .conn() + .query( + sql, + params![fts_query, min_trust, normalized_limit(limit) as i64], + ) + .await + } + .map_err(|e| db_error("fts_candidates", e))?; + + let mut scores = HashMap::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("fts_candidates", e))? + { + let rank = row + .get::(1) + .map_err(|e| db_error("fts_candidates", e))?; + scores.insert( + row.get::(0) + .map_err(|e| db_error("fts_candidates", e))?, + 1.0 / (1.0 + rank.abs()), + ); + } + Ok(scores) + } + + async fn entity_candidates( + &self, + query: &str, + query_tokens: &[String], + category: Option, + min_trust: f64, + limit: usize, + ) -> Result> { + let mut terms = Vec::new(); + let normalized_query = normalize_entity(query).to_ascii_lowercase(); + if !normalized_query.is_empty() { + terms.push(normalized_query); + } + terms.extend(query_tokens.iter().cloned()); + terms.sort(); + terms.dedup(); + if terms.is_empty() { + return Ok(Vec::new()); + } + + // Bind each term's exact and LIKE values as anonymous `?` placeholders in + // positional order. `escape_like` still governs wildcard semantics on the + // LIKE value, but the value is bound rather than interpolated. + let mut values: Vec = Vec::with_capacity(terms.len() * 2 + 3); + let predicates = terms + .iter() + .map(|term| { + values.push(libsql::Value::Text(term.clone())); + values.push(libsql::Value::Text(format!("%{}%", escape_like(term)))); + "(e.normalized_name = ? OR e.normalized_name LIKE ? ESCAPE '\\')".to_string() + }) + .collect::>() + .join(" OR "); + + let sql = if let Some(category) = category { + values.push(libsql::Value::Text(category.as_str().to_string())); + values.push(libsql::Value::Real(min_trust)); + values.push(libsql::Value::Integer(normalized_limit(limit) as i64)); + format!( + "SELECT DISTINCT f.fact_id + FROM memory_facts f + JOIN memory_fact_entities fe ON fe.fact_id = f.fact_id + JOIN memory_entities e ON e.entity_id = fe.entity_id + WHERE ({predicates}) + AND f.category = ? + AND f.trust_score >= ? + ORDER BY f.updated_at DESC, f.fact_id DESC + LIMIT ?" + ) + } else { + values.push(libsql::Value::Real(min_trust)); + values.push(libsql::Value::Integer(normalized_limit(limit) as i64)); + format!( + "SELECT DISTINCT f.fact_id + FROM memory_facts f + JOIN memory_fact_entities fe ON fe.fact_id = f.fact_id + JOIN memory_entities e ON e.entity_id = fe.entity_id + WHERE ({predicates}) + AND f.trust_score >= ? + ORDER BY f.updated_at DESC, f.fact_id DESC + LIMIT ?" + ) + }; + + let mut rows = self + .store + .conn() + .query(sql.as_str(), values) + .await + .map_err(|e| db_error("entity_candidates", e))?; + + let mut fact_ids = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("entity_candidates", e))? + { + fact_ids.push( + row.get::(0) + .map_err(|e| db_error("entity_candidates", e))?, + ); + } + Ok(fact_ids) + } + + async fn fact_ids_for_entity( + &self, + entity: &str, + category: Option, + min_trust: Option, + limit: usize, + ) -> Result> { + let normalized = normalize_entity(entity).to_ascii_lowercase(); + if normalized.is_empty() { + return Ok(Vec::new()); + } + + let sql = if category.is_some() { + "SELECT fe.fact_id + FROM memory_entities e + JOIN memory_fact_entities fe ON fe.entity_id = e.entity_id + JOIN memory_facts f ON f.fact_id = fe.fact_id + WHERE e.normalized_name = ?1 + AND f.category = ?2 + AND f.trust_score >= ?3 + ORDER BY f.updated_at DESC + LIMIT ?4" + } else { + "SELECT fe.fact_id + FROM memory_entities e + JOIN memory_fact_entities fe ON fe.entity_id = e.entity_id + JOIN memory_facts f ON f.fact_id = fe.fact_id + WHERE e.normalized_name = ?1 + AND f.trust_score >= ?2 + ORDER BY f.updated_at DESC + LIMIT ?3" + }; + let min_trust = min_trust.unwrap_or(DEFAULT_MIN_TRUST); + + let mut rows = if let Some(category) = category { + self.store + .conn() + .query( + sql, + params![ + normalized, + category.as_str(), + min_trust, + normalized_limit(limit) as i64 + ], + ) + .await + } else { + self.store + .conn() + .query( + sql, + params![normalized, min_trust, normalized_limit(limit) as i64], + ) + .await + } + .map_err(|e| db_error("fact_ids_for_entity", e))?; + + let mut fact_ids = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("fact_ids_for_entity", e))? + { + fact_ids.push( + row.get::(0) + .map_err(|e| db_error("fact_ids_for_entity", e))?, + ); + } + Ok(fact_ids) + } + + async fn results_for_fact_ids( + &self, + fact_ids: &[i64], + why: &str, + ) -> Result> { + if fact_ids.is_empty() { + return Ok(Vec::new()); + } + // One batched fetch, then iterate the ORIGINAL `fact_ids` order so the + // ordering callers rely on (probe/reason) is preserved exactly. + let facts = self.store.get_facts(fact_ids).await?; + let mut results = Vec::with_capacity(fact_ids.len()); + for fact_id in fact_ids { + if let Some(fact) = facts.get(fact_id).cloned() { + let trust_score = fact.trust_score; + results.push(FactSearchResult { + score: trust_score, + fts_score: 0.0, + jaccard_score: 0.0, + holographic_score: 1.0, + trust_score, + why: Some(why.to_string()), + fact, + }); + } + } + Ok(results) + } + + /// Holographic similarity between `query` and `fact`, using `stored_vector` + /// when present and otherwise encoding the fact's vector on the fly. This is + /// the pure form of the former `holographic_score`: callers preload vectors + /// in bulk via [`MemoryStore::fact_vectors`] and pass the result in here. + fn holographic_score_with( + &self, + query: &str, + fact: &FactRecord, + stored_vector: Option<&Vec>, + ) -> f64 { + let query_entities: Vec = tokenize(query); + let query_vector = self.encoder.encode_fact(query, &query_entities); + let similarity = if let Some(vector) = stored_vector { + self.encoder.similarity(&query_vector, vector) + } else { + let fact_vector = self.encoder.encode_fact(&fact.content, &fact.entities); + self.encoder.similarity(&query_vector, &fact_vector) + }; + f64::midpoint(similarity, 1.0).clamp(0.0, 1.0) + } +} + +fn build_fts_query(query: &str) -> Option { + let tokens = tokenize(query); + if tokens.is_empty() { + return None; + } + Some( + tokens + .into_iter() + .map(|token| format!("\"{}\"", token.replace('"', "\"\""))) + .collect::>() + .join(" OR "), + ) +} + +fn tokenize(text: &str) -> Vec { + let mut tokens = Vec::new(); + let mut current = String::new(); + for ch in text.chars() { + if ch.is_ascii_alphanumeric() || matches!(ch, '_' | '/' | ':' | '.') { + current.push(ch.to_ascii_lowercase()); + } else if !current.is_empty() { + push_token(&mut tokens, &mut current); + } + } + if !current.is_empty() { + push_token(&mut tokens, &mut current); + } + tokens.sort(); + tokens.dedup(); + tokens +} + +fn push_token(tokens: &mut Vec, current: &mut String) { + if current.len() >= 2 { + tokens.push(std::mem::take(current)); + } else { + current.clear(); + } +} + +fn fact_search_tokens(fact: &FactRecord) -> Vec { + let mut tokens = tokenize(&fact.content); + for tag in &fact.tags { + tokens.extend(tokenize(tag)); + } + for entity in &fact.entities { + tokens.extend(tokenize(entity)); + } + tokens.sort(); + tokens.dedup(); + tokens +} + +fn token_overlap(left: &[String], right: &[String]) -> usize { + let right_set: HashSet<&str> = right.iter().map(String::as_str).collect(); + left.iter() + .filter(|token| right_set.contains(token.as_str())) + .count() +} + +fn escape_like(value: &str) -> String { + value + .replace('\\', "\\\\") + .replace('%', "\\%") + .replace('_', "\\_") +} + +fn jaccard(left: &[String], right: &[String]) -> f64 { + if left.is_empty() || right.is_empty() { + return 0.0; + } + let left_set: HashSet<&str> = left.iter().map(String::as_str).collect(); + let right_set: HashSet<&str> = right.iter().map(String::as_str).collect(); + let intersection = left_set.intersection(&right_set).count(); + let union = left_set.union(&right_set).count(); + if union == 0 { + 0.0 + } else { + intersection as f64 / union as f64 + } +} + +fn combined_score( + fts: f64, + jaccard: f64, + holographic: f64, + trust: f64, + temporal_decay: f64, +) -> f64 { + let relevance = fts.mul_add( + FTS_SCORE_WEIGHT, + jaccard.mul_add(JACCARD_SCORE_WEIGHT, holographic * HOLOGRAPHIC_SCORE_WEIGHT), + ); + relevance * trust * temporal_decay.clamp(0.0, 1.0) +} + +fn temporal_decay_factor(updated_at: i64) -> f64 { + if updated_at <= 0 { + return 1.0; + } + let age_secs = current_timestamp().saturating_sub(updated_at).max(0) as f64; + let age_days = age_secs / 86_400.0; + 0.5_f64.powf(age_days / 365.0).clamp(0.10, 1.0) +} + +fn has_shared_entity(left: &FactRecord, right: &FactRecord) -> bool { + let right_entities: HashSet = right + .entities + .iter() + .map(|entity| entity.to_ascii_lowercase()) + .collect(); + left.entities + .iter() + .any(|entity| right_entities.contains(&entity.to_ascii_lowercase())) +} + +fn polarity_conflicts(left: &[String], right: &[String]) -> bool { + has_negative_marker(left) != has_negative_marker(right) +} + +fn has_negative_marker(tokens: &[String]) -> bool { + tokens.iter().any(|token| { + matches!( + token.as_str(), + "not" | "no" | "never" | "avoid" | "dont" | "don't" + ) + }) +} + +fn normalized_limit(limit: usize) -> usize { + if limit == 0 { + DEFAULT_LIMIT + } else { + limit.min(i64::MAX as usize) + } +} + +fn db_error(operation: &str, error: impl fmt::Display) -> TokenSaveError { + TokenSaveError::Database { + message: error.to_string(), + operation: operation.to_string(), + } +} diff --git a/src/memory/store.rs b/src/memory/store.rs new file mode 100644 index 0000000000..10987869bd --- /dev/null +++ b/src/memory/store.rs @@ -0,0 +1,1140 @@ +//! Persistence layer for memory facts, entities, vectors, and feedback. + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::fmt; + +use libsql::{params, Connection}; + +use super::encoding::HolographicEncoder; +use super::entities::{extract_entities, normalize_entity}; +use super::trust::{apply_feedback, clamp_trust, DEFAULT_MIN_TRUST}; +use super::types::{ + AddFactRequest, FactRecord, FeedbackAction, FeedbackRequest, FeedbackResult, MemoryCategory, + UpdateFactRequest, +}; +use crate::errors::{Result, TokenSaveError}; +use crate::tokensave::current_timestamp; + +const DEFAULT_LIMIT: usize = 50; +const ENTITY_BATCH_SIZE: usize = 500; +const MEMORY_SOURCE_DEFAULT: &str = "manual"; +const HRR_ALGEBRA: &str = "amari_fhrr"; + +pub struct MemoryStore<'a> { + conn: &'a Connection, + encoder: HolographicEncoder, +} + +impl<'a> MemoryStore<'a> { + pub const fn new(conn: &'a Connection) -> Self { + Self { + conn, + encoder: HolographicEncoder::new(), + } + } + + /// Runs `work` inside a `BEGIN IMMEDIATE` transaction, committing on success + /// and rolling back on error. The inner future is built before the + /// transaction opens, which is safe because async fns do no work until + /// polled — `work.await` is the first time any statement runs. + async fn with_immediate_tx( + &self, + operation: &str, + work: impl std::future::Future>, + ) -> Result { + self.conn + .execute("BEGIN IMMEDIATE", ()) + .await + .map_err(|e| db_error(operation, e))?; + match work.await { + Ok(value) => { + if let Err(error) = self.conn.execute("COMMIT", ()).await { + let _ = self.conn.execute("ROLLBACK", ()).await; + return Err(db_error(operation, error)); + } + Ok(value) + } + Err(error) => { + let _ = self.conn.execute("ROLLBACK", ()).await; + Err(error) + } + } + } + + pub async fn add_fact( + &self, + request: AddFactRequest, + default_trust: f64, + ) -> Result { + self.with_immediate_tx("add_fact", self.add_fact_inner(request, default_trust)) + .await + } + + async fn add_fact_inner( + &self, + request: AddFactRequest, + default_trust: f64, + ) -> Result { + let content = request.content.trim().to_string(); + if content.is_empty() { + return Err(db_message("add_fact", "fact content cannot be empty")); + } + + let now = current_timestamp(); + let entities = merge_entities(&content, &request.entities); + let tags_json = to_json_string(&request.tags, "add_fact")?; + let metadata_json = to_json_string(&request.metadata, "add_fact")?; + let vector = self.encode_vector(&content, &entities, "add_fact")?; + let source = request + .source + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| MEMORY_SOURCE_DEFAULT.to_string()); + + self.conn + .execute( + "INSERT OR IGNORE INTO memory_facts ( + content, category, tags, trust_score, created_at, + updated_at, source, metadata, hrr_vector, hrr_algebra, hrr_dim + ) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)", + params![ + content.as_str(), + request.category.as_str(), + tags_json, + clamp_trust(request.trust.unwrap_or(default_trust)), + now, + now, + source, + metadata_json, + vector, + HRR_ALGEBRA, + HolographicEncoder::DIMENSIONS as i64, + ], + ) + .await + .map_err(|e| db_error("add_fact", e))?; + + let Some(existing) = self.get_fact_by_content(&content).await? else { + return Err(db_message( + "add_fact", + "inserted or existing fact was not found by content", + )); + }; + let mut merged_entities = existing.entities.clone(); + let original_entities = merged_entities.clone(); + for entity in entities { + if !merged_entities + .iter() + .any(|stored| stored.eq_ignore_ascii_case(&entity)) + { + merged_entities.push(entity); + } + } + self.replace_fact_entities(existing.fact_id, &merged_entities) + .await?; + if merged_entities != original_entities { + self.update_fact_vector( + existing.fact_id, + &existing.content, + &merged_entities, + "add_fact", + ) + .await?; + } + let fact = self.get_fact(existing.fact_id).await?.ok_or_else(|| { + db_message( + "add_fact", + "inserted fact was not found when reading it back", + ) + })?; + self.mark_fact_banks_dirty(fact.category).await?; + Ok(fact) + } + + pub async fn update_fact(&self, request: UpdateFactRequest) -> Result { + self.with_immediate_tx("update_fact", self.update_fact_inner(request)) + .await + } + + async fn update_fact_inner(&self, request: UpdateFactRequest) -> Result { + let existing = self.get_fact(request.fact_id).await?.ok_or_else(|| { + db_message( + "update_fact", + format!("fact {} does not exist", request.fact_id), + ) + })?; + + let content = request.content.map_or_else( + || existing.content.clone(), + |value| value.trim().to_string(), + ); + if content.is_empty() { + return Err(db_message("update_fact", "fact content cannot be empty")); + } + + let category = request.category.unwrap_or(existing.category); + let tags = request.tags.unwrap_or(existing.tags); + let explicit_entities = request.entities.unwrap_or(existing.entities); + let entities = merge_entities(&content, &explicit_entities); + let trust = request.trust.map_or(existing.trust_score, clamp_trust); + let source = request.source.or(existing.source); + let metadata = request.metadata.unwrap_or(existing.metadata); + let tags_json = to_json_string(&tags, "update_fact")?; + let metadata_json = to_json_string(&metadata, "update_fact")?; + let vector = self.encode_vector(&content, &entities, "update_fact")?; + let now = current_timestamp(); + + self.conn + .execute( + "UPDATE memory_facts + SET content = ?1, + category = ?2, + tags = ?3, + trust_score = ?4, + source = ?5, + metadata = ?6, + hrr_vector = ?7, + hrr_algebra = ?8, + hrr_dim = ?9, + updated_at = ?10 + WHERE fact_id = ?11", + params![ + content, + category.as_str(), + tags_json, + trust, + source.unwrap_or_else(|| MEMORY_SOURCE_DEFAULT.to_string()), + metadata_json, + vector, + HRR_ALGEBRA, + HolographicEncoder::DIMENSIONS as i64, + now, + request.fact_id, + ], + ) + .await + .map_err(|e| db_error("update_fact", e))?; + + self.replace_fact_entities(request.fact_id, &entities) + .await?; + let updated = self.get_fact(request.fact_id).await?.ok_or_else(|| { + db_message( + "update_fact", + "updated fact was not found when reading it back", + ) + })?; + self.mark_fact_banks_dirty(existing.category).await?; + self.mark_fact_banks_dirty(updated.category).await?; + Ok(updated) + } + + pub async fn remove_fact(&self, fact_id: i64) -> Result { + self.with_immediate_tx("remove_fact", self.remove_fact_inner(fact_id)) + .await + } + + async fn remove_fact_inner(&self, fact_id: i64) -> Result { + let existing = self.get_fact(fact_id).await?; + let changed = self + .conn + .execute( + "DELETE FROM memory_facts WHERE fact_id = ?1", + params![fact_id], + ) + .await + .map_err(|e| db_error("remove_fact", e))?; + if changed > 0 { + if let Some(fact) = existing { + self.mark_fact_banks_dirty(fact.category).await?; + } + } + Ok(changed > 0) + } + + pub async fn list_facts( + &self, + category: Option, + min_trust: Option, + limit: usize, + ) -> Result> { + let min_trust = min_trust.unwrap_or(DEFAULT_MIN_TRUST); + let limit = normalized_limit(limit); + let sql = if category.is_some() { + "SELECT fact_id, content, category, tags, trust_score, source, + retrieval_count, helpful_count, unhelpful_count, + created_at, updated_at, last_retrieved_at, last_feedback_at, + metadata + FROM memory_facts + WHERE category = ?1 AND trust_score >= ?2 + ORDER BY updated_at DESC, fact_id DESC + LIMIT ?3" + } else { + "SELECT fact_id, content, category, tags, trust_score, source, + retrieval_count, helpful_count, unhelpful_count, + created_at, updated_at, last_retrieved_at, last_feedback_at, + metadata + FROM memory_facts + WHERE trust_score >= ?1 + ORDER BY updated_at DESC, fact_id DESC + LIMIT ?2" + }; + + let mut rows = if let Some(category) = category { + self.conn + .query(sql, params![category.as_str(), min_trust, limit as i64]) + .await + } else { + self.conn.query(sql, params![min_trust, limit as i64]).await + } + .map_err(|e| db_error("list_facts", e))?; + + let mut fact_ids = Vec::new(); + let mut facts = Vec::new(); + while let Some(row) = rows.next().await.map_err(|e| db_error("list_facts", e))? { + let fact = fact_from_row(&row, "list_facts", Vec::new())?; + fact_ids.push(fact.fact_id); + facts.push(fact); + } + + let mut entities_by_fact = self.load_entities_for_facts(&fact_ids).await?; + for fact in &mut facts { + fact.entities = entities_by_fact.remove(&fact.fact_id).unwrap_or_default(); + } + Ok(facts) + } + + pub async fn get_fact(&self, fact_id: i64) -> Result> { + let mut rows = self + .conn + .query( + "SELECT fact_id, content, category, tags, trust_score, source, + retrieval_count, helpful_count, unhelpful_count, + created_at, updated_at, last_retrieved_at, last_feedback_at, + metadata + FROM memory_facts + WHERE fact_id = ?1", + params![fact_id], + ) + .await + .map_err(|e| db_error("get_fact", e))?; + + let Some(row) = rows.next().await.map_err(|e| db_error("get_fact", e))? else { + return Ok(None); + }; + + Ok(Some(self.row_to_fact(&row, "get_fact").await?)) + } + + /// Bulk-loads facts by id, returning a map keyed by `fact_id`. Missing ids + /// are simply absent from the map. Entities are batch-loaded for the whole + /// set via [`Self::load_entities_for_facts`] rather than per fact, so this + /// replaces the per-id `get_fact` round-trips in the retrieval hot path. + /// + /// Ids are chunked at 256 per `IN (...)` statement to stay well clear of + /// `SQLite`'s 999-parameter limit. + pub async fn get_facts(&self, fact_ids: &[i64]) -> Result> { + const CHUNK: usize = 256; + let mut facts: HashMap = HashMap::new(); + for chunk in fact_ids.chunks(CHUNK) { + if chunk.is_empty() { + continue; + } + let placeholders = chunk.iter().map(|_| "?").collect::>().join(","); + let sql = format!( + "SELECT fact_id, content, category, tags, trust_score, source, + retrieval_count, helpful_count, unhelpful_count, + created_at, updated_at, last_retrieved_at, last_feedback_at, + metadata + FROM memory_facts + WHERE fact_id IN ({placeholders})" + ); + let values: Vec = + chunk.iter().map(|id| libsql::Value::Integer(*id)).collect(); + let mut rows = self + .conn + .query(&sql, values) + .await + .map_err(|e| db_error("get_facts", e))?; + while let Some(row) = rows.next().await.map_err(|e| db_error("get_facts", e))? { + let fact = fact_from_row(&row, "get_facts", Vec::new())?; + facts.insert(fact.fact_id, fact); + } + } + + if facts.is_empty() { + return Ok(facts); + } + let ids: Vec = facts.keys().copied().collect(); + let mut entities_by_fact = self.load_entities_for_facts(&ids).await?; + for fact in facts.values_mut() { + fact.entities = entities_by_fact.remove(&fact.fact_id).unwrap_or_default(); + } + Ok(facts) + } + + /// Bulk-loads stored HRR vectors by `fact_id`. Facts whose vector is NULL or + /// fails to decode are omitted from the map so callers fall back to encoding + /// the vector on the fly (preserving the per-fact fallback behaviour). + /// + /// Ids are chunked at 256 per `IN (...)` statement to stay well clear of + /// `SQLite`'s 999-parameter limit. + pub async fn fact_vectors(&self, fact_ids: &[i64]) -> Result>> { + const CHUNK: usize = 256; + let mut vectors: HashMap> = HashMap::new(); + for chunk in fact_ids.chunks(CHUNK) { + if chunk.is_empty() { + continue; + } + let placeholders = chunk.iter().map(|_| "?").collect::>().join(","); + let sql = format!( + "SELECT fact_id, hrr_vector FROM memory_facts WHERE fact_id IN ({placeholders})" + ); + let values: Vec = + chunk.iter().map(|id| libsql::Value::Integer(*id)).collect(); + let mut rows = self + .conn + .query(&sql, values) + .await + .map_err(|e| db_error("fact_vectors", e))?; + while let Some(row) = rows.next().await.map_err(|e| db_error("fact_vectors", e))? { + let fact_id = row.get::(0).map_err(|e| db_error("fact_vectors", e))?; + let value = row + .get::(1) + .map_err(|e| db_error("fact_vectors", e))?; + if let libsql::Value::Blob(bytes) = value { + if let Ok(vector) = HolographicEncoder::deserialize(&bytes) { + vectors.insert(fact_id, vector); + } + } + } + } + Ok(vectors) + } + + pub async fn increment_retrieval_counts(&self, fact_ids: &[i64]) -> Result<()> { + if fact_ids.is_empty() { + return Ok(()); + } + let now = current_timestamp(); + let mut counts = BTreeMap::new(); + for fact_id in fact_ids { + *counts.entry(*fact_id).or_insert(0_i64) += 1; + } + let ids: Vec = counts.keys().copied().collect(); + let id_list = sql_i64_list(&ids).ok_or_else(|| { + db_message( + "increment_retrieval_counts", + "retrieval count update had no fact ids", + ) + })?; + let increment_cases = counts + .iter() + .map(|(fact_id, count)| format!("WHEN {fact_id} THEN {count}")) + .collect::>() + .join(" "); + let sql = format!( + "UPDATE memory_facts + SET retrieval_count = retrieval_count + CASE fact_id {increment_cases} ELSE 0 END, + last_retrieved_at = ?1 + WHERE fact_id IN ({id_list})" + ); + self.conn + .execute(sql.as_str(), params![now]) + .await + .map_err(|e| db_error("increment_retrieval_counts", e))?; + Ok(()) + } + + pub async fn record_feedback_event(&self, request: FeedbackRequest) -> Result { + self.with_immediate_tx( + "record_feedback_event", + self.record_feedback_event_inner(request), + ) + .await + } + + async fn record_feedback_event_inner( + &self, + request: FeedbackRequest, + ) -> Result { + let existing = self.get_fact(request.fact_id).await?.ok_or_else(|| { + db_message( + "record_feedback_event", + format!("fact {} does not exist", request.fact_id), + ) + })?; + let old_trust = existing.trust_score; + let new_trust = apply_feedback(old_trust, request.action); + let delta = new_trust - old_trust; + let now = current_timestamp(); + let action = feedback_action_str(request.action); + let source = request + .source + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| "mcp".to_string()); + + self.conn + .execute( + "UPDATE memory_facts + SET trust_score = ?1, + helpful_count = helpful_count + ?2, + unhelpful_count = unhelpful_count + ?3, + last_feedback_at = ?4, + updated_at = ?4 + WHERE fact_id = ?5", + params![ + new_trust, + i64::from(request.action == FeedbackAction::Helpful), + i64::from(request.action == FeedbackAction::Unhelpful), + now, + request.fact_id, + ], + ) + .await + .map_err(|e| db_error("record_feedback_event", e))?; + + self.conn + .execute( + "INSERT INTO memory_feedback_events ( + fact_id, action, trust_delta, old_trust, new_trust, + created_at, source, note + ) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + params![ + request.fact_id, + action, + delta, + old_trust, + new_trust, + now, + source, + request.note, + ], + ) + .await + .map_err(|e| db_error("record_feedback_event", e))?; + + let event_id = self.last_insert_rowid("record_feedback_event").await?; + Ok(FeedbackResult { + event_id, + fact_id: request.fact_id, + action: request.action, + old_trust, + new_trust, + trust_delta: delta, + helpful_count: existing.helpful_count + + i64::from(request.action == FeedbackAction::Helpful), + unhelpful_count: existing.unhelpful_count + + i64::from(request.action == FeedbackAction::Unhelpful), + }) + } + + pub async fn compute_missing_vectors(&self, limit: usize) -> Result { + let limit = normalized_limit(limit); + let mut rows = self + .conn + .query( + "SELECT fact_id FROM memory_facts + WHERE hrr_vector IS NULL + OR hrr_algebra != ?1 + OR hrr_dim != ?2 + ORDER BY updated_at DESC + LIMIT ?3", + params![ + HRR_ALGEBRA, + HolographicEncoder::DIMENSIONS as i64, + limit as i64 + ], + ) + .await + .map_err(|e| db_error("compute_missing_vectors", e))?; + + let mut fact_ids = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("compute_missing_vectors", e))? + { + fact_ids.push( + row.get::(0) + .map_err(|e| db_error("compute_missing_vectors", e))?, + ); + } + + for fact_id in &fact_ids { + if let Some(fact) = self.get_fact(*fact_id).await? { + let vector = + self.encode_vector(&fact.content, &fact.entities, "compute_missing_vectors")?; + self.conn + .execute( + "UPDATE memory_facts + SET hrr_vector = ?1, hrr_algebra = ?2, hrr_dim = ?3, updated_at = ?4 + WHERE fact_id = ?5", + params![ + vector, + HRR_ALGEBRA, + HolographicEncoder::DIMENSIONS as i64, + current_timestamp(), + *fact_id, + ], + ) + .await + .map_err(|e| db_error("compute_missing_vectors", e))?; + } + } + + Ok(fact_ids.len()) + } + + pub async fn rebuild_bank( + &self, + bank_name: &str, + category: Option, + ) -> Result { + let (fact_count, vectors) = self.load_bank_vectors(category).await?; + if vectors.is_empty() { + self.conn + .execute( + "DELETE FROM memory_banks WHERE bank_name = ?1", + params![bank_name], + ) + .await + .map_err(|e| db_error("rebuild_bank", e))?; + return Ok(0); + } + + let averaged = average_vectors(&vectors); + let vector_bytes = HolographicEncoder::serialize(&averaged).map_err(|e| { + db_message( + "rebuild_bank", + format!("failed to serialize bank vector: {e}"), + ) + })?; + let normalized_name = normalize_bank_name(bank_name); + let now = current_timestamp(); + + self.conn + .execute( + "INSERT INTO memory_banks ( + bank_name, vector, hrr_algebra, hrr_dim, fact_count, updated_at + ) + VALUES (?1, ?2, ?3, ?4, ?5, ?6) + ON CONFLICT(bank_name) DO UPDATE SET + vector = excluded.vector, + hrr_algebra = excluded.hrr_algebra, + hrr_dim = excluded.hrr_dim, + fact_count = excluded.fact_count, + updated_at = excluded.updated_at", + params![ + normalized_name, + vector_bytes, + HRR_ALGEBRA, + HolographicEncoder::DIMENSIONS as i64, + fact_count as i64, + now, + ], + ) + .await + .map_err(|e| db_error("rebuild_bank", e))?; + + Ok(fact_count) + } + + pub async fn rebuild_all_banks(&self) -> Result { + let mut categories = Vec::new(); + let mut rows = self + .conn + .query("SELECT DISTINCT category FROM memory_facts", ()) + .await + .map_err(|e| db_error("rebuild_all_banks", e))?; + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("rebuild_all_banks", e))? + { + let category = row + .get::(0) + .map_err(|e| db_error("rebuild_all_banks", e))?; + categories.push(parse_category(&category, "rebuild_all_banks")?); + } + + let mut rebuilt = 0; + self.rebuild_bank("all", None).await?; + rebuilt += 1; + for category in categories { + self.rebuild_bank(category.as_str(), Some(category)).await?; + rebuilt += 1; + } + Ok(rebuilt) + } + + pub async fn rebuild_dirty_banks(&self) -> Result { + let mut rows = self + .conn + .query( + "SELECT bank_name FROM memory_bank_dirty ORDER BY bank_name", + (), + ) + .await + .map_err(|e| db_error("rebuild_dirty_banks", e))?; + let mut bank_names = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("rebuild_dirty_banks", e))? + { + bank_names.push( + row.get::(0) + .map_err(|e| db_error("rebuild_dirty_banks", e))?, + ); + } + + let mut rebuilt = 0; + for bank_name in bank_names { + if bank_name == "all" { + self.rebuild_bank("all", None).await?; + } else { + let category = parse_category(&bank_name, "rebuild_dirty_banks")?; + self.rebuild_bank(category.as_str(), Some(category)).await?; + } + self.conn + .execute( + "DELETE FROM memory_bank_dirty WHERE bank_name = ?1", + params![bank_name], + ) + .await + .map_err(|e| db_error("rebuild_dirty_banks", e))?; + rebuilt += 1; + } + Ok(rebuilt) + } + + pub(crate) fn conn(&self) -> &Connection { + self.conn + } + + async fn get_fact_by_content(&self, content: &str) -> Result> { + let mut rows = self + .conn + .query( + "SELECT fact_id FROM memory_facts WHERE content = ?1", + params![content], + ) + .await + .map_err(|e| db_error("get_fact_by_content", e))?; + let Some(row) = rows + .next() + .await + .map_err(|e| db_error("get_fact_by_content", e))? + else { + return Ok(None); + }; + let fact_id = row + .get::(0) + .map_err(|e| db_error("get_fact_by_content", e))?; + self.get_fact(fact_id).await + } + + async fn row_to_fact(&self, row: &libsql::Row, operation: &str) -> Result { + let fact_id = row.get::(0).map_err(|e| db_error(operation, e))?; + let entities = self.load_fact_entities(fact_id).await?; + fact_from_row(row, operation, entities) + } + + async fn load_entities_for_facts(&self, fact_ids: &[i64]) -> Result>> { + let mut entities: HashMap> = HashMap::new(); + for chunk in fact_ids.chunks(ENTITY_BATCH_SIZE) { + let Some(id_list) = sql_i64_list(chunk) else { + continue; + }; + let sql = format!( + "SELECT fe.fact_id, e.name + FROM memory_fact_entities fe + JOIN memory_entities e ON e.entity_id = fe.entity_id + WHERE fe.fact_id IN ({id_list}) + ORDER BY fe.fact_id, e.name" + ); + let mut rows = self + .conn + .query(sql.as_str(), ()) + .await + .map_err(|e| db_error("load_entities_for_facts", e))?; + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("load_entities_for_facts", e))? + { + let fact_id = row + .get::(0) + .map_err(|e| db_error("load_entities_for_facts", e))?; + let entity = row + .get::(1) + .map_err(|e| db_error("load_entities_for_facts", e))?; + entities.entry(fact_id).or_default().push(entity); + } + } + Ok(entities) + } + + async fn replace_fact_entities(&self, fact_id: i64, entities: &[String]) -> Result<()> { + self.conn + .execute( + "DELETE FROM memory_fact_entities WHERE fact_id = ?1", + params![fact_id], + ) + .await + .map_err(|e| db_error("replace_fact_entities", e))?; + + for entity in entities { + let entity_id = self.resolve_entity(entity).await?; + self.conn + .execute( + "INSERT OR IGNORE INTO memory_fact_entities (fact_id, entity_id) + VALUES (?1, ?2)", + params![fact_id, entity_id], + ) + .await + .map_err(|e| db_error("replace_fact_entities", e))?; + } + Ok(()) + } + + async fn resolve_entity(&self, entity: &str) -> Result { + let name = normalize_entity(entity); + let normalized = name.to_ascii_lowercase(); + let mut rows = self + .conn + .query( + "SELECT entity_id FROM memory_entities WHERE normalized_name = ?1", + params![normalized.as_str()], + ) + .await + .map_err(|e| db_error("resolve_entity", e))?; + if let Some(row) = rows + .next() + .await + .map_err(|e| db_error("resolve_entity", e))? + { + let entity_id = row + .get::(0) + .map_err(|e| db_error("resolve_entity", e))?; + return Ok(entity_id); + } + + self.conn + .execute( + "INSERT OR IGNORE INTO memory_entities ( + name, normalized_name, entity_type, aliases, created_at + ) + VALUES (?1, ?2, 'unknown', '[]', ?3)", + params![name, normalized.as_str(), current_timestamp(),], + ) + .await + .map_err(|e| db_error("resolve_entity", e))?; + let mut rows = self + .conn + .query( + "SELECT entity_id FROM memory_entities WHERE normalized_name = ?1", + params![normalized.as_str()], + ) + .await + .map_err(|e| db_error("resolve_entity", e))?; + let row = rows + .next() + .await + .map_err(|e| db_error("resolve_entity", e))? + .ok_or_else(|| db_message("resolve_entity", "entity insert/read returned no row"))?; + row.get::(0).map_err(|e| db_error("resolve_entity", e)) + } + + async fn load_fact_entities(&self, fact_id: i64) -> Result> { + let mut rows = self + .conn + .query( + "SELECT e.name + FROM memory_entities e + JOIN memory_fact_entities fe ON fe.entity_id = e.entity_id + WHERE fe.fact_id = ?1 + ORDER BY e.name", + params![fact_id], + ) + .await + .map_err(|e| db_error("load_fact_entities", e))?; + let mut entities = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("load_fact_entities", e))? + { + entities.push( + row.get::(0) + .map_err(|e| db_error("load_fact_entities", e))?, + ); + } + Ok(entities) + } + + async fn load_bank_vectors( + &self, + category: Option, + ) -> Result<(usize, Vec>)> { + let sql = if category.is_some() { + "SELECT hrr_vector + FROM memory_facts + WHERE category = ?1 AND trust_score >= ?2" + } else { + "SELECT hrr_vector + FROM memory_facts + WHERE trust_score >= ?1" + }; + + let mut rows = if let Some(category) = category { + self.conn.query(sql, params![category.as_str(), 0.0]).await + } else { + self.conn.query(sql, params![0.0]).await + } + .map_err(|e| db_error("load_bank_vectors", e))?; + + let mut fact_count = 0; + let mut vectors = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| db_error("load_bank_vectors", e))? + { + fact_count += 1; + let value = row + .get::(0) + .map_err(|e| db_error("load_bank_vectors", e))?; + if let Some(vector) = deserialize_vector_value(value, "load_bank_vectors")? { + vectors.push(vector); + } + } + Ok((fact_count, vectors)) + } + + async fn last_insert_rowid(&self, operation: &str) -> Result { + let mut rows = self + .conn + .query("SELECT last_insert_rowid()", ()) + .await + .map_err(|e| db_error(operation, e))?; + let row = rows + .next() + .await + .map_err(|e| db_error(operation, e))? + .ok_or_else(|| db_message(operation, "last_insert_rowid returned no rows"))?; + row.get::(0).map_err(|e| db_error(operation, e)) + } + + fn encode_vector( + &self, + content: &str, + entities: &[String], + operation: &str, + ) -> Result> { + let vector = self.encoder.encode_fact(content, entities); + HolographicEncoder::serialize(&vector) + .map_err(|e| db_message(operation, format!("failed to serialize vector: {e}"))) + } + + async fn update_fact_vector( + &self, + fact_id: i64, + content: &str, + entities: &[String], + operation: &str, + ) -> Result<()> { + let vector = self.encode_vector(content, entities, operation)?; + self.conn + .execute( + "UPDATE memory_facts + SET hrr_vector = ?1, + hrr_algebra = ?2, + hrr_dim = ?3, + updated_at = ?4 + WHERE fact_id = ?5", + params![ + vector, + HRR_ALGEBRA, + HolographicEncoder::DIMENSIONS as i64, + current_timestamp(), + fact_id, + ], + ) + .await + .map_err(|e| db_error(operation, e))?; + Ok(()) + } + + async fn mark_fact_banks_dirty(&self, category: MemoryCategory) -> Result<()> { + self.mark_bank_dirty("all").await?; + self.mark_bank_dirty(category.as_str()).await + } + + async fn mark_bank_dirty(&self, bank_name: &str) -> Result<()> { + self.conn + .execute( + "INSERT INTO memory_bank_dirty (bank_name, updated_at) + VALUES (?1, ?2) + ON CONFLICT(bank_name) DO UPDATE SET updated_at = excluded.updated_at", + params![bank_name, current_timestamp()], + ) + .await + .map_err(|e| db_error("mark_bank_dirty", e))?; + Ok(()) + } +} + +fn merge_entities(content: &str, explicit: &[String]) -> Vec { + let mut seen = HashSet::new(); + let mut entities = Vec::new(); + for entity in explicit.iter().cloned().chain(extract_entities(content)) { + let normalized = normalize_entity(&entity); + if normalized.is_empty() { + continue; + } + if seen.insert(normalized.to_ascii_lowercase()) { + entities.push(normalized); + } + } + entities +} + +fn to_json_string(value: &T, operation: &str) -> Result { + serde_json::to_string(value) + .map_err(|e| db_message(operation, format!("failed to serialize JSON: {e}"))) +} + +fn parse_json_array(value: &str, operation: &str) -> Result> { + serde_json::from_str(value) + .map_err(|e| db_message(operation, format!("failed to parse JSON array: {e}"))) +} + +fn parse_category(value: &str, operation: &str) -> Result { + value + .parse() + .map_err(|e| db_message(operation, format!("failed to parse category: {e}"))) +} + +fn fact_from_row(row: &libsql::Row, operation: &str, entities: Vec) -> Result { + let category = parse_category( + &row.get::(2).map_err(|e| db_error(operation, e))?, + operation, + )?; + let tags = parse_json_array( + &row.get::(3).map_err(|e| db_error(operation, e))?, + operation, + )?; + let metadata = + serde_json::from_str(&row.get::(13).map_err(|e| db_error(operation, e))?) + .map_err(|e| db_message(operation, format!("failed to parse metadata: {e}")))?; + + Ok(FactRecord { + fact_id: row.get::(0).map_err(|e| db_error(operation, e))?, + content: row.get::(1).map_err(|e| db_error(operation, e))?, + category, + tags, + entities, + trust_score: row.get::(4).map_err(|e| db_error(operation, e))?, + source: Some(row.get::(5).map_err(|e| db_error(operation, e))?), + retrieval_count: row.get::(6).map_err(|e| db_error(operation, e))?, + helpful_count: row.get::(7).map_err(|e| db_error(operation, e))?, + unhelpful_count: row.get::(8).map_err(|e| db_error(operation, e))?, + created_at: row.get::(9).map_err(|e| db_error(operation, e))?, + updated_at: row.get::(10).map_err(|e| db_error(operation, e))?, + last_retrieved_at: row + .get::>(11) + .map_err(|e| db_error(operation, e))?, + last_feedback_at: row + .get::>(12) + .map_err(|e| db_error(operation, e))?, + metadata, + }) +} + +fn deserialize_vector_value(value: libsql::Value, operation: &str) -> Result>> { + match value { + libsql::Value::Blob(bytes) => HolographicEncoder::deserialize(&bytes) + .map(Some) + .map_err(|e| db_message(operation, format!("failed to decode vector: {e}"))), + libsql::Value::Null => Ok(None), + _ => Err(db_message( + operation, + "hrr_vector contained a non-blob value", + )), + } +} + +fn sql_i64_list(ids: &[i64]) -> Option { + if ids.is_empty() { + None + } else { + Some( + ids.iter() + .map(ToString::to_string) + .collect::>() + .join(", "), + ) + } +} + +fn feedback_action_str(action: FeedbackAction) -> &'static str { + match action { + FeedbackAction::Helpful => "helpful", + FeedbackAction::Unhelpful => "unhelpful", + } +} + +fn normalized_limit(limit: usize) -> usize { + if limit == 0 { + DEFAULT_LIMIT + } else { + limit.min(i64::MAX as usize) + } +} + +fn average_vectors(vectors: &[Vec]) -> Vec { + if vectors.is_empty() { + return vec![0.0; HolographicEncoder::DIMENSIONS]; + } + + let mut average = vec![0.0; HolographicEncoder::DIMENSIONS]; + let mut count = 0.0; + for vector in vectors { + if vector.len() != HolographicEncoder::DIMENSIONS { + continue; + } + count += 1.0; + for (target, value) in average.iter_mut().zip(vector) { + *target += value; + } + } + if count > 0.0 { + for value in &mut average { + *value /= count; + } + } + average +} + +fn normalize_bank_name(bank_name: &str) -> String { + bank_name + .trim() + .to_ascii_lowercase() + .replace([' ', '-'], "_") +} + +fn db_error(operation: &str, error: impl fmt::Display) -> TokenSaveError { + TokenSaveError::Database { + message: error.to_string(), + operation: operation.to_string(), + } +} + +fn db_message(operation: &str, message: impl Into) -> TokenSaveError { + TokenSaveError::Database { + message: message.into(), + operation: operation.to_string(), + } +} diff --git a/src/memory/trust.rs b/src/memory/trust.rs new file mode 100644 index 0000000000..fad29f7c1c --- /dev/null +++ b/src/memory/trust.rs @@ -0,0 +1,52 @@ +//! Trust score helpers for bounded confidence, feedback, and aging. + +use super::types::FeedbackAction; + +pub const HELPFUL_DELTA: f64 = 0.05; +pub const UNHELPFUL_DELTA: f64 = -0.10; +pub const TRUST_MIN: f64 = 0.0; +pub const TRUST_MAX: f64 = 1.0; +pub const DEFAULT_TRUST: f64 = 0.5; +pub const DEFAULT_MIN_TRUST: f64 = 0.3; + +pub fn clamp_trust(score: f64) -> f64 { + score.clamp(TRUST_MIN, TRUST_MAX) +} + +pub fn apply_feedback(current_trust: f64, action: FeedbackAction) -> f64 { + let delta = match action { + FeedbackAction::Helpful => HELPFUL_DELTA, + FeedbackAction::Unhelpful => UNHELPFUL_DELTA, + }; + + clamp_trust(current_trust + delta) +} + +pub fn trust_bucket(score: f64) -> &'static str { + let clamped = clamp_trust(score); + if clamped < DEFAULT_MIN_TRUST { + "low" + } else if clamped < 0.75 { + "medium" + } else { + "high" + } +} + +pub fn trust_distribution(scores: &[f64]) -> (usize, usize, usize) { + scores.iter().fold((0, 0, 0), |(low, medium, high), score| { + match trust_bucket(*score) { + "low" => (low + 1, medium, high), + "medium" => (low, medium + 1, high), + _ => (low, medium, high + 1), + } + }) +} + +pub fn temporal_decay(current_trust: f64, age_days: f64) -> f64 { + let age = age_days.max(0.0); + let decay_weight = 1.0 - (-age / 180.0).exp(); + let decayed = current_trust.mul_add(1.0 - decay_weight, DEFAULT_TRUST * decay_weight); + + clamp_trust(decayed) +} diff --git a/src/memory/types.rs b/src/memory/types.rs new file mode 100644 index 0000000000..b949651637 --- /dev/null +++ b/src/memory/types.rs @@ -0,0 +1,195 @@ +use std::fmt; +use std::str::FromStr; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MemoryCategory { + General, + UserPref, + Project, + Tool, + Decision, + CodeArea, +} + +impl MemoryCategory { + pub const fn as_str(self) -> &'static str { + match self { + Self::General => "general", + Self::UserPref => "user_pref", + Self::Project => "project", + Self::Tool => "tool", + Self::Decision => "decision", + Self::CodeArea => "code_area", + } + } +} + +impl fmt::Display for MemoryCategory { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ParseMemoryCategoryError { + value: String, +} + +impl fmt::Display for ParseMemoryCategoryError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "unknown memory category: {}", self.value) + } +} + +impl std::error::Error for ParseMemoryCategoryError {} + +impl FromStr for MemoryCategory { + type Err = ParseMemoryCategoryError; + + fn from_str(value: &str) -> Result { + let normalized = value.trim().to_ascii_lowercase().replace(['-', ' '], "_"); + match normalized.as_str() { + "general" => Ok(Self::General), + "user_pref" | "user_preference" | "user_preferences" => Ok(Self::UserPref), + "project" => Ok(Self::Project), + "tool" => Ok(Self::Tool), + "decision" => Ok(Self::Decision), + "code_area" | "code" => Ok(Self::CodeArea), + _ => Err(ParseMemoryCategoryError { + value: value.to_string(), + }), + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FactRecord { + pub fact_id: i64, + pub content: String, + pub category: MemoryCategory, + pub tags: Vec, + pub entities: Vec, + pub trust_score: f64, + pub source: Option, + pub retrieval_count: i64, + pub helpful_count: i64, + pub unhelpful_count: i64, + pub created_at: i64, + pub updated_at: i64, + pub last_retrieved_at: Option, + pub last_feedback_at: Option, + pub metadata: Value, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct EntityRecord { + pub entity_id: i64, + pub name: String, + pub normalized_name: String, + pub entity_type: Option, + pub created_at: i64, + pub updated_at: i64, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FactSearchResult { + pub fact: FactRecord, + pub score: f64, + pub fts_score: f64, + pub jaccard_score: f64, + pub holographic_score: f64, + pub trust_score: f64, + pub why: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ContradictionResult { + pub existing_fact: FactRecord, + pub new_content: String, + pub score: f64, + pub why: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FeedbackAction { + Helpful, + Unhelpful, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FeedbackRequest { + pub fact_id: i64, + pub action: FeedbackAction, + #[serde(default)] + pub source: Option, + #[serde(default, alias = "reason")] + pub note: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FeedbackResult { + pub event_id: i64, + pub fact_id: i64, + pub action: FeedbackAction, + pub old_trust: f64, + pub new_trust: f64, + pub trust_delta: f64, + pub helpful_count: i64, + pub unhelpful_count: i64, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct MemoryStatus { + pub fact_count: usize, + pub entity_count: usize, + pub bank_count: usize, + pub algebra_name: String, + pub hrr_dim: usize, + pub estimated_capacity: usize, + pub trust_0_025_count: usize, + pub trust_025_050_count: usize, + pub trust_050_075_count: usize, + pub trust_075_100_count: usize, + pub below_default_recall_threshold_count: usize, + pub helpful_count: usize, + pub unhelpful_count: usize, + pub missing_vector_count: usize, + pub legacy_backfill_complete: bool, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct AddFactRequest { + pub content: String, + pub category: MemoryCategory, + pub source: Option, + pub tags: Vec, + pub entities: Vec, + pub trust: Option, + pub metadata: Value, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct SearchFactsRequest { + pub query: String, + pub category: Option, + pub limit: Option, + pub min_trust: Option, + pub include_why: bool, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct UpdateFactRequest { + pub fact_id: i64, + pub content: Option, + pub category: Option, + pub tags: Option>, + pub entities: Option>, + pub trust: Option, + pub source: Option, + pub metadata: Option, +} diff --git a/src/sessions/claude.rs b/src/sessions/claude.rs new file mode 100644 index 0000000000..ba87ce869e --- /dev/null +++ b/src/sessions/claude.rs @@ -0,0 +1,225 @@ +//! Claude Code transcript source. +//! +//! Claude Code appends one JSON object per line to +//! `~/.claude/projects//.jsonl` (with subagent transcripts +//! under `…//subagents/*.jsonl`). Each line carries a top-level `type` +//! (`"user"`/`"assistant"`/…), a `message` object (`role`, `content`, `model`, +//! `id`), an ISO-8601 `timestamp`, the session `cwd`, and `sessionId`/`uuid`. +//! +//! The accounting parser already reads these files for cost `turns`; this source +//! reuses the **same** append-only byte-offset machinery to also populate the +//! provider-neutral `session_messages` table. Files are scoped to the current +//! project by their recorded `cwd`, so a project only ingests its own sessions. + +use std::path::{Path, PathBuf}; + +use serde_json::Value; + +use crate::accounting::parser::parse_timestamp; +use crate::sessions::source::{ + collect_files_with_ext, paths_equal, stream_new_jsonl, title_from_messages, ParsedTranscript, + SessionDraft, StoredCursor, TranscriptSource, +}; +use crate::sessions::SessionMessageRecord; + +const PROVIDER: &str = "claude"; +/// `~/.claude/projects//<…>.jsonl` is at most a few levels deep. +const MAX_SCAN_DEPTH: u8 = 6; +/// `cwd` should appear on an early line; scan a few in case the first is a +/// `summary`/meta line without one. +const CWD_PROBE_LINES: usize = 8; + +/// Claude Code transcript locator + parser. +pub struct ClaudeSource { + projects_dir: PathBuf, +} + +impl ClaudeSource { + /// Source rooted at the real `~/.claude/projects`. Returns `None` when the + /// home directory cannot be resolved. + pub fn new() -> Option { + let home = dirs::home_dir()?; + Some(Self::with_home(&home)) + } + + /// Source rooted at `/.claude/projects` (used by tests). + pub fn with_home(home: &Path) -> Self { + Self { + projects_dir: home.join(".claude").join("projects"), + } + } +} + +impl TranscriptSource for ClaudeSource { + fn provider(&self) -> &'static str { + PROVIDER + } + + fn transcript_paths(&self, _project_root: &Path) -> Vec { + // Scan every project slug; `parse_new` filters by recorded `cwd` so each + // project only ingests its own sessions without us having to replicate + // Claude's slug-encoding scheme. + collect_files_with_ext(&self.projects_dir, "jsonl", MAX_SCAN_DEPTH) + } + + fn parse_new( + &self, + path: &Path, + prev: StoredCursor, + project_root: &Path, + max_new_bytes: Option, + ) -> Option { + // Cheap project scoping: a transcript belongs to exactly one cwd, so + // skip files that are not this project's without advancing the cursor. + match transcript_cwd(path) { + Some(cwd) if paths_equal(&cwd, project_root) => {} + _ => return None, + } + + let new = stream_new_jsonl(path, prev, max_new_bytes)?; + let session_id = path + .file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or("unknown") + .to_string(); + + let mut messages = Vec::new(); + for line in &new.lines { + if let Some(message) = message_from_line(&line.value, &session_id, path, line.offset) { + messages.push(message); + } + } + + let project = project_root.to_string_lossy().to_string(); + let draft = SessionDraft { + session_id, + project_key: project.clone(), + project_path: project, + title: title_from_messages(&messages), + metadata_json: serde_json::to_string(&serde_json::json!({ + "source": "claude_transcript", + })) + .ok(), + }; + + Some(ParsedTranscript { + draft, + messages, + new_cursor: new.new_cursor, + }) + } +} + +/// Reads the session `cwd` from an early line of a Claude transcript. +fn transcript_cwd(path: &Path) -> Option { + use std::io::BufRead; + let file = std::fs::File::open(path).ok()?; + let reader = std::io::BufReader::new(file); + for line in reader.lines().take(CWD_PROBE_LINES).map_while(Result::ok) { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + if let Ok(value) = serde_json::from_str::(trimmed) { + if let Some(cwd) = value.get("cwd").and_then(Value::as_str) { + if !cwd.is_empty() { + return Some(PathBuf::from(cwd)); + } + } + } + } + None +} + +/// Map one Claude transcript line to a provider-neutral message, or `None` for +/// lines that carry no conversational text (tool-result-only, meta lines, …). +fn message_from_line( + record: &Value, + session_id: &str, + path: &Path, + offset: i64, +) -> Option { + let kind = record.get("type").and_then(Value::as_str)?; + if kind != "user" && kind != "assistant" { + return None; + } + let message = record.get("message").unwrap_or(record); + let role = message + .get("role") + .and_then(Value::as_str) + .unwrap_or(kind) + .to_string(); + + let content = message.get("content").unwrap_or(message); + let (text, tool_names) = content_text_and_tools(content); + if text.trim().is_empty() { + return None; + } + + let message_id = message + .get("id") + .and_then(Value::as_str) + .or_else(|| record.get("uuid").and_then(Value::as_str)) + .filter(|id| !id.is_empty()) + .map_or_else(|| format!("{session_id}:{offset}"), ToString::to_string); + let model = message + .get("model") + .and_then(Value::as_str) + .map(str::to_string); + let timestamp = record + .get("timestamp") + .and_then(Value::as_str) + .and_then(parse_timestamp) + .map(|secs| secs as i64); + + Some(SessionMessageRecord { + provider: PROVIDER.to_string(), + message_id, + session_id: session_id.to_string(), + role, + timestamp, + ordinal: offset, + text, + kind: Some("message".to_string()), + model, + tool_names: (!tool_names.is_empty()).then(|| tool_names.join(",")), + source_path: Some(path.to_string_lossy().to_string()), + source_offset: Some(offset), + metadata_json: serde_json::to_string(&serde_json::json!({ + "source": "claude_transcript", + "raw_type": kind, + })) + .ok(), + }) +} + +/// Extract the concatenated text and tool-use names from a Claude `content` +/// field, which is either a plain string (user turns) or an array of typed +/// blocks (`text`, `tool_use`, `tool_result`, …) for assistant turns. +fn content_text_and_tools(content: &Value) -> (String, Vec) { + if let Some(text) = content.as_str() { + return (text.to_string(), Vec::new()); + } + let Some(items) = content.as_array() else { + return (String::new(), Vec::new()); + }; + + let mut texts = Vec::new(); + let mut tools = Vec::new(); + for item in items { + match item.get("type").and_then(Value::as_str) { + Some("text") => { + if let Some(text) = item.get("text").and_then(Value::as_str) { + texts.push(text.to_string()); + } + } + Some("tool_use") => { + if let Some(name) = item.get("name").and_then(Value::as_str) { + tools.push(name.to_string()); + } + } + _ => {} + } + } + (texts.join("\n\n"), tools) +} diff --git a/src/sessions/cline_like.rs b/src/sessions/cline_like.rs new file mode 100644 index 0000000000..42d234ef98 --- /dev/null +++ b/src/sessions/cline_like.rs @@ -0,0 +1,318 @@ +//! Cline/Roo Code/Kilo Code task-history transcript sources. +//! +//! These VS Code extension-family adapters persist each task in a directory with +//! JSON files such as: +//! +//! * `api_conversation_history.json` (or Roo's `api_messages.json`) - the +//! Anthropic-compatible conversation sent to/received from the model. +//! * `ui_messages.json` - webview-oriented messages. +//! * `task_metadata.json` / `history_item.json` - task metadata. +//! +//! The API conversation file is a **full-rewrite** JSON array, so the source uses +//! the shared `ContentHash` reader and deterministic `:` message +//! ids. To avoid mixing global VS Code extension history across projects, a task +//! is ingested only when its metadata contains a project/workspace/cwd path that +//! resolves to the current tokensave project root. + +use std::path::{Path, PathBuf}; + +use serde_json::Value; + +use crate::sessions::source::{ + paths_equal, read_changed_file, title_from_messages, ParsedTranscript, SessionDraft, + StoredCursor, TranscriptSource, +}; +use crate::sessions::SessionMessageRecord; + +/// One Cline-family provider configuration. +#[derive(Clone)] +pub struct ClineLikeSource { + provider: &'static str, + storage_roots: Vec, +} + +impl ClineLikeSource { + /// Cline VS Code extension storage: + /// `Code/User/globalStorage/saoudrizwan.claude-dev/tasks`. + pub fn cline() -> Option { + let home = dirs::home_dir()?; + Some(Self::cline_with_home(&home)) + } + + /// Roo Code VS Code extension storage: + /// `Code/User/globalStorage/rooveterinaryinc.roo-cline/tasks`. + pub fn roo_code() -> Option { + let home = dirs::home_dir()?; + Some(Self::roo_code_with_home(&home)) + } + + /// Kilo Code storage. Current docs mention both the VS Code extension root + /// and the CLI root (`~/.kilocode/cli/global/tasks`), so scan both. + pub fn kilo() -> Option { + let home = dirs::home_dir()?; + Some(Self::kilo_with_home(&home)) + } + + pub fn cline_with_home(home: &Path) -> Self { + Self { + provider: "cline", + storage_roots: vec![crate::agents::vscode_data_dir(home) + .join("User/globalStorage/saoudrizwan.claude-dev/tasks")], + } + } + + pub fn roo_code_with_home(home: &Path) -> Self { + Self { + provider: "roo-code", + storage_roots: vec![crate::agents::vscode_data_dir(home) + .join("User/globalStorage/rooveterinaryinc.roo-cline/tasks")], + } + } + + pub fn kilo_with_home(home: &Path) -> Self { + Self { + provider: "kilo", + storage_roots: vec![ + crate::agents::vscode_data_dir(home) + .join("User/globalStorage/kilocode.kilo-code/tasks"), + home.join(".kilocode/cli/global/tasks"), + ], + } + } +} + +impl TranscriptSource for ClineLikeSource { + fn provider(&self) -> &'static str { + self.provider + } + + fn transcript_paths(&self, _project_root: &Path) -> Vec { + let mut out = Vec::new(); + for root in &self.storage_roots { + let Ok(entries) = std::fs::read_dir(root) else { + continue; + }; + for entry in entries.flatten() { + let task_dir = entry.path(); + if !task_dir.is_dir() { + continue; + } + for name in ["api_conversation_history.json", "api_messages.json"] { + let path = task_dir.join(name); + if path.is_file() { + out.push(path); + } + } + } + } + out + } + + fn parse_new( + &self, + path: &Path, + prev: StoredCursor, + project_root: &Path, + _max_new_bytes: Option, + ) -> Option { + let task_dir = path.parent()?; + let metadata = read_task_metadata(task_dir)?; + if !metadata_belongs_to_project(&metadata, project_root) { + return None; + } + + let changed = read_changed_file(path, prev)?; + let document: Value = serde_json::from_str(&changed.contents).ok()?; + let entries = document.as_array()?; + let task_id = task_dir + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("unknown"); + + let mut messages = Vec::new(); + for (index, entry) in entries.iter().enumerate() { + if let Some(message) = message_from_entry(self.provider, entry, task_id, path, index) { + messages.push(message); + } + } + + let project = project_root.to_string_lossy().to_string(); + let draft = SessionDraft { + session_id: task_id.to_string(), + project_key: project.clone(), + project_path: project, + title: title_from_messages(&messages) + .or_else(|| metadata_task_title(&metadata).map(str::to_string)), + metadata_json: serde_json::to_string(&serde_json::json!({ + "source": format!("{}_task_history", self.provider), + })) + .ok(), + }; + + Some(ParsedTranscript { + draft, + messages, + new_cursor: changed.new_cursor, + }) + } +} + +fn read_task_metadata(task_dir: &Path) -> Option { + for name in ["task_metadata.json", "history_item.json", "history.json"] { + let path = task_dir.join(name); + if !path.is_file() { + continue; + } + if let Ok(contents) = std::fs::read_to_string(path) { + if let Ok(value) = serde_json::from_str::(&contents) { + return Some(value); + } + } + } + None +} + +fn metadata_belongs_to_project(metadata: &Value, project_root: &Path) -> bool { + metadata_project_paths(metadata) + .iter() + .any(|path| paths_equal(path, project_root)) +} + +fn metadata_project_paths(value: &Value) -> Vec { + let mut out = Vec::new(); + collect_metadata_project_paths(value, None, &mut out); + out +} + +fn collect_metadata_project_paths(value: &Value, key: Option<&str>, out: &mut Vec) { + match value { + Value::Object(map) => { + for (child_key, child_value) in map { + collect_metadata_project_paths(child_value, Some(child_key), out); + } + } + Value::Array(items) => { + for item in items { + collect_metadata_project_paths(item, key, out); + } + } + Value::String(s) => { + let key = key.unwrap_or_default().to_ascii_lowercase(); + let looks_like_project_path = key.contains("workspace") + || key.contains("project") + || key.contains("cwd") + || key.contains("workdir") + || key.contains("directory") + || key == "root"; + if looks_like_project_path && !s.is_empty() { + out.push(PathBuf::from(s)); + } + } + _ => {} + } +} + +fn metadata_task_title(metadata: &Value) -> Option<&str> { + metadata + .get("task") + .or_else(|| metadata.get("title")) + .or_else(|| metadata.get("summary")) + .and_then(Value::as_str) + .filter(|s| !s.is_empty()) +} + +fn message_from_entry( + provider: &str, + entry: &Value, + task_id: &str, + path: &Path, + index: usize, +) -> Option { + let role = match entry.get("role").and_then(Value::as_str)? { + "user" => "user", + "assistant" | "model" => "assistant", + _ => return None, + }; + let content = entry.get("content").unwrap_or(entry); + let (text, tool_names) = content_text_and_tools(content); + if text.trim().is_empty() { + return None; + } + let timestamp = entry + .get("ts") + .or_else(|| entry.get("timestamp")) + .or_else(|| entry.get("createdAt")) + .and_then(|value| { + value + .as_i64() + .or_else(|| value.as_str().and_then(|s| s.parse::().ok())) + }); + let model = entry + .get("model") + .and_then(Value::as_str) + .map(str::to_string); + let message_id = entry + .get("id") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map_or_else(|| format!("{task_id}:{index}"), ToString::to_string); + + Some(SessionMessageRecord { + provider: provider.to_string(), + message_id, + session_id: task_id.to_string(), + role: role.to_string(), + timestamp, + ordinal: index as i64, + text, + kind: Some("message".to_string()), + model, + tool_names: (!tool_names.is_empty()).then(|| tool_names.join(",")), + source_path: Some(path.to_string_lossy().to_string()), + source_offset: Some(index as i64), + metadata_json: serde_json::to_string(&serde_json::json!({ + "source": format!("{provider}_task_history"), + })) + .ok(), + }) +} + +fn content_text_and_tools(content: &Value) -> (String, Vec) { + if let Some(text) = content.as_str() { + return (text.to_string(), Vec::new()); + } + let Some(items) = content.as_array() else { + return ( + content + .get("text") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + Vec::new(), + ); + }; + + let mut texts = Vec::new(); + let mut tools = Vec::new(); + for item in items { + match item.get("type").and_then(Value::as_str) { + Some("text") | None => { + if let Some(text) = item.get("text").and_then(Value::as_str) { + texts.push(text.to_string()); + } + } + Some("tool_use") => { + if let Some(name) = item.get("name").and_then(Value::as_str) { + tools.push(name.to_string()); + } + } + Some("tool_result") => { + if let Some(text) = item.get("content").and_then(Value::as_str) { + texts.push(text.to_string()); + } + } + _ => {} + } + } + (texts.join("\n\n"), tools) +} diff --git a/src/sessions/codex.rs b/src/sessions/codex.rs new file mode 100644 index 0000000000..565f028643 --- /dev/null +++ b/src/sessions/codex.rs @@ -0,0 +1,210 @@ +//! Codex CLI transcript source. +//! +//! Codex appends one JSON object per line to +//! `~/.codex/sessions/YYYY/MM/DD/rollout-*.jsonl`. Each line is +//! `{"timestamp": "", "type": "", "payload": {…}}`. The relevant +//! kinds for conversation text are: +//! +//! * `session_meta` — first line; `payload.cwd`, session `id`, model info. +//! * `event_msg` with `payload.type == "user_message"` — a real user prompt +//! (`payload.message`). +//! * `event_msg` with `payload.type == "agent_message"` — a real assistant reply +//! (`payload.message`). +//! +//! `response_item` entries are intentionally skipped: they carry auto-injected +//! synthetic context and duplicate the `agent_message`/`user_message` turns, so +//! ingesting them would double-count the conversation. This append-only JSONL is +//! read with the shared byte-offset machinery and scoped to the current project +//! by `session_meta.cwd`. + +use std::path::{Path, PathBuf}; + +use serde_json::Value; + +use crate::accounting::parser::parse_timestamp; +use crate::sessions::source::{ + collect_files_with_ext, paths_equal, stream_new_jsonl, title_from_messages, ParsedTranscript, + SessionDraft, StoredCursor, TranscriptSource, +}; +use crate::sessions::SessionMessageRecord; + +const PROVIDER: &str = "codex"; +/// `~/.codex/sessions/YYYY/MM/DD/rollout-*.jsonl` → date dirs add depth. +const MAX_SCAN_DEPTH: u8 = 6; + +/// Session metadata read from a rollout's leading `session_meta` line. +struct CodexMeta { + cwd: PathBuf, + session_id: String, + model: Option, +} + +/// Codex CLI transcript locator + parser. +pub struct CodexSource { + sessions_dir: PathBuf, +} + +impl CodexSource { + /// Source rooted at the real `~/.codex/sessions`. Returns `None` when the + /// home directory cannot be resolved. + pub fn new() -> Option { + let home = dirs::home_dir()?; + Some(Self::with_home(&home)) + } + + /// Source rooted at `/.codex/sessions` (used by tests). + pub fn with_home(home: &Path) -> Self { + Self { + sessions_dir: home.join(".codex").join("sessions"), + } + } +} + +impl TranscriptSource for CodexSource { + fn provider(&self) -> &'static str { + PROVIDER + } + + fn transcript_paths(&self, _project_root: &Path) -> Vec { + collect_files_with_ext(&self.sessions_dir, "jsonl", MAX_SCAN_DEPTH) + } + + fn parse_new( + &self, + path: &Path, + prev: StoredCursor, + project_root: &Path, + max_new_bytes: Option, + ) -> Option { + // `session_meta` (line 1) is authoritative for cwd + session id; without + // it we cannot safely attribute the rollout to a project, so skip. + let meta = session_meta(path)?; + if !paths_equal(&meta.cwd, project_root) { + return None; + } + + let new = stream_new_jsonl(path, prev, max_new_bytes)?; + let mut messages = Vec::new(); + for line in &new.lines { + if let Some(message) = message_from_line(&line.value, &meta, path, line.offset) { + messages.push(message); + } + } + + let project = project_root.to_string_lossy().to_string(); + let draft = SessionDraft { + session_id: meta.session_id.clone(), + project_key: project.clone(), + project_path: project, + title: title_from_messages(&messages), + metadata_json: serde_json::to_string(&serde_json::json!({ + "source": "codex_rollout", + })) + .ok(), + }; + + Some(ParsedTranscript { + draft, + messages, + new_cursor: new.new_cursor, + }) + } +} + +/// Read the leading `session_meta` line of a rollout for cwd/session-id/model. +fn session_meta(path: &Path) -> Option { + use std::io::BufRead; + let file = std::fs::File::open(path).ok()?; + let reader = std::io::BufReader::new(file); + for line in reader.lines().take(4).map_while(Result::ok) { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let Ok(value) = serde_json::from_str::(trimmed) else { + continue; + }; + if value.get("type").and_then(Value::as_str) != Some("session_meta") { + continue; + } + let payload = value.get("payload").unwrap_or(&value); + let cwd = payload + .get("cwd") + .and_then(Value::as_str) + .filter(|cwd| !cwd.is_empty()) + .map(PathBuf::from)?; + let session_id = payload + .get("id") + .or_else(|| payload.get("session_id")) + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map_or_else( + || { + path.file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or("unknown") + .to_string() + }, + ToString::to_string, + ); + let model = payload + .get("model") + .or_else(|| payload.get("model_provider")) + .and_then(Value::as_str) + .map(str::to_string); + return Some(CodexMeta { + cwd, + session_id, + model, + }); + } + None +} + +/// Map one rollout line to a provider-neutral message, or `None` for non-message +/// events (`response_item`, tool calls, token counts, …). +fn message_from_line( + record: &Value, + meta: &CodexMeta, + path: &Path, + offset: i64, +) -> Option { + if record.get("type").and_then(Value::as_str) != Some("event_msg") { + return None; + } + let payload = record.get("payload")?; + let role = match payload.get("type").and_then(Value::as_str)? { + "user_message" => "user", + "agent_message" => "assistant", + _ => return None, + }; + let text = payload.get("message").and_then(Value::as_str).unwrap_or(""); + if text.trim().is_empty() { + return None; + } + + let timestamp = record + .get("timestamp") + .and_then(Value::as_str) + .and_then(parse_timestamp) + .map(|secs| secs as i64); + + Some(SessionMessageRecord { + provider: PROVIDER.to_string(), + message_id: format!("{}:{offset}", meta.session_id), + session_id: meta.session_id.clone(), + role: role.to_string(), + timestamp, + ordinal: offset, + text: text.to_string(), + kind: Some("message".to_string()), + model: meta.model.clone(), + tool_names: None, + source_path: Some(path.to_string_lossy().to_string()), + source_offset: Some(offset), + metadata_json: serde_json::to_string(&serde_json::json!({ + "source": "codex_rollout", + })) + .ok(), + }) +} diff --git a/src/sessions/cursor.rs b/src/sessions/cursor.rs new file mode 100644 index 0000000000..4cbb12339a --- /dev/null +++ b/src/sessions/cursor.rs @@ -0,0 +1,335 @@ +use std::path::{Path, PathBuf}; + +use serde_json::Value; + +use crate::global_db::GlobalDb; +use crate::sessions::source::{ + ingest_source, stream_new_jsonl, title_from_messages, ParsedTranscript, SessionDraft, + StoredCursor, TranscriptSource, +}; +use crate::sessions::SessionMessageRecord; + +const PROJECT_SESSION_DB_FILENAME: &str = "sessions.db"; + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct CursorTranscriptIngestStats { + pub sessions_upserted: u64, + pub messages_upserted: u64, +} + +pub fn project_session_db_path(project_root: &Path) -> PathBuf { + crate::config::get_tokensave_dir(project_root).join(PROJECT_SESSION_DB_FILENAME) +} + +pub async fn open_project_session_db(project_root: &Path) -> Option { + GlobalDb::open_at(&project_session_db_path(project_root)).await +} + +/// A Cursor hook event scoped to one transcript file. Cursor is hook-driven — +/// the transcript path, session id, and project all come from the event payload +/// rather than from a directory scan — so the source wraps the parsed event and +/// yields exactly that one path. +struct CursorEventSource { + event: Value, + transcript_path: PathBuf, +} + +impl TranscriptSource for CursorEventSource { + fn provider(&self) -> &'static str { + "cursor" + } + + fn transcript_paths(&self, _project_root: &Path) -> Vec { + vec![self.transcript_path.clone()] + } + + fn parse_new( + &self, + path: &Path, + prev: StoredCursor, + _project_root: &Path, + max_new_bytes: Option, + ) -> Option { + let new = stream_new_jsonl(path, prev, max_new_bytes)?; + let session_id = event_session_id(&self.event, path); + let mut messages = Vec::new(); + for line in &new.lines { + // The byte offset doubles as the message ordinal and source_offset, + // matching the original Cursor ingestion. + if let Some(message) = event_message( + &line.value, + &self.event, + &session_id, + path, + line.offset, + line.offset, + ) { + messages.push(message); + } + } + + // Defer the (filesystem-walking) project/title/metadata derivation until + // we actually have new messages; the driver ignores the draft otherwise. + let draft = if messages.is_empty() { + SessionDraft { + session_id, + project_key: String::new(), + project_path: String::new(), + title: None, + metadata_json: None, + } + } else { + let (project_key, project_path) = event_project(&self.event); + SessionDraft { + session_id, + project_key, + project_path, + title: title_from_messages(&messages), + metadata_json: serde_json::to_string(&session_metadata(&self.event)).ok(), + } + }; + + Some(ParsedTranscript { + draft, + messages, + new_cursor: new.new_cursor, + }) + } +} + +/// Ingest the Cursor transcript referenced by a hook payload into the +/// provider-neutral session/message tables for the provided database. Project +/// hooks should pass the project-local DB from [`open_project_session_db`]. +/// +/// Ingestion is **incremental**: it resumes from the byte offset recorded in the +/// DB's `parse_offsets` table (via the shared [`crate::sessions::source`] +/// driver), so each call only parses and upserts transcript lines appended since +/// the last run rather than re-reading the whole file. Repeated calls on an +/// unchanged file are a no-op. +pub async fn ingest_cursor_transcript_event( + event_json: &str, + db: &GlobalDb, +) -> CursorTranscriptIngestStats { + ingest_cursor_transcript_event_capped(event_json, db, None).await +} + +/// Like [`ingest_cursor_transcript_event`], but bounds how many newly-appended +/// bytes a single call will read. The Cursor `beforeSubmitPrompt` hot path passes +/// a small cap so it can never threaten the 5s hook budget; backlogs larger than +/// the cap are left for the lower-frequency `sessionStart` / `stop` hooks (which +/// pass `None` for an unbounded catch-up read). +pub async fn ingest_cursor_transcript_event_capped( + event_json: &str, + db: &GlobalDb, + max_new_bytes: Option, +) -> CursorTranscriptIngestStats { + let Ok(event) = serde_json::from_str::(event_json) else { + return CursorTranscriptIngestStats::default(); + }; + let Some(transcript_path) = event + .get("transcript_path") + .and_then(Value::as_str) + .filter(|path| !path.is_empty()) + .map(PathBuf::from) + else { + return CursorTranscriptIngestStats::default(); + }; + + // Cursor derives its project from the event, so the driver's project_root + // argument is unused by `CursorEventSource`; the transcript path's parent is + // a cheap, side-effect-free placeholder. + let project_root = transcript_path + .parent() + .map_or_else(|| transcript_path.clone(), Path::to_path_buf); + let source = CursorEventSource { + event, + transcript_path, + }; + let stats = ingest_source(db, &source, &project_root, max_new_bytes).await; + CursorTranscriptIngestStats { + sessions_upserted: stats.sessions_upserted, + messages_upserted: stats.messages_upserted, + } +} + +fn event_message( + record: &Value, + event: &Value, + session_id: &str, + transcript_path: &Path, + ordinal: i64, + source_offset: i64, +) -> Option { + let role = record + .get("role") + .and_then(Value::as_str) + .filter(|role| !role.is_empty())?; + let message = record.get("message").unwrap_or(record); + let content = message.get("content").unwrap_or(message); + let (text, tool_names) = content_text_and_tools(content); + if text.trim().is_empty() { + return None; + } + + let message_id = record + .get("id") + .or_else(|| message.get("id")) + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map_or_else( + || format!("{session_id}:{ordinal}"), + std::string::ToString::to_string, + ); + let model = record + .get("model") + .or_else(|| message.get("model")) + .or_else(|| event.get("model")) + .and_then(Value::as_str) + .map(str::to_string); + + Some(SessionMessageRecord { + provider: "cursor".to_string(), + message_id, + session_id: session_id.to_string(), + role: role.to_string(), + timestamp: record_timestamp(record).or_else(|| record_timestamp(event)), + ordinal, + text, + kind: content_kind(content).map(str::to_string), + model, + tool_names: (!tool_names.is_empty()).then(|| tool_names.join(",")), + source_path: Some(transcript_path.to_string_lossy().to_string()), + source_offset: Some(source_offset), + metadata_json: serde_json::to_string(&message_metadata(record)).ok(), + }) +} + +fn content_text_and_tools(content: &Value) -> (String, Vec) { + if let Some(text) = content.as_str() { + return (text.to_string(), Vec::new()); + } + let Some(items) = content.as_array() else { + return ( + content + .get("text") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + Vec::new(), + ); + }; + + let mut texts = Vec::new(); + let mut tools = Vec::new(); + for item in items { + match item.get("type").and_then(Value::as_str) { + Some("text") => { + if let Some(text) = item.get("text").and_then(Value::as_str) { + texts.push(text.to_string()); + } + } + Some("tool_use") => { + if let Some(name) = item.get("name").and_then(Value::as_str) { + tools.push(name.to_string()); + } + } + _ => {} + } + } + (texts.join("\n\n"), tools) +} + +fn content_kind(content: &Value) -> Option<&'static str> { + if content.is_array() { + Some("message") + } else if content.is_string() { + Some("text") + } else { + None + } +} + +fn event_session_id(event: &Value, transcript_path: &Path) -> String { + event + .get("session_id") + .or_else(|| event.get("conversation_id")) + .or_else(|| event.get("chat_id")) + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map_or_else( + || { + transcript_path + .file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or("unknown") + .to_string() + }, + str::to_string, + ) +} + +fn event_project(event: &Value) -> (String, String) { + let candidates = event_project_candidates(event); + let project = candidates + .iter() + .find_map(|candidate| crate::config::discover_project_root(candidate)) + .or_else(|| candidates.into_iter().next()) + .map_or_else( + || "unknown".to_string(), + |path| path.to_string_lossy().to_string(), + ); + (project.clone(), project) +} + +fn event_project_candidates(event: &Value) -> Vec { + let mut candidates = Vec::new(); + if let Some(roots) = event.get("workspace_roots").and_then(Value::as_array) { + for root in roots { + if let Some(path) = root.as_str().filter(|path| !path.is_empty()) { + candidates.push(PathBuf::from(path)); + } + } + } + if let Some(cwd) = event + .get("cwd") + .and_then(Value::as_str) + .filter(|path| !path.is_empty()) + { + candidates.push(PathBuf::from(cwd)); + } + if let Some(file_path) = event + .get("file_path") + .and_then(Value::as_str) + .filter(|path| !path.is_empty()) + { + candidates.push(PathBuf::from(file_path)); + } + candidates +} + +fn record_timestamp(value: &Value) -> Option { + value + .get("timestamp") + .or_else(|| value.get("created_at")) + .and_then(|timestamp| { + timestamp + .as_i64() + .or_else(|| timestamp.as_str().and_then(|s| s.parse::().ok())) + }) +} + +fn session_metadata(event: &Value) -> Value { + serde_json::json!({ + "source": "cursor_transcript", + "conversation_id": event.get("conversation_id").cloned(), + "hook_event_name": event.get("hook_event_name").cloned(), + "cursor_version": event.get("cursor_version").cloned(), + }) +} + +fn message_metadata(record: &Value) -> Value { + serde_json::json!({ + "source": "cursor_transcript", + "raw_type": record.get("type").cloned(), + }) +} diff --git a/src/sessions/mod.rs b/src/sessions/mod.rs index b0d8ddbf48..0f6ea5f2c2 100644 --- a/src/sessions/mod.rs +++ b/src/sessions/mod.rs @@ -1,5 +1,61 @@ +use std::path::Path; + use serde::{Deserialize, Serialize}; +use crate::global_db::GlobalDb; +use crate::sessions::source::{ingest_source, TranscriptIngestStats, TranscriptSource}; + +pub mod claude; +pub mod cline_like; +pub mod codex; +pub mod cursor; +pub mod source; +pub mod vibe; + +/// Ingest transcripts from every hookless, path-discoverable agent whose +/// sessions belong to `project_root`, into the project-local `sessions.db` +/// (`db`). This is the serve-side counterpart to the Cursor hooks: these agents +/// register no end-of-turn hook, so their transcripts are reconciled by the +/// startup catch-up sweep instead. Fail-open and incremental (unchanged files +/// are a no-op). +pub async fn ingest_global_sources(db: &GlobalDb, project_root: &Path) -> TranscriptIngestStats { + let mut sources: Vec> = Vec::new(); + if let Some(source) = claude::ClaudeSource::new() { + sources.push(Box::new(source)); + } + if let Some(source) = codex::CodexSource::new() { + sources.push(Box::new(source)); + } + if let Some(source) = vibe::VibeSource::new() { + sources.push(Box::new(source)); + } + if let Some(source) = cline_like::ClineLikeSource::cline() { + sources.push(Box::new(source)); + } + if let Some(source) = cline_like::ClineLikeSource::roo_code() { + sources.push(Box::new(source)); + } + if let Some(source) = cline_like::ClineLikeSource::kilo() { + sources.push(Box::new(source)); + } + ingest_sources(db, project_root, &sources).await +} + +/// Drive a set of sources against `db` for `project_root`. Separated from +/// [`ingest_global_sources`] so tests can supply sources rooted at a temporary +/// home directory instead of the real `~`. +pub(crate) async fn ingest_sources( + db: &GlobalDb, + project_root: &Path, + sources: &[Box], +) -> TranscriptIngestStats { + let mut stats = TranscriptIngestStats::default(); + for source in sources { + stats = stats.merge(ingest_source(db, source.as_ref(), project_root, None).await); + } + stats +} + /// Provider-neutral metadata for an indexed agent session. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct SessionRecord { diff --git a/src/sessions/source.rs b/src/sessions/source.rs new file mode 100644 index 0000000000..bb7a34a711 --- /dev/null +++ b/src/sessions/source.rs @@ -0,0 +1,590 @@ +//! Provider-neutral transcript ingestion framework. +//! +//! Every agent transcript — Cursor, Claude Code, Codex, Vibe, … — converges to +//! the same provider-neutral [`SessionMessageRecord`] rows in a per-project +//! `sessions.db`. This module factors the *incremental, fail-open* machinery +//! out of the original Cursor-specific implementation so any adapter can plug +//! in by implementing [`TranscriptSource`]. +//! +//! ## Incremental cursors +//! +//! Sources differ in how they store transcripts, so three cursor kinds are +//! supported, all persisted through the existing `parse_offsets` table +//! ([`GlobalDb::get_parse_offset`]/[`GlobalDb::set_parse_offset`]) keyed by file +//! path. The stored [`StoredCursor`] is `(position, mtime)` where `position` +//! means: +//! +//! * [`stream_new_jsonl`] — **`ByteOffset`**: append-only JSONL (Cursor, Claude, +//! Codex, …). `position` is the byte offset of the next unread line; we seek +//! there and stream only new lines. +//! * [`read_changed_file`] — **`ContentHash`**: full-file-rewrite JSON (Cline, +//! Roo Code, Kilo, …). `position` is a stable 64-bit prefix of the content +//! hash; combined with `mtime` it detects rewrites. On change the whole +//! document is re-parsed and re-upserted — idempotent `ON CONFLICT` upserts +//! make re-adding unchanged messages a no-op. +//! * [`read_new_rows`] — **`RowCursor`**: SQLite-backed stores (Zed, Copilot CLI +//! `session-store.db`). `position` is the last-seen `rowid`; we select rows +//! with a greater `rowid`. +//! +//! All three are fail-open: any I/O or parse error yields "nothing new" rather +//! than propagating, so ingestion never blocks an agent. + +use std::io::{BufRead, BufReader, Seek, SeekFrom}; +use std::path::{Path, PathBuf}; + +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::global_db::GlobalDb; +use crate::sessions::{SessionMessageRecord, SessionRecord}; + +/// Counters returned by an ingestion pass. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct TranscriptIngestStats { + pub sessions_upserted: u64, + pub messages_upserted: u64, +} + +impl TranscriptIngestStats { + /// Accumulate another pass's counters into this one. + #[must_use] + pub fn merge(self, other: Self) -> Self { + Self { + sessions_upserted: self + .sessions_upserted + .saturating_add(other.sessions_upserted), + messages_upserted: self + .messages_upserted + .saturating_add(other.messages_upserted), + } + } +} + +/// The incremental position persisted between ingestion runs. +/// +/// `position` is interpreted per cursor kind: a byte offset (`ByteOffset`), a +/// stable 64-bit content hash prefix (`ContentHash`), or a last-seen `rowid` +/// (`RowCursor`). `mtime` is the file modification time in epoch seconds, used +/// to detect rewrites and to skip unchanged files cheaply. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct StoredCursor { + pub position: u64, + pub mtime: u64, +} + +/// Provider-neutral session metadata an adapter derives while parsing. +/// +/// The driver merges this with any existing row so a session's original +/// `started_at`/`title` survive incremental appends. +pub struct SessionDraft { + pub session_id: String, + pub project_key: String, + pub project_path: String, + pub title: Option, + pub metadata_json: Option, +} + +/// The result of parsing only the *new* portion of one transcript file. +pub struct ParsedTranscript { + pub draft: SessionDraft, + pub messages: Vec, + pub new_cursor: StoredCursor, +} + +/// A pluggable transcript provider. +/// +/// Implementors locate their transcript files for a project and parse only the +/// content appended/changed since the last run. The shared [`ingest_source`] +/// driver handles offset persistence and idempotent session/message upserts. +/// +/// `Send + Sync` is required so boxed sources can be driven from detached +/// background tasks (e.g. the serve-side startup sweep). +pub trait TranscriptSource: Send + Sync { + /// Stable provider id stored on every session/message row (e.g. `"claude"`). + fn provider(&self) -> &'static str; + + /// Candidate transcript files to consider for `project_root`. May scan + /// per-project and/or OS-specific global directories. Non-existent paths + /// are tolerated by the driver. + fn transcript_paths(&self, project_root: &Path) -> Vec; + + /// Parse only the new content of `path` given the previously stored cursor. + /// + /// Returns `None` to mean "ingest nothing and do not advance the cursor" + /// (unreadable file, hot-path byte cap exceeded, or the transcript does not + /// belong to `project_root`). Returns `Some` with a possibly-empty message + /// list otherwise; an empty list still advances the cursor (e.g. only + /// non-message lines were appended). + fn parse_new( + &self, + path: &Path, + prev: StoredCursor, + project_root: &Path, + max_new_bytes: Option, + ) -> Option; +} + +/// Drive a single source to completion against `db`, ingesting every transcript +/// it locates for `project_root`. Fail-open: per-file errors are swallowed. +/// +/// `max_new_bytes` bounds how much newly-appended content a byte-offset source +/// will read in one call (used to keep per-prompt hot paths inside budget); +/// pass `None` for an unbounded catch-up. +pub async fn ingest_source( + db: &GlobalDb, + source: &dyn TranscriptSource, + project_root: &Path, + max_new_bytes: Option, +) -> TranscriptIngestStats { + let mut stats = TranscriptIngestStats::default(); + for path in source.transcript_paths(project_root) { + stats = stats.merge(ingest_one(db, source, &path, project_root, max_new_bytes).await); + } + stats +} + +/// Ingest one transcript file: load the prior cursor, parse new content, persist +/// the advanced cursor, then upsert the session (merging preserved fields) and +/// its new messages. +async fn ingest_one( + db: &GlobalDb, + source: &dyn TranscriptSource, + path: &Path, + project_root: &Path, + max_new_bytes: Option, +) -> TranscriptIngestStats { + let path_str = path.to_string_lossy().to_string(); + let (prev_position, prev_mtime) = db.get_parse_offset(&path_str).await.unwrap_or((0, 0)); + let prev = StoredCursor { + position: prev_position, + mtime: prev_mtime, + }; + let Some(parsed) = source.parse_new(path, prev, project_root, max_new_bytes) else { + return TranscriptIngestStats::default(); + }; + + // Persist progress (even with zero messages) so the next run only sees + // genuinely new content. + db.set_parse_offset( + &path_str, + parsed.new_cursor.position, + parsed.new_cursor.mtime, + ) + .await; + + if parsed.messages.is_empty() { + return TranscriptIngestStats::default(); + } + + let provider = source.provider(); + let draft = parsed.draft; + let existing = db.get_session(provider, &draft.session_id).await; + // Preserve the session's original start time and title across appends; only + // advance ended_at to the latest message seen. + let started_at = existing + .as_ref() + .and_then(|session| session.started_at) + .or_else(|| { + parsed + .messages + .first() + .and_then(|message| message.timestamp) + }); + let title = existing + .as_ref() + .and_then(|session| session.title.clone()) + .or(draft.title); + let ended_at = parsed + .messages + .last() + .and_then(|message| message.timestamp) + .or_else(|| existing.as_ref().and_then(|session| session.ended_at)); + + let session = SessionRecord { + provider: provider.to_string(), + session_id: draft.session_id, + project_key: draft.project_key, + project_path: draft.project_path, + title, + started_at, + ended_at, + transcript_path: Some(path.to_string_lossy().to_string()), + metadata_json: draft.metadata_json, + }; + + let sessions_upserted = u64::from(db.upsert_session(&session).await); + let mut messages_upserted = 0_u64; + for message in &parsed.messages { + if db.upsert_session_message(message).await { + messages_upserted = messages_upserted.saturating_add(1); + } + } + TranscriptIngestStats { + sessions_upserted, + messages_upserted, + } +} + +/// One newly-read JSONL line: its starting byte offset and decoded value. +pub struct JsonlLine { + pub offset: i64, + pub value: Value, +} + +/// New JSONL content read from a file, plus the advanced cursor. +pub struct NewJsonl { + pub lines: Vec, + pub new_cursor: StoredCursor, +} + +/// **`ByteOffset`** reader for append-only JSONL. +/// +/// Seeks to `prev.position` (when the file has only grown and its mtime has not +/// regressed) and streams complete, newline-terminated lines, decoding each as +/// JSON. Blank and undecodable lines still advance the offset (so they are not +/// re-read) but are omitted from `lines`. A trailing line without a newline is a +/// partial write and is left unconsumed for the next call. +/// +/// Returns `None` when the file cannot be stat-ed/opened, or when +/// `max_new_bytes` is set and the unread tail exceeds it (so a hot path can defer +/// a large backlog to a lower-frequency caller without advancing the cursor). +pub fn stream_new_jsonl( + path: &Path, + prev: StoredCursor, + max_new_bytes: Option, +) -> Option { + let meta = std::fs::metadata(path).ok()?; + let file_size = meta.len(); + let mtime = file_mtime_secs(&meta); + + // Resume from the saved offset only when the file has grown (or stayed) and + // its mtime has not regressed; otherwise treat it as truncated/rewritten and + // restart from the beginning. + let resume = prev.position > 0 && file_size >= prev.position && mtime >= prev.mtime; + let seek_to = if resume { prev.position } else { 0 }; + + if seek_to >= file_size { + // Nothing new; refresh mtime so we stop re-stat-ing an idle file. + return Some(NewJsonl { + lines: Vec::new(), + new_cursor: StoredCursor { + position: seek_to, + mtime, + }, + }); + } + + if let Some(cap) = max_new_bytes { + if file_size.saturating_sub(seek_to) > cap { + return None; + } + } + + let file = std::fs::File::open(path).ok()?; + let mut reader = BufReader::new(file); + if seek_to > 0 && reader.seek(SeekFrom::Start(seek_to)).is_err() { + return None; + } + + let mut lines = Vec::new(); + let mut offset = seek_to; + let mut line = String::new(); + loop { + line.clear(); + match reader.read_line(&mut line) { + Ok(0) | Err(_) => break, + Ok(n) => { + // A line without a trailing newline is a partial write at EOF: + // stop without consuming it so the next call re-reads it whole. + if !line.ends_with('\n') { + break; + } + let line_offset = offset; + offset = offset.saturating_add(n as u64); + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + if let Ok(value) = serde_json::from_str::(trimmed) { + lines.push(JsonlLine { + offset: line_offset as i64, + value, + }); + } + } + } + } + + Some(NewJsonl { + lines, + new_cursor: StoredCursor { + position: offset, + mtime, + }, + }) +} + +/// Full contents of a changed file plus the advanced cursor. +pub struct ChangedFile { + pub contents: String, + pub new_cursor: StoredCursor, +} + +/// **`ContentHash`** reader for full-file-rewrite JSON. +/// +/// Detects a change via `(content_hash64, mtime)` versus the stored cursor and, +/// on change, returns the whole file so the caller can re-derive every message +/// with deterministic ids. Idempotent upserts make re-adding unchanged messages +/// a no-op. Returns `None` when the file cannot be read or is unchanged since +/// the last run. +pub fn read_changed_file(path: &Path, prev: StoredCursor) -> Option { + let meta = std::fs::metadata(path).ok()?; + let mtime = file_mtime_secs(&meta); + let contents = std::fs::read_to_string(path).ok()?; + let hash = content_hash64(&contents); + + // Unchanged since last run (we have read it before and neither content hash + // nor mtime moved) -> nothing to do. + if prev.position == hash && prev.mtime == mtime && (prev.position != 0 || prev.mtime != 0) { + return None; + } + + Some(ChangedFile { + contents, + new_cursor: StoredCursor { + position: hash, + mtime, + }, + }) +} + +/// Mapped rows read past the stored cursor, plus the advanced cursor. +pub struct NewRows { + pub items: Vec, + pub new_cursor: StoredCursor, +} + +/// **`RowCursor`** reader for SQLite-backed transcript stores (Zed, Copilot CLI +/// `session-store.db`). +/// +/// Selects rows whose rowid is greater than `prev.position` (the last-seen +/// rowid), ordered ascending, mapping each through `map_row` *during* iteration +/// (libsql rows must not outlive the cursor) and advancing the stored cursor to +/// the maximum rowid seen. `select_sql` must select the rowid as its first +/// column and accept a single `?` bound to the previous rowid, e.g. +/// `"SELECT rowid, role, text FROM turns WHERE rowid > ? ORDER BY rowid"`. +/// Fail-open: any query error yields `None`; `map_row` returning `None` skips +/// that row while still advancing the cursor. +pub async fn read_new_rows( + conn: &libsql::Connection, + select_sql: &str, + prev: StoredCursor, + mut map_row: impl FnMut(i64, &libsql::Row) -> Option, +) -> Option> { + let mut result_rows = conn + .query(select_sql, libsql::params![prev.position as i64]) + .await + .ok()?; + + let mut items = Vec::new(); + let mut max_rowid = prev.position; + while let Ok(Some(row)) = result_rows.next().await { + let Ok(rowid) = row.get::(0) else { + continue; + }; + if rowid as u64 > max_rowid { + max_rowid = rowid as u64; + } + if let Some(item) = map_row(rowid, &row) { + items.push(item); + } + } + + Some(NewRows { + items, + new_cursor: StoredCursor { + position: max_rowid, + // Row stores have no single file mtime; the rowid alone is the + // monotonic cursor, so mtime is left as a sentinel. + mtime: 0, + }, + }) +} + +/// Recursively collect files with the given extension under `dir`, bounded by +/// `max_depth` to avoid runaway traversal. Returns an empty vec when `dir` is +/// missing or unreadable. Used by global-store adapters (Claude, Codex) whose +/// transcripts live in nested date/slug directories. +pub(crate) fn collect_files_with_ext(dir: &Path, ext: &str, max_depth: u8) -> Vec { + let mut out = Vec::new(); + collect_files_inner(dir, ext, max_depth, 0, &mut out); + out +} + +fn collect_files_inner(dir: &Path, ext: &str, max_depth: u8, depth: u8, out: &mut Vec) { + if depth > max_depth { + return; + } + let Ok(entries) = std::fs::read_dir(dir) else { + return; + }; + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + collect_files_inner(&path, ext, max_depth, depth + 1, out); + } else if path.extension().and_then(|e| e.to_str()) == Some(ext) { + out.push(path); + } + } +} + +/// Compare two paths for equality, canonicalizing when possible so that +/// symlinks/`..`/trailing differences do not cause false mismatches. Falls back +/// to a literal comparison when canonicalization fails (e.g. a path that no +/// longer exists). +pub(crate) fn paths_equal(a: &Path, b: &Path) -> bool { + match (a.canonicalize(), b.canonicalize()) { + (Ok(a), Ok(b)) => a == b, + _ => a == b, + } +} + +/// File modification time in epoch seconds, or 0 when unavailable. +fn file_mtime_secs(meta: &std::fs::Metadata) -> u64 { + meta.modified() + .ok() + .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()) + .map_or(0, |d| d.as_secs()) +} + +/// Stable 64-bit content hash prefix suitable for the existing integer +/// `parse_offsets.byte_offset` column. +fn content_hash64(contents: &str) -> u64 { + let mut hasher = Sha256::new(); + hasher.update(contents.as_bytes()); + let digest = hasher.finalize(); + let mut bytes = [0_u8; 8]; + bytes.copy_from_slice(&digest[..8]); + u64::from_be_bytes(bytes) +} + +/// Collapse whitespace and clip to a short preview suitable for a session title. +pub(crate) fn preview_title(text: &str) -> String { + const MAX_TITLE_CHARS: usize = 80; + let collapsed = text.split_whitespace().collect::>().join(" "); + if collapsed.chars().count() <= MAX_TITLE_CHARS { + collapsed + } else { + collapsed.chars().take(MAX_TITLE_CHARS).collect() + } +} + +/// Build a session title from the first user message, if any. +pub(crate) fn title_from_messages(messages: &[SessionMessageRecord]) -> Option { + messages + .iter() + .find(|message| message.role == "user") + .map(|message| preview_title(&message.text)) +} + +#[cfg(test)] +#[allow(clippy::unwrap_used)] +mod tests { + use super::*; + use std::io::Write; + + #[test] + fn stream_new_jsonl_reads_only_appended_lines() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("t.jsonl"); + std::fs::write(&path, "{\"a\":1}\n{\"a\":2}\n").unwrap(); + + let first = stream_new_jsonl(&path, StoredCursor::default(), None).unwrap(); + assert_eq!(first.lines.len(), 2); + + // Re-reading from the advanced cursor yields nothing. + let again = stream_new_jsonl(&path, first.new_cursor, None).unwrap(); + assert_eq!(again.lines.len(), 0); + + // Appending one line yields only that line on the next read. + let mut f = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + f.write_all(b"{\"a\":3}\n").unwrap(); + drop(f); + let third = stream_new_jsonl(&path, again.new_cursor, None).unwrap(); + assert_eq!(third.lines.len(), 1); + assert_eq!(third.lines[0].value["a"], 3); + } + + #[test] + fn stream_new_jsonl_defers_partial_final_line_and_respects_cap() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("t.jsonl"); + std::fs::write(&path, "{\"a\":1}\n{\"a\":2}").unwrap(); // second line unterminated + + let read = stream_new_jsonl(&path, StoredCursor::default(), None).unwrap(); + assert_eq!(read.lines.len(), 1, "partial final line must be deferred"); + + // A cap smaller than the unread tail defers the whole read (no cursor advance). + assert!(stream_new_jsonl(&path, StoredCursor::default(), Some(1)).is_none()); + } + + #[test] + fn read_changed_file_detects_change_and_noops_when_unchanged() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("chat.json"); + std::fs::write(&path, "[{\"role\":\"user\"}]").unwrap(); + + let changed = read_changed_file(&path, StoredCursor::default()).unwrap(); + assert!(changed.contents.contains("user")); + // Unchanged file → None. + assert!(read_changed_file(&path, changed.new_cursor).is_none()); + } + + #[tokio::test] + async fn read_new_rows_tracks_last_rowid() { + // A synthetic SQLite-backed source exercises the RowCursor kind. + let db = libsql::Builder::new_local(":memory:") + .build() + .await + .unwrap(); + let conn = db.connect().unwrap(); + conn.execute("CREATE TABLE turns (role TEXT, text TEXT)", ()) + .await + .unwrap(); + conn.execute( + "INSERT INTO turns (role, text) VALUES ('user', 'hello'), ('assistant', 'hi')", + (), + ) + .await + .unwrap(); + + let sql = "SELECT rowid, role, text FROM turns WHERE rowid > ? ORDER BY rowid"; + let map = |_rowid: i64, row: &libsql::Row| row.get::(2).ok(); + let first = read_new_rows(&conn, sql, StoredCursor::default(), map) + .await + .unwrap(); + assert_eq!(first.items, vec!["hello".to_string(), "hi".to_string()]); + assert_eq!(first.new_cursor.position, 2); + + // No new rows past the advanced cursor. + let again = read_new_rows(&conn, sql, first.new_cursor, map) + .await + .unwrap(); + assert_eq!(again.items.len(), 0); + + conn.execute( + "INSERT INTO turns (role, text) VALUES ('user', 'again')", + (), + ) + .await + .unwrap(); + let third = read_new_rows(&conn, sql, again.new_cursor, map) + .await + .unwrap(); + assert_eq!(third.items, vec!["again".to_string()]); + assert_eq!(third.new_cursor.position, 3); + } +} diff --git a/src/sessions/vibe.rs b/src/sessions/vibe.rs new file mode 100644 index 0000000000..c603c7fd0d --- /dev/null +++ b/src/sessions/vibe.rs @@ -0,0 +1,238 @@ +//! Mistral Vibe transcript source. +//! +//! Vibe stores sessions under `$VIBE_HOME/logs/session/` or +//! `~/.vibe/logs/session/`. Each session directory contains: +//! +//! * `meta.json` - cumulative metadata, including session id, active model, and +//! the working directory (`environment.working_directory` in current releases). +//! * `messages.jsonl` - append-only line-delimited LLM messages. +//! +//! This source uses the shared **`ByteOffset`** reader for `messages.jsonl` and +//! scopes sessions to a tokensave project by matching the working directory in +//! `meta.json` to `project_root`. + +use std::path::{Path, PathBuf}; + +use serde_json::Value; + +use crate::sessions::source::{ + collect_files_with_ext, paths_equal, stream_new_jsonl, title_from_messages, ParsedTranscript, + SessionDraft, StoredCursor, TranscriptSource, +}; +use crate::sessions::SessionMessageRecord; + +const PROVIDER: &str = "vibe"; +const MAX_SCAN_DEPTH: u8 = 4; + +/// Vibe session locator + parser. +pub struct VibeSource { + session_root: PathBuf, +} + +impl VibeSource { + /// Source rooted at the real Vibe home. Returns `None` when the home + /// directory cannot be resolved. + pub fn new() -> Option { + let home = dirs::home_dir()?; + Some(Self::with_home(&home)) + } + + /// Source rooted at `/.vibe/logs/session` (used by tests). This does + /// not read `VIBE_HOME`; tests can pass the desired base explicitly. + pub fn with_home(home: &Path) -> Self { + Self::with_vibe_home(&home.join(".vibe")) + } + + /// Source rooted at `/logs/session`. + pub fn with_vibe_home(vibe_home: &Path) -> Self { + Self { + session_root: vibe_home.join("logs").join("session"), + } + } +} + +impl TranscriptSource for VibeSource { + fn provider(&self) -> &'static str { + PROVIDER + } + + fn transcript_paths(&self, _project_root: &Path) -> Vec { + collect_files_with_ext(&self.session_root, "jsonl", MAX_SCAN_DEPTH) + .into_iter() + .filter(|path| { + path.file_name().and_then(|name| name.to_str()) == Some("messages.jsonl") + }) + .collect() + } + + fn parse_new( + &self, + path: &Path, + prev: StoredCursor, + project_root: &Path, + max_new_bytes: Option, + ) -> Option { + let meta_path = path.parent()?.join("meta.json"); + let meta = read_meta(&meta_path)?; + if !paths_equal(&meta.working_directory, project_root) { + return None; + } + + let new = stream_new_jsonl(path, prev, max_new_bytes)?; + let mut messages = Vec::new(); + for line in &new.lines { + if let Some(message) = message_from_line(&line.value, &meta, path, line.offset) { + messages.push(message); + } + } + + let project = project_root.to_string_lossy().to_string(); + let draft = SessionDraft { + session_id: meta.session_id, + project_key: project.clone(), + project_path: project, + title: title_from_messages(&messages), + metadata_json: serde_json::to_string(&serde_json::json!({ + "source": "vibe_messages", + })) + .ok(), + }; + + Some(ParsedTranscript { + draft, + messages, + new_cursor: new.new_cursor, + }) + } +} + +struct VibeMeta { + session_id: String, + working_directory: PathBuf, + model: Option, +} + +fn read_meta(path: &Path) -> Option { + let value: Value = serde_json::from_str(&std::fs::read_to_string(path).ok()?).ok()?; + let session_id = value + .get("session_id") + .or_else(|| value.get("id")) + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map_or_else( + || { + path.parent() + .and_then(Path::file_name) + .and_then(|name| name.to_str()) + .unwrap_or("unknown") + .to_string() + }, + ToString::to_string, + ); + let working_directory = value + .pointer("/environment/working_directory") + .or_else(|| value.pointer("/environment/workdir")) + .or_else(|| value.pointer("/config/working_directory")) + .or_else(|| value.pointer("/config/workdir")) + .or_else(|| value.get("working_directory")) + .or_else(|| value.get("cwd")) + .and_then(Value::as_str) + .filter(|path| !path.is_empty()) + .map(PathBuf::from)?; + let model = value + .pointer("/config/active_model") + .or_else(|| value.get("active_model")) + .or_else(|| value.get("model")) + .and_then(Value::as_str) + .map(str::to_string); + + Some(VibeMeta { + session_id, + working_directory, + model, + }) +} + +fn message_from_line( + record: &Value, + meta: &VibeMeta, + path: &Path, + offset: i64, +) -> Option { + let role = record + .get("role") + .or_else(|| record.pointer("/message/role")) + .and_then(Value::as_str) + .filter(|role| matches!(*role, "user" | "assistant" | "model"))?; + let normalized_role = if role == "model" { "assistant" } else { role }; + let content = record + .get("content") + .or_else(|| record.pointer("/message/content")) + .unwrap_or(record); + let (text, tool_names) = content_text_and_tools(content); + if text.trim().is_empty() { + return None; + } + let timestamp = record + .get("timestamp") + .or_else(|| record.get("created_at")) + .and_then(|value| { + value + .as_i64() + .or_else(|| value.as_str().and_then(|s| s.parse::().ok())) + }); + + Some(SessionMessageRecord { + provider: PROVIDER.to_string(), + message_id: format!("{}:{offset}", meta.session_id), + session_id: meta.session_id.clone(), + role: normalized_role.to_string(), + timestamp, + ordinal: offset, + text, + kind: Some("message".to_string()), + model: meta.model.clone(), + tool_names: (!tool_names.is_empty()).then(|| tool_names.join(",")), + source_path: Some(path.to_string_lossy().to_string()), + source_offset: Some(offset), + metadata_json: serde_json::to_string(&serde_json::json!({ + "source": "vibe_messages", + })) + .ok(), + }) +} + +fn content_text_and_tools(content: &Value) -> (String, Vec) { + if let Some(text) = content.as_str() { + return (text.to_string(), Vec::new()); + } + let Some(items) = content.as_array() else { + return ( + content + .get("text") + .or_else(|| content.get("message")) + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + Vec::new(), + ); + }; + + let mut texts = Vec::new(); + let mut tools = Vec::new(); + for item in items { + if let Some(text) = item.get("text").and_then(Value::as_str) { + texts.push(text.to_string()); + } + if let Some(name) = item + .get("tool_call") + .or_else(|| item.get("functionCall")) + .or_else(|| item.get("function_call")) + .and_then(|call| call.get("name")) + .and_then(Value::as_str) + { + tools.push(name.to_string()); + } + } + (texts.join("\n\n"), tools) +} diff --git a/src/tokensave.rs b/src/tokensave.rs index d6dbef7545..5c0ebbbab3 100644 --- a/src/tokensave.rs +++ b/src/tokensave.rs @@ -17,6 +17,13 @@ use crate::db::Database; use crate::errors::{Result, TokenSaveError}; use crate::extraction::LanguageRegistry; use crate::graph::{GraphQueryManager, GraphTraverser}; +use crate::memory::retrieval::FactRetriever; +use crate::memory::store::MemoryStore; +use crate::memory::trust::{DEFAULT_MIN_TRUST, DEFAULT_TRUST}; +use crate::memory::types::{ + AddFactRequest, ContradictionResult, FactRecord, FactSearchResult, FeedbackRequest, + FeedbackResult, MemoryCategory, MemoryStatus, SearchFactsRequest, UpdateFactRequest, +}; use crate::resolution::ReferenceResolver; use crate::sync; use crate::types::*; @@ -152,38 +159,6 @@ pub struct TokenSave { fallback_warning: Option, } -/// A decision recorded by an agent during a session. -#[derive(Debug, Clone, serde::Serialize)] -pub struct DecisionRecord { - /// Row id. - pub id: i64, - /// The decision text. - pub text: String, - /// Optional rationale for the decision. - pub reason: Option, - /// UNIX timestamp (seconds) when the decision was recorded. - pub created_at: i64, - /// File paths relevant to this decision. - pub files: Vec, - /// Arbitrary tags for categorisation. - pub tags: Vec, -} - -/// A code area (file path) that an agent has touched during a session. -#[derive(Debug, Clone, serde::Serialize)] -pub struct CodeAreaRecord { - /// Row id. - pub id: i64, - /// Relative file path. - pub path: String, - /// Optional human-readable description of the area. - pub description: Option, - /// UNIX timestamp (seconds) of the most recent touch. - pub last_touched_at: i64, - /// How many times this path has been touched. - pub touch_count: u32, -} - /// Result of a full indexing operation. pub struct IndexResult { /// Number of files scanned and indexed. @@ -2855,189 +2830,280 @@ impl TokenSave { // Session memory // --------------------------------------------------------------------------- -const MAX_RECALL_LIMIT: usize = 200; -const MAX_CODE_AREAS_LIMIT: usize = 200; +const MAX_FACT_LIMIT: usize = 200; +const DEFAULT_FACT_LIMIT: usize = 20; + +fn memory_database_error(operation: &str, message: impl std::fmt::Display) -> TokenSaveError { + TokenSaveError::Database { + message: format!("{operation} failed: {message}"), + operation: operation.to_string(), + } +} + +fn fact_result_ids(results: &[FactSearchResult]) -> Vec { + results.iter().map(|result| result.fact.fact_id).collect() +} + +fn fact_ids(facts: &[FactRecord]) -> Vec { + facts.iter().map(|fact| fact.fact_id).collect() +} impl TokenSave { - /// Record an agent decision. Returns the new row id. - pub async fn record_decision( - &self, - text: &str, - reason: Option<&str>, - files: &[String], - tags: &[String], - ) -> crate::errors::Result { - debug_assert!(!text.is_empty(), "decision text must not be empty"); - let files_json = - serde_json::to_string(files).map_err(|e| crate::errors::TokenSaveError::Database { - message: format!("record_decision files serialization failed: {e}"), - operation: "record_decision".to_string(), - })?; - let tags_json = - serde_json::to_string(tags).map_err(|e| crate::errors::TokenSaveError::Database { - message: format!("record_decision tags serialization failed: {e}"), - operation: "record_decision".to_string(), - })?; - let now = current_timestamp(); - let conn = self.db.conn(); - conn.execute( - "INSERT INTO memory_decisions (text, reason, created_at, files, tags) \ - VALUES (?1, ?2, ?3, ?4, ?5)", - libsql::params![text, reason, now, files_json, tags_json], - ) - .await - .map_err(|e| crate::errors::TokenSaveError::Database { - message: format!("record_decision insert failed: {e}"), - operation: "record_decision".to_string(), - })?; - Ok(conn.last_insert_rowid()) + /// Add or replace a fact in the holographic memory store. + pub async fn add_fact(&self, request: AddFactRequest) -> Result { + MemoryStore::new(self.db.conn()) + .add_fact(request, DEFAULT_TRUST) + .await + } + + /// Search facts by lexical overlap, entity metadata, category, and trust. + pub async fn search_facts(&self, request: SearchFactsRequest) -> Result> { + let mut results = FactRetriever::new(self.db.conn()) + .search( + &request.query, + request.category, + request.min_trust, + request.limit.unwrap_or(DEFAULT_FACT_LIMIT), + ) + .await?; + if !request.include_why { + for result in &mut results { + result.why = None; + } + } + MemoryStore::new(self.db.conn()) + .increment_retrieval_counts(&fact_result_ids(&results)) + .await?; + Ok(results) } - /// Recall decisions. With `query`, runs FTS5 MATCH against text+reason. - /// Without `query`, returns newest-first. - pub async fn session_recall( + pub async fn probe_entity( &self, - query: Option<&str>, - since: Option, + entity: &str, + category: Option, + min_trust: Option, limit: usize, - ) -> crate::errors::Result> { - let limit = limit.clamp(1, MAX_RECALL_LIMIT) as i64; - let conn = self.db.conn(); + ) -> Result> { + let results = FactRetriever::new(self.db.conn()) + .probe(entity, category, min_trust, limit) + .await?; + MemoryStore::new(self.db.conn()) + .increment_retrieval_counts(&fact_result_ids(&results)) + .await?; + Ok(results) + } - let db_err = |e: libsql::Error| crate::errors::TokenSaveError::Database { - message: format!("session_recall query failed: {e}"), - operation: "session_recall".to_string(), - }; + pub async fn related_facts( + &self, + entity: &str, + category: Option, + min_trust: Option, + limit: usize, + ) -> Result> { + let retriever = FactRetriever::new(self.db.conn()); + let related_entities = retriever.related(entity, limit).await?; + let mut seen = std::collections::HashSet::new(); + let mut results = Vec::new(); + for related in related_entities { + for result in retriever + .probe(&related.name, category, min_trust, limit.saturating_mul(2)) + .await? + { + if seen.insert(result.fact.fact_id) { + results.push(result); + if results.len() >= limit.clamp(1, MAX_FACT_LIMIT) { + break; + } + } + } + if results.len() >= limit.clamp(1, MAX_FACT_LIMIT) { + break; + } + } + MemoryStore::new(self.db.conn()) + .increment_retrieval_counts(&fact_result_ids(&results)) + .await?; + Ok(results) + } - let mut rows = match (query, since) { - (Some(q), Some(ts)) => conn - .query( - "SELECT d.id, d.text, d.reason, d.created_at, d.files, d.tags \ - FROM memory_decisions d \ - JOIN memory_decisions_fts f ON f.rowid = d.id \ - WHERE memory_decisions_fts MATCH ?1 AND d.created_at >= ?2 \ - ORDER BY d.created_at DESC LIMIT ?3", - libsql::params![q, ts, limit], - ) - .await - .map_err(db_err)?, - (Some(q), None) => conn - .query( - "SELECT d.id, d.text, d.reason, d.created_at, d.files, d.tags \ - FROM memory_decisions d \ - JOIN memory_decisions_fts f ON f.rowid = d.id \ - WHERE memory_decisions_fts MATCH ?1 \ - ORDER BY d.created_at DESC LIMIT ?2", - libsql::params![q, limit], - ) - .await - .map_err(db_err)?, - (None, Some(ts)) => conn - .query( - "SELECT id, text, reason, created_at, files, tags \ - FROM memory_decisions WHERE created_at >= ?1 \ - ORDER BY created_at DESC LIMIT ?2", - libsql::params![ts, limit], - ) - .await - .map_err(db_err)?, - (None, None) => conn - .query( - "SELECT id, text, reason, created_at, files, tags \ - FROM memory_decisions ORDER BY created_at DESC LIMIT ?1", - libsql::params![limit], - ) - .await - .map_err(db_err)?, - }; + pub async fn reason_facts( + &self, + entities: &[String], + category: Option, + min_trust: Option, + limit: usize, + ) -> Result> { + let results = FactRetriever::new(self.db.conn()) + .reason(entities, category, min_trust, limit) + .await?; + MemoryStore::new(self.db.conn()) + .increment_retrieval_counts(&fact_result_ids(&results)) + .await?; + Ok(results) + } - let row_err = |e: libsql::Error| crate::errors::TokenSaveError::Database { - message: format!("session_recall row read failed: {e}"), - operation: "session_recall".to_string(), - }; - let json_err = |e: serde_json::Error| crate::errors::TokenSaveError::Database { - message: format!("session_recall JSON parse failed: {e}"), - operation: "session_recall".to_string(), - }; + pub async fn contradict_facts( + &self, + category: Option, + threshold: f64, + limit: usize, + ) -> Result> { + let retriever = FactRetriever::new(self.db.conn()); + if let Some(category) = category { + return retriever.contradict(category, threshold, limit).await; + } let mut out = Vec::new(); - while let Some(row) = rows.next().await.map_err(row_err)? { - let files_json: String = row.get(4).map_err(row_err)?; - let tags_json: String = row.get(5).map_err(row_err)?; - out.push(DecisionRecord { - id: row.get(0).map_err(row_err)?, - text: row.get(1).map_err(row_err)?, - reason: row.get::>(2).map_err(row_err)?, - created_at: row.get(3).map_err(row_err)?, - files: serde_json::from_str(&files_json).map_err(json_err)?, - tags: serde_json::from_str(&tags_json).map_err(json_err)?, - }); + for category in [ + MemoryCategory::General, + MemoryCategory::UserPref, + MemoryCategory::Project, + MemoryCategory::Tool, + MemoryCategory::Decision, + MemoryCategory::CodeArea, + ] { + out.extend(retriever.contradict(category, threshold, limit).await?); + if out.len() >= limit.clamp(1, MAX_FACT_LIMIT) { + out.truncate(limit.clamp(1, MAX_FACT_LIMIT)); + break; + } } Ok(out) } - /// Record (or update) a code area the agent worked in. Increments `touch_count` - /// on re-touch. Description is set on first insert; subsequent `None` values - /// preserve the existing description. - pub async fn record_code_area( - &self, - path: &str, - description: Option<&str>, - ) -> crate::errors::Result<()> { - debug_assert!(!path.is_empty(), "code area path must not be empty"); - let now = current_timestamp(); - let conn = self.db.conn(); - conn.execute( - "INSERT INTO memory_code_areas (path, description, last_touched_at, touch_count) \ - VALUES (?1, ?2, ?3, 1) \ - ON CONFLICT(path) DO UPDATE SET \ - description = COALESCE(excluded.description, memory_code_areas.description), \ - last_touched_at = excluded.last_touched_at, \ - touch_count = memory_code_areas.touch_count + 1", - libsql::params![path, description, now], - ) - .await - .map_err(|e| crate::errors::TokenSaveError::Database { - message: format!("record_code_area upsert failed: {e}"), - operation: "record_code_area".to_string(), - })?; - Ok(()) + pub async fn update_fact(&self, request: UpdateFactRequest) -> Result { + MemoryStore::new(self.db.conn()).update_fact(request).await + } + + pub async fn remove_fact(&self, fact_id: i64) -> Result { + MemoryStore::new(self.db.conn()).remove_fact(fact_id).await } - /// List code areas, most-recently-touched first. - pub async fn list_code_areas( + pub async fn list_facts( &self, + category: Option, + min_trust: Option, limit: usize, - ) -> crate::errors::Result> { - let limit = limit.clamp(1, MAX_CODE_AREAS_LIMIT) as i64; + ) -> Result> { + let facts = MemoryStore::new(self.db.conn()) + .list_facts(category, min_trust, limit) + .await?; + MemoryStore::new(self.db.conn()) + .increment_retrieval_counts(&fact_ids(&facts)) + .await?; + Ok(facts) + } + + pub async fn get_fact(&self, fact_id: i64) -> Result> { + MemoryStore::new(self.db.conn()).get_fact(fact_id).await + } + + pub async fn record_fact_feedback(&self, request: FeedbackRequest) -> Result { + MemoryStore::new(self.db.conn()) + .record_feedback_event(request) + .await + } + + pub async fn memory_status(&self) -> Result { + let operation = "memory_status"; let conn = self.db.conn(); - let mut rows = conn + MemoryStore::new(conn).rebuild_dirty_banks().await?; + let mut fact_rows = conn + .query("SELECT trust_score FROM memory_facts", ()) + .await + .map_err(|e| memory_database_error(operation, e))?; + let row_err = |e: libsql::Error| memory_database_error(operation, e); + let mut trust_0_025_count = 0_usize; + let mut trust_025_050_count = 0_usize; + let mut trust_050_075_count = 0_usize; + let mut trust_075_100_count = 0_usize; + let mut below_default_recall_threshold_count = 0_usize; + let mut fact_count = 0_usize; + while let Some(row) = fact_rows.next().await.map_err(row_err)? { + fact_count += 1; + let trust_score = row.get::(0).map_err(row_err)?; + if trust_score < DEFAULT_MIN_TRUST { + below_default_recall_threshold_count += 1; + } + if trust_score < 0.25 { + trust_0_025_count += 1; + } else if trust_score < 0.50 { + trust_025_050_count += 1; + } else if trust_score < 0.75 { + trust_050_075_count += 1; + } else { + trust_075_100_count += 1; + } + } + let mut entity_rows = conn + .query("SELECT COUNT(*) FROM memory_entities", ()) + .await + .map_err(|e| memory_database_error(operation, e))?; + let entity_count = entity_rows + .next() + .await + .map_err(row_err)? + .map_or(Ok(0_i64), |row| row.get(0).map_err(row_err))?; + let mut bank_rows = conn + .query("SELECT COUNT(*) FROM memory_banks", ()) + .await + .map_err(|e| memory_database_error(operation, e))?; + let bank_count = bank_rows + .next() + .await + .map_err(row_err)? + .map_or(Ok(0_i64), |row| row.get(0).map_err(row_err))?; + let mut aggregate_rows = conn .query( - "SELECT id, path, description, last_touched_at, touch_count \ - FROM memory_code_areas ORDER BY last_touched_at DESC LIMIT ?1", - libsql::params![limit], + "SELECT COALESCE(SUM(helpful_count), 0), + COALESCE(SUM(unhelpful_count), 0), + COALESCE(SUM(CASE WHEN hrr_vector IS NULL THEN 1 ELSE 0 END), 0) + FROM memory_facts", + (), ) .await - .map_err(|e| crate::errors::TokenSaveError::Database { - message: format!("list_code_areas query failed: {e}"), - operation: "list_code_areas".to_string(), - })?; - let row_err = |e: libsql::Error| crate::errors::TokenSaveError::Database { - message: format!("list_code_areas row read failed: {e}"), - operation: "list_code_areas".to_string(), + .map_err(|e| memory_database_error(operation, e))?; + let Some(aggregate_row) = aggregate_rows.next().await.map_err(row_err)? else { + return Err(memory_database_error( + operation, + "memory aggregate query returned no rows", + )); }; - - let mut out = Vec::new(); - while let Some(row) = rows.next().await.map_err(row_err)? { - out.push(CodeAreaRecord { - id: row.get(0).map_err(row_err)?, - path: row.get(1).map_err(row_err)?, - description: row.get::>(2).map_err(row_err)?, - last_touched_at: row.get(3).map_err(row_err)?, - touch_count: row.get::(4).map_err(row_err)? as u32, - }); - } - Ok(out) + let helpful_count = aggregate_row.get::(0).map_err(row_err)?; + let unhelpful_count = aggregate_row.get::(1).map_err(row_err)?; + let missing_vector_count = aggregate_row.get::(2).map_err(row_err)?; + let mut backfill_rows = conn + .query( + "SELECT COUNT(*) FROM memory_facts + WHERE json_extract(metadata, '$.holographic_memory_backfill_v1') = 1", + (), + ) + .await + .map_err(|e| memory_database_error(operation, e))?; + let backfilled_count = backfill_rows + .next() + .await + .map_err(row_err)? + .map_or(Ok(0_i64), |row| row.get(0).map_err(row_err))?; + let hrr_dim = 2048_usize; + let estimated_capacity = (hrr_dim as f64 / (hrr_dim as f64).ln()).round() as usize; + Ok(MemoryStatus { + fact_count, + entity_count: entity_count as usize, + bank_count: bank_count as usize, + algebra_name: "amari_fhrr".to_string(), + hrr_dim, + estimated_capacity, + trust_0_025_count, + trust_025_050_count, + trust_050_075_count, + trust_075_100_count, + below_default_recall_threshold_count, + helpful_count: helpful_count as usize, + unhelpful_count: unhelpful_count as usize, + missing_vector_count: missing_vector_count as usize, + legacy_backfill_complete: backfilled_count > 0, + }) } } diff --git a/src/tool_command.rs b/src/tool_command.rs index 00729def20..2adeb0fc7a 100644 --- a/src/tool_command.rs +++ b/src/tool_command.rs @@ -371,7 +371,7 @@ fn print_tool_list(defs: &[ToolDefinition]) { let is_always = def .meta .as_ref() - .and_then(|m| m.pointer("/anthropic/alwaysLoad")) + .and_then(|m| m.get("anthropic/alwaysLoad")) .and_then(Value::as_bool) .unwrap_or(false); if is_always { @@ -446,9 +446,9 @@ fn group_for(def: &ToolDefinition) -> &'static str { || n == "tokensave_insert_at_symbol" { "edit" - } else if n == "tokensave_record_decision" - || n == "tokensave_record_code_area" - || n == "tokensave_session_recall" + } else if n == "tokensave_fact_store" + || n == "tokensave_fact_feedback" + || n == "tokensave_memory_status" || n == "tokensave_session_start" || n == "tokensave_session_end" { diff --git a/tests/agent_test.rs b/tests/agent_test.rs index 51ece4fd22..90584937ca 100644 --- a/tests/agent_test.rs +++ b/tests/agent_test.rs @@ -324,6 +324,18 @@ fn test_local_install_cursor_writes_project_config_only() { "Cursor workspaceOpen hook should run a catch-up sync" ); + let stop_hooks = hooks["hooks"]["stop"] + .as_array() + .expect("stop hooks should be an array"); + assert!( + stop_hooks.iter().any(|hook| { + hook["command"] + .as_str() + .is_some_and(|command| command.contains("hook-cursor-stop")) + }), + "Cursor stop hook should ingest the session transcript at end of turn" + ); + assert!( !home.path().join(".cursor/mcp.json").exists(), "local install must not write the global Cursor config" @@ -946,6 +958,18 @@ fn test_local_install_supported_agents_write_project_paths() { path.extension().and_then(|ext| ext.to_str()), Some("md" | "mdc") ); + if is_instruction_file { + assert!( + body.contains("tokensave_fact_store"), + "{agent} local instruction file {} should mention fact memory tools", + path.display() + ); + assert!( + body.contains("tokensave_message_search"), + "{agent} local instruction file {} should mention transcript message search", + path.display() + ); + } let is_cursor_permissions = agent == "cursor" && relative == ".cursor/permissions.json"; if !is_instruction_file && !is_cursor_permissions { let expected = expected_tokensave_bin(); diff --git a/tests/claude_transcript_ingest_test.rs b/tests/claude_transcript_ingest_test.rs new file mode 100644 index 0000000000..51e765af43 --- /dev/null +++ b/tests/claude_transcript_ingest_test.rs @@ -0,0 +1,146 @@ +use std::io::Write; + +use tempfile::TempDir; +use tokensave::sessions::claude::ClaudeSource; +use tokensave::sessions::cursor::open_project_session_db; +use tokensave::sessions::source::ingest_source; + +/// Builds an initialized project dir and returns (home, project_root). +fn setup(tmp: &TempDir) -> (std::path::PathBuf, std::path::PathBuf) { + let home = tmp.path().join("home"); + let project = tmp.path().join("project"); + std::fs::create_dir_all(&project).unwrap(); + std::fs::create_dir(project.join(".tokensave")).unwrap(); + std::fs::write(project.join(".tokensave/tokensave.db"), "").unwrap(); + (home, project) +} + +/// Writes a Claude Code transcript (one JSON object per line) for `session` whose +/// recorded `cwd` is `project`. +fn write_claude_transcript( + home: &std::path::Path, + project: &std::path::Path, + session: &str, +) -> std::path::PathBuf { + let dir = home.join(".claude/projects/-some-slug"); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join(format!("{session}.jsonl")); + let cwd = project.to_string_lossy(); + let contents = format!( + "{}\n{}\n", + serde_json::json!({ + "type": "user", + "cwd": cwd, + "sessionId": session, + "uuid": "u1", + "timestamp": "2026-01-01T00:00:00.000Z", + "message": {"role": "user", "content": "Investigate the billing pipeline regression"} + }), + serde_json::json!({ + "type": "assistant", + "cwd": cwd, + "sessionId": session, + "uuid": "u2", + "timestamp": "2026-01-01T00:00:05.000Z", + "message": { + "id": "msg_claude_1", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [ + {"type": "text", "text": "The billing pipeline regression is fixed."}, + {"type": "tool_use", "name": "tokensave_context", "input": {}} + ] + } + }), + ); + std::fs::write(&path, contents).unwrap(); + path +} + +#[tokio::test] +async fn claude_transcript_populates_searchable_messages() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + write_claude_transcript(&home, &project, "claude-sess"); + + let db = open_project_session_db(&project).await.unwrap(); + let source = ClaudeSource::with_home(&home); + + let stats = ingest_source(&db, &source, &project, None).await; + assert_eq!(stats.messages_upserted, 2); + assert_eq!(stats.sessions_upserted, 1); + + let results = db + .search_session_messages( + "claude", + Some(project.to_string_lossy().as_ref()), + "billing pipeline", + 10, + ) + .await; + assert_eq!(results.len(), 2); + assert!(results + .iter() + .any(|hit| hit.message.tool_names.as_deref() == Some("tokensave_context"))); + assert!(results + .iter() + .any(|hit| hit.message.model.as_deref() == Some("claude-opus-4-8"))); +} + +#[tokio::test] +async fn claude_transcript_ingest_is_incremental() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + let path = write_claude_transcript(&home, &project, "claude-sess"); + + let db = open_project_session_db(&project).await.unwrap(); + let source = ClaudeSource::with_home(&home); + + let first = ingest_source(&db, &source, &project, None).await; + assert_eq!(first.messages_upserted, 2); + // Re-ingesting the unchanged file is a no-op. + let second = ingest_source(&db, &source, &project, None).await; + assert_eq!(second.messages_upserted, 0); + + // Appending one line ingests only that line. + let mut f = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + writeln!( + f, + "{}", + serde_json::json!({ + "type": "user", + "cwd": project.to_string_lossy(), + "sessionId": "claude-sess", + "uuid": "u3", + "timestamp": "2026-01-01T00:01:00.000Z", + "message": {"role": "user", "content": "Add a regression test for billing"} + }) + ) + .unwrap(); + drop(f); + + let third = ingest_source(&db, &source, &project, None).await; + assert_eq!(third.messages_upserted, 1); +} + +#[tokio::test] +async fn claude_transcript_for_other_project_is_skipped() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + let other = tmp.path().join("other-project"); + std::fs::create_dir_all(&other).unwrap(); + // Transcript records a cwd that is NOT the project we ingest for. + write_claude_transcript(&home, &other, "claude-other"); + + let db = open_project_session_db(&project).await.unwrap(); + let source = ClaudeSource::with_home(&home); + + let stats = ingest_source(&db, &source, &project, None).await; + assert_eq!( + stats.messages_upserted, 0, + "a transcript whose cwd is a different project must be skipped" + ); +} diff --git a/tests/cline_like_transcript_ingest_test.rs b/tests/cline_like_transcript_ingest_test.rs new file mode 100644 index 0000000000..cafbdbd415 --- /dev/null +++ b/tests/cline_like_transcript_ingest_test.rs @@ -0,0 +1,175 @@ +use tempfile::TempDir; +use tokensave::sessions::cline_like::ClineLikeSource; +use tokensave::sessions::cursor::open_project_session_db; +use tokensave::sessions::source::ingest_source; + +fn setup(tmp: &TempDir) -> (std::path::PathBuf, std::path::PathBuf) { + let home = tmp.path().join("home"); + let project = tmp.path().join("project"); + std::fs::create_dir_all(&project).unwrap(); + std::fs::create_dir(project.join(".tokensave")).unwrap(); + std::fs::write(project.join(".tokensave/tokensave.db"), "").unwrap(); + (home, project) +} + +fn vscode_storage_root(home: &std::path::Path, extension_id: &str) -> std::path::PathBuf { + tokensave::agents::vscode_data_dir(home) + .join("User/globalStorage") + .join(extension_id) + .join("tasks") +} + +fn write_task( + root: &std::path::Path, + project: &std::path::Path, + task_id: &str, +) -> std::path::PathBuf { + let dir = root.join(task_id); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write( + dir.join("task_metadata.json"), + serde_json::to_string_pretty(&serde_json::json!({ + "task": "Investigate the billing pipeline regression", + "workspacePath": project + })) + .unwrap(), + ) + .unwrap(); + let api = dir.join("api_conversation_history.json"); + std::fs::write( + &api, + serde_json::to_string_pretty(&serde_json::json!([ + { + "role": "user", + "content": "Investigate the billing pipeline regression", + "ts": 1_800_000_000_i64 + }, + { + "role": "assistant", + "model": "claude-sonnet-4.6", + "content": [ + {"type": "text", "text": "The billing pipeline regression is fixed."}, + {"type": "tool_use", "name": "read_file"} + ], + "ts": 1_800_000_010_i64 + } + ])) + .unwrap(), + ) + .unwrap(); + api +} + +async fn assert_provider_ingests( + provider: &str, + source: ClineLikeSource, + db: &tokensave::global_db::GlobalDb, + project: &std::path::Path, +) { + let stats = ingest_source(db, &source, project, None).await; + assert_eq!(stats.messages_upserted, 2); + + let results = db + .search_session_messages( + provider, + Some(project.to_string_lossy().as_ref()), + "billing pipeline", + 10, + ) + .await; + assert_eq!(results.len(), 2); + assert!(results + .iter() + .any(|hit| hit.message.tool_names.as_deref() == Some("read_file"))); + + // ContentHash: unchanged full-rewrite file is a no-op. + assert_eq!( + ingest_source(db, &source, project, None) + .await + .messages_upserted, + 0 + ); +} + +#[tokio::test] +async fn cline_task_history_populates_searchable_messages() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + write_task( + &vscode_storage_root(&home, "saoudrizwan.claude-dev"), + &project, + "cline-task", + ); + + let db = open_project_session_db(&project).await.unwrap(); + assert_provider_ingests( + "cline", + ClineLikeSource::cline_with_home(&home), + &db, + &project, + ) + .await; +} + +#[tokio::test] +async fn roo_code_task_history_populates_searchable_messages() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + write_task( + &vscode_storage_root(&home, "rooveterinaryinc.roo-cline"), + &project, + "roo-task", + ); + + let db = open_project_session_db(&project).await.unwrap(); + assert_provider_ingests( + "roo-code", + ClineLikeSource::roo_code_with_home(&home), + &db, + &project, + ) + .await; +} + +#[tokio::test] +async fn kilo_task_history_populates_searchable_messages() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + write_task( + &vscode_storage_root(&home, "kilocode.kilo-code"), + &project, + "kilo-task", + ); + + let db = open_project_session_db(&project).await.unwrap(); + assert_provider_ingests( + "kilo", + ClineLikeSource::kilo_with_home(&home), + &db, + &project, + ) + .await; +} + +#[tokio::test] +async fn cline_like_task_for_other_project_is_skipped() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + let other = tmp.path().join("other-project"); + std::fs::create_dir_all(&other).unwrap(); + write_task( + &vscode_storage_root(&home, "saoudrizwan.claude-dev"), + &other, + "other-task", + ); + + let db = open_project_session_db(&project).await.unwrap(); + let stats = ingest_source( + &db, + &ClineLikeSource::cline_with_home(&home), + &project, + None, + ) + .await; + assert_eq!(stats.messages_upserted, 0); +} diff --git a/tests/codex_transcript_ingest_test.rs b/tests/codex_transcript_ingest_test.rs new file mode 100644 index 0000000000..a062935cb0 --- /dev/null +++ b/tests/codex_transcript_ingest_test.rs @@ -0,0 +1,128 @@ +use std::io::Write; + +use tempfile::TempDir; +use tokensave::sessions::codex::CodexSource; +use tokensave::sessions::cursor::open_project_session_db; +use tokensave::sessions::source::ingest_source; + +fn setup(tmp: &TempDir) -> (std::path::PathBuf, std::path::PathBuf) { + let home = tmp.path().join("home"); + let project = tmp.path().join("project"); + std::fs::create_dir_all(&project).unwrap(); + std::fs::create_dir(project.join(".tokensave")).unwrap(); + std::fs::write(project.join(".tokensave/tokensave.db"), "").unwrap(); + (home, project) +} + +/// Writes a Codex rollout JSONL whose `session_meta.cwd` is `project`. Includes a +/// `response_item` line that must be ignored (it duplicates the agent_message). +fn write_codex_rollout( + home: &std::path::Path, + project: &std::path::Path, + session: &str, +) -> std::path::PathBuf { + let dir = home.join(".codex/sessions/2026/01/01"); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join(format!("rollout-2026-01-01T00-00-00-{session}.jsonl")); + let contents = format!( + "{}\n{}\n{}\n{}\n", + serde_json::json!({ + "timestamp": "2026-01-01T00:00:00.000Z", + "type": "session_meta", + "payload": {"id": session, "cwd": project.to_string_lossy(), "model": "gpt-5.5"} + }), + serde_json::json!({ + "timestamp": "2026-01-01T00:00:01.000Z", + "type": "event_msg", + "payload": {"type": "user_message", "message": "Investigate the billing pipeline regression"} + }), + serde_json::json!({ + "timestamp": "2026-01-01T00:00:02.000Z", + "type": "event_msg", + "payload": {"type": "agent_message", "message": "The billing pipeline regression is fixed."} + }), + serde_json::json!({ + "timestamp": "2026-01-01T00:00:02.500Z", + "type": "response_item", + "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "duplicate"}]} + }), + ); + std::fs::write(&path, contents).unwrap(); + path +} + +#[tokio::test] +async fn codex_rollout_populates_user_and_agent_messages_only() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + write_codex_rollout(&home, &project, "codex-sess"); + + let db = open_project_session_db(&project).await.unwrap(); + let source = CodexSource::with_home(&home); + + let stats = ingest_source(&db, &source, &project, None).await; + // user_message + agent_message; the response_item duplicate is skipped. + assert_eq!(stats.messages_upserted, 2); + assert_eq!(stats.sessions_upserted, 1); + + let results = db + .search_session_messages( + "codex", + Some(project.to_string_lossy().as_ref()), + "billing pipeline", + 10, + ) + .await; + assert_eq!(results.len(), 2); + assert!(results.iter().any(|hit| hit.message.role == "user")); + assert!(results.iter().any(|hit| hit.message.role == "assistant")); + assert!(results + .iter() + .all(|hit| hit.message.model.as_deref() == Some("gpt-5.5"))); +} + +#[tokio::test] +async fn codex_rollout_ingest_is_incremental() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + let path = write_codex_rollout(&home, &project, "codex-sess"); + + let db = open_project_session_db(&project).await.unwrap(); + let source = CodexSource::with_home(&home); + + assert_eq!( + ingest_source(&db, &source, &project, None) + .await + .messages_upserted, + 2 + ); + assert_eq!( + ingest_source(&db, &source, &project, None) + .await + .messages_upserted, + 0 + ); + + let mut f = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + writeln!( + f, + "{}", + serde_json::json!({ + "timestamp": "2026-01-01T00:01:00.000Z", + "type": "event_msg", + "payload": {"type": "agent_message", "message": "Added a regression test."} + }) + ) + .unwrap(); + drop(f); + + assert_eq!( + ingest_source(&db, &source, &project, None) + .await + .messages_upserted, + 1 + ); +} diff --git a/tests/cursor_transcript_ingest_test.rs b/tests/cursor_transcript_ingest_test.rs new file mode 100644 index 0000000000..a1e6499d04 --- /dev/null +++ b/tests/cursor_transcript_ingest_test.rs @@ -0,0 +1,211 @@ +use std::io::Write; + +use tempfile::TempDir; +use tokensave::sessions::cursor::{ + ingest_cursor_transcript_event, ingest_cursor_transcript_event_capped, open_project_session_db, + project_session_db_path, +}; + +fn init_project(tmp: &TempDir) -> std::path::PathBuf { + let project = tmp.path().join("project"); + std::fs::create_dir_all(&project).unwrap(); + std::fs::create_dir(project.join(".tokensave")).unwrap(); + std::fs::write(project.join(".tokensave/tokensave.db"), "").unwrap(); + project +} + +#[tokio::test] +async fn cursor_transcript_ingest_populates_searchable_messages() { + let tmp = TempDir::new().unwrap(); + let project = tmp.path().join("project"); + std::fs::create_dir_all(&project).unwrap(); + std::fs::create_dir(project.join(".tokensave")).unwrap(); + std::fs::write(project.join(".tokensave/tokensave.db"), "").unwrap(); + + let transcript = tmp.path().join("cursor-session.jsonl"); + std::fs::write( + &transcript, + r#"{"role":"user","message":{"content":[{"type":"text","text":"Please check billing ingestion from Cursor transcripts."}]}} +{"role":"assistant","message":{"content":[{"type":"text","text":"The billing ingestion plan is ready."},{"type":"tool_use","name":"tokensave_context","input":{"task":"billing ingestion"}}]}} +"#, + ) + .unwrap(); + + let db = open_project_session_db(&project).await.unwrap(); + let event = serde_json::json!({ + "session_id": "cursor-session", + "conversation_id": "conversation-1", + "transcript_path": transcript, + "cwd": project, + "model": "gpt-5.5" + }); + + let stats = ingest_cursor_transcript_event(&event.to_string(), &db).await; + assert_eq!(stats.sessions_upserted, 1); + assert_eq!(stats.messages_upserted, 2); + assert!(project_session_db_path(&project).exists()); + + let results = db + .search_session_messages( + "cursor", + Some(project.to_string_lossy().as_ref()), + "billing ingestion", + 10, + ) + .await; + assert_eq!(results.len(), 2); + assert_eq!(results[0].session.session_id, "cursor-session"); + assert_eq!( + results[0].session.transcript_path.as_deref(), + transcript.to_str() + ); + assert!(results + .iter() + .any(|hit| hit.message.tool_names.as_deref() == Some("tokensave_context"))); +} + +#[tokio::test] +async fn cursor_transcript_ingest_is_idempotent() { + let tmp = TempDir::new().unwrap(); + let project = init_project(&tmp); + + let transcript = tmp.path().join("cursor-session.jsonl"); + std::fs::write( + &transcript, + r#"{"role":"user","message":{"content":[{"type":"text","text":"Remember the Cursor transcript parser decision."}]}} +"#, + ) + .unwrap(); + + let db = open_project_session_db(&project).await.unwrap(); + let event = serde_json::json!({ + "session_id": "cursor-session", + "transcript_path": transcript, + "workspace_roots": [project] + }); + + // Ingestion is now incremental: the first call ingests the message and + // records a parse offset, so a second call over the *unchanged* file is a + // no-op rather than re-upserting the same row. + let first = ingest_cursor_transcript_event(&event.to_string(), &db).await; + let second = ingest_cursor_transcript_event(&event.to_string(), &db).await; + assert_eq!(first.messages_upserted, 1); + assert_eq!(second.messages_upserted, 0); + + let results = db + .search_session_messages("cursor", None, "parser decision", 10) + .await; + assert_eq!(results.len(), 1); +} + +#[tokio::test] +async fn cursor_transcript_ingest_reads_only_appended_lines() { + let tmp = TempDir::new().unwrap(); + let project = init_project(&tmp); + + let transcript = tmp.path().join("cursor-session.jsonl"); + std::fs::write( + &transcript, + r#"{"role":"user","message":{"content":[{"type":"text","text":"First message about incremental ingestion."}]}} +"#, + ) + .unwrap(); + + let db = open_project_session_db(&project).await.unwrap(); + let event = serde_json::json!({ + "session_id": "cursor-session", + "transcript_path": transcript, + "workspace_roots": [project] + }); + + let first = ingest_cursor_transcript_event(&event.to_string(), &db).await; + assert_eq!(first.messages_upserted, 1); + + // Append a new line; only the appended line should be parsed/upserted. + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&transcript) + .unwrap(); + file.write_all( + b"{\"role\":\"assistant\",\"message\":{\"content\":[{\"type\":\"text\",\"text\":\"Second message about incremental ingestion.\"}]}}\n", + ) + .unwrap(); + drop(file); + + let second = ingest_cursor_transcript_event(&event.to_string(), &db).await; + assert_eq!(second.messages_upserted, 1); + + let results = db + .search_session_messages("cursor", None, "incremental ingestion", 10) + .await; + assert_eq!(results.len(), 2); +} + +#[tokio::test] +async fn cursor_transcript_ingest_cap_defers_large_backlog() { + let tmp = TempDir::new().unwrap(); + let project = init_project(&tmp); + + let transcript = tmp.path().join("cursor-session.jsonl"); + let large_text = "x".repeat(2048); + std::fs::write( + &transcript, + format!( + "{{\"role\":\"user\",\"message\":{{\"content\":[{{\"type\":\"text\",\"text\":\"{large_text}\"}}]}}}}\n" + ), + ) + .unwrap(); + + let db = open_project_session_db(&project).await.unwrap(); + let event = serde_json::json!({ + "session_id": "cursor-session", + "transcript_path": transcript, + "workspace_roots": [project] + }); + + let capped = ingest_cursor_transcript_event_capped(&event.to_string(), &db, Some(128)).await; + assert_eq!(capped.messages_upserted, 0); + + let uncapped = ingest_cursor_transcript_event(&event.to_string(), &db).await; + assert_eq!(uncapped.messages_upserted, 1); +} + +#[tokio::test] +async fn cursor_transcript_ingest_defers_partial_final_line() { + let tmp = TempDir::new().unwrap(); + let project = init_project(&tmp); + + let transcript = tmp.path().join("cursor-session.jsonl"); + // A complete first line followed by a partial (un-terminated) second line, + // as can happen mid-flush while Cursor is still writing the transcript. + let complete = "{\"role\":\"user\",\"message\":{\"content\":[{\"type\":\"text\",\"text\":\"Complete line about partial handling.\"}]}}\n"; + let partial = "{\"role\":\"assistant\",\"message\":{\"content\":[{\"type\":\"text\",\"text\":\"Partial line about partial handling.\"}]}}"; + std::fs::write(&transcript, format!("{complete}{partial}")).unwrap(); + + let db = open_project_session_db(&project).await.unwrap(); + let event = serde_json::json!({ + "session_id": "cursor-session", + "transcript_path": transcript, + "workspace_roots": [project] + }); + + // The partial final line is left unconsumed. + let first = ingest_cursor_transcript_event(&event.to_string(), &db).await; + assert_eq!(first.messages_upserted, 1); + + // Once the trailing newline arrives, the previously-partial line is ingested. + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&transcript) + .unwrap(); + file.write_all(b"\n").unwrap(); + drop(file); + + let second = ingest_cursor_transcript_event(&event.to_string(), &db).await; + assert_eq!(second.messages_upserted, 1); + + let results = db + .search_session_messages("cursor", None, "partial handling", 10) + .await; + assert_eq!(results.len(), 2); +} diff --git a/tests/db_test.rs b/tests/db_test.rs index 29a911c469..873774d185 100644 --- a/tests/db_test.rs +++ b/tests/db_test.rs @@ -455,14 +455,14 @@ async fn test_migrate_v7_adds_and_backfills_attrs_start_line() { .expect("migrate failed"); assert!(migrated, "expected v7 migration to run"); - // user_version is now LATEST (= 10). + // user_version is now LATEST (= 12). let mut rows = conn .query("PRAGMA user_version", ()) .await .expect("read version"); let row = rows.next().await.expect("row").expect("some row"); let version: i64 = row.get(0).expect("version"); - assert_eq!(version, 10); + assert_eq!(version, 12); // attrs_start_line is backfilled from start_line for both rows. // Row a: start_line=42 -> attrs_start_line=42. diff --git a/tests/mcp_handler_test.rs b/tests/mcp_handler_test.rs index d061176803..d95a4987e4 100644 --- a/tests/mcp_handler_test.rs +++ b/tests/mcp_handler_test.rs @@ -6,7 +6,9 @@ use serde_json::{json, Value}; use std::fs; use tempfile::TempDir; -use tokensave::mcp::handle_tool_call; +use tokensave::mcp::{get_tool_definitions, handle_tool_call}; +use tokensave::sessions::cursor::open_project_session_db; +use tokensave::sessions::{SessionMessageRecord, SessionRecord}; use tokensave::tokensave::TokenSave; // --------------------------------------------------------------------------- @@ -75,6 +77,13 @@ fn extract_text(value: &Value) -> &str { .unwrap_or("") } +fn expect_tool_error(result: tokensave::errors::Result) -> String { + match result { + Ok(_) => panic!("expected tool call to fail"), + Err(err) => format!("{err}"), + } +} + /// Searches for `name` via the search handler and returns the first matching /// node id whose name field equals `name`. async fn find_node_id(cg: &TokenSave, name: &str) -> String { @@ -3108,94 +3117,536 @@ async fn test_by_qualified_name_requires_param() { assert!(format!("{err}").contains("qualified_name")); } -// --------------------------------------------------------------------------- -// Memory handler tests (record_decision, record_code_area, session_recall) -// --------------------------------------------------------------------------- - #[tokio::test] -async fn test_handle_record_decision() { +async fn memory_fact_store_add_search_update_remove_and_wrappers() { let (cg, _dir) = setup_project().await; - let result = handle_tool_call( + + let added = handle_tool_call( &cg, - "tokensave_record_decision", - json!({"text": "use JWT", "reason": "legal flagged sessions"}), + "tokensave_fact_store", + json!({ + "action": "add", + "content": "Project Phoenix uses Amari Memory in src/memory/types.rs", + "category": "project", + "entity": "Project Phoenix", + "entities": ["Amari Memory"], + "tags": ["memory", "holographic"], + "source": "mcp-test", + "metadata": {"plan": "holographic"} + }), None, None, ) .await .unwrap(); - let text = extract_text(&result.value); - let output: Value = serde_json::from_str(text).unwrap(); + let added: Value = serde_json::from_str(extract_text(&added.value)).unwrap(); + let fact_id = added["fact"]["fact_id"] + .as_i64() + .expect("fact_store add should return numeric id"); + assert!(added["fact"].get("id").is_none()); + assert!(added["fact"].get("trust").is_none()); + assert!(added["fact"]["trust_score"].as_f64().is_some()); + assert_eq!(added["action"], "add"); + assert_eq!(added["fact"]["category"], "project"); + assert_eq!(added["fact"]["source"], "mcp-test"); + + let search = handle_tool_call( + &cg, + "tokensave_fact_store", + json!({ + "action": "search", + "query": "Amari Memory", + "category": "project", + "min_trust": 0.1, + "limit": 5 + }), + None, + None, + ) + .await + .unwrap(); + let search: Value = serde_json::from_str(extract_text(&search.value)).unwrap(); + assert_eq!(search["action"], "search"); + assert_eq!(search["count"].as_u64(), Some(1)); + assert_eq!(search["results"], search["facts"]); assert!( - output.get("id").is_some(), - "response should contain 'id', got: {output}" - ); + search["facts"] + .as_array() + .unwrap() + .iter() + .any(|hit| hit["fact"]["fact_id"].as_i64() == Some(fact_id)), + "search results should include added fact: {search}" + ); + + for (action, payload) in [ + ("probe", json!({"entity": "Project Phoenix"})), + ("related", json!({"entity": "Amari Memory"})), + ( + "reason", + json!({"entities": ["Project Phoenix", "Amari Memory"]}), + ), + ( + "contradict", + json!({"category": "project", "threshold": 0.8}), + ), + ("list", json!({"category": "project", "min_trust": 0.1})), + ] { + let mut args = payload; + args["action"] = json!(action); + let result = handle_tool_call(&cg, "tokensave_fact_store", args, None, None) + .await + .unwrap(); + let output: Value = serde_json::from_str(extract_text(&result.value)).unwrap(); + assert_eq!(output["action"], action, "{action} should echo action"); + assert!( + output["results"].is_array(), + "{action} should include results array: {output}" + ); + assert!( + output["count"].is_number(), + "{action} should include count: {output}" + ); + if action == "related" { + assert!( + output["count"].as_u64().unwrap_or_default() > 0, + "related should return facts connected through adjacent entities: {output}" + ); + } + } + + let updated = handle_tool_call( + &cg, + "tokensave_fact_store", + json!({ + "action": "update", + "fact_id": fact_id, + "content": "Project Phoenix uses deterministic Amari Memory", + "entities": ["Project Phoenix", "Amari Memory"], + "metadata": {"updated": true} + }), + None, + None, + ) + .await + .unwrap(); + let updated: Value = serde_json::from_str(extract_text(&updated.value)).unwrap(); assert_eq!( - output["status"].as_str().unwrap(), - "recorded", - "status should be 'recorded', got: {output}" + updated["fact"]["content"], + "Project Phoenix uses deterministic Amari Memory" ); + assert_eq!(updated["count"].as_u64(), Some(1)); + + let removed = handle_tool_call( + &cg, + "tokensave_fact_store", + json!({"action": "remove", "fact_id": fact_id.to_string()}), + None, + None, + ) + .await + .unwrap(); + let removed: Value = serde_json::from_str(extract_text(&removed.value)).unwrap(); + assert_eq!(removed["removed"], true); } #[tokio::test] -async fn test_handle_record_code_area() { +async fn memory_recall_updates_retrieval_count() { let (cg, _dir) = setup_project().await; - let result = handle_tool_call( + let added = handle_tool_call( &cg, - "tokensave_record_code_area", - json!({"path": "src/auth.rs", "description": "OAuth provider"}), + "tokensave_fact_store", + json!({ + "action": "add", + "content": "Retrieval counters move after search", + "entity": "Counter Entity" + }), None, None, ) .await .unwrap(); - let text = extract_text(&result.value); - let output: Value = serde_json::from_str(text).unwrap(); - assert_eq!( - output["status"].as_str().unwrap(), - "recorded", - "status should be 'recorded', got: {output}" + let added: Value = serde_json::from_str(extract_text(&added.value)).unwrap(); + let fact_id = added["fact"]["fact_id"].as_i64().unwrap(); + + handle_tool_call( + &cg, + "tokensave_fact_store", + json!({"action": "search", "query": "Retrieval counters", "limit": 5}), + None, + None, + ) + .await + .unwrap(); + + let status = handle_tool_call( + &cg, + "tokensave_fact_store", + json!({"action": "list", "min_trust": 0.0, "limit": 10}), + None, + None, + ) + .await + .unwrap(); + let status: Value = serde_json::from_str(extract_text(&status.value)).unwrap(); + let fact = status["results"] + .as_array() + .unwrap() + .iter() + .find(|fact| fact["fact_id"].as_i64() == Some(fact_id)) + .unwrap(); + assert!( + fact["retrieval_count"].as_i64().unwrap_or_default() > 0, + "returned facts should increment retrieval_count: {status}" ); } #[tokio::test] -async fn test_handle_session_recall_returns_recorded_decision() { +async fn memory_fact_store_update_trust_delta_uses_direct_fact_lookup() { let (cg, _dir) = setup_project().await; - // Seed a decision first - handle_tool_call( + let first = handle_tool_call( + &cg, + "tokensave_fact_store", + json!({ + "action": "add", + "content": "First fact should remain updateable after many later facts", + "trust": 0.4 + }), + None, + None, + ) + .await + .unwrap(); + let first: Value = serde_json::from_str(extract_text(&first.value)).unwrap(); + let first_id = first["fact"]["fact_id"].as_i64().unwrap(); + + for i in 0..205 { + handle_tool_call( + &cg, + "tokensave_fact_store", + json!({ + "action": "add", + "content": format!("Later fact {i} should not hide the first fact"), + }), + None, + None, + ) + .await + .unwrap(); + } + + let updated = handle_tool_call( &cg, - "tokensave_record_decision", - json!({"text": "use JWT", "reason": "legal flagged sessions"}), + "tokensave_fact_store", + json!({ + "action": "update", + "fact_id": first_id, + "trust_delta": 0.2 + }), None, None, ) .await .unwrap(); - // Recall and verify the seeded decision appears + let updated: Value = serde_json::from_str(extract_text(&updated.value)).unwrap(); + assert_eq!(updated["fact"]["fact_id"].as_i64(), Some(first_id)); + assert!( + (updated["fact"]["trust_score"].as_f64().unwrap() - 0.6).abs() < 0.000_001, + "trust_delta should apply through direct fact lookup: {updated}" + ); +} + +#[tokio::test] +async fn memory_feedback_and_status_include_trust_fields() { + let (cg, _dir) = setup_project().await; + let added = handle_tool_call( + &cg, + "tokensave_fact_store", + json!({ + "action": "add", + "content": "Helpful memory fact for feedback", + "category": "general" + }), + None, + None, + ) + .await + .unwrap(); + let added: Value = serde_json::from_str(extract_text(&added.value)).unwrap(); + let fact_id = added["fact"]["fact_id"].as_i64().unwrap(); + assert!(added["fact"].get("id").is_none()); + assert!(added["fact"].get("trust").is_none()); + assert!(added["fact"]["trust_score"].as_f64().is_some()); + + let helpful = handle_tool_call( + &cg, + "tokensave_fact_feedback", + json!({"fact_id": fact_id, "helpful": true, "source": "mcp-test", "note": "matched"}), + None, + None, + ) + .await + .unwrap(); + let helpful: Value = serde_json::from_str(extract_text(&helpful.value)).unwrap(); + assert!(helpful["feedback"]["event_id"].as_i64().unwrap() > 0); + assert_eq!(helpful["feedback"]["fact_id"], fact_id); + assert_eq!(helpful["feedback"]["action"], "helpful"); + assert_eq!(helpful["feedback"]["old_trust"], 0.5); + assert!(helpful["feedback"]["new_trust"].as_f64().unwrap() > 0.5); + assert!(helpful["feedback"]["trust_delta"].as_f64().unwrap() > 0.0); + assert_eq!(helpful["feedback"]["helpful_count"], 1); + assert_eq!(helpful["feedback"]["unhelpful_count"], 0); + + let unhelpful = handle_tool_call( + &cg, + "tokensave_fact_feedback", + json!({"fact_id": fact_id, "unhelpful": true}), + None, + None, + ) + .await + .unwrap(); + let unhelpful: Value = serde_json::from_str(extract_text(&unhelpful.value)).unwrap(); + assert_eq!(unhelpful["feedback"]["action"], "unhelpful"); + assert!( + unhelpful["feedback"]["new_trust"].as_f64().unwrap() + < helpful["feedback"]["new_trust"].as_f64().unwrap() + ); + assert_eq!(unhelpful["feedback"]["helpful_count"], 1); + assert_eq!(unhelpful["feedback"]["unhelpful_count"], 1); + + let status = handle_tool_call(&cg, "tokensave_memory_status", json!({}), None, None) + .await + .unwrap(); + let status: Value = serde_json::from_str(extract_text(&status.value)).unwrap(); + assert_eq!(status["status"], "ok"); + assert!(status["memory"]["fact_count"].as_u64().unwrap() >= 1); + assert!(status["memory"].get("trust_0_025_count").is_some()); + assert!(status["memory"].get("trust_025_050_count").is_some()); + assert!(status["memory"].get("trust_050_075_count").is_some()); + assert!(status["memory"].get("trust_075_100_count").is_some()); + assert!(status["memory"].get("helpful_count").is_some()); + assert!(status["memory"].get("unhelpful_count").is_some()); + assert!(status["memory"].get("missing_vector_count").is_some()); +} + +#[tokio::test] +async fn memory_tools_validate_malformed_inputs() { + let (cg, _dir) = setup_project().await; + + let missing_action = handle_tool_call(&cg, "tokensave_fact_store", json!({}), None, None).await; + assert!(expect_tool_error(missing_action).contains("action")); + + let bad_action = handle_tool_call( + &cg, + "tokensave_fact_store", + json!({"action": "teleport"}), + None, + None, + ) + .await; + assert!(expect_tool_error(bad_action).contains("unknown fact_store action")); + + let bad_category = handle_tool_call( + &cg, + "tokensave_fact_store", + json!({"action": "list", "category": "definitely-not-a-category"}), + None, + None, + ) + .await; + assert!(expect_tool_error(bad_category).contains("category")); + + let missing_feedback_action = handle_tool_call( + &cg, + "tokensave_fact_feedback", + json!({"fact_id": 123}), + None, + None, + ) + .await; + assert!(expect_tool_error(missing_feedback_action).contains("helpful")); +} + +#[tokio::test] +async fn message_search_reads_project_local_session_db() { + let (cg, _dir) = setup_project().await; + let db = open_project_session_db(cg.project_root()) + .await + .expect("project-local session db should open"); + let session = SessionRecord { + provider: "cursor".to_string(), + session_id: "cursor-session".to_string(), + project_key: cg.project_root().to_string_lossy().to_string(), + project_path: cg.project_root().to_string_lossy().to_string(), + title: Some("Cursor transcript".to_string()), + started_at: Some(1), + ended_at: None, + transcript_path: Some("cursor-session.jsonl".to_string()), + metadata_json: None, + }; + assert!(db.upsert_session(&session).await); + assert!( + db.upsert_session_message(&SessionMessageRecord { + provider: "cursor".to_string(), + message_id: "cursor-message".to_string(), + session_id: "cursor-session".to_string(), + role: "user".to_string(), + timestamp: Some(2), + ordinal: 1, + text: "Project-local transcript search is working.".to_string(), + kind: Some("message".to_string()), + model: Some("test-model".to_string()), + tool_names: None, + source_path: Some("cursor-session.jsonl".to_string()), + source_offset: Some(0), + metadata_json: None, + }) + .await + ); + let result = handle_tool_call( &cg, - "tokensave_session_recall", - json!({"query": "JWT"}), + "tokensave_message_search", + json!({"query": "transcript search", "provider": "cursor", "limit": 5}), None, None, ) .await .unwrap(); - let text = extract_text(&result.value); - let output: Value = serde_json::from_str(text).unwrap(); - let decisions = output["decisions"] - .as_array() - .expect("decisions should be an array"); + let parsed: Value = serde_json::from_str(extract_text(&result.value)).unwrap(); + assert_eq!(parsed["status"], "ok"); + assert_eq!(parsed["count"], 1); + assert_eq!( + parsed["results"][0]["message"]["message_id"], + "cursor-message" + ); + assert_eq!( + parsed["results"][0]["session"]["project_key"], + cg.project_root().to_string_lossy().to_string() + ); +} + +#[test] +fn memory_tool_definitions_include_hermes_payload_fields() { + let tools = get_tool_definitions(); + let tool_names: std::collections::HashSet<_> = + tools.iter().map(|tool| tool.name.as_str()).collect(); + let fact_store = tools + .iter() + .find(|tool| tool.name == "tokensave_fact_store") + .expect("tokensave_fact_store definition"); + let feedback = tools + .iter() + .find(|tool| tool.name == "tokensave_fact_feedback") + .expect("tokensave_fact_feedback definition"); + let status = tools + .iter() + .find(|tool| tool.name == "tokensave_memory_status") + .expect("tokensave_memory_status definition"); + + assert_eq!( + fact_store.annotations.as_ref().unwrap()["readOnlyHint"], + false + ); + assert_eq!( + feedback.annotations.as_ref().unwrap()["readOnlyHint"], + false + ); + assert_eq!(status.annotations.as_ref().unwrap()["readOnlyHint"], false); + + for field in [ + "action", + "content", + "query", + "entity", + "entities", + "fact_id", + "category", + "tags", + "min_trust", + "trust", + "trust_delta", + "threshold", + "limit", + "source", + "metadata", + "note", + ] { + assert!( + fact_store.input_schema["properties"].get(field).is_some(), + "fact_store schema missing Hermes field {field}" + ); + } + assert_eq!( + feedback.input_schema["required"], + serde_json::json!(["fact_id"]) + ); + assert_eq!( + fact_store.input_schema["properties"]["trust"]["type"], + "number" + ); + assert_eq!(fact_store.input_schema["properties"]["trust"]["minimum"], 0); + assert_eq!(fact_store.input_schema["properties"]["trust"]["maximum"], 1); + assert!( - !decisions.is_empty(), - "recall should return at least one decision after seeding" + !tool_names.contains("tokensave_record_decision"), + "unshipped legacy decision tool should not be exposed" ); - let found = decisions + assert!( + !tool_names.contains("tokensave_record_code_area"), + "unshipped legacy code-area tool should not be exposed" + ); + assert!( + !tool_names.contains("tokensave_session_recall"), + "unshipped legacy recall tool should not be exposed" + ); +} + +#[test] +fn message_search_provider_schema_matches_ingested_providers() { + let tools = get_tool_definitions(); + let message_search = tools .iter() - .any(|d| d["text"].as_str().unwrap_or("").contains("JWT")); + .find(|tool| tool.name == "tokensave_message_search") + .expect("tokensave_message_search definition"); + + assert_eq!( + message_search.input_schema["properties"]["provider"]["enum"], + serde_json::json!(["cursor", "claude", "codex", "vibe", "cline", "roo-code", "kilo"]) + ); +} + +#[tokio::test] +async fn memory_status_repairs_dirty_banks_before_reporting() { + let (cg, _dir) = setup_project().await; + handle_tool_call( + &cg, + "tokensave_fact_store", + json!({ + "action": "add", + "content": "Status should repair dirty holographic banks", + "category": "project", + "entity": "Holographic Banks" + }), + None, + None, + ) + .await + .unwrap(); + + let status = handle_tool_call(&cg, "tokensave_memory_status", json!({}), None, None) + .await + .unwrap(); + let status: Value = serde_json::from_str(extract_text(&status.value)).unwrap(); + assert_eq!(status["status"], "ok"); assert!( - found, - "seeded 'JWT' decision should appear in recall results" + status["memory"]["bank_count"].as_u64().unwrap_or_default() >= 2, + "memory_status should rebuild all and category banks before reporting: {status}" + ); + assert_eq!( + status["memory"]["missing_vector_count"].as_u64(), + Some(0), + "status-triggered bank repair should not leave missing vectors" ); } @@ -4637,9 +5088,15 @@ async fn mcp_server_owns_watcher_and_refreshes_token_map_on_change() { "find_stale_files should detect newly written b.rs" ); server.cg().sync_if_stale_silent(&stale).await.unwrap(); - server.refresh_file_token_map().await; - - let after_count = server.file_token_map_snapshot().len(); + let mut after_count = initial_count; + for _ in 0..10 { + server.refresh_file_token_map().await; + after_count = server.file_token_map_snapshot().len(); + if after_count > initial_count { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } assert!( after_count > initial_count, "lazy sync should have refreshed map ({initial_count} -> {after_count})" diff --git a/tests/mcp_test.rs b/tests/mcp_test.rs index d9bd82727d..f755e3066d 100644 --- a/tests/mcp_test.rs +++ b/tests/mcp_test.rs @@ -97,9 +97,9 @@ fn test_tool_definitions_count() { // the external `ast-grep` binary is on PATH — hide-when-missing so // agents never receive a tool that will instantly fail. let expected = if tokensave::mcp::tools::ast_grep_available() { - 76 + 77 } else { - 75 + 76 }; assert_eq!(tools.len(), expected); } diff --git a/tests/memory_test.rs b/tests/memory_test.rs index 40df9830f2..842a5798b2 100644 --- a/tests/memory_test.rs +++ b/tests/memory_test.rs @@ -1,4 +1,16 @@ use tempfile::TempDir; +use tokensave::db::Database; +use tokensave::memory::encoding::HolographicEncoder; +use tokensave::memory::entities::{extract_entities, normalize_entity}; +use tokensave::memory::retrieval::FactRetriever; +use tokensave::memory::store::MemoryStore; +use tokensave::memory::trust::{ + apply_feedback, clamp_trust, temporal_decay, trust_bucket, trust_distribution, DEFAULT_TRUST, +}; +use tokensave::memory::types::{ + AddFactRequest, FactRecord, FeedbackAction, FeedbackRequest, MemoryCategory, + SearchFactsRequest, UpdateFactRequest, +}; use tokensave::tokensave::TokenSave; async fn make_project() -> (TempDir, TokenSave) { @@ -8,85 +20,865 @@ async fn make_project() -> (TempDir, TokenSave) { (tmp, cg) } +async fn make_memory_store() -> (Database, TempDir) { + let tmp = TempDir::new().unwrap(); + let db_path = tmp.path().join("tokensave.db"); + let (db, _) = Database::initialize(&db_path).await.unwrap(); + (db, tmp) +} + +fn fact_request(content: &str, category: MemoryCategory, trust: f64) -> AddFactRequest { + AddFactRequest { + content: content.to_string(), + category, + source: Some("test".to_string()), + tags: Vec::new(), + entities: Vec::new(), + trust: Some(trust), + metadata: serde_json::json!({}), + } +} + +async fn dirty_bank_names(db: &Database) -> Vec { + let mut rows = db + .conn() + .query( + "SELECT bank_name FROM memory_bank_dirty ORDER BY bank_name", + (), + ) + .await + .unwrap(); + let mut names = Vec::new(); + while let Some(row) = rows.next().await.unwrap() { + names.push(row.get::(0).unwrap()); + } + names +} + +async fn memory_bank_count(db: &Database) -> i64 { + let mut rows = db + .conn() + .query("SELECT COUNT(*) FROM memory_banks", ()) + .await + .unwrap(); + rows.next().await.unwrap().unwrap().get(0).unwrap() +} + +async fn memory_bank_fact_count(db: &Database, bank_name: &str) -> Option { + let mut rows = db + .conn() + .query( + "SELECT fact_count FROM memory_banks WHERE bank_name = ?1", + libsql::params![bank_name], + ) + .await + .unwrap(); + rows.next() + .await + .unwrap() + .map(|row| row.get::(0).unwrap()) +} + +async fn fact_hrr_vector(db: &Database, fact_id: i64) -> Vec { + let mut rows = db + .conn() + .query( + "SELECT hrr_vector FROM memory_facts WHERE fact_id = ?1", + libsql::params![fact_id], + ) + .await + .unwrap(); + let row = rows.next().await.unwrap().unwrap(); + let bytes = row.get::>(0).unwrap(); + HolographicEncoder::deserialize(&bytes).unwrap() +} + +#[test] +fn core_memory_types_use_stable_json_strings() { + assert_eq!(MemoryCategory::UserPref.to_string(), "user_pref"); + assert_eq!( + "code_area".parse::().unwrap(), + MemoryCategory::CodeArea + ); + + let fact = FactRecord { + fact_id: 42, + content: "Prefer Rust-native memory".to_string(), + category: MemoryCategory::Decision, + tags: vec!["memory".to_string()], + entities: vec!["Rust-native memory".to_string()], + trust_score: 0.7, + source: Some("test".to_string()), + retrieval_count: 3, + helpful_count: 1, + unhelpful_count: 0, + created_at: 1, + updated_at: 2, + last_retrieved_at: Some(3), + last_feedback_at: Some(4), + metadata: serde_json::json!({"scope": "core"}), + }; + + let json = serde_json::to_string(&fact).unwrap(); + assert!(json.contains(r#""fact_id":42"#)); + assert!(json.contains(r#""trust_score":0.7"#)); + assert!(!json.contains(r#""id":"#)); + assert!(!json.contains(r#""trust":"#)); + assert!(json.contains(r#""category":"decision""#)); + let round_trip: FactRecord = serde_json::from_str(&json).unwrap(); + assert_eq!(round_trip, fact); +} + +#[test] +fn memory_request_types_round_trip_through_json() { + let add = AddFactRequest { + content: "Use amari-holographic for fact vectors".to_string(), + category: MemoryCategory::Project, + source: Some("plan".to_string()), + tags: vec!["hrr".to_string()], + entities: vec!["amari-holographic".to_string()], + trust: Some(0.8), + metadata: serde_json::json!({"phase": "core"}), + }; + let search = SearchFactsRequest { + query: "fact vectors".to_string(), + category: Some(MemoryCategory::Project), + limit: Some(5), + min_trust: Some(0.4), + include_why: true, + }; + let update = UpdateFactRequest { + fact_id: 7, + content: Some("Use deterministic fact vectors".to_string()), + category: Some(MemoryCategory::Decision), + tags: Some(vec!["reviewed".to_string()]), + entities: Some(vec!["deterministic fact vectors".to_string()]), + trust: Some(0.9), + source: Some("review".to_string()), + metadata: Some(serde_json::json!({"reviewed": true})), + }; + let feedback = FeedbackRequest { + fact_id: 7, + action: FeedbackAction::Helpful, + source: Some("test".to_string()), + note: Some("matched project context".to_string()), + }; + + assert_eq!( + serde_json::from_value::(serde_json::to_value(add.clone()).unwrap()) + .unwrap(), + add + ); + assert_eq!( + serde_json::from_value::(serde_json::to_value(search.clone()).unwrap()) + .unwrap(), + search + ); + assert_eq!( + serde_json::from_value::(serde_json::to_value(update.clone()).unwrap()) + .unwrap(), + update + ); + assert_eq!( + serde_json::from_value::(serde_json::to_value(feedback.clone()).unwrap()) + .unwrap(), + feedback + ); +} + +#[test] +fn trust_feedback_clamps_buckets_and_decays() { + assert_eq!(clamp_trust(-0.2), 0.0); + assert_eq!(clamp_trust(1.2), 1.0); + assert!((apply_feedback(DEFAULT_TRUST, FeedbackAction::Helpful) - 0.55).abs() < f64::EPSILON); + assert!((apply_feedback(DEFAULT_TRUST, FeedbackAction::Unhelpful) - 0.4).abs() < f64::EPSILON); + assert_eq!(trust_bucket(0.2), "low"); + assert_eq!(trust_bucket(0.5), "medium"); + assert_eq!(trust_bucket(0.8), "high"); + assert_eq!(trust_distribution(&[0.2, 0.31, 0.6, 0.8]), (1, 2, 1)); + assert!(temporal_decay(0.9, 30.0) < 0.9); + assert!(temporal_decay(0.1, 30.0) > 0.1); +} + +#[test] +fn entity_extraction_finds_expected_patterns_and_dedupes() { + let entities = extract_entities( + r#"Project Phoenix uses "holographic memory" aka Amari Memory, also known as Fact Lens in src/memory/types.rs via HolographicEncoder::encode_fact and tokensave_search. Project Phoenix keeps RustNative::Memory nearby."#, + ); + + assert_eq!( + entities, + vec![ + "Project Phoenix", + "holographic memory", + "Amari Memory", + "Fact Lens", + "src/memory/types.rs", + "HolographicEncoder::encode_fact", + "tokensave_search", + "RustNative::Memory", + ] + ); +} + +#[test] +fn entity_extraction_handles_alias_paths_tools_and_whitespace_edges() { + assert_eq!( + normalize_entity(" Project\tPhoenix\nCore "), + "Project Phoenix Core" + ); + + let entities = extract_entities( + r#"Implement Project Phoenix AKA Firebird via src\memory\mod.rs and /etc/config. Then use TOKENSAVE-SEARCH with .gitignore. Project Phoenix appears again."#, + ); + + assert!(entities.contains(&"Project Phoenix".to_string())); + assert!(entities.contains(&"Firebird".to_string())); + assert!(entities.contains(&"src\\memory\\mod.rs".to_string())); + assert!(entities.contains(&"/etc/config".to_string())); + assert!(entities.contains(&".gitignore".to_string())); + assert!(entities.contains(&"tokensave_search".to_string())); + assert_eq!( + entities + .iter() + .filter(|entity| entity.eq_ignore_ascii_case("Project Phoenix")) + .count(), + 1 + ); +} + +#[test] +fn holographic_encoding_is_deterministic_and_round_trips() { + let encoder = HolographicEncoder; + assert_eq!(HolographicEncoder::ROLE_CONTENT, "__hrr_role_content__"); + assert_eq!(HolographicEncoder::ROLE_ENTITY, "__hrr_role_entity__"); + assert_eq!( + encoder.encode_text("Prefer Rust-native memory"), + encoder.encode_text("Prefer Rust-native memory") + ); + let first = encoder.encode_fact( + "Prefer Rust-native memory", + &["Project Phoenix".to_string()], + ); + let same = encoder.encode_fact( + "Prefer Rust-native memory", + &["Project Phoenix".to_string()], + ); + let different = encoder.encode_fact("Prefer Python memory", &["Project Phoenix".to_string()]); + let reordered = encoder.encode_fact( + "Prefer Rust-native memory", + &["SQLite".to_string(), "Project Phoenix".to_string()], + ); + let reordered_same = encoder.encode_fact( + "Prefer Rust-native memory", + &["Project Phoenix".to_string(), "SQLite".to_string()], + ); + + assert_eq!(first, same); + assert_eq!(reordered, reordered_same); + assert_eq!(first.len(), HolographicEncoder::DIMENSIONS); + assert!(first.iter().all(|value| (-1.0..=1.0).contains(value))); + assert!(encoder.similarity(&first, &same) > 0.999_999); + assert!(encoder.similarity(&first, &different) < 0.95); + assert_ne!( + encoder.encode_text("Prefer Rust-native memory"), + encoder.encode_fact("Prefer Rust-native memory", &[]) + ); + assert_eq!( + encoder.encode_fact("Prefer Rust-native memory", &["SQLite".to_string()]), + encoder.encode_fact("Prefer Rust-native memory", &["sqlite".to_string()]) + ); + assert_eq!(encoder.similarity(&[], &first), 0.0); + + let bytes = HolographicEncoder::serialize(&first).unwrap(); + let decoded = HolographicEncoder::deserialize(&bytes).unwrap(); + assert_eq!(decoded, first); + assert!(HolographicEncoder::deserialize(b"not bincode").is_err()); +} + #[tokio::test] -async fn record_decision_persists_and_recalls() { - let (_tmp, cg) = make_project().await; +async fn memory_store_marks_and_rebuilds_dirty_banks() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); - let id = cg - .record_decision( - "use JWT for auth", - Some("session tokens flagged by legal"), - &["src/auth.rs".to_string()], - &["security".to_string(), "decision".to_string()], + let fact = store + .add_fact( + fact_request( + "Project facts should dirty project banks", + MemoryCategory::Project, + 0.8, + ), + DEFAULT_TRUST, ) .await .unwrap(); - assert!(id > 0); + assert_eq!(dirty_bank_names(&db).await, vec!["all", "project"]); + + assert_eq!(store.rebuild_dirty_banks().await.unwrap(), 2); + assert!(dirty_bank_names(&db).await.is_empty()); + assert_eq!(memory_bank_fact_count(&db, "all").await, Some(1)); + assert_eq!(memory_bank_fact_count(&db, "project").await, Some(1)); + + store + .update_fact(UpdateFactRequest { + fact_id: fact.fact_id, + content: Some("Decision facts should replace project bank membership".to_string()), + category: Some(MemoryCategory::Decision), + tags: None, + entities: None, + trust: None, + source: None, + metadata: None, + }) + .await + .unwrap(); + assert_eq!( + dirty_bank_names(&db).await, + vec!["all", "decision", "project"] + ); + + assert_eq!(store.rebuild_dirty_banks().await.unwrap(), 3); + assert!(dirty_bank_names(&db).await.is_empty()); + assert_eq!(memory_bank_fact_count(&db, "all").await, Some(1)); + assert_eq!(memory_bank_fact_count(&db, "decision").await, Some(1)); + assert_eq!(memory_bank_fact_count(&db, "project").await, None); + + assert!(store.remove_fact(fact.fact_id).await.unwrap()); + assert_eq!(dirty_bank_names(&db).await, vec!["all", "decision"]); + + assert_eq!(store.rebuild_dirty_banks().await.unwrap(), 2); + assert!(dirty_bank_names(&db).await.is_empty()); + assert_eq!(memory_bank_count(&db).await, 0); +} + +#[tokio::test] +async fn memory_store_add_list_get_and_deduplicates_by_content() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + + let mut request = fact_request( + "Use SQLite-backed holographic memory", + MemoryCategory::Decision, + 0.72, + ); + request.tags = vec!["storage".to_string()]; + request.entities = vec!["SQLite".to_string()]; + + let first = store + .add_fact(request.clone(), DEFAULT_TRUST) + .await + .unwrap(); + let duplicate = store.add_fact(request, DEFAULT_TRUST).await.unwrap(); + + assert_eq!(duplicate.fact_id, first.fact_id); + assert_eq!(first.tags, vec!["storage"]); + assert_eq!(first.entities, vec!["SQLite"]); + + let fetched = store.get_fact(first.fact_id).await.unwrap().unwrap(); + assert_eq!(fetched, first); + + let listed = store + .list_facts(Some(MemoryCategory::Decision), Some(0.7), 10) + .await + .unwrap(); + assert_eq!(listed.len(), 1); + assert_eq!(listed[0].fact_id, first.fact_id); - let hits = cg.session_recall(Some("JWT"), None, 10).await.unwrap(); - assert_eq!(hits.len(), 1); - assert_eq!(hits[0].text, "use JWT for auth"); + assert!(store.remove_fact(first.fact_id).await.unwrap()); + assert!(store.get_fact(first.fact_id).await.unwrap().is_none()); + assert!(!store.remove_fact(first.fact_id).await.unwrap()); +} + +#[tokio::test] +async fn memory_store_refreshes_vector_when_duplicate_add_merges_entities() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + let encoder = HolographicEncoder; + let content = "persist duplicate vector content"; + + let mut first_request = fact_request(content, MemoryCategory::Project, 0.8); + first_request.entities = vec!["FirstEntity".to_string()]; + let first = store.add_fact(first_request, DEFAULT_TRUST).await.unwrap(); assert_eq!( - hits[0].reason.as_deref(), - Some("session tokens flagged by legal") + fact_hrr_vector(&db, first.fact_id).await, + encoder.encode_fact(content, &["FirstEntity".to_string()]) + ); + + let mut duplicate_request = fact_request(content, MemoryCategory::Project, 0.8); + duplicate_request.entities = vec!["SecondEntity".to_string()]; + let duplicate = store + .add_fact(duplicate_request, DEFAULT_TRUST) + .await + .unwrap(); + + assert_eq!(duplicate.fact_id, first.fact_id); + assert!(duplicate.entities.contains(&"FirstEntity".to_string())); + assert!(duplicate.entities.contains(&"SecondEntity".to_string())); + assert_eq!( + fact_hrr_vector(&db, first.fact_id).await, + encoder.encode_fact( + content, + &["FirstEntity".to_string(), "SecondEntity".to_string()] + ) ); - assert_eq!(hits[0].files, vec!["src/auth.rs"]); - assert_eq!(hits[0].tags, vec!["security", "decision"]); } #[tokio::test] -async fn session_recall_orders_newest_first_when_no_query() { +async fn memory_store_links_explicit_and_extracted_entities_and_updates_fields() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + + let mut request = fact_request( + r#"Project Phoenix stores facts in src/memory/store.rs via HolographicEncoder::encode_fact"#, + MemoryCategory::Project, + 0.6, + ); + request.entities = vec!["Manual Entity".to_string(), "Project Phoenix".to_string()]; + + let fact = store.add_fact(request, DEFAULT_TRUST).await.unwrap(); + assert!(fact.entities.contains(&"Manual Entity".to_string())); + assert!(fact.entities.contains(&"Project Phoenix".to_string())); + assert!(fact.entities.contains(&"src/memory/store.rs".to_string())); + assert!(fact + .entities + .contains(&"HolographicEncoder::encode_fact".to_string())); + + let updated = store + .update_fact(UpdateFactRequest { + fact_id: fact.fact_id, + content: Some("Use deterministic HRR banks for Project Phoenix".to_string()), + category: Some(MemoryCategory::Decision), + tags: Some(vec!["updated".to_string()]), + entities: Some(vec!["Project Phoenix".to_string(), "HRR banks".to_string()]), + trust: Some(0.88), + source: Some("review".to_string()), + metadata: Some(serde_json::json!({"reviewed": true})), + }) + .await + .unwrap(); + + assert_eq!(updated.category, MemoryCategory::Decision); + assert_eq!(updated.tags, vec!["updated"]); + assert_eq!(updated.source.as_deref(), Some("review")); + assert!((updated.trust_score - 0.88).abs() < f64::EPSILON); + assert_eq!(updated.metadata, serde_json::json!({"reviewed": true})); + assert!(updated.entities.contains(&"HRR banks".to_string())); +} + +#[tokio::test] +async fn memory_store_persists_vectors_and_rebuilds_missing_vectors_and_banks() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + + let fact = store + .add_fact( + fact_request( + "Persist an HRR vector for each fact", + MemoryCategory::Project, + 0.8, + ), + DEFAULT_TRUST, + ) + .await + .unwrap(); + let fact_without_vector = store + .add_fact( + fact_request( + "Bank rebuild still counts facts while skipping missing vectors", + MemoryCategory::Project, + 0.8, + ), + DEFAULT_TRUST, + ) + .await + .unwrap(); + + let mut rows = db + .conn() + .query( + "SELECT length(hrr_vector) FROM memory_facts WHERE fact_id = ?1", + libsql::params![fact.fact_id], + ) + .await + .unwrap(); + let row = rows.next().await.unwrap().unwrap(); + let vector_len: i64 = row.get(0).unwrap(); + assert!(vector_len > 0); + + db.conn() + .execute( + "UPDATE memory_facts SET hrr_vector = NULL, hrr_dim = 8 WHERE fact_id = ?1", + libsql::params![fact.fact_id], + ) + .await + .unwrap(); + + assert_eq!(store.compute_missing_vectors(10).await.unwrap(), 1); + assert_eq!(store.compute_missing_vectors(10).await.unwrap(), 0); + let mut rows = db + .conn() + .query( + "SELECT hrr_dim FROM memory_facts WHERE fact_id = ?1", + libsql::params![fact.fact_id], + ) + .await + .unwrap(); + let hrr_dim: i64 = rows.next().await.unwrap().unwrap().get(0).unwrap(); + assert_eq!(hrr_dim, HolographicEncoder::DIMENSIONS as i64); + + db.conn() + .execute( + "UPDATE memory_facts SET hrr_vector = NULL WHERE fact_id = ?1", + libsql::params![fact_without_vector.fact_id], + ) + .await + .unwrap(); + assert_eq!( + store + .rebuild_bank("project", Some(MemoryCategory::Project)) + .await + .unwrap(), + 2 + ); + assert!(store.rebuild_all_banks().await.unwrap() >= 1); + store.remove_fact(fact.fact_id).await.unwrap(); + store + .remove_fact(fact_without_vector.fact_id) + .await + .unwrap(); + assert_eq!( + store + .rebuild_bank("project", Some(MemoryCategory::Project)) + .await + .unwrap(), + 0 + ); +} + +#[tokio::test] +async fn memory_store_records_feedback_audit_and_retrieval_counts() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + let fact = store + .add_fact( + fact_request( + "Feedback adjusts trust with an audit trail", + MemoryCategory::General, + 0.5, + ), + DEFAULT_TRUST, + ) + .await + .unwrap(); + let other_fact = store + .add_fact( + fact_request( + "Batch retrieval count updates preserve duplicate IDs", + MemoryCategory::General, + 0.5, + ), + DEFAULT_TRUST, + ) + .await + .unwrap(); + + store + .increment_retrieval_counts(&[fact.fact_id, other_fact.fact_id, fact.fact_id]) + .await + .unwrap(); + let retrieved = store.get_fact(fact.fact_id).await.unwrap().unwrap(); + assert_eq!(retrieved.retrieval_count, 2); + assert!(retrieved.last_retrieved_at.is_some()); + assert_eq!( + retrieved.updated_at, fact.updated_at, + "retrieval is a read event and must not change updated_at ordering" + ); + let other_retrieved = store.get_fact(other_fact.fact_id).await.unwrap().unwrap(); + assert_eq!(other_retrieved.retrieval_count, 1); + assert!(other_retrieved.last_retrieved_at.is_some()); + + let helpful = store + .record_feedback_event(FeedbackRequest { + fact_id: fact.fact_id, + action: FeedbackAction::Helpful, + source: Some("test".to_string()), + note: Some("useful".to_string()), + }) + .await + .unwrap(); + assert!(helpful.event_id > 0); + assert_eq!(helpful.fact_id, fact.fact_id); + assert_eq!(helpful.action, FeedbackAction::Helpful); + assert!((helpful.old_trust - 0.5).abs() < f64::EPSILON); + assert!((helpful.new_trust - 0.55).abs() < f64::EPSILON); + assert!((helpful.trust_delta - 0.05).abs() < f64::EPSILON); + assert_eq!(helpful.helpful_count, 1); + assert_eq!(helpful.unhelpful_count, 0); + + let unhelpful = store + .record_feedback_event(FeedbackRequest { + fact_id: fact.fact_id, + action: FeedbackAction::Unhelpful, + source: None, + note: None, + }) + .await + .unwrap(); + assert!((unhelpful.old_trust - 0.55).abs() < f64::EPSILON); + assert!((unhelpful.new_trust - 0.45).abs() < f64::EPSILON); + assert_eq!(unhelpful.helpful_count, 1); + assert_eq!(unhelpful.unhelpful_count, 1); + + let updated = store.get_fact(fact.fact_id).await.unwrap().unwrap(); + assert_eq!(updated.helpful_count, 1); + assert_eq!(updated.unhelpful_count, 1); + assert!(updated.last_feedback_at.is_some()); +} + +#[tokio::test] +async fn memory_status_reports_exact_bucket_and_feedback_counts() { let (_tmp, cg) = make_project().await; + let trusts = [0.24, 0.25, 0.50, 0.75]; + let mut fact_ids = Vec::new(); + for trust in trusts { + let fact = cg + .add_fact(AddFactRequest { + content: format!("bucket fact {trust}"), + category: MemoryCategory::General, + source: Some("test".to_string()), + tags: Vec::new(), + entities: Vec::new(), + trust: Some(trust), + metadata: serde_json::json!({}), + }) + .await + .unwrap(); + fact_ids.push(fact.fact_id); + } - cg.record_decision("first", None, &[], &[]).await.unwrap(); - // current_timestamp() is second-granularity, so we need a >1s gap to guarantee - // the two decisions have distinct created_at values for a deterministic ordering. - tokio::time::sleep(std::time::Duration::from_millis(1100)).await; - cg.record_decision("second", None, &[], &[]).await.unwrap(); + cg.record_fact_feedback(FeedbackRequest { + fact_id: fact_ids[1], + action: FeedbackAction::Helpful, + source: Some("test".to_string()), + note: None, + }) + .await + .unwrap(); + cg.record_fact_feedback(FeedbackRequest { + fact_id: fact_ids[2], + action: FeedbackAction::Unhelpful, + source: Some("test".to_string()), + note: None, + }) + .await + .unwrap(); - let hits = cg.session_recall(None, None, 10).await.unwrap(); - assert_eq!(hits.len(), 2); - assert_eq!(hits[0].text, "second"); - assert_eq!(hits[1].text, "first"); + let status = cg.memory_status().await.unwrap(); + assert_eq!(status.fact_count, 4); + assert_eq!(status.trust_0_025_count, 1); + assert_eq!(status.trust_025_050_count, 2); + assert_eq!(status.trust_050_075_count, 0); + assert_eq!(status.trust_075_100_count, 1); + assert_eq!(status.below_default_recall_threshold_count, 1); + assert_eq!(status.helpful_count, 1); + assert_eq!(status.unhelpful_count, 1); + assert_eq!(status.missing_vector_count, 0); } #[tokio::test] -async fn record_code_area_upserts_touch_count() { +async fn memory_status_handles_empty_fact_store() { let (_tmp, cg) = make_project().await; + let status = cg.memory_status().await.unwrap(); + assert_eq!(status.fact_count, 0); + assert_eq!(status.missing_vector_count, 0); +} + +#[tokio::test] +async fn fact_retriever_search_sanitizes_fts_chars_and_trust_weights_ordering() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + let retriever = FactRetriever::new(db.conn()); - cg.record_code_area("src/auth.rs", Some("OAuth provider")) + store + .add_fact( + fact_request( + "Rust HRR auth memory is preferred", + MemoryCategory::Decision, + 0.9, + ), + DEFAULT_TRUST, + ) + .await + .unwrap(); + store + .add_fact( + fact_request( + "Rust HRR auth memory is experimental", + MemoryCategory::Decision, + 0.2, + ), + DEFAULT_TRUST, + ) .await .unwrap(); - cg.record_code_area("src/auth.rs", None).await.unwrap(); - cg.record_code_area("src/auth.rs", None).await.unwrap(); - let areas = cg.list_code_areas(10).await.unwrap(); - assert_eq!(areas.len(), 1); - assert_eq!(areas[0].path, "src/auth.rs"); - assert_eq!(areas[0].touch_count, 3); - assert_eq!(areas[0].description.as_deref(), Some("OAuth provider")); + let results = retriever + .search( + "Rust (HRR) + auth?", + Some(MemoryCategory::Decision), + None, + 10, + ) + .await + .unwrap(); + + assert_eq!(results.len(), 1); + assert!(results[0].score > 0.0); + assert!(results[0].fts_score >= 0.0); + assert!(results[0].jaccard_score > 0.0); + assert!(results[0].holographic_score >= 0.0); + assert_eq!(results[0].trust_score, results[0].fact.trust_score); + assert!(results[0].why.as_deref().unwrap_or("").contains("trust")); + assert_eq!(results[0].fact.content, "Rust HRR auth memory is preferred"); } #[tokio::test] -async fn session_recall_filters_by_since() { - let (_tmp, cg) = make_project().await; - cg.record_decision("old decision", None, &[], &[]) +async fn fact_retriever_search_includes_old_entity_only_matches() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + let retriever = FactRetriever::new(db.conn()); + + let mut matching = fact_request( + "Older durable fact without the query words", + MemoryCategory::Project, + 0.9, + ); + matching.entities = vec!["EntityNeedle".to_string()]; + store.add_fact(matching, DEFAULT_TRUST).await.unwrap(); + + for i in 0..125 { + let mut unrelated = fact_request( + &format!("Newer unrelated project fact {i}"), + MemoryCategory::Project, + 0.9, + ); + unrelated.entities = vec![format!("UnrelatedEntity{i}")]; + store.add_fact(unrelated, DEFAULT_TRUST).await.unwrap(); + } + + let results = retriever + .search("EntityNeedle", Some(MemoryCategory::Project), Some(0.3), 5) .await .unwrap(); - // Force a > 1s gap so created_at values differ deterministically. - tokio::time::sleep(std::time::Duration::from_millis(1100)).await; - let cutoff = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs() as i64; - // Force a > 1s gap on the new side too, otherwise the new record could share - // its created_at with `cutoff` (second-granularity). - tokio::time::sleep(std::time::Duration::from_millis(1100)).await; - cg.record_decision("new decision", None, &[], &[]) + + assert!( + results + .iter() + .any(|result| result.fact.content == "Older durable fact without the query words"), + "search should include facts found only through stored entities" + ); +} + +#[tokio::test] +async fn fact_retriever_probe_related_reason_and_contradiction() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + let retriever = FactRetriever::new(db.conn()); + + let mut first = fact_request( + "Project Phoenix uses SQLite memory", + MemoryCategory::Decision, + 0.8, + ); + first.entities = vec!["Project Phoenix".to_string(), "SQLite".to_string()]; + store.add_fact(first, DEFAULT_TRUST).await.unwrap(); + + let mut second = fact_request( + "Project Phoenix uses HRR banks", + MemoryCategory::Decision, + 0.8, + ); + second.entities = vec!["Project Phoenix".to_string(), "HRR banks".to_string()]; + store.add_fact(second, DEFAULT_TRUST).await.unwrap(); + + let mut third = fact_request( + "Do not use SQLite memory for Project Phoenix", + MemoryCategory::Decision, + 0.8, + ); + third.entities = vec!["Project Phoenix".to_string(), "SQLite".to_string()]; + store.add_fact(third, DEFAULT_TRUST).await.unwrap(); + + let probe = retriever + .probe("Project Phoenix", None, Some(0.0), 10) + .await + .unwrap(); + assert_eq!(probe.len(), 3); + + let related = retriever.related("Project Phoenix", 10).await.unwrap(); + let related_names: Vec<_> = related.into_iter().map(|entity| entity.name).collect(); + assert!(related_names.contains(&"SQLite".to_string())); + assert!(related_names.contains(&"HRR banks".to_string())); + + let reason = retriever + .reason( + &["Project Phoenix".to_string(), "SQLite".to_string()], + None, + Some(0.0), + 10, + ) .await .unwrap(); + assert_eq!(reason.len(), 2); - let hits = cg.session_recall(None, Some(cutoff), 10).await.unwrap(); - assert_eq!(hits.len(), 1); - assert_eq!(hits[0].text, "new decision"); + let contradictions = retriever + .contradict(MemoryCategory::Decision, 0.2, 10) + .await + .unwrap(); + assert!(contradictions.iter().any(|result| result + .existing_fact + .content + .contains("uses SQLite") + && result.new_content.contains("Do not use SQLite"))); +} + +#[tokio::test] +async fn fact_retriever_reason_applies_entity_predicates_before_limit() { + let (db, _tmp) = make_memory_store().await; + let store = MemoryStore::new(db.conn()); + let retriever = FactRetriever::new(db.conn()); + + let mut matching = fact_request( + "Older fact links Project Phoenix and SQLite", + MemoryCategory::Decision, + 0.9, + ); + matching.entities = vec!["Project Phoenix".to_string(), "SQLite".to_string()]; + store.add_fact(matching, DEFAULT_TRUST).await.unwrap(); + + for i in 0..125 { + let mut unrelated = fact_request( + &format!("Newer unrelated fact {i}"), + MemoryCategory::Decision, + 0.9, + ); + unrelated.entities = vec![format!("Unrelated {i}")]; + store.add_fact(unrelated, DEFAULT_TRUST).await.unwrap(); + } + + let results = retriever + .reason( + &["Project Phoenix".to_string(), "SQLite".to_string()], + Some(MemoryCategory::Decision), + Some(0.3), + 10, + ) + .await + .unwrap(); + assert!( + results + .iter() + .any(|result| result.fact.content.contains("Older fact links")), + "reason should find matching facts before applying the result cap" + ); } diff --git a/tests/migration_test.rs b/tests/migration_test.rs index 838a50ef85..9c2ba5737c 100644 --- a/tests/migration_test.rs +++ b/tests/migration_test.rs @@ -79,6 +79,96 @@ async fn index_exists(conn: &Connection, index_name: &str) -> bool { .is_some() } +/// Checks whether a trigger exists in sqlite_master. +async fn trigger_exists(conn: &Connection, trigger_name: &str) -> bool { + let mut rows = conn + .query( + "SELECT name FROM sqlite_master WHERE type='trigger' AND name=?1", + libsql::params![trigger_name], + ) + .await + .expect("failed to query sqlite_master"); + rows.next() + .await + .expect("failed to read sqlite_master row") + .is_some() +} + +/// Returns the first column from the first row as i64. +async fn scalar_i64(conn: &Connection, sql: &str) -> i64 { + let mut rows = conn.query(sql, ()).await.expect("failed to query scalar"); + let row = rows + .next() + .await + .expect("failed to read scalar row") + .expect("scalar query should return a row"); + row.get(0).expect("failed to read scalar value") +} + +async fn assert_backfilled_memory_has_vectors_and_banks( + conn: &Connection, + category: &str, + expected_fact_count: i64, +) { + assert_eq!( + scalar_i64( + conn, + &format!( + "SELECT COUNT(*) FROM memory_facts + WHERE category = '{category}' + AND hrr_vector IS NOT NULL + AND length(hrr_vector) > 0 + AND hrr_algebra = 'amari_fhrr' + AND hrr_dim = 2048" + ) + ) + .await, + expected_fact_count, + "all backfilled {category} facts should have serialized HRR vectors" + ); + assert_eq!( + scalar_i64( + conn, + "SELECT COUNT(*) FROM memory_facts + WHERE hrr_vector IS NULL OR hrr_algebra != 'amari_fhrr' OR hrr_dim != 2048" + ) + .await, + 0, + "v11 migration should leave no backfilled facts missing vectors" + ); + assert_eq!( + scalar_i64( + conn, + "SELECT COUNT(*) FROM memory_banks + WHERE bank_name = 'all' + AND vector IS NOT NULL + AND length(vector) > 0 + AND hrr_algebra = 'amari_fhrr' + AND hrr_dim = 2048" + ) + .await, + 1, + "v11 migration should build the global memory bank" + ); + assert_eq!( + scalar_i64( + conn, + &format!( + "SELECT COUNT(*) FROM memory_banks + WHERE bank_name = '{category}' + AND vector IS NOT NULL + AND length(vector) > 0 + AND hrr_algebra = 'amari_fhrr' + AND hrr_dim = 2048 + AND fact_count = {expected_fact_count}" + ) + ) + .await, + 1, + "v11 migration should build the {category} memory bank" + ); +} + /// Checks whether a column exists on a table via PRAGMA table_info. async fn column_exists(conn: &Connection, table: &str, column: &str) -> bool { let mut rows = conn @@ -97,6 +187,29 @@ async fn column_exists(conn: &Connection, table: &str, column: &str) -> bool { false } +/// Returns the declared SQLite type and primary-key ordinal for a column. +async fn column_type_and_pk(conn: &Connection, table: &str, column: &str) -> (String, i64) { + let mut rows = conn + .query(&format!("PRAGMA table_info({table})"), ()) + .await + .expect("failed to query table_info"); + while let Some(row) = rows.next().await.expect("failed to read table_info row") { + let name = row + .get_str(1) + .expect("failed to read column name") + .to_string(); + if name == column { + return ( + row.get_str(2) + .expect("failed to read column type") + .to_string(), + row.get(5).expect("failed to read primary key ordinal"), + ); + } + } + panic!("{table}.{column} not found"); +} + /// Creates the V1 schema (tables, FTS, indexes — no metadata, no complexity columns). async fn create_v1_schema(conn: &Connection) { conn.execute_batch( @@ -234,11 +347,54 @@ async fn apply_v4(conn: &Connection) { set_user_version(conn, 4).await; } +/// Creates a latest pre-v11 schema with legacy memory tables but no holographic tables. +async fn create_v10_schema_for_v11_tests(conn: &Connection) { + create_schema(conn) + .await + .expect("failed to create baseline schema"); + conn.execute_batch( + "DROP TRIGGER IF EXISTS memory_facts_fts_insert; + DROP TRIGGER IF EXISTS memory_facts_fts_delete; + DROP TRIGGER IF EXISTS memory_facts_fts_update; + DROP TABLE IF EXISTS memory_facts_fts; + DROP TABLE IF EXISTS memory_feedback_events; + DROP TABLE IF EXISTS memory_fact_entities; + DROP TABLE IF EXISTS memory_banks; + DROP TABLE IF EXISTS memory_entities; + DROP TABLE IF EXISTS memory_facts; + + CREATE TABLE IF NOT EXISTS memory_decisions ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + text TEXT NOT NULL, + reason TEXT, + created_at INTEGER NOT NULL, + files TEXT NOT NULL DEFAULT '[]', + tags TEXT NOT NULL DEFAULT '[]' + ); + + CREATE TABLE IF NOT EXISTS memory_code_areas ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + path TEXT NOT NULL, + description TEXT, + last_touched_at INTEGER NOT NULL, + touch_count INTEGER NOT NULL DEFAULT 1 + ); + + CREATE UNIQUE INDEX IF NOT EXISTS idx_memory_code_areas_path + ON memory_code_areas(path); + CREATE INDEX IF NOT EXISTS idx_memory_decisions_created_at + ON memory_decisions(created_at);", + ) + .await + .expect("failed to remove v11 tables"); + set_user_version(conn, 10).await; +} + // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- -/// create_schema on a fresh database sets user_version to 5 and creates all tables. +/// create_schema on a fresh database sets user_version to latest and creates all tables. #[tokio::test] async fn test_create_schema_fresh_db() { let (conn, _db, _dir) = create_raw_db().await; @@ -247,7 +403,7 @@ async fn test_create_schema_fresh_db() { .await .expect("create_schema should succeed"); - assert_eq!(get_user_version(&conn).await, 10); + assert_eq!(get_user_version(&conn).await, 12); assert!(table_exists(&conn, "nodes").await); assert!(table_exists(&conn, "edges").await); assert!(table_exists(&conn, "files").await); @@ -255,6 +411,15 @@ async fn test_create_schema_fresh_db() { assert!(table_exists(&conn, "vectors").await); assert!(table_exists(&conn, "metadata").await); assert!(table_exists(&conn, "nodes_fts").await); + assert!(!table_exists(&conn, "memory_decisions").await); + assert!(!table_exists(&conn, "memory_code_areas").await); + assert!(table_exists(&conn, "memory_facts").await); + assert!(table_exists(&conn, "memory_entities").await); + assert!(table_exists(&conn, "memory_fact_entities").await); + assert!(table_exists(&conn, "memory_banks").await); + assert!(table_exists(&conn, "memory_bank_dirty").await); + assert!(table_exists(&conn, "memory_feedback_events").await); + assert!(table_exists(&conn, "memory_facts_fts").await); } /// create_schema is idempotent — calling it twice does not error. @@ -269,7 +434,7 @@ async fn test_create_schema_idempotent() { .await .expect("second create_schema should succeed"); - assert_eq!(get_user_version(&conn).await, 10); + assert_eq!(get_user_version(&conn).await, 12); } /// migrate returns false when already at the latest version. @@ -287,7 +452,7 @@ async fn test_migrate_already_latest_returns_false() { !migrated, "migrate should return false when already at latest" ); - assert_eq!(get_user_version(&conn).await, 10); + assert_eq!(get_user_version(&conn).await, 12); } /// migrate from v0 (completely empty database) applies all migrations to latest. @@ -306,7 +471,7 @@ async fn test_migrate_from_v0() { migrated, "migrate should return true when migrations were applied" ); - assert_eq!(get_user_version(&conn).await, 10); + assert_eq!(get_user_version(&conn).await, 12); // All expected tables should exist assert!(table_exists(&conn, "nodes").await); @@ -347,7 +512,7 @@ async fn test_migrate_from_v1() { .expect("migrate from v1 should succeed"); assert!(migrated); - assert_eq!(get_user_version(&conn).await, 10); + assert_eq!(get_user_version(&conn).await, 12); // V2: metadata table assert!(table_exists(&conn, "metadata").await); @@ -383,7 +548,7 @@ async fn test_migrate_from_v2() { .expect("migrate from v2 should succeed"); assert!(migrated); - assert_eq!(get_user_version(&conn).await, 10); + assert_eq!(get_user_version(&conn).await, 12); // V3 columns assert!(column_exists(&conn, "nodes", "branches").await); @@ -413,7 +578,7 @@ async fn test_migrate_from_v3() { .expect("migrate from v3 should succeed"); assert!(migrated); - assert_eq!(get_user_version(&conn).await, 10); + assert_eq!(get_user_version(&conn).await, 12); // V4 columns assert!(column_exists(&conn, "nodes", "unsafe_blocks").await); @@ -441,7 +606,7 @@ async fn test_migrate_from_v4() { .expect("migrate from v4 should succeed"); assert!(migrated); - assert_eq!(get_user_version(&conn).await, 10); + assert_eq!(get_user_version(&conn).await, 12); assert!(index_exists(&conn, "idx_edges_unique").await); } @@ -583,7 +748,7 @@ async fn test_database_initialize_creates_latest_version() { .expect("failed to read row") .expect("should have row"); let version: i64 = row.get(0).expect("failed to read version"); - assert_eq!(version, 10); + assert_eq!(version, 12); } /// Database::open on an already-current database does not re-migrate. @@ -650,7 +815,7 @@ async fn test_database_open_migrates_v1_to_latest() { .expect("failed to read row") .expect("should have row"); let version: i64 = row.get(0).expect("failed to read version"); - assert_eq!(version, 10); + assert_eq!(version, 12); } /// After create_schema, all v5 columns on nodes exist. @@ -767,31 +932,86 @@ async fn test_fts_triggers_exist_after_migration() { } #[tokio::test] -async fn test_v8_creates_memory_tables() { +async fn test_latest_schema_omits_legacy_memory_tables() { let (conn, _db, _dir) = create_raw_db().await; create_schema(&conn).await.unwrap(); - // memory_decisions table exists with expected columns + assert!(!table_exists(&conn, "memory_decisions").await); + assert!(!table_exists(&conn, "memory_code_areas").await); + assert!(!table_exists(&conn, "memory_decisions_fts").await); + assert!(table_exists(&conn, "memory_facts").await); + assert!(table_exists(&conn, "memory_entities").await); +} + +#[tokio::test] +async fn test_v7_to_latest_upgrade_path() { + let (conn, _db, _dir) = create_raw_db().await; + + create_schema(&conn).await.unwrap(); + conn.execute("PRAGMA user_version = 7", ()).await.unwrap(); + // Drop the v8+ tables to simulate a true v7 starting state + conn.execute("DROP TABLE IF EXISTS memory_decisions_fts", ()) + .await + .unwrap(); + conn.execute("DROP TABLE IF EXISTS memory_decisions", ()) + .await + .unwrap(); + conn.execute("DROP TABLE IF EXISTS memory_code_areas", ()) + .await + .unwrap(); + conn.execute("DROP TABLE IF EXISTS read_cache", ()) + .await + .unwrap(); + + let did_migrate = migrate(&conn).await.unwrap(); + assert!(did_migrate, "expected migrate() to return true"); + + let mut rows = conn.query("PRAGMA user_version", ()).await.unwrap(); + let row = rows.next().await.unwrap().unwrap(); + let v: i64 = row.get(0).unwrap(); + assert_eq!(v, 12); + let mut rows = conn .query( - "SELECT name FROM pragma_table_info('memory_decisions') ORDER BY cid", + "SELECT name FROM sqlite_master WHERE type='table' AND name IN \ + ('memory_decisions','memory_code_areas','memory_decisions_fts','read_cache') ORDER BY name", (), ) .await .unwrap(); - let mut cols = Vec::new(); + let mut names = Vec::new(); while let Some(row) = rows.next().await.unwrap() { - cols.push(row.get::(0).unwrap()); + names.push(row.get::(0).unwrap()); } - assert_eq!( - cols, - vec!["id", "text", "reason", "created_at", "files", "tags"] + assert_eq!(names, vec!["read_cache"]); +} + +/// V9 adds the `read_cache` table used by `tokensave_read`. +#[tokio::test] +async fn test_migrate_v9_adds_read_cache() { + let (conn, _db, _dir) = create_raw_db().await; + migrate(&conn).await.expect("migrate should succeed"); + + assert!( + table_exists(&conn, "read_cache").await, + "v9 migration should create the read_cache table" + ); + assert!( + index_exists(&conn, "idx_read_cache_session").await, + "v9 migration should create idx_read_cache_session" ); +} + +#[tokio::test] +async fn test_v11_create_schema_has_holographic_memory_schema() { + let (conn, _db, _dir) = create_raw_db().await; + create_schema(&conn) + .await + .expect("create_schema should succeed"); - // memory_code_areas table exists let mut rows = conn .query( - "SELECT name FROM pragma_table_info('memory_code_areas') ORDER BY cid", + "SELECT name FROM pragma_table_info('memory_facts') ORDER BY cid", (), ) .await @@ -803,114 +1023,526 @@ async fn test_v8_creates_memory_tables() { assert_eq!( cols, vec![ - "id", - "path", - "description", - "last_touched_at", - "touch_count" + "fact_id", + "content", + "category", + "tags", + "trust_score", + "retrieval_count", + "helpful_count", + "unhelpful_count", + "created_at", + "updated_at", + "last_retrieved_at", + "last_feedback_at", + "source", + "metadata", + "hrr_vector", + "hrr_algebra", + "hrr_dim", ] ); + assert_eq!( + column_type_and_pk(&conn, "memory_facts", "fact_id").await, + ("INTEGER".to_string(), 1) + ); + assert_eq!( + column_type_and_pk(&conn, "memory_entities", "entity_id").await, + ("INTEGER".to_string(), 1) + ); + assert_eq!( + column_type_and_pk(&conn, "memory_banks", "bank_id").await, + ("INTEGER".to_string(), 1) + ); + assert_eq!( + column_type_and_pk(&conn, "memory_fact_entities", "fact_id").await, + ("INTEGER".to_string(), 1) + ); + assert_eq!( + column_type_and_pk(&conn, "memory_fact_entities", "entity_id").await, + ("INTEGER".to_string(), 2) + ); + assert_eq!( + column_type_and_pk(&conn, "memory_feedback_events", "fact_id").await, + ("INTEGER".to_string(), 0) + ); + + for table in [ + "memory_entities", + "memory_fact_entities", + "memory_banks", + "memory_feedback_events", + "memory_facts_fts", + ] { + assert!(table_exists(&conn, table).await, "{table} should exist"); + } + + for index in [ + "idx_memory_facts_category", + "idx_memory_facts_updated_at", + "idx_memory_entities_type", + "idx_memory_fact_entities_entity_id", + "idx_memory_feedback_events_fact_id", + ] { + assert!(index_exists(&conn, index).await, "{index} should exist"); + } + + for trigger in [ + "memory_facts_fts_insert", + "memory_facts_fts_delete", + "memory_facts_fts_update", + ] { + assert!( + trigger_exists(&conn, trigger).await, + "{trigger} should exist" + ); + } + + conn.execute( + "INSERT INTO memory_facts (content, category) VALUES ('Default values matter', 'test')", + (), + ) + .await + .expect("minimal memory_facts insert should use defaults"); + let fact_id = scalar_i64(&conn, "SELECT fact_id FROM memory_facts").await; + assert!(fact_id > 0); - // FTS table exists let mut rows = conn .query( - "SELECT name FROM sqlite_master WHERE type='table' AND name='memory_decisions_fts'", - (), + "SELECT tags, trust_score, retrieval_count, helpful_count, unhelpful_count, source, metadata, hrr_algebra, hrr_dim FROM memory_facts WHERE fact_id=?1", + libsql::params![fact_id], ) .await .unwrap(); - assert!( - rows.next().await.unwrap().is_some(), - "memory_decisions_fts missing" - ); + let row = rows.next().await.unwrap().unwrap(); + assert_eq!(row.get::(0).unwrap(), "[]"); + assert_eq!(row.get::(1).unwrap(), 0.5); + assert_eq!(row.get::(2).unwrap(), 0); + assert_eq!(row.get::(3).unwrap(), 0); + assert_eq!(row.get::(4).unwrap(), 0); + assert_eq!(row.get::(5).unwrap(), "manual"); + assert_eq!(row.get::(6).unwrap(), "{}"); + assert_eq!(row.get::(7).unwrap(), "amari_fhrr"); + assert_eq!(row.get::(8).unwrap(), 2048); +} + +#[tokio::test] +async fn test_v10_to_v11_backfills_and_drops_legacy_memory_tables() { + let (conn, _db, _dir) = create_raw_db().await; + create_v10_schema_for_v11_tests(&conn).await; + + assert_eq!(get_user_version(&conn).await, 10); + assert!(table_exists(&conn, "memory_decisions").await); + assert!(table_exists(&conn, "memory_code_areas").await); + assert!(!table_exists(&conn, "memory_facts").await); + + let did_migrate = migrate(&conn).await.expect("v10 to v11 should migrate"); + + assert!(did_migrate); + assert_eq!(get_user_version(&conn).await, 12); + assert!(!table_exists(&conn, "memory_decisions").await); + assert!(!table_exists(&conn, "memory_code_areas").await); + assert!(table_exists(&conn, "memory_facts").await); + assert!(table_exists(&conn, "memory_entities").await); + assert!(table_exists(&conn, "memory_fact_entities").await); + assert!(table_exists(&conn, "memory_banks").await); + assert!(table_exists(&conn, "memory_feedback_events").await); + assert!(table_exists(&conn, "memory_facts_fts").await); +} + +#[tokio::test] +async fn test_v11_to_v12_adds_memory_bank_dirty_table() { + let (conn, _db, _dir) = create_raw_db().await; + create_schema(&conn).await.unwrap(); + conn.execute("DROP TABLE IF EXISTS memory_bank_dirty", ()) + .await + .unwrap(); + set_user_version(&conn, 11).await; + + assert_eq!(get_user_version(&conn).await, 11); + assert!(!table_exists(&conn, "memory_bank_dirty").await); + + let did_migrate = migrate(&conn).await.expect("v11 to v12 should migrate"); + + assert!(did_migrate); + assert_eq!(get_user_version(&conn).await, 12); + assert!(table_exists(&conn, "memory_bank_dirty").await); +} + +#[tokio::test] +async fn test_v11_feedback_events_enforce_action_and_cascade_with_facts() { + let (conn, _db, _dir) = create_raw_db().await; + create_schema(&conn).await.unwrap(); + + conn.execute( + "INSERT INTO memory_facts (content, category) VALUES ('Feedback fact', 'test')", + (), + ) + .await + .expect("failed to insert memory fact"); + let fact_id = scalar_i64(&conn, "SELECT fact_id FROM memory_facts").await; + conn.execute( + "INSERT INTO memory_feedback_events (fact_id, action, trust_delta, old_trust, new_trust, note) + VALUES (?1, 'helpful', 0.1, 0.5, 0.6, 'worked')", + libsql::params![fact_id], + ) + .await + .expect("valid feedback action should insert"); + + let invalid = conn + .execute( + "INSERT INTO memory_feedback_events (fact_id, action, trust_delta, old_trust, new_trust) + VALUES (?1, 'neutral', 0.0, 0.5, 0.5)", + libsql::params![fact_id], + ) + .await; + assert!(invalid.is_err(), "invalid feedback action should fail"); - // All three FTS triggers exist let mut rows = conn .query( - "SELECT name FROM sqlite_master WHERE type='trigger' \ - AND name IN ('memory_decisions_fts_insert', 'memory_decisions_fts_delete', 'memory_decisions_fts_update') \ - ORDER BY name", - (), + "SELECT source FROM memory_feedback_events WHERE fact_id=?1", + libsql::params![fact_id], ) .await .unwrap(); - let mut trigger_names = Vec::new(); - while let Some(row) = rows.next().await.unwrap() { - trigger_names.push(row.get::(0).unwrap()); - } + let row = rows.next().await.unwrap().unwrap(); + assert_eq!(row.get::(0).unwrap(), "mcp"); + + conn.execute( + "DELETE FROM memory_facts WHERE fact_id=?1", + libsql::params![fact_id], + ) + .await + .expect("deleting memory fact should cascade"); assert_eq!( - trigger_names, - vec![ - "memory_decisions_fts_delete", - "memory_decisions_fts_insert", - "memory_decisions_fts_update", - ] + { + let mut rows = conn + .query( + "SELECT COUNT(*) FROM memory_feedback_events WHERE fact_id=?1", + libsql::params![fact_id], + ) + .await + .unwrap(); + let row = rows.next().await.unwrap().unwrap(); + row.get::(0).unwrap() + }, + 0 ); } #[tokio::test] -async fn test_v7_to_latest_upgrade_path() { +async fn test_v11_memory_facts_fts_triggers_track_insert_update_delete() { let (conn, _db, _dir) = create_raw_db().await; - create_schema(&conn).await.unwrap(); - conn.execute("PRAGMA user_version = 7", ()).await.unwrap(); - // Drop the v8+ tables to simulate a true v7 starting state - conn.execute("DROP TABLE IF EXISTS memory_decisions_fts", ()) + + conn.execute( + "INSERT INTO memory_facts (content, category, tags) + VALUES ('Use orbital retrieval for context', 'test', '[\"retrieval\"]')", + (), + ) + .await + .expect("failed to insert memory fact"); + let fact_id = scalar_i64(&conn, "SELECT fact_id FROM memory_facts").await; + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(*) FROM memory_facts_fts WHERE memory_facts_fts MATCH 'orbital'" + ) + .await, + 1 + ); + + conn.execute( + "UPDATE memory_facts SET content='Use semantic banana storage', tags='[\"banana\"]' WHERE fact_id=?1", + libsql::params![fact_id], + ) + .await + .expect("failed to update memory fact"); + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(*) FROM memory_facts_fts WHERE memory_facts_fts MATCH 'orbital'" + ) + .await, + 0 + ); + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(*) FROM memory_facts_fts WHERE memory_facts_fts MATCH 'banana'" + ) + .await, + 1 + ); + + conn.execute( + "DELETE FROM memory_facts WHERE fact_id=?1", + libsql::params![fact_id], + ) + .await + .expect("failed to delete memory fact"); + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(*) FROM memory_facts_fts WHERE memory_facts_fts MATCH 'banana'" + ) + .await, + 0 + ); +} + +#[tokio::test] +async fn test_v11_backfills_legacy_memory_decisions_as_facts() { + let (conn, _db, _dir) = create_raw_db().await; + create_v10_schema_for_v11_tests(&conn).await; + conn.execute( + "INSERT INTO memory_decisions (text, reason, created_at, files, tags) + VALUES ('Prefer libsql migrations', 'Keeps install path simple', 1234, '[\"src/db/migrations.rs\"]', '[\"db\",\"memory\"]')", + (), + ) + .await + .expect("failed to insert legacy decision"); + + migrate(&conn).await.expect("v11 migration should backfill"); + + let mut rows = conn + .query( + "SELECT fact_id, content, tags, metadata FROM memory_facts WHERE category='decision'", + (), + ) .await .unwrap(); - conn.execute("DROP TABLE IF EXISTS memory_decisions", ()) + let row = rows.next().await.unwrap().unwrap(); + let fact_id = row.get::(0).unwrap(); + let content = row.get::(1).unwrap(); + let tags = row.get::(2).unwrap(); + let metadata = row.get::(3).unwrap(); + + assert!(fact_id > 0); + assert!(content.contains("Prefer libsql migrations")); + assert!(content.contains("Keeps install path simple")); + assert_eq!(tags, "[\"db\",\"memory\"]"); + assert!(!metadata.contains("legacy-decision-")); + assert!(metadata.contains("holographic_memory_backfill_v1")); + assert!(metadata.contains("memory_decisions")); + assert!(metadata.contains("\"legacy_id\":1")); + assert!(metadata.contains("\"decision_text\":\"Prefer libsql migrations\"")); + assert!(metadata.contains("src/db/migrations.rs")); + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(*) + FROM memory_fact_entities fe + JOIN memory_entities e ON e.entity_id = fe.entity_id + WHERE fe.fact_id = 1 + AND e.normalized_name IN ('src/db/migrations.rs', 'db', 'memory')" + ) + .await, + 3 + ); + assert_backfilled_memory_has_vectors_and_banks(&conn, "decision", 1).await; +} + +#[tokio::test] +async fn test_v11_backfills_legacy_memory_code_areas_as_facts() { + let (conn, _db, _dir) = create_raw_db().await; + create_v10_schema_for_v11_tests(&conn).await; + conn.execute( + "INSERT INTO memory_code_areas (path, description, last_touched_at, touch_count) + VALUES ('src/db/migrations.rs', 'Schema migration code', 5678, 3)", + (), + ) + .await + .expect("failed to insert legacy code area"); + + migrate(&conn).await.expect("v11 migration should backfill"); + + let mut rows = conn + .query( + "SELECT fact_id, content, tags, metadata FROM memory_facts WHERE category='code_area'", + (), + ) .await .unwrap(); - conn.execute("DROP TABLE IF EXISTS memory_code_areas", ()) + let row = rows.next().await.unwrap().unwrap(); + let fact_id = row.get::(0).unwrap(); + let content = row.get::(1).unwrap(); + let tags = row.get::(2).unwrap(); + let metadata = row.get::(3).unwrap(); + + assert!(fact_id > 0); + assert!(content.contains("src/db/migrations.rs")); + assert!(content.contains("Schema migration code")); + assert!(tags.contains("code_area")); + assert!(tags.contains("src/db/migrations.rs")); + assert!(!metadata.contains("legacy-code-area-")); + assert!(metadata.contains("holographic_memory_backfill_v1")); + assert!(metadata.contains("memory_code_areas")); + assert!(metadata.contains("\"legacy_id\":1")); + assert!(metadata.contains("touch_count")); + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(*) + FROM memory_fact_entities fe + JOIN memory_entities e ON e.entity_id = fe.entity_id + WHERE fe.fact_id = 1 + AND e.normalized_name = 'src/db/migrations.rs'" + ) + .await, + 1 + ); + assert_backfilled_memory_has_vectors_and_banks(&conn, "code_area", 1).await; +} + +#[tokio::test] +async fn test_v11_backfill_is_idempotent_when_migration_reruns() { + let (conn, _db, _dir) = create_raw_db().await; + create_v10_schema_for_v11_tests(&conn).await; + conn.execute( + "INSERT INTO memory_decisions (text, reason, created_at, tags) + VALUES ('Avoid duplicate facts', 'Content has a unique constraint', 1000, '[\"dedupe\"]')", + (), + ) + .await + .expect("failed to insert legacy decision"); + conn.execute( + "INSERT INTO memory_code_areas (path, description, last_touched_at) + VALUES ('src/memory.rs', 'Legacy memory facade', 1000)", + (), + ) + .await + .expect("failed to insert legacy code area"); + + migrate(&conn) .await - .unwrap(); - conn.execute("DROP TABLE IF EXISTS read_cache", ()) + .expect("first v11 migration should succeed"); + assert_eq!( + scalar_i64(&conn, "SELECT COUNT(*) FROM memory_facts").await, + 2 + ); + + set_user_version(&conn, 10).await; + migrate(&conn) .await - .unwrap(); + .expect("rerunning v11 migration should succeed"); - let did_migrate = migrate(&conn).await.unwrap(); - assert!(did_migrate, "expected migrate() to return true"); + assert_eq!( + scalar_i64(&conn, "SELECT COUNT(*) FROM memory_facts").await, + 2 + ); +} - let mut rows = conn.query("PRAGMA user_version", ()).await.unwrap(); - let row = rows.next().await.unwrap().unwrap(); - let v: i64 = row.get(0).unwrap(); - assert_eq!(v, 10); +#[tokio::test] +async fn test_v11_backfill_handles_malformed_and_blank_legacy_json() { + let (conn, _db, _dir) = create_raw_db().await; + create_v10_schema_for_v11_tests(&conn).await; + conn.execute( + "INSERT INTO memory_decisions (text, reason, created_at, files, tags) + VALUES ('Bad JSON is normalized', '', 1000, '[invalid json', 'not-an-array')", + (), + ) + .await + .expect("failed to insert bad-json legacy decision"); + conn.execute( + "INSERT INTO memory_code_areas (path, description, last_touched_at, touch_count) + VALUES ('src/blank.rs', '', 1001, 1)", + (), + ) + .await + .expect("failed to insert blank legacy code area"); + + migrate(&conn) + .await + .expect("v11 migration should tolerate malformed legacy JSON"); let mut rows = conn .query( - "SELECT name FROM sqlite_master WHERE type='table' AND name IN \ - ('memory_decisions','memory_code_areas','memory_decisions_fts','read_cache') ORDER BY name", + "SELECT content, tags, metadata FROM memory_facts WHERE category='decision'", (), ) .await .unwrap(); - let mut names = Vec::new(); - while let Some(row) = rows.next().await.unwrap() { - names.push(row.get::(0).unwrap()); - } + let row = rows.next().await.unwrap().unwrap(); + let content = row.get::(0).unwrap(); + let tags = row.get::(1).unwrap(); + let metadata = row.get::(2).unwrap(); + assert!(content.contains("Bad JSON is normalized")); + assert!(!content.contains("Reason:")); + assert_eq!(tags, "[]"); + assert!(metadata.contains("\"files\":[]")); + assert!(metadata.contains("\"tags\":[]")); assert_eq!( - names, - vec![ - "memory_code_areas", - "memory_decisions", - "memory_decisions_fts", - "read_cache", - ] + scalar_i64( + &conn, + "SELECT COUNT(*) + FROM memory_fact_entities fe + JOIN memory_facts f ON f.fact_id = fe.fact_id + WHERE f.category = 'decision'" + ) + .await, + 0 ); + + let mut rows = conn + .query( + "SELECT content FROM memory_facts WHERE category='code_area'", + (), + ) + .await + .unwrap(); + let content = rows + .next() + .await + .unwrap() + .unwrap() + .get::(0) + .unwrap(); + assert!(content.contains("src/blank.rs")); + assert!(!content.contains("\n\n\n")); } -/// V9 adds the `read_cache` table used by `tokensave_read`. #[tokio::test] -async fn test_migrate_v9_adds_read_cache() { +async fn test_v11_backfill_preserves_duplicate_legacy_content() { let (conn, _db, _dir) = create_raw_db().await; - migrate(&conn).await.expect("migrate should succeed"); + create_v10_schema_for_v11_tests(&conn).await; + for tag in ["rust", "performance"] { + conn.execute( + "INSERT INTO memory_decisions (text, reason, created_at, files, tags) + VALUES ('Use Rust', 'same reason', 1000, '[]', json_array(?1))", + libsql::params![tag], + ) + .await + .expect("failed to insert duplicate legacy decision"); + } - assert!( - table_exists(&conn, "read_cache").await, - "v9 migration should create the read_cache table" + migrate(&conn).await.expect("v11 migration should backfill"); + + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(*) FROM memory_facts WHERE category='decision'" + ) + .await, + 2 ); - assert!( - index_exists(&conn, "idx_read_cache_session").await, - "v9 migration should create idx_read_cache_session" + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(DISTINCT content) FROM memory_facts WHERE category='decision'" + ) + .await, + 2 + ); + assert_eq!( + scalar_i64( + &conn, + "SELECT COUNT(*) + FROM memory_fact_entities fe + JOIN memory_entities e ON e.entity_id = fe.entity_id + WHERE e.normalized_name IN ('rust', 'performance')" + ) + .await, + 2 ); } diff --git a/tests/session_global_db_test.rs b/tests/session_global_db_test.rs index 29b9ac8f2f..e27a0377ba 100644 --- a/tests/session_global_db_test.rs +++ b/tests/session_global_db_test.rs @@ -110,6 +110,28 @@ async fn upsert_session_message_round_trips_and_updates() { assert_eq!(fetched.source_offset, Some(99)); } +#[tokio::test] +async fn upsert_session_message_truncates_oversized_text_deterministically() { + let tmp = TempDir::new().unwrap(); + let db = open_isolated_db(&tmp).await; + let session = sample_session("cursor", "session-1", "project-a"); + db.upsert_session(&session).await; + + let oversized = "x".repeat(300_000); + let message = sample_message("cursor", "message-1", "session-1", &oversized); + assert!(db.upsert_session_message(&message).await); + + let fetched = db + .get_session_message("cursor", "message-1") + .await + .expect("message should exist"); + assert!(fetched.text.len() < oversized.len()); + assert!( + fetched.text.ends_with("[truncated by tokensave]"), + "truncated message should be explicitly marked" + ); +} + #[tokio::test] async fn search_session_messages_uses_fts_and_filters_provider_project() { let tmp = TempDir::new().unwrap(); diff --git a/tests/vibe_transcript_ingest_test.rs b/tests/vibe_transcript_ingest_test.rs new file mode 100644 index 0000000000..157c68fa73 --- /dev/null +++ b/tests/vibe_transcript_ingest_test.rs @@ -0,0 +1,131 @@ +use std::io::Write; + +use tempfile::TempDir; +use tokensave::sessions::cursor::open_project_session_db; +use tokensave::sessions::source::ingest_source; +use tokensave::sessions::vibe::VibeSource; + +fn setup(tmp: &TempDir) -> (std::path::PathBuf, std::path::PathBuf) { + let home = tmp.path().join("home"); + let project = tmp.path().join("project"); + std::fs::create_dir_all(&project).unwrap(); + std::fs::create_dir(project.join(".tokensave")).unwrap(); + std::fs::write(project.join(".tokensave/tokensave.db"), "").unwrap(); + (home, project) +} + +fn write_vibe_session( + home: &std::path::Path, + project: &std::path::Path, + session_id: &str, +) -> std::path::PathBuf { + let dir = home + .join(".vibe/logs/session") + .join(format!("session_20260608_010000_{session_id}")); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write( + dir.join("meta.json"), + serde_json::to_string_pretty(&serde_json::json!({ + "session_id": session_id, + "environment": {"working_directory": project}, + "config": {"active_model": "mistral-medium-3.5"} + })) + .unwrap(), + ) + .unwrap(); + let messages = dir.join("messages.jsonl"); + std::fs::write( + &messages, + format!( + "{}\n{}\n", + serde_json::json!({ + "role": "user", + "content": "Investigate the billing pipeline regression", + "timestamp": 1_800_000_000_i64 + }), + serde_json::json!({ + "role": "assistant", + "content": [ + {"text": "The billing pipeline regression is fixed."}, + {"tool_call": {"name": "read_file"}} + ], + "timestamp": 1_800_000_010_i64 + }), + ), + ) + .unwrap(); + messages +} + +#[tokio::test] +async fn vibe_messages_populate_searchable_session_messages() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + write_vibe_session(&home, &project, "vibe-sess"); + + let db = open_project_session_db(&project).await.unwrap(); + let source = VibeSource::with_home(&home); + let stats = ingest_source(&db, &source, &project, None).await; + assert_eq!(stats.messages_upserted, 2); + + let results = db + .search_session_messages( + "vibe", + Some(project.to_string_lossy().as_ref()), + "billing pipeline", + 10, + ) + .await; + assert_eq!(results.len(), 2); + assert!(results + .iter() + .any(|hit| hit.message.tool_names.as_deref() == Some("read_file"))); + assert!(results + .iter() + .all(|hit| hit.message.model.as_deref() == Some("mistral-medium-3.5"))); +} + +#[tokio::test] +async fn vibe_messages_are_incremental() { + let tmp = TempDir::new().unwrap(); + let (home, project) = setup(&tmp); + let messages = write_vibe_session(&home, &project, "vibe-sess"); + + let db = open_project_session_db(&project).await.unwrap(); + let source = VibeSource::with_home(&home); + assert_eq!( + ingest_source(&db, &source, &project, None) + .await + .messages_upserted, + 2 + ); + assert_eq!( + ingest_source(&db, &source, &project, None) + .await + .messages_upserted, + 0 + ); + + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&messages) + .unwrap(); + writeln!( + file, + "{}", + serde_json::json!({ + "role": "assistant", + "content": "Added the regression test.", + "timestamp": 1_800_000_020_i64 + }) + ) + .unwrap(); + drop(file); + + assert_eq!( + ingest_source(&db, &source, &project, None) + .await + .messages_upserted, + 1 + ); +}