From 5e781e9ae4ade5c4ce291b62470c304e0321856f Mon Sep 17 00:00:00 2001 From: Kent Overstreet Date: Thu, 2 Apr 2026 02:52:59 -0400 Subject: [PATCH] Fix budget counting: remove stale refresh_context_message refresh_context_message was injecting personality into conversation messages (assuming fixed positions that no longer exist). Replaced with refresh_context_state which just re-measures and publishes. conv_tokens now subtracts mem_tokens since memory tool results are in the conversation message list. Co-Authored-By: Proof of Concept --- src/agent/runner.rs | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/src/agent/runner.rs b/src/agent/runner.rs index 992765a..aee7cdf 100644 --- a/src/agent/runner.rs +++ b/src/agent/runner.rs @@ -195,8 +195,9 @@ impl Agent { .map(|e| count(&e.content)).sum(); let mem_tokens: usize = self.context.loaded_nodes.iter() .map(|node| count(&node.render())).sum(); - let conv_tokens: usize = self.context.messages.iter() + let total_conv: usize = self.context.messages.iter() .map(|m| crate::agent::context::msg_token_count(&self.tokenizer, m)).sum(); + let conv_tokens = total_conv.saturating_sub(mem_tokens); self.context_budget = ContextBudget { identity_tokens: id_tokens, @@ -485,7 +486,7 @@ impl Agent { // Re-render the context message so the model sees the updated stack if !result.starts_with("Error:") { - self.refresh_context_message(); + self.refresh_context_state(); } return; } @@ -827,14 +828,9 @@ impl Agent { dbg_log!("[journal] context.journal now has {} entries", self.context.journal.len()); } - /// Re-render the context message in self.context.messages from live ContextState. /// Called after any change to context state (working stack, etc). - fn refresh_context_message(&mut self) { - let rendered = self.context.render_context_message(); - // The context message is the first user message (index 1, after system prompt) - if self.context.messages.len() >= 2 && self.context.messages[1].role == Role::User { - self.context.messages[1] = Message::user(rendered); - } + fn refresh_context_state(&mut self) { + self.measure_budget(); self.publish_context_state(); self.save_working_stack(); }