From f4664ca06f4b571af0ce72be08eb1c461e008bf8 Mon Sep 17 00:00:00 2001
From: Kent Overstreet <kent.overstreet@linux.dev>
Date: Mon, 6 Apr 2026 18:36:33 -0400
Subject: [PATCH] Cache context budget instead of recomputing every frame

budget() called tiktoken on every UI tick, which was the main CPU hog
during rapid key input. Move the cached ContextBudget onto ContextState
and recompute only when entries actually change (push_entry, compact,
restore_from_log).

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
---
 src/agent/context.rs | 13 ++++++++-----
 src/agent/mod.rs     | 14 ++++++++------
 src/user/chat.rs     |  2 +-
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/agent/context.rs b/src/agent/context.rs
index 3ca6b1b..6042fea 100644
--- a/src/agent/context.rs
+++ b/src/agent/context.rs
@@ -230,13 +230,15 @@ 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.
-    pub fn budget(&self, count_str: &dyn Fn(&str) -> usize,
+    /// 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 {
+                   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();
@@ -246,13 +248,14 @@ impl ContextState {
             let tokens = count_msg(entry.api_message());
             if entry.is_memory() { mem += tokens } else { conv += tokens }
         }
-        ContextBudget {
+        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 {
diff --git a/src/agent/mod.rs b/src/agent/mod.rs
index 91f7451..7cf7c40 100644
--- a/src/agent/mod.rs
+++ b/src/agent/mod.rs
@@ -215,6 +215,7 @@ 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);
@@ -297,6 +298,7 @@ impl Agent {
             }
         }
         self.context.entries.push(entry);
+        self.recompute_budget();
         self.changed.notify_one();
     }
 
@@ -318,11 +320,11 @@ impl Agent {
         self.changed.notify_one();
     }
 
-    pub fn budget(&self) -> ContextBudget {
+    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.budget(&count_str, &count_msg, window)
+        self.context.recompute_budget(&count_str, &count_msg, window)
     }
 
     /// Send a user message and run the agent loop until the model
@@ -1040,8 +1042,8 @@ impl Agent {
         dbglog!("[compact] entries: {} → {} (mem: {} → {}, conv: {} → {})",
             before, after, before_mem, after_mem, before_conv, after_conv);
 
-        let budget = self.budget();
-        dbglog!("[compact] budget: {}", budget.status_string());
+        self.recompute_budget();
+        dbglog!("[compact] budget: {}", self.context.budget.status_string());
 
         self.load_startup_journal();
         self.generation += 1;
@@ -1073,8 +1075,8 @@ impl Agent {
         self.context.entries = all;
         self.compact();
         // Estimate prompt tokens from budget so status bar isn't 0 on startup
-        let b = self.budget();
-        self.last_prompt_tokens = b.used() as u32;
+        self.recompute_budget();
+        self.last_prompt_tokens = self.context.budget.used() as u32;
         true
     }
 
diff --git a/src/user/chat.rs b/src/user/chat.rs
index ad8c8fe..0510662 100644
--- a/src/user/chat.rs
+++ b/src/user/chat.rs
@@ -828,7 +828,7 @@ impl ScreenView for InteractScreen {
             agent.expire_activities();
             app.status.prompt_tokens = agent.last_prompt_tokens();
             app.status.model = agent.model().to_string();
-            app.status.context_budget = agent.budget().status_string();
+            app.status.context_budget = agent.context.budget.status_string();
             app.activity = agent.activities.last()
                 .map(|a| a.label.clone())
                 .unwrap_or_default();