Unify budget and context state — single source of truth

Kill ContextBudget and recompute_budget entirely. Budget percentages,
used token counts, and compaction threshold checks now all derive from
the ContextSection tree built by context_state_summary(). This
eliminates the stale-budget bug where the cached budget diverged from
actual context contents.

Also: remove MindCommand::Turn — user input flows through
shared_mind.input exclusively. Mind::start_turn() atomically moves
text from pending input into the agent's context and spawns the turn.
Kill /retry. Make Agent::turn() take no input parameter.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-06 20:34:51 -04:00
parent f63c341f94
commit c22b8c3a6f
4 changed files with 83 additions and 117 deletions

View file

@ -24,7 +24,7 @@ use tiktoken_rs::CoreBPE;
use api::{ApiClient, ToolCall};
use api::types::{ContentPart, Message, MessageContent, Role};
use context::{ConversationEntry, ContextState, ContextBudget};
use context::{ConversationEntry, ContextState};
use tools::{summarize_args, working_stack};
use crate::mind::log::ConversationLog;
@ -214,7 +214,6 @@ impl Agent {
journal: Vec::new(),
working_stack: Vec::new(),
entries: Vec::new(),
budget: ContextBudget::default(),
};
let session_id = format!("consciousness-{}", chrono::Utc::now().format("%Y%m%d-%H%M%S"));
let agent_cycles = crate::subconscious::subconscious::AgentCycleState::new(&session_id);
@ -284,7 +283,7 @@ impl Agent {
}
/// Push a conversation message — stamped and logged.
fn push_message(&mut self, mut msg: Message) {
pub fn push_message(&mut self, mut msg: Message) {
msg.stamp();
let entry = ConversationEntry::Message(msg);
self.push_entry(entry);
@ -297,7 +296,7 @@ impl Agent {
}
}
self.context.entries.push(entry);
self.recompute_budget();
self.changed.notify_one();
}
@ -319,12 +318,6 @@ impl Agent {
self.changed.notify_one();
}
pub fn recompute_budget(&mut self) -> &ContextBudget {
let count_str = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
let count_msg = |m: &Message| crate::agent::context::msg_token_count(&self.tokenizer, m);
let window = crate::agent::context::context_window();
self.context.recompute_budget(&count_str, &count_msg, window)
}
/// Send a user message and run the agent loop until the model
/// produces a text response (no more tool calls). Streams text
@ -334,7 +327,6 @@ impl Agent {
/// lock is never held across I/O (API streaming, tool dispatch).
pub async fn turn(
agent: Arc<tokio::sync::Mutex<Agent>>,
user_input: &str,
) -> Result<TurnResult> {
// --- Pre-loop setup (lock 1): agent cycle, memories, user input ---
let active_tools = {
@ -385,14 +377,13 @@ impl Agent {
}
}
// Re-acquire to apply results and push user input
{
// Re-acquire to apply background tool results
if !bg_results.is_empty() {
let mut me = agent.lock().await;
let mut bg_ds = DispatchState::new();
for (call, output) in bg_results {
me.apply_tool_result(&call, output, &mut bg_ds);
}
me.push_message(Message::user(user_input));
}
tools
@ -1041,13 +1032,14 @@ impl Agent {
dbglog!("[compact] entries: {} → {} (mem: {} → {}, conv: {} → {})",
before, after, before_mem, after_mem, before_conv, after_conv);
self.recompute_budget();
dbglog!("[compact] budget: {}", self.context.budget.status_string());
self.load_startup_journal();
self.generation += 1;
self.last_prompt_tokens = 0;
self.publish_context_state();
let sections = self.shared_context.read().map(|s| s.clone()).unwrap_or_default();
dbglog!("[compact] budget: {}", context::sections_budget_string(&sections));
}
/// Restore from the conversation log. Builds the context window
@ -1073,9 +1065,9 @@ impl Agent {
all.len(), mem_count, conv_count);
self.context.entries = all;
self.compact();
// Estimate prompt tokens from budget so status bar isn't 0 on startup
self.recompute_budget();
self.last_prompt_tokens = self.context.budget.used() as u32;
// Estimate prompt tokens from sections so status bar isn't 0 on startup
let sections = self.shared_context.read().map(|s| s.clone()).unwrap_or_default();
self.last_prompt_tokens = context::sections_used(&sections) as u32;
true
}