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();
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();
}