From cf1c64f9361e2f5660ce3e77cc3054b8ef041ce8 Mon Sep 17 00:00:00 2001 From: Kent Overstreet Date: Tue, 7 Apr 2026 19:02:58 -0400 Subject: [PATCH] Split context_state_summary: ContextBudget for compaction, UI-only for display MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit context_state_summary() was used for both compaction decisions (just needs token counts) and debug screen display (needs full tree with labels). Split into: - Agent::context_budget() -> ContextBudget: cheap token counting by category, used by compact(), restore_from_log(), mind event loop - ContextBudget::format(): replaces sections_budget_string() which fragily pattern-matched on section name strings - context_state_summary(): now UI-only, formatting code stays here Also extracted entry_sections() as shared helper with include_memories param — false for context_state_summary (memories have own section), true for conversation_sections_from() (subconscious screen shows all). Co-Authored-By: Proof of Concept --- src/agent/context.rs | 50 +++++++-------- src/agent/mod.rs | 145 ++++++++++++++++++++++++++++--------------- src/mind/mod.rs | 18 ++++-- src/user/chat.rs | 3 +- 4 files changed, 130 insertions(+), 86 deletions(-) diff --git a/src/agent/context.rs b/src/agent/context.rs index defff60..96ef79d 100644 --- a/src/agent/context.rs +++ b/src/agent/context.rs @@ -336,33 +336,29 @@ impl ContextState { } } -/// Total tokens used across all context sections. -pub fn sections_used(sections: &[ContextSection]) -> usize { - sections.iter().map(|s| s.tokens).sum() +/// Token budget per context category — cheap to compute, no formatting. +pub struct ContextBudget { + pub system: usize, + pub identity: usize, + pub journal: usize, + pub memory: usize, + pub conversation: usize, } -/// 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 = 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)) +impl ContextBudget { + pub fn total(&self) -> usize { + self.system + self.identity + self.journal + self.memory + self.conversation + } + + pub fn format(&self) -> String { + let window = context_window(); + if window == 0 { return String::new(); } + let used = self.total(); + let free = window.saturating_sub(used); + let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / window).max(1) }; + format!("sys:{}% id:{}% jnl:{}% mem:{}% conv:{}% free:{}%", + pct(self.system), pct(self.identity), pct(self.journal), + pct(self.memory), pct(self.conversation), pct(free)) + } } + diff --git a/src/agent/mod.rs b/src/agent/mod.rs index a5cf5a5..a6c6dde 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -655,6 +655,28 @@ impl Agent { self.push_message(Message::tool_result(&call.id, &output)); } + /// Token budget by category — cheap, no formatting. Used for compaction decisions. + pub fn context_budget(&self) -> context::ContextBudget { + let count = |m: &Message| context::msg_token_count(&self.tokenizer, m); + + let system = count(&Message::system(&self.context.system_prompt)); + let identity = count(&Message::user(&self.context.render_context_message())); + let journal_rendered = context::render_journal(&self.context.journal); + let journal = if journal_rendered.is_empty() { 0 } else { + count(&Message::user(&journal_rendered)) + }; + let memory: usize = self.context.entries.iter() + .filter(|e| e.is_memory()) + .map(|e| count(e.message())) + .sum(); + let conversation: usize = self.context.entries.iter() + .filter(|e| !e.is_memory() && !e.is_log()) + .map(|e| count(e.api_message())) + .sum(); + + context::ContextBudget { system, identity, journal, memory, conversation } + } + /// Build context state summary for the debug screen. pub fn context_state_summary(&self) -> Vec { let count_msg = |m: &Message| context::msg_token_count(&self.tokenizer, m); @@ -723,52 +745,8 @@ impl Agent { }); } - // Conversation — non-memory entries only (memories counted above) - let conv_children: Vec = self.context.entries.iter().enumerate() - .filter(|(_, e)| !e.is_memory() && !e.is_log()) - .map(|(i, entry)| { - let m = entry.message(); - let text = m.content.as_ref() - .map(|c| c.as_text().to_string()) - .unwrap_or_default(); - let tool_info = m.tool_calls.as_ref().map(|tc| { - tc.iter() - .map(|c| c.function.name.clone()) - .collect::>() - .join(", ") - }); - let label = if entry.is_memory() { - if let ConversationEntry::Memory { key, .. } = entry { - format!("[memory: {}]", key) - } else { unreachable!() } - } else { - match &tool_info { - Some(tools) => format!("[tool_call: {}]", tools), - None => { - let preview: String = text.chars().take(60).collect(); - let preview = preview.replace('\n', " "); - if text.len() > 60 { format!("{}...", preview) } else { preview } - } - } - }; - let tokens = count_msg(entry.api_message()); - let cfg = crate::config::get(); - let role_name = if entry.is_memory() { "mem".to_string() } else { - match m.role { - Role::Assistant => cfg.assistant_name.clone(), - Role::User => cfg.user_name.clone(), - Role::Tool => "tool".to_string(), - Role::System => "system".to_string(), - } - }; - ContextSection { - name: format!("[{}] {}: {}", i, role_name, label), - tokens, - content: text, - children: Vec::new(), - } - }) - .collect(); + // Conversation — memories excluded (counted in their own section above) + let conv_children = self.entry_sections(&count_msg, 0, false); let conv_tokens: usize = conv_children.iter().map(|c| c.tokens).sum(); sections.push(ContextSection { name: format!("Conversation ({} messages)", conv_children.len()), @@ -780,6 +758,72 @@ impl Agent { sections } + /// Build ContextSection nodes for conversation entries starting at `from`. + /// When `include_memories` is false, memory entries are excluded (they get + /// their own section in context_state_summary to avoid double-counting). + fn entry_sections( + &self, + count_msg: &dyn Fn(&Message) -> usize, + from: usize, + include_memories: bool, + ) -> Vec { + let cfg = crate::config::get(); + self.context.entries.iter().enumerate() + .skip(from) + .filter(|(_, e)| !e.is_log() && (include_memories || !e.is_memory())) + .map(|(i, entry)| { + let m = entry.message(); + let text = m.content.as_ref() + .map(|c| c.as_text().to_string()) + .unwrap_or_default(); + + let (role_name, label) = if let ConversationEntry::Memory { key, score, .. } = entry { + let label = match score { + Some(s) => format!("[memory: {} score:{:.1}]", key, s), + None => format!("[memory: {}]", key), + }; + ("mem".to_string(), label) + } else { + let tool_info = m.tool_calls.as_ref().map(|tc| { + tc.iter() + .map(|c| c.function.name.clone()) + .collect::>() + .join(", ") + }); + let label = match &tool_info { + Some(tools) => format!("[tool_call: {}]", tools), + None => { + let preview: String = text.chars().take(60).collect(); + let preview = preview.replace('\n', " "); + if text.len() > 60 { format!("{}...", preview) } else { preview } + } + }; + let role_name = match m.role { + Role::Assistant => cfg.assistant_name.clone(), + Role::User => cfg.user_name.clone(), + Role::Tool => "tool".to_string(), + Role::System => "system".to_string(), + }; + (role_name, label) + }; + + ContextSection { + name: format!("[{}] {}: {}", i, role_name, label), + tokens: count_msg(entry.api_message()), + content: text, + children: Vec::new(), + } + }) + .collect() + } + + /// Context sections for entries from `from` onward — used by the + /// subconscious debug screen to show forked agent conversations. + pub fn conversation_sections_from(&self, from: usize) -> Vec { + let count_msg = |m: &Message| context::msg_token_count(&self.tokenizer, m); + self.entry_sections(&count_msg, from, true) + } + /// Load recent journal entries at startup for orientation. /// Uses the same budget logic as compaction but with empty conversation. /// Only parses the tail of the journal file (last 64KB) for speed. @@ -959,8 +1003,8 @@ impl Agent { self.generation += 1; self.last_prompt_tokens = 0; - let sections = self.context_state_summary(); - dbglog!("[compact] budget: {}", context::sections_budget_string(§ions)); + let budget = self.context_budget(); + dbglog!("[compact] budget: {}", budget.format()); } /// Restore from the conversation log. Builds the context window @@ -986,9 +1030,8 @@ impl Agent { all.len(), mem_count, conv_count); self.context.entries = all; self.compact(); - // Estimate prompt tokens from sections so status bar isn't 0 on startup - self.last_prompt_tokens = context::sections_used( - &self.context_state_summary()) as u32; + // Estimate prompt tokens so status bar isn't 0 on startup + self.last_prompt_tokens = self.context_budget().total() as u32; true } diff --git a/src/mind/mod.rs b/src/mind/mod.rs index 3bbca79..b167ad8 100644 --- a/src/mind/mod.rs +++ b/src/mind/mod.rs @@ -249,7 +249,12 @@ impl Mind { /// Initialize — restore log, start daemons and background agents. pub async fn subconscious_snapshots(&self) -> Vec { - self.subconscious.lock().await.snapshots() + let store = crate::store::Store::cached().await.ok(); + let store_guard = match &store { + Some(s) => Some(s.lock().await), + None => None, + }; + self.subconscious.lock().await.snapshots(store_guard.as_deref()) } pub async fn subconscious_walked(&self) -> Vec { @@ -262,6 +267,10 @@ impl Mind { ag.restore_from_log(); ag.changed.notify_one(); drop(ag); + + // Load persistent subconscious state + let state_path = self.config.session_dir.join("subconscious-state.json"); + self.subconscious.lock().await.set_state_path(state_path); } pub fn turn_watch(&self) -> tokio::sync::watch::Receiver { @@ -276,8 +285,7 @@ impl Mind { MindCommand::Compact => { let threshold = compaction_threshold(&self.config.app) as usize; let mut ag = self.agent.lock().await; - let sections = ag.context_state_summary(); - if crate::agent::context::sections_used(§ions) > threshold { + if ag.context_budget().total() > threshold { ag.compact(); ag.notify("compacted"); } @@ -374,9 +382,7 @@ impl Mind { // Compact if over budget before sending let threshold = compaction_threshold(&self.config.app) as usize; - let used = crate::agent::context::sections_used( - &ag.context_state_summary()); - if used > threshold { + if ag.context_budget().total() > threshold { ag.compact(); ag.notify("compacted"); } diff --git a/src/user/chat.rs b/src/user/chat.rs index b863667..55c14ec 100644 --- a/src/user/chat.rs +++ b/src/user/chat.rs @@ -863,8 +863,7 @@ impl ScreenView for InteractScreen { agent.expire_activities(); app.status.prompt_tokens = agent.last_prompt_tokens(); app.status.model = agent.model().to_string(); - let sections = agent.context_state_summary(); - app.status.context_budget = crate::agent::context::sections_budget_string(§ions); + app.status.context_budget = agent.context_budget().format(); app.activity = agent.activities.last() .map(|a| a.label.clone()) .unwrap_or_default();