Restore trim_conversation: dedup memories, evict to budget, snap boundary

Ported the old trim_entries logic to the new AstNode types:
- Phase 1: Dedup Memory nodes by key (keep last), drop DMN entries
- Phase 2: While over budget, evict lowest-scored memory (if memories
  > 50% of conv tokens) or oldest conversation entry
- Phase 3: Snap to User message boundary at start

Called from compact() which runs on startup and on /compact.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-08 21:20:50 -04:00
parent 7237baba11
commit bbffc2213e
2 changed files with 84 additions and 0 deletions

View file

@ -795,6 +795,88 @@ impl ContextState {
self.section_mut(section).clear();
}
/// Dedup and trim conversation entries to fit within the context budget.
///
/// Phase 1: Drop duplicate memories (keep last) and DMN entries.
/// Phase 2: While over budget, drop lowest-scored memory (if memories
/// are > 50% of conversation tokens) or oldest conversation entry.
/// Phase 3: Snap to user message boundary at start.
pub fn trim_conversation(&mut self) {
let max_tokens = context_budget_tokens();
let fixed = self.system.iter().map(|n| n.tokens()).sum::<usize>()
+ self.identity.iter().map(|n| n.tokens()).sum::<usize>()
+ self.journal.iter().map(|n| n.tokens()).sum::<usize>();
// Phase 1: dedup memories by key (keep last), drop DMN
let mut seen_keys: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
let mut drop = std::collections::HashSet::new();
for (i, node) in self.conversation.iter().enumerate() {
if let AstNode::Leaf(leaf) = node {
match leaf.body() {
NodeBody::Dmn(_) => { drop.insert(i); }
NodeBody::Memory { key, .. } => {
if let Some(prev) = seen_keys.insert(key.clone(), i) {
drop.insert(prev);
}
}
_ => {}
}
}
}
if !drop.is_empty() {
let mut i = 0;
self.conversation.retain(|_| { let keep = !drop.contains(&i); i += 1; keep });
}
// Phase 2: while over budget, evict
loop {
let total: usize = self.conversation.iter().map(|n| n.tokens()).sum();
if fixed + total <= max_tokens { break; }
let mt: usize = self.conversation.iter()
.filter(|n| matches!(n, AstNode::Leaf(l) if matches!(l.body(), NodeBody::Memory { .. })))
.map(|n| n.tokens()).sum();
let ct = total - mt;
if mt > ct {
// Memories > 50% — drop lowest-scored
if let Some(i) = self.lowest_scored_memory() {
self.conversation.remove(i);
continue;
}
}
// Drop oldest non-memory entry
if let Some(i) = self.conversation.iter().position(|n|
!matches!(n, AstNode::Leaf(l) if matches!(l.body(), NodeBody::Memory { .. })))
{
self.conversation.remove(i);
} else {
break;
}
}
// Phase 3: snap to user message boundary
while let Some(first) = self.conversation.first() {
if matches!(first, AstNode::Branch { role: Role::User, .. }) { break; }
self.conversation.remove(0);
}
}
fn lowest_scored_memory(&self) -> Option<usize> {
self.conversation.iter().enumerate()
.filter_map(|(i, n)| {
if let AstNode::Leaf(l) = n {
if let NodeBody::Memory { score: Some(s), .. } = l.body() {
return Some((i, *s));
}
}
None
})
.min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
.map(|(i, _)| i)
}
/// Push a child node into a branch at `index` in `section`.
pub fn push_child(&mut self, section: Section, index: usize, child: AstNode) {
let node = &mut self.section_mut(section)[index];

View file

@ -555,6 +555,8 @@ impl Agent {
self.load_startup_journal().await;
self.context.lock().await.trim_conversation();
let mut st = self.state.lock().await;
st.generation += 1;
st.last_prompt_tokens = 0;