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 <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-02 02:52:59 -04:00
parent a0aacfc552
commit 5e781e9ae4

View file

@ -195,8 +195,9 @@ impl Agent {
.map(|e| count(&e.content)).sum(); .map(|e| count(&e.content)).sum();
let mem_tokens: usize = self.context.loaded_nodes.iter() let mem_tokens: usize = self.context.loaded_nodes.iter()
.map(|node| count(&node.render())).sum(); .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(); .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 { self.context_budget = ContextBudget {
identity_tokens: id_tokens, identity_tokens: id_tokens,
@ -485,7 +486,7 @@ impl Agent {
// Re-render the context message so the model sees the updated stack // Re-render the context message so the model sees the updated stack
if !result.starts_with("Error:") { if !result.starts_with("Error:") {
self.refresh_context_message(); self.refresh_context_state();
} }
return; return;
} }
@ -827,14 +828,9 @@ impl Agent {
dbg_log!("[journal] context.journal now has {} entries", self.context.journal.len()); 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). /// Called after any change to context state (working stack, etc).
fn refresh_context_message(&mut self) { fn refresh_context_state(&mut self) {
let rendered = self.context.render_context_message(); self.measure_budget();
// 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);
}
self.publish_context_state(); self.publish_context_state();
self.save_working_stack(); self.save_working_stack();
} }