Unify budget and context state — single source of truth

Kill ContextBudget and recompute_budget entirely. Budget percentages,
used token counts, and compaction threshold checks now all derive from
the ContextSection tree built by context_state_summary(). This
eliminates the stale-budget bug where the cached budget diverged from
actual context contents.

Also: remove MindCommand::Turn — user input flows through
shared_mind.input exclusively. Mind::start_turn() atomically moves
text from pending input into the agent's context and spawns the turn.
Kill /retry. Make Agent::turn() take no input parameter.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-06 20:34:51 -04:00
parent f63c341f94
commit c22b8c3a6f
4 changed files with 83 additions and 117 deletions

View file

@ -230,34 +230,9 @@ pub struct ContextState {
/// Conversation entries — messages and memory, interleaved in order.
/// Does NOT include system prompt, personality, or journal.
pub entries: Vec<ConversationEntry>,
/// Cached token budget — recomputed when entries change, not every frame.
pub budget: ContextBudget,
}
impl ContextState {
/// Compute the context budget from typed sources and cache the result.
pub fn recompute_budget(&mut self, count_str: &dyn Fn(&str) -> usize,
count_msg: &dyn Fn(&Message) -> usize,
window_tokens: usize) -> &ContextBudget {
let id = count_str(&self.system_prompt)
+ self.personality.iter().map(|(_, c)| count_str(c)).sum::<usize>();
let jnl: usize = self.journal.iter().map(|e| count_str(&e.content)).sum();
let mut mem = 0;
let mut conv = 0;
for entry in &self.entries {
let tokens = count_msg(entry.api_message());
if entry.is_memory() { mem += tokens } else { conv += tokens }
}
self.budget = ContextBudget {
identity_tokens: id,
memory_tokens: mem,
journal_tokens: jnl,
conversation_tokens: conv,
window_tokens,
};
&self.budget
}
pub fn render_context_message(&self) -> String {
let mut parts: Vec<String> = self.personality.iter()
.map(|(name, content)| format!("## {}\n\n{}", name, content))
@ -281,28 +256,33 @@ impl ContextState {
}
}
#[derive(Debug, Clone, Default)]
pub struct ContextBudget {
pub identity_tokens: usize,
pub memory_tokens: usize,
pub journal_tokens: usize,
pub conversation_tokens: usize,
pub window_tokens: usize,
/// Total tokens used across all context sections.
pub fn sections_used(sections: &[ContextSection]) -> usize {
sections.iter().map(|s| s.tokens).sum()
}
impl ContextBudget {
pub fn used(&self) -> usize {
self.identity_tokens + self.memory_tokens + self.journal_tokens + self.conversation_tokens
}
pub fn free(&self) -> usize {
self.window_tokens.saturating_sub(self.used())
}
pub fn status_string(&self) -> String {
let total = self.window_tokens;
if total == 0 { return String::new(); }
let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / total).max(1) };
format!("id:{}% mem:{}% jnl:{}% conv:{}% free:{}%",
pct(self.identity_tokens), pct(self.memory_tokens),
pct(self.journal_tokens), pct(self.conversation_tokens), pct(self.free()))
}
/// 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))
}