budget: count personality + loaded nodes as memory tokens

mem% was always 0 because memory_tokens was hardcoded to 0. Now
counts personality context + loaded nodes from memory tool calls.
Also calls measure_budget + publish_context_state after memory tool
dispatch so the debug screen updates immediately.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
ProofOfConcept 2026-03-25 02:27:25 -04:00
parent a865285313
commit 79672cbe53

View file

@ -169,8 +169,9 @@ impl Agent {
/// Measure context window usage by category. Uses the BPE tokenizer
/// for direct token counting (no chars/4 approximation).
fn measure_budget(&mut self) {
let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
let mut id_tokens: usize = 0;
let mem_tokens: usize = 0;
let mut jnl_tokens: usize = 0;
let mut conv_tokens: usize = 0;
let mut in_conversation = false;
@ -192,8 +193,8 @@ impl Agent {
} else if text.starts_with("Your context was just rebuilt") {
jnl_tokens += tokens;
} else if jnl_tokens == 0 && conv_tokens == 0 {
// Static identity context (before any journal/conversation)
id_tokens += tokens;
// First user message is personality/memory context
// Count it as memory, not identity
} else {
in_conversation = true;
conv_tokens += tokens;
@ -206,6 +207,14 @@ impl Agent {
}
}
// Memory = personality context + loaded nodes from tool calls
let mut mem_tokens: usize = self.context.personality.iter()
.map(|(_, content)| count(content))
.sum();
for node in &self.context.loaded_nodes {
mem_tokens += count(&node.render());
}
self.context_budget = ContextBudget {
identity_tokens: id_tokens,
memory_tokens: mem_tokens,
@ -487,6 +496,8 @@ impl Agent {
if output.text.starts_with("Error:") {
ds.tool_errors += 1;
}
self.measure_budget();
self.publish_context_state();
return;
}