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

@ -81,8 +81,6 @@ impl Clone for MindState {
/// What should happen after a state transition.
pub enum MindCommand {
/// Start a turn with this input
Turn(String, StreamTarget),
/// Run compaction check
Compact,
/// Run memory scoring
@ -113,10 +111,12 @@ impl MindState {
}
}
/// Consume pending input, return a Turn command if ready.
fn take_pending_input(&mut self) -> MindCommand {
/// Consume pending user input if no turn is active.
/// Returns the text to send; caller is responsible for pushing it
/// into the Agent's context and starting the turn.
fn take_pending_input(&mut self) -> Option<String> {
if self.turn_active || self.input.is_empty() {
return MindCommand::None;
return None;
}
let text = self.input.join("\n");
self.input.clear();
@ -124,7 +124,7 @@ impl MindState {
self.consecutive_errors = 0;
self.last_user_input = Instant::now();
self.dmn = dmn::State::Engaged;
MindCommand::Turn(text, StreamTarget::Conversation)
Some(text)
}
/// Process turn completion, return model switch name if requested.
@ -158,17 +158,17 @@ impl MindState {
}
}
/// DMN tick — returns a Turn action with the DMN prompt, or None.
fn dmn_tick(&mut self) -> MindCommand {
/// DMN tick — returns a prompt and target if we should run a turn.
fn dmn_tick(&mut self) -> Option<(String, StreamTarget)> {
if matches!(self.dmn, dmn::State::Paused | dmn::State::Off) {
return MindCommand::None;
return None;
}
self.dmn_turns += 1;
if self.dmn_turns > self.max_dmn_turns {
self.dmn = dmn::State::Resting { since: Instant::now() };
self.dmn_turns = 0;
return MindCommand::None;
return None;
}
let dmn_ctx = dmn::DmnContext {
@ -177,7 +177,7 @@ impl MindState {
last_turn_had_tools: self.last_turn_had_tools,
};
let prompt = self.dmn.prompt(&dmn_ctx);
MindCommand::Turn(prompt, StreamTarget::Autonomous)
Some((prompt, StreamTarget::Autonomous))
}
fn interrupt(&mut self) {
@ -261,9 +261,10 @@ impl Mind {
match cmd {
MindCommand::None => {}
MindCommand::Compact => {
let threshold = compaction_threshold(&self.config.app);
let threshold = compaction_threshold(&self.config.app) as usize;
let mut ag = self.agent.lock().await;
if ag.last_prompt_tokens() > threshold {
let sections = ag.shared_context.read().map(|s| s.clone()).unwrap_or_default();
if crate::agent::context::sections_used(&sections) > threshold {
ag.compact();
ag.notify("compacted");
}
@ -305,16 +306,6 @@ impl Mind {
new_log, shared_ctx, shared_tools,
);
}
MindCommand::Turn(input, target) => {
self.shared.lock().unwrap().turn_active = true;
let _ = self.turn_watch.send(true);
let agent = self.agent.clone();
let result_tx = self.turn_tx.clone();
self.shared.lock().unwrap().turn_handle = Some(tokio::spawn(async move {
let result = Agent::turn(agent, &input).await;
let _ = result_tx.send((result, target)).await;
}));
}
}
}
}
@ -344,6 +335,25 @@ impl Mind {
});
}
/// Push user/DMN message into the agent's context and spawn a turn.
/// The text moves from pending_input to ContextState atomically —
/// by the time this returns, the message is in context and the turn
/// is running.
async fn start_turn(&self, text: &str, target: StreamTarget) {
{
let mut ag = self.agent.lock().await;
ag.push_message(crate::agent::api::types::Message::user(text));
}
self.shared.lock().unwrap().turn_active = true;
let _ = self.turn_watch.send(true);
let agent = self.agent.clone();
let result_tx = self.turn_tx.clone();
self.shared.lock().unwrap().turn_handle = Some(tokio::spawn(async move {
let result = Agent::turn(agent).await;
let _ = result_tx.send((result, target)).await;
}));
}
pub async fn shutdown(&self) {
if let Some(handle) = self.shared.lock().unwrap().turn_handle.take() { handle.abort(); }
}
@ -404,12 +414,17 @@ impl Mind {
_ = tokio::time::sleep(timeout), if !turn_active => {
let tick = self.shared.lock().unwrap().dmn_tick();
cmds.push(tick);
if let Some((prompt, target)) = tick {
self.start_turn(&prompt, target).await;
}
}
}
// Always check for pending input
cmds.push(self.shared.lock().unwrap().take_pending_input());
// Check for pending user input → push to agent context and start turn
let pending = self.shared.lock().unwrap().take_pending_input();
if let Some(text) = pending {
self.start_turn(&text, StreamTarget::Conversation).await;
}
self.run_commands(cmds).await;
}