kent/consciousness
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Split context_state_summary: ContextBudget for compaction, UI-only for display

Browse source 
context_state_summary() was used for both compaction decisions (just
needs token counts) and debug screen display (needs full tree with
labels). Split into:

- Agent::context_budget() -> ContextBudget: cheap token counting by
  category, used by compact(), restore_from_log(), mind event loop
- ContextBudget::format(): replaces sections_budget_string() which
  fragily pattern-matched on section name strings
- context_state_summary(): now UI-only, formatting code stays here

Also extracted entry_sections() as shared helper with include_memories
param — false for context_state_summary (memories have own section),
true for conversation_sections_from() (subconscious screen shows all).

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-07 19:02:58 -04:00
parent 9e49398689
commit cf1c64f936
4 changed files with 130 additions and 86 deletions
Download patch file Download diff file Expand all files Collapse all files

50
src/agent/context.rs
View file

@ -336,33 +336,29 @@ impl ContextState {
}
}
/// Total tokens used across all context sections.
pub fn sections_used(sections: &[ContextSection]) -> usize {
sections.iter().map(|s| s.tokens).sum()
/// Token budget per context category — cheap to compute, no formatting.
pub struct ContextBudget {
pub system: usize,
pub identity: usize,
pub journal: usize,
pub memory: usize,
pub conversation: usize,
}
/// Budget status string derived from context sections.
pub fn sections_budget_string(sections: &[ContextSection]) -> String {
let window = context_window();
if window == 0 { return String::new(); }
let used: usize = sections.iter().map(|s| s.tokens).sum();
let free = window.saturating_sub(used);
let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / window).max(1) };
let parts: Vec<String> = sections.iter()
.map(|s| {
// Short label from section name
let label = match s.name.as_str() {
n if n.starts_with("System") => "sys",
n if n.starts_with("Personality") => "id",
n if n.starts_with("Journal") => "jnl",
n if n.starts_with("Working") => "stack",
n if n.starts_with("Memory") => "mem",
n if n.starts_with("Conversation") => "conv",
_ => return String::new(),
};
format!("{}:{}%", label, pct(s.tokens))
})
.filter(|s| !s.is_empty())
.collect();
format!("{} free:{}%", parts.join(" "), pct(free))
impl ContextBudget {
pub fn total(&self) -> usize {
self.system + self.identity + self.journal + self.memory + self.conversation
}
pub fn format(&self) -> String {
let window = context_window();
if window == 0 { return String::new(); }
let used = self.total();
let free = window.saturating_sub(used);
let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / window).max(1) };
format!("sys:{}% id:{}% jnl:{}% mem:{}% conv:{}% free:{}%",
pct(self.system), pct(self.identity), pct(self.journal),
pct(self.memory), pct(self.conversation), pct(free))
}
}

145
src/agent/mod.rs
View file

@ -655,6 +655,28 @@ impl Agent {
self.push_message(Message::tool_result(&call.id, &output));
}
/// Token budget by category — cheap, no formatting. Used for compaction decisions.
pub fn context_budget(&self) -> context::ContextBudget {
let count = |m: &Message| context::msg_token_count(&self.tokenizer, m);
let system = count(&Message::system(&self.context.system_prompt));
let identity = count(&Message::user(&self.context.render_context_message()));
let journal_rendered = context::render_journal(&self.context.journal);
let journal = if journal_rendered.is_empty() { 0 } else {
count(&Message::user(&journal_rendered))
};
let memory: usize = self.context.entries.iter()
.filter(|e| e.is_memory())
.map(|e| count(e.message()))
.sum();
let conversation: usize = self.context.entries.iter()
.filter(|e| !e.is_memory() && !e.is_log())
.map(|e| count(e.api_message()))
.sum();
context::ContextBudget { system, identity, journal, memory, conversation }
}
/// Build context state summary for the debug screen.
pub fn context_state_summary(&self) -> Vec<ContextSection> {
let count_msg = |m: &Message| context::msg_token_count(&self.tokenizer, m);
@ -723,52 +745,8 @@ impl Agent {
});
}
// Conversation — non-memory entries only (memories counted above)
let conv_children: Vec<ContextSection> = self.context.entries.iter().enumerate()
.filter(|(_, e)| !e.is_memory() && !e.is_log())
.map(|(i, entry)| {
let m = entry.message();
let text = m.content.as_ref()
.map(|c| c.as_text().to_string())
.unwrap_or_default();
let tool_info = m.tool_calls.as_ref().map(|tc| {
tc.iter()
.map(|c| c.function.name.clone())
.collect::<Vec<_>>()
.join(", ")
});
let label = if entry.is_memory() {
if let ConversationEntry::Memory { key, .. } = entry {
format!("[memory: {}]", key)
} else { unreachable!() }
} else {
match &tool_info {
Some(tools) => format!("[tool_call: {}]", tools),
None => {
let preview: String = text.chars().take(60).collect();
let preview = preview.replace('\n', " ");
if text.len() > 60 { format!("{}...", preview) } else { preview }
}
}
};
let tokens = count_msg(entry.api_message());
let cfg = crate::config::get();
let role_name = if entry.is_memory() { "mem".to_string() } else {
match m.role {
Role::Assistant => cfg.assistant_name.clone(),
Role::User => cfg.user_name.clone(),
Role::Tool => "tool".to_string(),
Role::System => "system".to_string(),
}
};
ContextSection {
name: format!("[{}] {}: {}", i, role_name, label),
tokens,
content: text,
children: Vec::new(),
}
})
.collect();
// Conversation — memories excluded (counted in their own section above)
let conv_children = self.entry_sections(&count_msg, 0, false);
let conv_tokens: usize = conv_children.iter().map(|c| c.tokens).sum();
sections.push(ContextSection {
name: format!("Conversation ({} messages)", conv_children.len()),
@ -780,6 +758,72 @@ impl Agent {
sections
}
/// Build ContextSection nodes for conversation entries starting at `from`.
/// When `include_memories` is false, memory entries are excluded (they get
/// their own section in context_state_summary to avoid double-counting).
fn entry_sections(
&self,
count_msg: &dyn Fn(&Message) -> usize,
from: usize,
include_memories: bool,
) -> Vec<ContextSection> {
let cfg = crate::config::get();
self.context.entries.iter().enumerate()
.skip(from)
.filter(|(_, e)| !e.is_log() && (include_memories || !e.is_memory()))
.map(|(i, entry)| {
let m = entry.message();
let text = m.content.as_ref()
.map(|c| c.as_text().to_string())
.unwrap_or_default();
let (role_name, label) = if let ConversationEntry::Memory { key, score, .. } = entry {
let label = match score {
Some(s) => format!("[memory: {} score:{:.1}]", key, s),
None => format!("[memory: {}]", key),
};
("mem".to_string(), label)
} else {
let tool_info = m.tool_calls.as_ref().map(|tc| {
tc.iter()
.map(|c| c.function.name.clone())
.collect::<Vec<_>>()
.join(", ")
});
let label = match &tool_info {
Some(tools) => format!("[tool_call: {}]", tools),
None => {
let preview: String = text.chars().take(60).collect();
let preview = preview.replace('\n', " ");
if text.len() > 60 { format!("{}...", preview) } else { preview }
}
};
let role_name = match m.role {
Role::Assistant => cfg.assistant_name.clone(),
Role::User => cfg.user_name.clone(),
Role::Tool => "tool".to_string(),
Role::System => "system".to_string(),
};
(role_name, label)
};
ContextSection {
name: format!("[{}] {}: {}", i, role_name, label),
tokens: count_msg(entry.api_message()),
content: text,
children: Vec::new(),
}
})
.collect()
}
/// Context sections for entries from `from` onward — used by the
/// subconscious debug screen to show forked agent conversations.
pub fn conversation_sections_from(&self, from: usize) -> Vec<ContextSection> {
let count_msg = |m: &Message| context::msg_token_count(&self.tokenizer, m);
self.entry_sections(&count_msg, from, true)
}
/// Load recent journal entries at startup for orientation.
/// Uses the same budget logic as compaction but with empty conversation.
/// Only parses the tail of the journal file (last 64KB) for speed.
@ -959,8 +1003,8 @@ impl Agent {
self.generation += 1;
self.last_prompt_tokens = 0;
let sections = self.context_state_summary();
dbglog!("[compact] budget: {}", context::sections_budget_string(&sections));
let budget = self.context_budget();
dbglog!("[compact] budget: {}", budget.format());
}
/// Restore from the conversation log. Builds the context window
@ -986,9 +1030,8 @@ impl Agent {
all.len(), mem_count, conv_count);
self.context.entries = all;
self.compact();
// Estimate prompt tokens from sections so status bar isn't 0 on startup
self.last_prompt_tokens = context::sections_used(
&self.context_state_summary()) as u32;
// Estimate prompt tokens so status bar isn't 0 on startup
self.last_prompt_tokens = self.context_budget().total() as u32;
true
}