kent/consciousness
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Unify budget and context state — single source of truth

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

74
src/agent/context.rs
View file

@ -230,34 +230,9 @@ pub struct ContextState {
/// Conversation entries — messages and memory, interleaved in order.
/// Does NOT include system prompt, personality, or journal.
pub entries: Vec<ConversationEntry>,
/// Cached token budget — recomputed when entries change, not every frame.
pub budget: ContextBudget,
}
impl ContextState {
/// Compute the context budget from typed sources and cache the result.
pub fn recompute_budget(&mut self, count_str: &dyn Fn(&str) -> usize,
count_msg: &dyn Fn(&Message) -> usize,
window_tokens: usize) -> &ContextBudget {
let id = count_str(&self.system_prompt)
+ self.personality.iter().map(|(_, c)| count_str(c)).sum::<usize>();
let jnl: usize = self.journal.iter().map(|e| count_str(&e.content)).sum();
let mut mem = 0;
let mut conv = 0;
for entry in &self.entries {
let tokens = count_msg(entry.api_message());
if entry.is_memory() { mem += tokens } else { conv += tokens }
}
self.budget = ContextBudget {
identity_tokens: id,
memory_tokens: mem,
journal_tokens: jnl,
conversation_tokens: conv,
window_tokens,
};
&self.budget
}
pub fn render_context_message(&self) -> String {
let mut parts: Vec<String> = self.personality.iter()
.map(|(name, content)| format!("## {}\n\n{}", name, content))
@ -281,28 +256,33 @@ impl ContextState {
}
}
#[derive(Debug, Clone, Default)]
pub struct ContextBudget {
pub identity_tokens: usize,
pub memory_tokens: usize,
pub journal_tokens: usize,
pub conversation_tokens: usize,
pub window_tokens: usize,
/// Total tokens used across all context sections.
pub fn sections_used(sections: &[ContextSection]) -> usize {
sections.iter().map(|s| s.tokens).sum()
}
impl ContextBudget {
pub fn used(&self) -> usize {
self.identity_tokens + self.memory_tokens + self.journal_tokens + self.conversation_tokens
}
pub fn free(&self) -> usize {
self.window_tokens.saturating_sub(self.used())
}
pub fn status_string(&self) -> String {
let total = self.window_tokens;
if total == 0 { return String::new(); }
let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / total).max(1) };
format!("id:{}% mem:{}% jnl:{}% conv:{}% free:{}%",
pct(self.identity_tokens), pct(self.memory_tokens),
pct(self.journal_tokens), pct(self.conversation_tokens), pct(self.free()))
}
/// Budget status string derived from context sections.
pub fn sections_budget_string(sections: &[ContextSection]) -> String {
let window = context_window();
if window == 0 { return String::new(); }
let used: usize = sections.iter().map(|s| s.tokens).sum();
let free = window.saturating_sub(used);
let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / window).max(1) };
let parts: Vec<String> = sections.iter()
.map(|s| {
// Short label from section name
let label = match s.name.as_str() {
n if n.starts_with("System") => "sys",
n if n.starts_with("Personality") => "id",
n if n.starts_with("Journal") => "jnl",
n if n.starts_with("Working") => "stack",
n if n.starts_with("Memory") => "mem",
n if n.starts_with("Conversation") => "conv",
_ => return String::new(),
};
format!("{}:{}%", label, pct(s.tokens))
})
.filter(|s| !s.is_empty())
.collect();
format!("{} free:{}%", parts.join(" "), pct(free))
}

30
src/agent/mod.rs
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
}