From 4580f5dadec3969c2ef37fe0a2ca789ea216ad4c Mon Sep 17 00:00:00 2001 From: Kent Overstreet Date: Thu, 2 Apr 2026 02:32:26 -0400 Subject: [PATCH] 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 --- src/agent/runner.rs | 35 ++++++++--------------------------- 1 file changed, 8 insertions(+), 27 deletions(-) diff --git a/src/agent/runner.rs b/src/agent/runner.rs index 7308504..97546a5 100644 --- a/src/agent/runner.rs +++ b/src/agent/runner.rs @@ -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::(); 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,