diff --git a/src/agent/context.rs b/src/agent/context.rs index 79ca030..a571203 100644 --- a/src/agent/context.rs +++ b/src/agent/context.rs @@ -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::() + + self.identity.iter().map(|n| n.tokens()).sum::() + + self.journal.iter().map(|n| n.tokens()).sum::(); + + // Phase 1: dedup memories by key (keep last), drop DMN + let mut seen_keys: std::collections::HashMap = 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 { + 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]; diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 8390ec4..b9b1d9b 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -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;