diff --git a/src/agent/runner.rs b/src/agent/runner.rs index 5bc3ce4..cba6296 100644 --- a/src/agent/runner.rs +++ b/src/agent/runner.rs @@ -62,8 +62,6 @@ pub struct Agent { pub reasoning_effort: String, /// Persistent conversation log — append-only record of all messages. conversation_log: Option, - /// Current context window budget breakdown. - pub context_budget: ContextBudget, /// BPE tokenizer for token counting (cl100k_base — close enough /// for Claude and Qwen budget allocation, ~85-90% count accuracy). tokenizer: CoreBPE, @@ -116,7 +114,6 @@ impl Agent { process_tracker: ProcessTracker::new(), reasoning_effort: "none".to_string(), conversation_log, - context_budget: ContextBudget::default(), tokenizer, context, shared_context, @@ -126,7 +123,6 @@ impl Agent { agent.load_startup_journal(); agent.load_working_stack(); - agent.measure_budget(); agent.publish_context_state(); agent } @@ -183,28 +179,11 @@ impl Agent { /// Push a context-only message (system prompt, identity context, /// journal summaries). Not logged — these are reconstructed on /// every startup/compaction. - /// Measure context window usage by category. Uses the BPE tokenizer - /// for direct token counting (no chars/4 approximation). - fn measure_budget(&mut self) { - let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len(); - - let id_tokens = count(&self.context.system_prompt) - + self.context.personality.iter() - .map(|(_, content)| count(content)).sum::(); - let jnl_tokens: usize = self.context.journal.iter() - .map(|e| count(&e.content)).sum(); - let mem_tokens: usize = self.context.loaded_nodes.iter() - .map(|node| count(&node.render())).sum(); - let conv_tokens: usize = self.context.messages.iter() - .map(|m| crate::agent::context::msg_token_count(&self.tokenizer, m)).sum(); - - self.context_budget = ContextBudget { - identity_tokens: id_tokens, - memory_tokens: mem_tokens, - journal_tokens: jnl_tokens, - conversation_tokens: conv_tokens, - window_tokens: crate::agent::context::model_context_window(&self.client.model), - }; + pub fn budget(&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::model_context_window(&self.client.model); + self.context.budget(&count_str, &count_msg, window) } /// Send a user message and run the agent loop until the model @@ -372,7 +351,7 @@ impl Agent { if let Some(usage) = &usage { self.last_prompt_tokens = usage.prompt_tokens; - self.measure_budget(); + self.publish_context_state(); let _ = ui_tx.send(UiMessage::StatusUpdate(StatusInfo { dmn_state: String::new(), // filled by main loop @@ -382,7 +361,7 @@ impl Agent { completion_tokens: usage.completion_tokens, model: self.client.model.clone(), turn_tools: 0, // tracked by TUI from ToolCall messages - context_budget: self.context_budget.status_string(), + context_budget: self.budget().status_string(), })); } @@ -547,7 +526,7 @@ impl Agent { if output.text.starts_with("Error:") { ds.tool_errors += 1; } - self.measure_budget(); + self.publish_context_state(); return; } @@ -826,7 +805,7 @@ impl Agent { /// Called after any change to context state (working stack, etc). fn refresh_context_state(&mut self) { - self.measure_budget(); + self.publish_context_state(); self.save_working_stack(); } @@ -849,8 +828,16 @@ impl Agent { /// Push the current context summary to the shared state for the TUI to read. fn publish_context_state(&self) { + let summary = self.context_state_summary(); + if let Ok(mut dbg) = std::fs::OpenOptions::new().create(true).append(true) + .open("/tmp/poc-journal-debug.log") { + use std::io::Write; + for s in &summary { + let _ = writeln!(dbg, "[publish] {} ({} tokens, {} children)", s.name, s.tokens, s.children.len()); + } + } if let Ok(mut state) = self.shared_context.write() { - *state = self.context_state_summary(); + *state = summary; } } @@ -978,7 +965,7 @@ impl Agent { self.context.journal = journal::parse_journal_text(&journal); self.context.messages = messages; self.last_prompt_tokens = 0; - self.measure_budget(); + self.publish_context_state(); } @@ -1041,7 +1028,7 @@ impl Agent { self.context.messages = messages; dbglog!("[restore] built context window: {} messages", self.context.messages.len()); self.last_prompt_tokens = 0; - self.measure_budget(); + self.publish_context_state(); true } diff --git a/src/agent/types.rs b/src/agent/types.rs index 1e04425..dbd8f4f 100644 --- a/src/agent/types.rs +++ b/src/agent/types.rs @@ -342,6 +342,57 @@ pub const WORKING_STACK_INSTRUCTIONS: &str = "/home/kent/.consciousness/config/w pub const WORKING_STACK_FILE: &str = "/home/kent/.consciousness/working-stack.json"; impl ContextState { + /// Compute the context budget from typed sources. + pub fn budget(&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::(); + let jnl: usize = self.journal.iter().map(|e| count_str(&e.content)).sum(); + let (mem, conv) = self.split_memory_conversation(count_msg); + ContextBudget { + identity_tokens: id, + memory_tokens: mem, + journal_tokens: jnl, + conversation_tokens: conv, + window_tokens, + } + } + + /// Split conversation messages into memory tool interactions and + /// everything else. Returns (memory_tokens, conversation_tokens). + pub fn split_memory_conversation(&self, count: &dyn Fn(&Message) -> usize) -> (usize, usize) { + // Collect tool_call_ids that belong to memory tools + let mut memory_call_ids: std::collections::HashSet = std::collections::HashSet::new(); + for msg in &self.messages { + if let Some(ref calls) = msg.tool_calls { + for call in calls { + if call.function.name.starts_with("memory_") + || call.function.name.starts_with("journal_") { + memory_call_ids.insert(call.id.clone()); + } + } + } + } + + let mut mem_tokens = 0; + let mut conv_tokens = 0; + for msg in &self.messages { + let tokens = count(msg); + let is_memory = match &msg.tool_call_id { + Some(id) => memory_call_ids.contains(id), + None => msg.tool_calls.as_ref().map_or(false, |calls| + calls.iter().all(|c| memory_call_ids.contains(&c.id))), + }; + if is_memory { + mem_tokens += tokens; + } else { + conv_tokens += tokens; + } + } + (mem_tokens, conv_tokens) + } + pub fn render_context_message(&self) -> String { let mut parts: Vec = self.personality.iter() .map(|(name, content)| format!("## {}\n\n{}", name, content)) diff --git a/src/bin/poc-agent.rs b/src/bin/poc-agent.rs index 8d399f1..09e481f 100644 --- a/src/bin/poc-agent.rs +++ b/src/bin/poc-agent.rs @@ -964,7 +964,7 @@ async fn run(cli: cli::CliArgs) -> Result<()> { completion_tokens: 0, model: agent_guard.model().to_string(), turn_tools: 0, - context_budget: agent_guard.context_budget.status_string(), + context_budget: agent_guard.budget().status_string(), })); }