diff --git a/src/agent/runner.rs b/src/agent/runner.rs
index 7308504..97546a5 100644
--- a/src/agent/runner.rs
+++ b/src/agent/runner.rs
@@ -185,35 +185,16 @@ impl Agent {
     fn measure_budget(&mut self) {
         let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
 
-        let mut id_tokens: usize = 0;
+        let id_tokens = count(&self.context.system_prompt)
+            + self.context.personality.iter()
+                .map(|(_, content)| count(content)).sum::<usize>();
         let jnl_tokens: usize = self.context.journal.iter()
-            .map(|e| count(&e.content))
-            .sum();
-        let mut conv_tokens: usize = 0;
-        let mut in_conversation = false;
-
-        for msg in &self.messages {
-            let tokens = crate::agent::context::msg_token_count(&self.tokenizer, msg);
-
-            if in_conversation {
-                conv_tokens += tokens;
-                continue;
-            }
-
-            if in_conversation {
-                conv_tokens += tokens;
-            } else if msg.role == Role::System || (!in_conversation && conv_tokens == 0) {
-                id_tokens += tokens;
-            } else {
-                in_conversation = true;
-                conv_tokens += tokens;
-            }
-        }
-
-        // Memory = nodes loaded during the session via tool calls
+            .map(|e| count(&e.content)).sum();
         let mem_tokens: usize = self.context.loaded_nodes.iter()
-            .map(|node| count(&node.render()))
-            .sum();
+            .map(|node| count(&node.render())).sum();
+        let total: usize = self.messages.iter()
+            .map(|m| crate::agent::context::msg_token_count(&self.tokenizer, m)).sum();
+        let conv_tokens = total.saturating_sub(id_tokens + jnl_tokens);
 
         self.context_budget = ContextBudget {
             identity_tokens: id_tokens,