feat: memory v2 — categories, dream cleanup, self-skeptical (Claude Code inspired)#61
feat: memory v2 — categories, dream cleanup, self-skeptical (Claude Code inspired)#61kienbui1995 merged 2 commits intomainfrom
Conversation
…pact Inspired by Claude Code leaked memory architecture: 1. Memory categories: project, user, feedback, reference - Fact struct gains 'category' field (serde default: project) - memory_write tool accepts category parameter - to_prompt_section groups by category - Auto-memory detects category from content patterns 2. Self-skeptical prompt: - 'Treat as hints — verify against actual code before acting' - Injected in to_prompt_section header 3. Dream cleanup (compact): - memory.compact() deduplicates by key, keeps newest - /memory compact available via TUI - auto_compact_on_start(150) runs on session init 4. Auto-categorized detection: - 'convention is' / 'always use' → feedback - 'running on port' / 'located at' → reference - 'prefers' / 'user wants' → user - Default → project 274 tests, 0 fail.
📝 WalkthroughWalkthroughAdds memory categorization (project, user, feedback, reference), persists categories with facts, auto-compacts duplicate facts on startup when above a threshold, and updates auto-save and tool schema to detect/submit categories. Changes
Sequence DiagramsequenceDiagram
participant CLI as CLI (mc-cli)
participant Memory as MemoryStore
participant Runtime as Runtime
participant AutoSave as auto_save_memory
CLI->>Memory: load_from_disk()
activate Memory
Memory-->>CLI: MemoryStore
deactivate Memory
CLI->>Memory: auto_compact_on_start(150)
activate Memory
Memory->>Memory: compact() -- dedupe by key, keep newest
Memory-->>CLI: compacted MemoryStore
deactivate Memory
CLI->>Runtime: set_memory(compacted_store)
Runtime->>AutoSave: process output -> detect fact?
activate AutoSave
AutoSave->>AutoSave: determine (is_fact, category)
AutoSave-->>Runtime: (is_fact, category)
deactivate AutoSave
Runtime->>Memory: set_with_category(key, value, category)
activate Memory
Memory->>Memory: store Fact {key, value, category, updated_at}
Memory-->>Runtime: ok
deactivate Memory
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces categorized memory storage for facts, allowing them to be grouped into 'project', 'user', 'feedback', and 'reference' categories. It updates the auto-detection logic to assign these categories, enhances prompt generation to display categorized sections, and adds a memory compaction feature to deduplicate entries. Feedback includes moving synchronous disk writes outside of loops to improve performance, optimizing the prompt section generation to avoid multiple iterations, aligning documentation with the compaction implementation, and using integer comparisons for timestamps to ensure reliability.
| memory.set_with_category(&key, trimmed, category); | ||
| let _ = memory.save(); |
There was a problem hiding this comment.
memory.save() is called inside the loop for every detected fact. Since save() performs a synchronous disk write, this can cause significant performance degradation if multiple facts are detected in a single response. It should be called once after the loop finishes to minimize disk I/O.
| for cat in &["project", "user", "feedback", "reference"] { | ||
| let facts: Vec<_> = self.facts.iter().filter(|f| f.category == *cat).collect(); | ||
| if !facts.is_empty() { | ||
| section.push_str(&format!("\n### {}\n", capitalize(cat))); | ||
| for f in facts { | ||
| section.push_str(&format!("- {}: {}\n", f.key, f.value)); | ||
| } | ||
| } | ||
| } | ||
| // Uncategorized | ||
| let other: Vec<_> = self | ||
| .facts | ||
| .iter() | ||
| .map(|f| format!("- {}: {}", f.key, f.value)) | ||
| .filter(|f| { | ||
| !["project", "user", "feedback", "reference"].contains(&f.category.as_str()) | ||
| }) | ||
| .collect(); | ||
| format!("\n\n## Project Memory\n{}", lines.join("\n")) | ||
| for f in other { | ||
| section.push_str(&format!("- {}: {}\n", f.key, f.value)); | ||
| } |
There was a problem hiding this comment.
The to_prompt_section method iterates over the facts list multiple times (once for each hardcoded category plus once for uncategorized facts). This is inefficient and creates a maintenance burden whenever a new category is added. Consider grouping the facts in a single pass using a HashMap or similar structure, and deriving the category list from a central definition to ensure consistency with the tool specifications.
| /// Dream cleanup: remove duplicates, resolve contradictions (keep newest), | ||
| /// remove entries with same key prefix. Call on session start or /memory compact. |
There was a problem hiding this comment.
The docstring for compact mentions removing entries with the same key prefix, but the implementation only deduplicates by exact key. If prefix-based deduplication is intended (e.g., to handle auto_ timestamped keys), the logic should be updated; otherwise, the docstring should be corrected to reflect the actual behavior.
| for (i, f) in self.facts.iter().enumerate() { | ||
| seen.entry(f.key.clone()) | ||
| .and_modify(|(idx, ts): &mut (usize, String)| { | ||
| if f.updated_at > *ts { |
There was a problem hiding this comment.
Comparing updated_at as a String lexicographically is unreliable for numeric timestamps if their string lengths differ (e.g., "100" < "99" is true). While Unix timestamps in seconds currently have the same number of digits, it is safer to parse them as integers before comparison to avoid future bugs when the timestamp length increases.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
mc/crates/mc-core/src/memory.rs (1)
324-336:⚠️ Potential issue | 🟡 MinorMove test inside
mod testsblock.This test is defined outside the
#[cfg(test)] mod testsblock (which ends at line 322). It will compile into non-test builds, increasing binary size unnecessarily.🧪 Move test into mod tests
-#[test] -fn handle_write_delete() { - let path = std::env::temp_dir().join(format!("mc-mem-del-{}", std::process::id())); - let mut store = MemoryStore::load(&path, 100); - store.handle_write(&serde_json::json!({"key": "temp", "value": "data"})); - assert!(store - .handle_read(&serde_json::json!({"key": "temp"})) - .contains("data")); - store.handle_write(&serde_json::json!({"key": "temp", "delete": true})); - assert!(!store - .handle_read(&serde_json::json!({"key": "temp"})) - .contains("data")); - std::fs::remove_file(path).ok(); -}Move this test inside the
mod testsblock above, using thetmp_path()helper for consistency:#[test] fn handle_write_delete() { let path = tmp_path(); let mut store = MemoryStore::load(&path, 100); store.handle_write(&serde_json::json!({"key": "temp", "value": "data"})); assert!(store .handle_read(&serde_json::json!({"key": "temp"})) .contains("data")); store.handle_write(&serde_json::json!({"key": "temp", "delete": true})); assert!(!store .handle_read(&serde_json::json!({"key": "temp"})) .contains("data")); fs::remove_file(path).ok(); }As per coding guidelines: "Unit tests must use
#[cfg(test)] mod testsat bottom of each file."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@mc/crates/mc-core/src/memory.rs` around lines 324 - 336, The test function handle_write_delete is declared outside the #[cfg(test)] mod tests block; move the entire handle_write_delete test into the existing test module (inside the mod tests block) so it only compiles for tests, and update the temp file usage to use the tmp_path() helper for consistency; keep the test body using MemoryStore::load, store.handle_write, store.handle_read, and remove the file with fs::remove_file(path).ok() as suggested.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@mc/crates/mc-core/src/memory.rs`:
- Around line 324-336: The test function handle_write_delete is declared outside
the #[cfg(test)] mod tests block; move the entire handle_write_delete test into
the existing test module (inside the mod tests block) so it only compiles for
tests, and update the temp file usage to use the tmp_path() helper for
consistency; keep the test body using MemoryStore::load, store.handle_write,
store.handle_read, and remove the file with fs::remove_file(path).ok() as
suggested.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 05303b9b-e201-4aa2-9455-a48710ed5b09
📒 Files selected for processing (4)
mc/crates/mc-cli/src/main.rsmc/crates/mc-core/src/memory.rsmc/crates/mc-core/src/runtime.rsmc/crates/mc-tools/src/spec.rs
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
mc/crates/mc-core/src/memory.rs (2)
226-231: Surface startup-compaction save failures.If
save()fails here, memory is compacted in RAM but the JSON file stays stale, so the same duplicates come back next launch with no clue why. Returning aResultor bubbling the error to the caller would make this diagnosable.💡 Possible direction
- pub fn auto_compact_on_start(&mut self, threshold: usize) { + pub fn auto_compact_on_start(&mut self, threshold: usize) -> Result<(), std::io::Error> { if self.facts.len() > threshold { let removed = self.compact(); if removed > 0 { - let _ = self.save(); + self.save()?; } } + Ok(()) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@mc/crates/mc-core/src/memory.rs` around lines 226 - 231, auto_compact_on_start currently compacts in-memory (calls compact()) and swallows save() failures, leaving on-disk JSON stale; change auto_compact_on_start to return Result<(), E> (or Result<(), anyhow::Error>) instead of (), call save() only when removed > 0 and propagate its error (use the ? operator) so callers see and can handle disk write failures; update any callers of auto_compact_on_start to handle the Result and adjust tests accordingly.
323-334: Please add regression coverage for the new memory-v2 paths.This test still only exercises plain write/delete. The risky additions here are defaulting missing
category, preserving categories acrossset(), invalid-category handling, andcompact()semantics, so a few focused cases there would make regressions much easier to catch.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@mc/crates/mc-core/src/memory.rs` around lines 323 - 334, Add regression tests that exercise the new memory-v2 behavior: extend or add tests around MemoryStore::load, handle_write, handle_read and the underlying set()/compact() logic to cover (1) defaulting a missing "category" on write/read (ensure writes without category land in the default category and reads without category find them), (2) preserving an existing category when calling set() on an existing key (write key with category A, update value without category and assert category A is kept), (3) invalid-category handling (attempt writes/reads with an invalid category and assert the expected error or no-op behavior), and (4) compact() semantics (create multiple versions across categories, call compact() and assert it removes expired/older entries per the memory-v2 rules). Reference MemoryStore::load, handle_write, handle_read, set, and compact in your new tests.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@mc/crates/mc-core/src/memory.rs`:
- Line 68: In set(), don't hard-code "project" when delegating to
set_with_category; first check whether a fact for key already exists (e.g. via
self.facts.get(key) or the module's existing lookup helper) and if it does use
that fact's f.category when calling set_with_category(key, value, category),
otherwise default to "project"; update the call site that currently does
self.set_with_category(key, value, "project") to compute category =
existing_fact.map(|f| f.category.clone()).unwrap_or("project".into()) and pass
that through to preserve existing categories on updates.
- Around line 165-168: The code currently accepts any string for the local
variable `category` (derived from input.get("category")...) and must reject
unknown values before persisting; replace the unvalidated unwrap_or("project")
logic by parsing/validating the input string against the canonical enum of
allowed categories (e.g. via a FromStr/try_from or explicit match against the
four allowed variants) and return an error (or validation failure) when the
value is not one of the permitted options; apply the same validation fix to the
other occurrence of this pattern around the second instance noted (lines
~181-182) so only valid enum categories are persisted and invalid inputs are
rejected.
- Around line 207-210: In compact(), inside the .and_modify closure that updates
the (idx, ts) tuple, change the comparison from strict greater (f.updated_at >
*ts) to greater-or-equal (f.updated_at >= *ts) so that when updated_at
timestamps are equal (same-second duplicates) the later duplicate (current index
i) wins; update the conditional and leave the assignments to *idx = i and *ts =
f.updated_at.clone() unchanged.
---
Nitpick comments:
In `@mc/crates/mc-core/src/memory.rs`:
- Around line 226-231: auto_compact_on_start currently compacts in-memory (calls
compact()) and swallows save() failures, leaving on-disk JSON stale; change
auto_compact_on_start to return Result<(), E> (or Result<(), anyhow::Error>)
instead of (), call save() only when removed > 0 and propagate its error (use
the ? operator) so callers see and can handle disk write failures; update any
callers of auto_compact_on_start to handle the Result and adjust tests
accordingly.
- Around line 323-334: Add regression tests that exercise the new memory-v2
behavior: extend or add tests around MemoryStore::load, handle_write,
handle_read and the underlying set()/compact() logic to cover (1) defaulting a
missing "category" on write/read (ensure writes without category land in the
default category and reads without category find them), (2) preserving an
existing category when calling set() on an existing key (write key with category
A, update value without category and assert category A is kept), (3)
invalid-category handling (attempt writes/reads with an invalid category and
assert the expected error or no-op behavior), and (4) compact() semantics
(create multiple versions across categories, call compact() and assert it
removes expired/older entries per the memory-v2 rules). Reference
MemoryStore::load, handle_write, handle_read, set, and compact in your new
tests.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: abe1b12a-3750-45b5-b36e-f97b71db7137
📒 Files selected for processing (1)
mc/crates/mc-core/src/memory.rs
|
|
||
| /// Set. | ||
| pub fn set(&mut self, key: &str, value: &str) { | ||
| self.set_with_category(key, value, "project"); |
There was a problem hiding this comment.
Preserve the current category when set() updates an existing key.
set_with_category() overwrites f.category on update, so hard-coding "project" here silently reclassifies any existing user/feedback/reference fact written through the legacy API. Only default new keys to "project".
💡 Proposed fix
pub fn set(&mut self, key: &str, value: &str) {
- self.set_with_category(key, value, "project");
+ let category = self
+ .get(key)
+ .map_or_else(|| "project".to_string(), |fact| fact.category.clone());
+ self.set_with_category(key, value, &category);
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self.set_with_category(key, value, "project"); | |
| pub fn set(&mut self, key: &str, value: &str) { | |
| let category = self | |
| .get(key) | |
| .map_or_else(|| "project".to_string(), |fact| fact.category.clone()); | |
| self.set_with_category(key, value, &category); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@mc/crates/mc-core/src/memory.rs` at line 68, In set(), don't hard-code
"project" when delegating to set_with_category; first check whether a fact for
key already exists (e.g. via self.facts.get(key) or the module's existing lookup
helper) and if it does use that fact's f.category when calling
set_with_category(key, value, category), otherwise default to "project"; update
the call site that currently does self.set_with_category(key, value, "project")
to compute category = existing_fact.map(|f|
f.category.clone()).unwrap_or("project".into()) and pass that through to
preserve existing categories on updates.
| let category = input | ||
| .get("category") | ||
| .and_then(|v| v.as_str()) | ||
| .unwrap_or("project"); |
There was a problem hiding this comment.
Reject unknown categories before persisting them.
This accepts any string, but mc/crates/mc-tools/src/spec.rs:117-130 defines memory_write.category as a four-value enum. A typo from the model or a direct caller becomes invalid stored state, and to_prompt_section() will render it as a trailing uncategorized bullet.
💡 Proposed fix
let category = input
.get("category")
.and_then(|v| v.as_str())
.unwrap_or("project");
if key.is_empty() {
return "Error: key is required".into();
}
if let Some(delete) = input.get("delete").and_then(serde_json::Value::as_bool) {
if delete {
return if self.delete(key) {
format!("Deleted: {key}")
} else {
format!("Key not found: {key}")
};
}
}
+ if !matches!(category, "project" | "user" | "feedback" | "reference") {
+ return format!("Error: invalid category: {category}");
+ }
self.set_with_category(key, value, category);Also applies to: 181-182
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@mc/crates/mc-core/src/memory.rs` around lines 165 - 168, The code currently
accepts any string for the local variable `category` (derived from
input.get("category")...) and must reject unknown values before persisting;
replace the unvalidated unwrap_or("project") logic by parsing/validating the
input string against the canonical enum of allowed categories (e.g. via a
FromStr/try_from or explicit match against the four allowed variants) and return
an error (or validation failure) when the value is not one of the permitted
options; apply the same validation fix to the other occurrence of this pattern
around the second instance noted (lines ~181-182) so only valid enum categories
are persisted and invalid inputs are rejected.
| .and_modify(|(idx, ts): &mut (usize, String)| { | ||
| if f.updated_at > *ts { | ||
| *idx = i; | ||
| *ts = f.updated_at.clone(); |
There was a problem hiding this comment.
Break same-second ties toward the later duplicate.
updated_at only has second resolution, so duplicate rows written in the same second compare equal here. With the strict > check, the older entry wins, which breaks the "keep newest" contract of compact().
💡 Proposed fix
.and_modify(|(idx, ts): &mut (usize, String)| {
- if f.updated_at > *ts {
+ if f.updated_at > *ts || (f.updated_at == *ts && i > *idx) {
*idx = i;
*ts = f.updated_at.clone();
}
})📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| .and_modify(|(idx, ts): &mut (usize, String)| { | |
| if f.updated_at > *ts { | |
| *idx = i; | |
| *ts = f.updated_at.clone(); | |
| .and_modify(|(idx, ts): &mut (usize, String)| { | |
| if f.updated_at > *ts || (f.updated_at == *ts && i > *idx) { | |
| *idx = i; | |
| *ts = f.updated_at.clone(); |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@mc/crates/mc-core/src/memory.rs` around lines 207 - 210, In compact(), inside
the .and_modify closure that updates the (idx, ts) tuple, change the comparison
from strict greater (f.updated_at > *ts) to greater-or-equal (f.updated_at >=
*ts) so that when updated_at timestamps are equal (same-second duplicates) the
later duplicate (current index i) wins; update the conditional and leave the
assignments to *idx = i and *ts = f.updated_at.clone() unchanged.
1 participant
Memory v2
Inspired by Claude Code's leaked 4-layer memory architecture.
Changes
Self-skeptical (from Claude Code)
274 tests, 0 fail.
Summary by CodeRabbit
New Features
Improvements
Bug Fixes / Maintenance