measure_budget: count from typed sources, not message scanning

Identity tokens from system_prompt + personality vec. Journal
from journal entries vec. Memory from loaded_nodes. Conversation
is the remainder. No string prefix matching.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-02 02:32:26 -04:00
parent 4bdc7ae112
commit 4580f5dade

View file

@ -185,35 +185,16 @@ impl Agent {
fn measure_budget(&mut self) {
let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
let mut id_tokens: usize = 0;
let id_tokens = count(&self.context.system_prompt)
+ self.context.personality.iter()
.map(|(_, content)| count(content)).sum::<usize>();
let jnl_tokens: usize = self.context.journal.iter()
.map(|e| count(&e.content))
.sum();
let mut conv_tokens: usize = 0;
let mut in_conversation = false;
for msg in &self.messages {
let tokens = crate::agent::context::msg_token_count(&self.tokenizer, msg);
if in_conversation {
conv_tokens += tokens;
continue;
}
if in_conversation {
conv_tokens += tokens;
} else if msg.role == Role::System || (!in_conversation && conv_tokens == 0) {
id_tokens += tokens;
} else {
in_conversation = true;
conv_tokens += tokens;
}
}
// Memory = nodes loaded during the session via tool calls
.map(|e| count(&e.content)).sum();
let mem_tokens: usize = self.context.loaded_nodes.iter()
.map(|node| count(&node.render()))
.sum();
.map(|node| count(&node.render())).sum();
let total: usize = self.messages.iter()
.map(|m| crate::agent::context::msg_token_count(&self.tokenizer, m)).sum();
let conv_tokens = total.saturating_sub(id_tokens + jnl_tokens);
self.context_budget = ContextBudget {
identity_tokens: id_tokens,