WIP: ContextEntry/ContextSection data structures for incremental token counting

New types — not yet wired to callers:

- ContextEntry: wraps ConversationEntry with cached token count and
  timestamp
- ContextSection: named group of entries with cached token total.
  Private entries/tokens, read via entries()/tokens().
  Mutation via push(entry), set(index, entry), del(index).
- ContextState: system/identity/journal/conversation sections + working_stack
- ConversationEntry::System variant for system prompt entries

Token counting happens once at push time. Sections maintain their
totals incrementally via push/set/del. No more recomputing from
scratch on every budget check.

Does not compile — callers need updating.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-07 20:15:31 -04:00
parent 776ac527f1
commit 62996e27d7
10 changed files with 450 additions and 403 deletions

View file

@ -518,12 +518,12 @@ impl Subconscious {
pub async fn trigger(&mut self, agent: &Arc<tokio::sync::Mutex<Agent>>) {
let (conversation_bytes, memory_keys) = {
let ag = agent.lock().await;
let bytes = ag.context.entries.iter()
.filter(|e| !e.is_log() && !e.is_memory())
.map(|e| e.message().content_text().len() as u64)
let bytes = ag.context.conversation.entries().iter()
.filter(|ce| !ce.entry.is_log() && !ce.entry.is_memory())
.map(|ce| ce.entry.message().content_text().len() as u64)
.sum::<u64>();
let keys: Vec<String> = ag.context.entries.iter().filter_map(|e| {
if let ConversationEntry::Memory { key, .. } = e {
let keys: Vec<String> = ag.context.conversation.entries().iter().filter_map(|ce| {
if let ConversationEntry::Memory { key, .. } = &ce.entry {
Some(key.clone())
} else { None }
}).collect();
@ -550,7 +550,7 @@ impl Subconscious {
let mut forked = conscious.fork(auto.tools.clone());
forked.provenance = format!("agent:{}", auto.name);
let fork_point = forked.context.entries.len();
let fork_point = forked.context.conversation.len();
let shared_forked = Arc::new(tokio::sync::Mutex::new(forked));
self.agents[idx].forked_agent = Some(shared_forked.clone());

View file

@ -31,7 +31,9 @@ pub use dmn::{SubconsciousSnapshot, Subconscious};
use crate::agent::context::ConversationEntry;
/// Load persisted memory scores from disk and apply to Memory entries.
fn load_memory_scores(entries: &mut [ConversationEntry], path: &std::path::Path) {
use crate::agent::context::ContextSection;
fn load_memory_scores(section: &mut ContextSection, path: &std::path::Path) {
let data = match std::fs::read_to_string(path) {
Ok(d) => d,
Err(_) => return,
@ -41,10 +43,10 @@ fn load_memory_scores(entries: &mut [ConversationEntry], path: &std::path::Path)
Err(_) => return,
};
let mut applied = 0;
for entry in entries.iter_mut() {
if let ConversationEntry::Memory { key, score, .. } = entry {
for i in 0..section.len() {
if let ConversationEntry::Memory { key, .. } = &section.entries()[i].entry {
if let Some(&s) = scores.get(key.as_str()) {
*score = Some(s);
section.set_score(i, Some(s));
applied += 1;
}
}
@ -55,10 +57,10 @@ fn load_memory_scores(entries: &mut [ConversationEntry], path: &std::path::Path)
}
/// Save all memory scores to disk.
fn save_memory_scores(entries: &[ConversationEntry], path: &std::path::Path) {
let scores: std::collections::BTreeMap<String, f64> = entries.iter()
.filter_map(|e| {
if let ConversationEntry::Memory { key, score: Some(s), .. } = e {
fn save_memory_scores(section: &ContextSection, path: &std::path::Path) {
let scores: std::collections::BTreeMap<String, f64> = section.entries().iter()
.filter_map(|ce| {
if let ConversationEntry::Memory { key, score: Some(s), .. } = &ce.entry {
Some((key.clone(), *s))
} else {
None
@ -313,7 +315,7 @@ impl Mind {
// Restore persisted memory scores
let scores_path = self.config.session_dir.join("memory-scores.json");
load_memory_scores(&mut ag.context.entries, &scores_path);
load_memory_scores(&mut ag.context.conversation, &scores_path);
ag.changed.notify_one();
drop(ag);
@ -403,16 +405,16 @@ impl Mind {
if let Ok(ref scores) = result {
// Write scores onto Memory entries
for (key, weight) in scores {
for entry in &mut ag.context.entries {
if let crate::agent::context::ConversationEntry::Memory {
key: k, score, ..
} = entry {
if k == key { *score = Some(*weight); }
for i in 0..ag.context.conversation.len() {
if let ConversationEntry::Memory { key: k, .. } = &ag.context.conversation.entries()[i].entry {
if k == key {
ag.context.conversation.set_score(i, Some(*weight));
}
}
}
}
// Persist all scores to disk
save_memory_scores(&ag.context.entries, &scores_path);
save_memory_scores(&ag.context.conversation, &scores_path);
}
}
let _ = bg_tx.send(BgEvent::ScoringDone);