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

WIP: ContextEntry/ContextSection data structures for incremental token counting

Browse source 
New types — not yet wired to callers:

- ContextEntry: wraps ConversationEntry with cached token count and
  timestamp
- ContextSection: named group of entries with cached token total.
  Private entries/tokens, read via entries()/tokens().
  Mutation via push(entry), set(index, entry), del(index).
- ContextState: system/identity/journal/conversation sections + working_stack
- ConversationEntry::System variant for system prompt entries

Token counting happens once at push time. Sections maintain their
totals incrementally via push/set/del. No more recomputing from
scratch on every budget check.

Does not compile — callers need updating.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-07 20:15:31 -04:00
parent 776ac527f1
commit 62996e27d7
10 changed files with 450 additions and 403 deletions
Download patch file Download diff file Expand all files Collapse all files

225
src/agent/context.rs
View file

@ -10,20 +10,130 @@ use serde::{Deserialize, Serialize};
use tiktoken_rs::CoreBPE; use tiktoken_rs::CoreBPE;
use crate::agent::tools::working_stack; use crate::agent::tools::working_stack;
/// A section of the context window, possibly with children. // --- Context state types ---
/// Conversation entry — either a regular message or memory content.
/// Memory entries preserve the original message for KV cache round-tripping.
#[derive(Debug, Clone, PartialEq)]
pub enum ConversationEntry {
/// System prompt or system-level instruction.
System(Message),
Message(Message),
Memory { key: String, message: Message, score: Option<f64> },
/// DMN heartbeat/autonomous prompt — evicted aggressively during compaction.
Dmn(Message),
/// Debug/status log line — written to conversation log for tracing,
/// skipped on read-back.
Log(String),
}
/// Entry in the context window — wraps a ConversationEntry with cached metadata.
#[derive(Debug, Clone)]
pub struct ContextEntry {
pub entry: ConversationEntry,
/// Cached token count (0 for Log entries).
pub tokens: usize,
/// When this entry was added to the context.
pub timestamp: Option<DateTime<Utc>>,
}
/// A named section of the context window with cached token total.
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct ContextSection { pub struct ContextSection {
pub name: String, pub name: String,
pub tokens: usize, /// Cached sum of entry tokens.
pub content: String, tokens: usize,
pub children: Vec<ContextSection>, entries: Vec<ContextEntry>,
} }
/// A single journal entry with its timestamp and content. impl ContextSection {
#[derive(Debug, Clone)] pub fn new(name: impl Into<String>) -> Self {
pub struct JournalEntry { Self { name: name.into(), tokens: 0, entries: Vec::new() }
pub timestamp: DateTime<Utc>, }
pub content: String,
pub fn entries(&self) -> &[ContextEntry] { &self.entries }
pub fn tokens(&self) -> usize { self.tokens }
pub fn len(&self) -> usize { self.entries.len() }
pub fn is_empty(&self) -> bool { self.entries.is_empty() }
/// Push an entry, updating the cached token total.
pub fn push(&mut self, entry: ContextEntry) {
self.tokens += entry.tokens;
self.entries.push(entry);
}
/// Replace an entry at `index`, adjusting the token total.
pub fn set(&mut self, index: usize, entry: ContextEntry) {
self.tokens -= self.entries[index].tokens;
self.tokens += entry.tokens;
self.entries[index] = entry;
}
/// Remove an entry at `index`, adjusting the token total.
pub fn del(&mut self, index: usize) -> ContextEntry {
let removed = self.entries.remove(index);
self.tokens -= removed.tokens;
removed
}
/// Replace the message inside an entry, recomputing its token count.
pub fn set_message(&mut self, index: usize, tokenizer: &CoreBPE, msg: Message) {
let old_tokens = self.entries[index].tokens;
*self.entries[index].entry.message_mut() = msg;
let new_tokens = msg_token_count(tokenizer, self.entries[index].entry.api_message());
self.entries[index].tokens = new_tokens;
self.tokens = self.tokens - old_tokens + new_tokens;
}
/// Set the score on a Memory entry. No token change.
pub fn set_score(&mut self, index: usize, score: Option<f64>) {
if let ConversationEntry::Memory { score: s, .. } = &mut self.entries[index].entry {
*s = score;
}
}
/// Bulk replace all entries, recomputing token total.
pub fn set_entries(&mut self, entries: Vec<ContextEntry>) {
self.tokens = entries.iter().map(|e| e.tokens).sum();
self.entries = entries;
}
/// Dedup and trim entries to fit within context budget.
pub fn trim(&mut self, budget: &ContextBudget, tokenizer: &CoreBPE) {
let result = trim_entries(&self.entries, tokenizer, budget);
self.entries = result;
self.tokens = self.entries.iter().map(|e| e.tokens).sum();
}
/// Clear all entries.
pub fn clear(&mut self) {
self.entries.clear();
self.tokens = 0;
}
}
#[derive(Clone)]
pub struct ContextState {
pub system: ContextSection,
pub identity: ContextSection,
pub journal: ContextSection,
pub conversation: ContextSection,
/// Working stack — separate from identity because it's managed
/// by its own tool, not loaded from personality files.
pub working_stack: Vec<String>,
}
impl ContextState {
/// Total tokens across all sections.
pub fn total_tokens(&self) -> usize {
self.system.tokens() + self.identity.tokens()
+ self.journal.tokens() + self.conversation.tokens()
}
/// All sections as a slice for iteration.
pub fn sections(&self) -> [&ContextSection; 4] {
[&self.system, &self.identity, &self.journal, &self.conversation]
}
} }
/// Context window size in tokens (from config). /// Context window size in tokens (from config).
@ -44,41 +154,39 @@ fn context_budget_tokens() -> usize {
/// corresponding assistant tool_call message). /// corresponding assistant tool_call message).
/// 2. Trim: drop oldest entries until the conversation fits, snapping /// 2. Trim: drop oldest entries until the conversation fits, snapping
/// to user message boundaries. /// to user message boundaries.
pub fn trim_entries( fn trim_entries(
entries: &[ConversationEntry], entries: &[ContextEntry],
tokenizer: &CoreBPE, _tokenizer: &CoreBPE,
budget: &ContextBudget, budget: &ContextBudget,
) -> Vec<ConversationEntry> { ) -> Vec<ContextEntry> {
let fixed_tokens = budget.system + budget.identity + budget.journal; let fixed_tokens = budget.system + budget.identity + budget.journal;
// --- Phase 1: dedup memory entries by key (keep last) --- // --- Phase 1: dedup memory entries by key (keep last) ---
let mut seen_keys: std::collections::HashMap<&str, usize> = std::collections::HashMap::new(); let mut seen_keys: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
let mut drop_indices: std::collections::HashSet<usize> = std::collections::HashSet::new(); let mut drop_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();
for (i, entry) in entries.iter().enumerate() { for (i, ce) in entries.iter().enumerate() {
if let ConversationEntry::Memory { key, .. } = entry { if let ConversationEntry::Memory { key, .. } = &ce.entry {
if let Some(prev) = seen_keys.insert(key.as_str(), i) { if let Some(prev) = seen_keys.insert(key.as_str(), i) {
drop_indices.insert(prev); drop_indices.insert(prev);
} }
} }
} }
let deduped: Vec<ConversationEntry> = entries.iter().enumerate() let deduped: Vec<ContextEntry> = entries.iter().enumerate()
.filter(|(i, _)| !drop_indices.contains(i)) .filter(|(i, _)| !drop_indices.contains(i))
.map(|(_, e)| e.clone()) .map(|(_, e)| e.clone())
.collect(); .collect();
// --- Phase 2: trim to fit context budget --- // --- Phase 2: trim to fit context budget ---
let max_tokens = context_budget_tokens(); let max_tokens = context_budget_tokens();
let count_msg = |m: &Message| msg_token_count(tokenizer, m);
let msg_costs: Vec<usize> = deduped.iter() let msg_costs: Vec<usize> = deduped.iter().map(|e| e.tokens).collect();
.map(|e| if e.is_log() { 0 } else { count_msg(e.api_message()) }).collect();
let entry_total: usize = msg_costs.iter().sum(); let entry_total: usize = msg_costs.iter().sum();
let total: usize = fixed_tokens + entry_total; let total: usize = fixed_tokens + entry_total;
let mem_tokens: usize = deduped.iter().zip(&msg_costs) let mem_tokens: usize = deduped.iter()
.filter(|(e, _)| e.is_memory()) .filter(|ce| ce.entry.is_memory())
.map(|(_, &c)| c).sum(); .map(|ce| ce.tokens).sum();
let conv_tokens: usize = entry_total - mem_tokens; let conv_tokens: usize = entry_total - mem_tokens;
dbglog!("[trim] max_tokens={} fixed={} mem={} conv={} total={} entries={}", dbglog!("[trim] max_tokens={} fixed={} mem={} conv={} total={} entries={}",
@ -90,7 +198,7 @@ pub fn trim_entries(
let mut cur_mem = mem_tokens; let mut cur_mem = mem_tokens;
for i in 0..deduped.len() { for i in 0..deduped.len() {
if deduped[i].is_dmn() { if deduped[i].entry.is_dmn() {
drop[i] = true; drop[i] = true;
trimmed -= msg_costs[i]; trimmed -= msg_costs[i];
} }
@ -99,14 +207,14 @@ pub fn trim_entries(
// Phase 2b: if memories > 50% of context, evict lowest-scored first // Phase 2b: if memories > 50% of context, evict lowest-scored first
if cur_mem > conv_tokens && trimmed > max_tokens { if cur_mem > conv_tokens && trimmed > max_tokens {
let mut mem_indices: Vec<usize> = (0..deduped.len()) let mut mem_indices: Vec<usize> = (0..deduped.len())
.filter(|&i| !drop[i] && deduped[i].is_memory()) .filter(|&i| !drop[i] && deduped[i].entry.is_memory())
.collect(); .collect();
mem_indices.sort_by(|&a, &b| { mem_indices.sort_by(|&a, &b| {
let sa = match &deduped[a] { let sa = match &deduped[a].entry {
ConversationEntry::Memory { score, .. } => score.unwrap_or(0.0), ConversationEntry::Memory { score, .. } => score.unwrap_or(0.0),
_ => 0.0, _ => 0.0,
}; };
let sb = match &deduped[b] { let sb = match &deduped[b].entry {
ConversationEntry::Memory { score, .. } => score.unwrap_or(0.0), ConversationEntry::Memory { score, .. } => score.unwrap_or(0.0),
_ => 0.0, _ => 0.0,
}; };
@ -130,16 +238,16 @@ pub fn trim_entries(
} }
// Walk forward to include complete conversation boundaries // Walk forward to include complete conversation boundaries
let mut result: Vec<ConversationEntry> = Vec::new(); let mut result: Vec<ContextEntry> = Vec::new();
let mut skipping = true; let mut skipping = true;
for (i, entry) in deduped.into_iter().enumerate() { for (i, ce) in deduped.into_iter().enumerate() {
if skipping { if skipping {
if drop[i] { continue; } if drop[i] { continue; }
// Snap to user message boundary // Snap to user message boundary
if entry.message().role != Role::User { continue; } if ce.entry.message().role != Role::User { continue; }
skipping = false; skipping = false;
} }
result.push(entry); result.push(ce);
} }
dbglog!("[trim] result={} trimmed_total={}", result.len(), trimmed); dbglog!("[trim] result={} trimmed_total={}", result.len(), trimmed);
@ -186,28 +294,13 @@ pub fn is_stream_error(err: &anyhow::Error) -> bool {
err.to_string().contains("model stream error") err.to_string().contains("model stream error")
} }
// --- Context state types ---
/// Conversation entry — either a regular message or memory content.
/// Memory entries preserve the original message for KV cache round-tripping.
#[derive(Debug, Clone, PartialEq)]
pub enum ConversationEntry {
Message(Message),
Memory { key: String, message: Message, score: Option<f64> },
/// DMN heartbeat/autonomous prompt — evicted aggressively during compaction.
Dmn(Message),
/// Debug/status log line — written to conversation log for tracing,
/// skipped on read-back.
Log(String),
}
// Custom serde: serialize Memory with a "memory_key" field added to the message, // Custom serde: serialize Memory with a "memory_key" field added to the message,
// plain messages serialize as-is. This keeps the conversation log readable. // plain messages serialize as-is. This keeps the conversation log readable.
impl Serialize for ConversationEntry { impl Serialize for ConversationEntry {
fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> { fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
use serde::ser::SerializeMap; use serde::ser::SerializeMap;
match self { match self {
Self::Message(m) | Self::Dmn(m) => m.serialize(s), Self::System(m) | Self::Message(m) | Self::Dmn(m) => m.serialize(s),
Self::Memory { key, message, score } => { Self::Memory { key, message, score } => {
let json = serde_json::to_value(message).map_err(serde::ser::Error::custom)?; let json = serde_json::to_value(message).map_err(serde::ser::Error::custom)?;
let mut map = s.serialize_map(None)?; let mut map = s.serialize_map(None)?;
@ -259,7 +352,7 @@ impl ConversationEntry {
/// Panics on Log entries (which should be filtered before API calls). /// Panics on Log entries (which should be filtered before API calls).
pub fn api_message(&self) -> &Message { pub fn api_message(&self) -> &Message {
match self { match self {
Self::Message(m) | Self::Dmn(m) => m, Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
Self::Memory { message, .. } => message, Self::Memory { message, .. } => message,
Self::Log(_) => panic!("Log entries have no API message"), Self::Log(_) => panic!("Log entries have no API message"),
} }
@ -281,7 +374,7 @@ impl ConversationEntry {
/// Panics on Log entries. /// Panics on Log entries.
pub fn message(&self) -> &Message { pub fn message(&self) -> &Message {
match self { match self {
Self::Message(m) | Self::Dmn(m) => m, Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
Self::Memory { message, .. } => message, Self::Memory { message, .. } => message,
Self::Log(_) => panic!("Log entries have no message"), Self::Log(_) => panic!("Log entries have no message"),
} }
@ -291,38 +384,36 @@ impl ConversationEntry {
/// Panics on Log entries. /// Panics on Log entries.
pub fn message_mut(&mut self) -> &mut Message { pub fn message_mut(&mut self) -> &mut Message {
match self { match self {
Self::Message(m) | Self::Dmn(m) => m, Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
Self::Memory { message, .. } => message, Self::Memory { message, .. } => message,
Self::Log(_) => panic!("Log entries have no message"), Self::Log(_) => panic!("Log entries have no message"),
} }
} }
} }
#[derive(Clone)] impl ContextState {
pub struct ContextState { /// Render journal entries into a single text block.
pub system_prompt: String, pub fn render_journal(&self) -> String {
pub personality: Vec<(String, String)>, if self.journal.is_empty() { return String::new(); }
pub journal: Vec<JournalEntry>,
pub working_stack: Vec<String>,
/// Conversation entries — messages and memory, interleaved in order.
/// Does NOT include system prompt, personality, or journal.
pub entries: Vec<ConversationEntry>,
}
pub fn render_journal(entries: &[JournalEntry]) -> String {
if entries.is_empty() { return String::new(); }
let mut text = String::from("[Earlier — from your journal]\n\n"); let mut text = String::from("[Earlier — from your journal]\n\n");
for entry in entries { for e in self.journal.entries() {
use std::fmt::Write; use std::fmt::Write;
writeln!(text, "## {}\n{}\n", entry.timestamp.format("%Y-%m-%dT%H:%M"), entry.content).ok(); if let Some(ts) = &e.timestamp {
writeln!(text, "## {}\n{}\n",
ts.format("%Y-%m-%dT%H:%M"),
e.entry.message().content_text()).ok();
} else {
text.push_str(&e.entry.message().content_text());
text.push_str("\n\n");
}
} }
text text
} }
impl ContextState { /// Render identity files + working stack into a single user message.
pub fn render_context_message(&self) -> String { pub fn render_context_message(&self) -> String {
let mut parts: Vec<String> = self.personality.iter() let mut parts: Vec<String> = self.identity.entries().iter()
.map(|(name, content)| format!("## {}\n\n{}", name, content)) .map(|e| e.entry.message().content_text().to_string())
.collect(); .collect();
let instructions = std::fs::read_to_string(working_stack::instructions_path()).unwrap_or_default(); let instructions = std::fs::read_to_string(working_stack::instructions_path()).unwrap_or_default();
let mut stack_section = instructions; let mut stack_section = instructions;

378
src/agent/mod.rs
View file

@ -24,7 +24,7 @@ use tiktoken_rs::CoreBPE;
use api::{ApiClient, ToolCall}; use api::{ApiClient, ToolCall};
use api::{ContentPart, Message, MessageContent, Role}; use api::{ContentPart, Message, MessageContent, Role};
use context::{ConversationEntry, ContextState}; use context::{ConversationEntry, ContextEntry, ContextState};
use tools::{summarize_args, working_stack}; use tools::{summarize_args, working_stack};
use crate::mind::log::ConversationLog; use crate::mind::log::ConversationLog;
@ -195,12 +195,27 @@ impl Agent {
let tokenizer = tiktoken_rs::cl100k_base() let tokenizer = tiktoken_rs::cl100k_base()
.expect("failed to load cl100k_base tokenizer"); .expect("failed to load cl100k_base tokenizer");
let mut system = ContextSection::new("System prompt");
system.push(ContextEntry {
entry: ConversationEntry::System(Message::system(&system_prompt)),
tokens: context::msg_token_count(&tokenizer, &Message::system(&system_prompt)),
timestamp: None,
});
let mut identity = ContextSection::new("Identity");
for (_name, content) in &personality {
let msg = Message::user(content);
identity.push(ContextEntry {
tokens: context::msg_token_count(&tokenizer, &msg),
entry: ConversationEntry::Message(msg),
timestamp: None,
});
}
let context = ContextState { let context = ContextState {
system_prompt: system_prompt.clone(), system,
personality, identity,
journal: Vec::new(), journal: ContextSection::new("Journal"),
conversation: ContextSection::new("Conversation"),
working_stack: Vec::new(), working_stack: Vec::new(),
entries: Vec::new(),
}; };
let session_id = format!("consciousness-{}", chrono::Utc::now().format("%Y%m%d-%H%M%S")); let session_id = format!("consciousness-{}", chrono::Utc::now().format("%Y%m%d-%H%M%S"));
let mut agent = Self { let mut agent = Self {
@ -274,18 +289,24 @@ impl Agent {
/// System prompt + personality context + journal + conversation messages. /// System prompt + personality context + journal + conversation messages.
pub fn assemble_api_messages(&self) -> Vec<Message> { pub fn assemble_api_messages(&self) -> Vec<Message> {
let mut msgs = Vec::new(); let mut msgs = Vec::new();
msgs.push(Message::system(&self.context.system_prompt)); // System section
for e in self.context.system.entries() {
msgs.push(e.entry.api_message().clone());
}
// Identity — render personality files + working stack into one user message
let ctx = self.context.render_context_message(); let ctx = self.context.render_context_message();
if !ctx.is_empty() { if !ctx.is_empty() {
msgs.push(Message::user(ctx)); msgs.push(Message::user(ctx));
} }
let jnl = context::render_journal(&self.context.journal); // Journal — render into one user message
let jnl = self.context.render_journal();
if !jnl.is_empty() { if !jnl.is_empty() {
msgs.push(Message::user(jnl)); msgs.push(Message::user(jnl));
} }
msgs.extend(self.context.entries.iter() // Conversation entries
.filter(|e| !e.is_log()) msgs.extend(self.context.conversation.entries().iter()
.map(|e| e.api_message().clone())); .filter(|e| !e.entry.is_log())
.map(|e| e.entry.api_message().clone()));
msgs msgs
} }
@ -303,50 +324,64 @@ impl Agent {
eprintln!("warning: failed to log entry: {:#}", e); eprintln!("warning: failed to log entry: {:#}", e);
} }
} }
self.context.entries.push(entry); let tokens = if entry.is_log() { 0 } else {
context::msg_token_count(&self.tokenizer, entry.api_message())
};
self.context.conversation.push(ContextEntry {
entry, tokens, timestamp: Some(chrono::Utc::now()),
});
self.changed.notify_one(); self.changed.notify_one();
} }
fn streaming_entry(&mut self) -> Option<&mut Message> { /// Find the index of the in-progress streaming entry (unstamped assistant message).
for entry in self.context.entries.iter_mut().rev() { fn streaming_index(&self) -> Option<usize> {
let m = entry.message_mut(); self.context.conversation.entries().iter().rposition(|ce| {
if m.role == Role::Assistant { let m = ce.entry.message();
return if m.timestamp.is_none() { Some(m) } else { None } m.role == Role::Assistant && m.timestamp.is_none()
} })
}
None
} }
/// Append streaming text to the last entry (creating a partial /// Append streaming text to the last entry (creating a partial
/// assistant entry if needed). Called by collect_stream per token batch. /// assistant entry if needed). Called by collect_stream per token batch.
fn append_streaming(&mut self, text: &str) { fn append_streaming(&mut self, text: &str) {
if let Some(m) = self.streaming_entry() { if let Some(idx) = self.streaming_index() {
m.append_content(text); let mut msg = self.context.conversation.entries()[idx].entry.message().clone();
msg.append_content(text);
self.context.conversation.set_message(idx, &self.tokenizer, msg);
} else { } else {
// No streaming entry — create without timestamp so finalize can find it let msg = Message {
self.context.entries.push(ConversationEntry::Message(Message {
role: Role::Assistant, role: Role::Assistant,
content: Some(MessageContent::Text(text.to_string())), content: Some(MessageContent::Text(text.to_string())),
tool_calls: None, tool_calls: None,
tool_call_id: None, tool_call_id: None,
name: None, name: None,
timestamp: None, timestamp: None,
})); };
let tokens = context::msg_token_count(&self.tokenizer, &msg);
self.context.conversation.push(ContextEntry {
entry: ConversationEntry::Message(msg),
tokens,
timestamp: None,
});
} }
self.changed.notify_one(); self.changed.notify_one();
} }
/// Finalize the streaming entry with the complete response message. /// Finalize the streaming entry with the complete response message.
/// Finds the unstamped assistant entry, updates it in place, and logs it. /// Finds the unstamped assistant entry, replaces it via set() with proper token count.
fn finalize_streaming(&mut self, msg: Message) { fn finalize_streaming(&mut self, msg: Message) {
if let Some(m) = self.streaming_entry() { if let Some(i) = self.streaming_index() {
*m = msg.clone(); let mut stamped = msg.clone();
m.stamp(); stamped.stamp();
let tokens = context::msg_token_count(&self.tokenizer, &stamped);
self.context.conversation.set(i, ContextEntry {
entry: ConversationEntry::Message(stamped),
tokens,
timestamp: Some(chrono::Utc::now()),
});
} else { } else {
// No streaming entry found — push as new (this logs via push_message)
self.push_message(msg.clone()); self.push_message(msg.clone());
} }
@ -655,173 +690,32 @@ impl Agent {
self.push_message(Message::tool_result(&call.id, &output)); self.push_message(Message::tool_result(&call.id, &output));
} }
/// Token budget by category — cheap, no formatting. Used for compaction decisions. /// Token budget by category — just reads cached section totals.
pub fn context_budget(&self) -> context::ContextBudget { pub fn context_budget(&self) -> context::ContextBudget {
let count = |m: &Message| context::msg_token_count(&self.tokenizer, m); let memory: usize = self.context.conversation.entries().iter()
.filter(|e| e.entry.is_memory())
let system = count(&Message::system(&self.context.system_prompt)); .map(|e| e.tokens)
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(); .sum();
let conversation: usize = self.context.entries.iter() let conv_total = self.context.conversation.tokens();
.filter(|e| !e.is_memory() && !e.is_log()) context::ContextBudget {
.map(|e| count(e.api_message())) system: self.context.system.tokens(),
.sum(); identity: self.context.identity.tokens(),
journal: self.context.journal.tokens(),
context::ContextBudget { system, identity, journal, memory, conversation } memory,
conversation: conv_total - memory,
}
} }
/// Build context state summary for the debug screen. /// Context state sections — just returns references to the live data.
pub fn context_state_summary(&self) -> Vec<ContextSection> { pub fn context_sections(&self) -> [&ContextSection; 4] {
let count_msg = |m: &Message| context::msg_token_count(&self.tokenizer, m); self.context.sections()
let mut sections = Vec::new();
// System prompt — counted as the actual message sent
let system_msg = Message::system(&self.context.system_prompt);
sections.push(ContextSection {
name: "System prompt".into(),
tokens: count_msg(&system_msg),
content: self.context.system_prompt.clone(),
children: Vec::new(),
});
// Context message (personality + working stack) — counted as the
// single user message that assemble_api_messages sends
let context_rendered = self.context.render_context_message();
let context_msg = Message::user(&context_rendered);
sections.push(ContextSection {
name: format!("Identity ({} files + stack)", self.context.personality.len()),
tokens: count_msg(&context_msg),
content: context_rendered,
children: Vec::new(),
});
// Journal — counted as the single rendered message sent
let journal_rendered = context::render_journal(&self.context.journal);
let journal_msg = Message::user(&journal_rendered);
sections.push(ContextSection {
name: format!("Journal ({} entries)", self.context.journal.len()),
tokens: if journal_rendered.is_empty() { 0 } else { count_msg(&journal_msg) },
content: journal_rendered,
children: Vec::new(),
});
// Memory nodes — extracted from Memory entries in the conversation
let memory_entries: Vec<&ConversationEntry> = self.context.entries.iter()
.filter(|e| e.is_memory())
.collect();
if !memory_entries.is_empty() {
let node_children: Vec<ContextSection> = memory_entries.iter()
.map(|entry| {
let (key, score) = match entry {
ConversationEntry::Memory { key, score, .. } => (key.as_str(), *score),
_ => unreachable!(),
};
let label = match score {
Some(s) => format!("{} (score:{:.1})", key, s),
None => key.to_string(),
};
ContextSection {
name: label,
tokens: count_msg(entry.message()),
content: String::new(),
children: Vec::new(),
}
})
.collect();
let node_tokens: usize = node_children.iter().map(|c| c.tokens).sum();
sections.push(ContextSection {
name: format!("Memory nodes ({} loaded)", memory_entries.len()),
tokens: node_tokens,
content: String::new(),
children: node_children,
});
} }
// Conversation — memories excluded (counted in their own section above) /// Conversation entries from `from` onward — used by the
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()),
tokens: conv_tokens,
content: String::new(),
children: conv_children,
});
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<ContextSection> {
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::<Vec<_>>()
.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. /// subconscious debug screen to show forked agent conversations.
pub fn conversation_sections_from(&self, from: usize) -> Vec<ContextSection> { pub fn conversation_entries_from(&self, from: usize) -> &[ContextEntry] {
let count_msg = |m: &Message| context::msg_token_count(&self.tokenizer, m); let entries = self.context.conversation.entries();
self.entry_sections(&count_msg, from, true) if from < entries.len() { &entries[from..] } else { &[] }
} }
/// Load recent journal entries at startup for orientation. /// Load recent journal entries at startup for orientation.
@ -876,35 +770,38 @@ impl Agent {
dbg_log!("[journal] cutoff_idx={}", cutoff_idx); dbg_log!("[journal] cutoff_idx={}", cutoff_idx);
// Walk backwards from cutoff, accumulating entries within 15% of context // Walk backwards from cutoff, accumulating entries within 15% of context
let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
let context_window = crate::agent::context::context_window(); let context_window = crate::agent::context::context_window();
let journal_budget = context_window * 15 / 100; let journal_budget = context_window * 15 / 100;
dbg_log!("[journal] budget={} tokens ({}*15%)", journal_budget, context_window); dbg_log!("[journal] budget={} tokens ({}*15%)", journal_budget, context_window);
let mut entries = Vec::new(); let mut journal_entries = Vec::new();
let mut total_tokens = 0; let mut total_tokens = 0;
for node in journal_nodes[..cutoff_idx].iter().rev() { for node in journal_nodes[..cutoff_idx].iter().rev() {
let tokens = count(&node.content); let msg = Message::user(&node.content);
if total_tokens + tokens > journal_budget && !entries.is_empty() { let tokens = context::msg_token_count(&self.tokenizer, &msg);
if total_tokens + tokens > journal_budget && !journal_entries.is_empty() {
break; break;
} }
entries.push(context::JournalEntry { journal_entries.push(ContextEntry {
timestamp: chrono::DateTime::from_timestamp(node.created_at, 0) entry: ConversationEntry::Message(msg),
.unwrap_or_default(), tokens,
content: node.content.clone(), timestamp: chrono::DateTime::from_timestamp(node.created_at, 0),
}); });
total_tokens += tokens; total_tokens += tokens;
} }
entries.reverse(); journal_entries.reverse();
dbg_log!("[journal] loaded {} entries, {} tokens", entries.len(), total_tokens); dbg_log!("[journal] loaded {} entries, {} tokens", journal_entries.len(), total_tokens);
if entries.is_empty() { if journal_entries.is_empty() {
dbg_log!("[journal] no entries!"); dbg_log!("[journal] no entries!");
return; return;
} }
self.context.journal = entries; self.context.journal.clear();
for entry in journal_entries {
self.context.journal.push(entry);
}
dbg_log!("[journal] context.journal now has {} entries", self.context.journal.len()); dbg_log!("[journal] context.journal now has {} entries", self.context.journal.len());
} }
@ -923,10 +820,10 @@ impl Agent {
/// The tool result message before each image records what was loaded. /// The tool result message before each image records what was loaded.
pub fn age_out_images(&mut self) { pub fn age_out_images(&mut self) {
// Find image entries newest-first, skip 1 (caller is about to add another) // Find image entries newest-first, skip 1 (caller is about to add another)
let to_age: Vec<usize> = self.context.entries.iter().enumerate() let to_age: Vec<usize> = self.context.conversation.entries().iter().enumerate()
.rev() .rev()
.filter(|(_, e)| { .filter(|(_, ce)| {
if let Some(MessageContent::Parts(parts)) = &e.message().content { if let Some(MessageContent::Parts(parts)) = &ce.entry.message().content {
parts.iter().any(|p| matches!(p, ContentPart::ImageUrl { .. })) parts.iter().any(|p| matches!(p, ContentPart::ImageUrl { .. }))
} else { false } } else { false }
}) })
@ -935,7 +832,9 @@ impl Agent {
.collect(); .collect();
for i in to_age { for i in to_age {
let msg = self.context.entries[i].message_mut(); // Build replacement entry with image data stripped
let old = &self.context.conversation.entries()[i];
let msg = old.entry.message();
if let Some(MessageContent::Parts(parts)) = &msg.content { if let Some(MessageContent::Parts(parts)) = &msg.content {
let mut replacement = String::new(); let mut replacement = String::new();
for part in parts { for part in parts {
@ -950,7 +849,14 @@ impl Agent {
} }
} }
} }
msg.content = Some(MessageContent::Text(replacement)); let mut new_msg = msg.clone();
new_msg.content = Some(MessageContent::Text(replacement));
let tokens = context::msg_token_count(&self.tokenizer, &new_msg);
self.context.conversation.set(i, ContextEntry {
entry: ConversationEntry::Message(new_msg),
tokens,
timestamp: old.timestamp,
});
} }
} }
self.generation += 1; self.generation += 1;
@ -968,16 +874,30 @@ impl Agent {
// Reload identity from config // Reload identity from config
match crate::config::reload_for_model(&self.app_config, &self.prompt_file) { match crate::config::reload_for_model(&self.app_config, &self.prompt_file) {
Ok((system_prompt, personality)) => { Ok((system_prompt, personality)) => {
self.context.system_prompt = system_prompt; self.context.system.clear();
self.context.personality = personality; self.context.system.push(ContextEntry {
entry: ConversationEntry::System(Message::system(&system_prompt)),
tokens: context::msg_token_count(&self.tokenizer, &Message::system(&system_prompt)),
timestamp: None,
});
self.context.identity.clear();
for (_name, content) in &personality {
let msg = Message::user(content);
self.context.identity.push(ContextEntry {
tokens: context::msg_token_count(&self.tokenizer, &msg),
entry: ConversationEntry::Message(msg),
timestamp: None,
});
}
} }
Err(e) => { Err(e) => {
eprintln!("warning: failed to reload identity: {:#}", e); eprintln!("warning: failed to reload identity: {:#}", e);
} }
} }
let before = self.context.entries.len(); let before = self.context.conversation.len();
let before_mem = self.context.entries.iter().filter(|e| e.is_memory()).count(); let before_mem = self.context.conversation.entries().iter()
.filter(|e| e.entry.is_memory()).count();
let before_conv = before - before_mem; let before_conv = before - before_mem;
// Age out images before trimming — they're huge in the request payload // Age out images before trimming — they're huge in the request payload
@ -988,15 +908,11 @@ impl Agent {
// Dedup memory, trim to budget // Dedup memory, trim to budget
let budget = self.context_budget(); let budget = self.context_budget();
let entries = self.context.entries.clone(); self.context.conversation.trim(&budget, &self.tokenizer);
self.context.entries = crate::agent::context::trim_entries(
&entries,
&self.tokenizer,
&budget,
);
let after = self.context.entries.len(); let after = self.context.conversation.len();
let after_mem = self.context.entries.iter().filter(|e| e.is_memory()).count(); let after_mem = self.context.conversation.entries().iter()
.filter(|e| e.entry.is_memory()).count();
let after_conv = after - after_mem; let after_conv = after - after_mem;
dbglog!("[compact] entries: {} → {} (mem: {} → {}, conv: {} → {})", dbglog!("[compact] entries: {} → {} (mem: {} → {}, conv: {} → {})",
@ -1022,14 +938,26 @@ impl Agent {
}; };
// Load extra — compact() will dedup, trim, reload identity + journal // Load extra — compact() will dedup, trim, reload identity + journal
let all: Vec<_> = entries.into_iter() let all: Vec<ContextEntry> = entries.into_iter()
.filter(|e| !e.is_log() && e.message().role != Role::System) .filter(|e| !e.is_log() && e.message().role != Role::System)
.map(|e| {
let tokens = if e.is_log() { 0 } else {
context::msg_token_count(&self.tokenizer, e.api_message())
};
let timestamp = if e.is_log() { None } else {
e.message().timestamp.as_ref().and_then(|ts| {
chrono::DateTime::parse_from_rfc3339(ts).ok()
.map(|dt| dt.with_timezone(&chrono::Utc))
})
};
ContextEntry { entry: e, tokens, timestamp }
})
.collect(); .collect();
let mem_count = all.iter().filter(|e| e.is_memory()).count(); let mem_count = all.iter().filter(|e| e.entry.is_memory()).count();
let conv_count = all.len() - mem_count; let conv_count = all.len() - mem_count;
dbglog!("[restore] loaded {} entries from log (mem: {}, conv: {})", dbglog!("[restore] loaded {} entries from log (mem: {}, conv: {})",
all.len(), mem_count, conv_count); all.len(), mem_count, conv_count);
self.context.entries = all; self.context.conversation.set_entries(all);
self.compact(); self.compact();
// Estimate prompt tokens so status bar isn't 0 on startup // Estimate prompt tokens so status bar isn't 0 on startup
self.last_prompt_tokens = self.context_budget().total() as u32; self.last_prompt_tokens = self.context_budget().total() as u32;
@ -1046,9 +974,9 @@ impl Agent {
&self.client.model &self.client.model
} }
/// Get the conversation entries for persistence. /// Get the conversation entries.
pub fn entries(&self) -> &[ConversationEntry] { pub fn entries(&self) -> &[ContextEntry] {
&self.context.entries self.context.conversation.entries()
} }
/// Mutable access to conversation entries (for /retry). /// Mutable access to conversation entries (for /retry).

3
src/agent/oneshot.rs
View file

@ -93,8 +93,7 @@ impl Backend {
match self { match self {
Backend::Standalone { messages, .. } => messages.push(msg), Backend::Standalone { messages, .. } => messages.push(msg),
Backend::Forked(agent) => { Backend::Forked(agent) => {
agent.lock().await.context.entries.push( agent.lock().await.push_message(msg);
super::context::ConversationEntry::Message(msg));
} }
} }
} }

12
src/mind/dmn.rs
View file

@ -518,12 +518,12 @@ impl Subconscious {
pub async fn trigger(&mut self, agent: &Arc<tokio::sync::Mutex<Agent>>) { pub async fn trigger(&mut self, agent: &Arc<tokio::sync::Mutex<Agent>>) {
let (conversation_bytes, memory_keys) = { let (conversation_bytes, memory_keys) = {
let ag = agent.lock().await; let ag = agent.lock().await;
let bytes = ag.context.entries.iter() let bytes = ag.context.conversation.entries().iter()
.filter(|e| !e.is_log() && !e.is_memory()) .filter(|ce| !ce.entry.is_log() && !ce.entry.is_memory())
.map(|e| e.message().content_text().len() as u64) .map(|ce| ce.entry.message().content_text().len() as u64)
.sum::<u64>(); .sum::<u64>();
let keys: Vec<String> = ag.context.entries.iter().filter_map(|e| { let keys: Vec<String> = ag.context.conversation.entries().iter().filter_map(|ce| {
if let ConversationEntry::Memory { key, .. } = e { if let ConversationEntry::Memory { key, .. } = &ce.entry {
Some(key.clone()) Some(key.clone())
} else { None } } else { None }
}).collect(); }).collect();
@ -550,7 +550,7 @@ impl Subconscious {
let mut forked = conscious.fork(auto.tools.clone()); let mut forked = conscious.fork(auto.tools.clone());
forked.provenance = format!("agent:{}", auto.name); forked.provenance = format!("agent:{}", auto.name);
let fork_point = forked.context.entries.len(); let fork_point = forked.context.conversation.len();
let shared_forked = Arc::new(tokio::sync::Mutex::new(forked)); let shared_forked = Arc::new(tokio::sync::Mutex::new(forked));
self.agents[idx].forked_agent = Some(shared_forked.clone()); self.agents[idx].forked_agent = Some(shared_forked.clone());

32
src/mind/mod.rs
View file

@ -31,7 +31,9 @@ pub use dmn::{SubconsciousSnapshot, Subconscious};
use crate::agent::context::ConversationEntry; use crate::agent::context::ConversationEntry;
/// Load persisted memory scores from disk and apply to Memory entries. /// Load persisted memory scores from disk and apply to Memory entries.
fn load_memory_scores(entries: &mut [ConversationEntry], path: &std::path::Path) { use crate::agent::context::ContextSection;
fn load_memory_scores(section: &mut ContextSection, path: &std::path::Path) {
let data = match std::fs::read_to_string(path) { let data = match std::fs::read_to_string(path) {
Ok(d) => d, Ok(d) => d,
Err(_) => return, Err(_) => return,
@ -41,10 +43,10 @@ fn load_memory_scores(entries: &mut [ConversationEntry], path: &std::path::Path)
Err(_) => return, Err(_) => return,
}; };
let mut applied = 0; let mut applied = 0;
for entry in entries.iter_mut() { for i in 0..section.len() {
if let ConversationEntry::Memory { key, score, .. } = entry { if let ConversationEntry::Memory { key, .. } = &section.entries()[i].entry {
if let Some(&s) = scores.get(key.as_str()) { if let Some(&s) = scores.get(key.as_str()) {
*score = Some(s); section.set_score(i, Some(s));
applied += 1; applied += 1;
} }
} }
@ -55,10 +57,10 @@ fn load_memory_scores(entries: &mut [ConversationEntry], path: &std::path::Path)
} }
/// Save all memory scores to disk. /// Save all memory scores to disk.
fn save_memory_scores(entries: &[ConversationEntry], path: &std::path::Path) { fn save_memory_scores(section: &ContextSection, path: &std::path::Path) {
let scores: std::collections::BTreeMap<String, f64> = entries.iter() let scores: std::collections::BTreeMap<String, f64> = section.entries().iter()
.filter_map(|e| { .filter_map(|ce| {
if let ConversationEntry::Memory { key, score: Some(s), .. } = e { if let ConversationEntry::Memory { key, score: Some(s), .. } = &ce.entry {
Some((key.clone(), *s)) Some((key.clone(), *s))
} else { } else {
None None
@ -313,7 +315,7 @@ impl Mind {
// Restore persisted memory scores // Restore persisted memory scores
let scores_path = self.config.session_dir.join("memory-scores.json"); let scores_path = self.config.session_dir.join("memory-scores.json");
load_memory_scores(&mut ag.context.entries, &scores_path); load_memory_scores(&mut ag.context.conversation, &scores_path);
ag.changed.notify_one(); ag.changed.notify_one();
drop(ag); drop(ag);
@ -403,16 +405,16 @@ impl Mind {
if let Ok(ref scores) = result { if let Ok(ref scores) = result {
// Write scores onto Memory entries // Write scores onto Memory entries
for (key, weight) in scores { for (key, weight) in scores {
for entry in &mut ag.context.entries { for i in 0..ag.context.conversation.len() {
if let crate::agent::context::ConversationEntry::Memory { if let ConversationEntry::Memory { key: k, .. } = &ag.context.conversation.entries()[i].entry {
key: k, score, .. if k == key {
} = entry { ag.context.conversation.set_score(i, Some(*weight));
if k == key { *score = Some(*weight); } }
} }
} }
} }
// Persist all scores to disk // Persist all scores to disk
save_memory_scores(&ag.context.entries, &scores_path); save_memory_scores(&ag.context.conversation, &scores_path);
} }
} }
let _ = bg_tx.send(BgEvent::ScoringDone); let _ = bg_tx.send(BgEvent::ScoringDone);

52
src/subconscious/learn.rs
View file

@ -16,7 +16,7 @@
use crate::agent::api::ApiClient; use crate::agent::api::ApiClient;
use crate::agent::api::*; use crate::agent::api::*;
use crate::agent::context::{ConversationEntry, ContextState}; use crate::agent::context::{ConversationEntry, ContextEntry, ContextState};
const SCORE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120); const SCORE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
@ -39,19 +39,22 @@ fn build_messages(
range: std::ops::Range<usize>, range: std::ops::Range<usize>,
filter: Filter, filter: Filter,
) -> Vec<serde_json::Value> { ) -> Vec<serde_json::Value> {
let mut msgs = vec![ let mut msgs = Vec::new();
serde_json::json!({"role": "system", "content": &context.system_prompt}), for e in context.system.entries() {
]; msgs.push(serde_json::json!({"role": "system", "content": e.entry.message().content_text()}));
}
let ctx = context.render_context_message(); let ctx = context.render_context_message();
if !ctx.is_empty() { if !ctx.is_empty() {
msgs.push(serde_json::json!({"role": "user", "content": ctx})); msgs.push(serde_json::json!({"role": "user", "content": ctx}));
} }
let entries = context.conversation.entries();
for i in range { for i in range {
let entry = &context.entries[i]; let ce = &entries[i];
let entry = &ce.entry;
let skip = match &filter { let skip = match &filter {
Filter::None => false, Filter::None => false,
Filter::SkipIndex(idx) => i == *idx, Filter::SkipIndex(idx) => i == *idx,
Filter::SkipKey(key) => matches!(entry, ConversationEntry::Memory { key: k, .. } if k == key), Filter::SkipKey(key) => matches!(entry, ConversationEntry::Memory { key: k, .. } if k == *key),
Filter::SkipAllMemories => entry.is_memory(), Filter::SkipAllMemories => entry.is_memory(),
}; };
if skip { continue; } if skip { continue; }
@ -175,16 +178,16 @@ pub async fn score_memories(
context: &ContextState, context: &ContextState,
client: &ApiClient, client: &ApiClient,
) -> anyhow::Result<MemoryScore> { ) -> anyhow::Result<MemoryScore> {
let mut memory_keys: Vec<String> = context.entries.iter() let mut memory_keys: Vec<String> = context.conversation.entries().iter()
.filter_map(|e| match e { .filter_map(|ce| match &ce.entry {
ConversationEntry::Memory { key, .. } => Some(key.clone()), ConversationEntry::Memory { key, .. } => Some(key.clone()),
_ => None, _ => None,
}) })
.collect(); .collect();
memory_keys.dedup(); memory_keys.dedup();
let response_indices: Vec<usize> = context.entries.iter().enumerate() let response_indices: Vec<usize> = context.conversation.entries().iter().enumerate()
.filter(|(_, e)| e.message().role == Role::Assistant) .filter(|(_, ce)| ce.entry.message().role == Role::Assistant)
.map(|(i, _)| i) .map(|(i, _)| i)
.collect(); .collect();
@ -198,7 +201,7 @@ pub async fn score_memories(
let http = http_client(); let http = http_client();
let range = 0..context.entries.len(); let range = 0..context.conversation.entries().len();
let baseline = call_score(&http, client, &build_messages(context, range.clone(), Filter::None)).await?; let baseline = call_score(&http, client, &build_messages(context, range.clone(), Filter::None)).await?;
@ -242,10 +245,10 @@ pub async fn score_memories(
/// Find the entry index after `start` that contains the Nth assistant response. /// Find the entry index after `start` that contains the Nth assistant response.
/// Returns (end_index, true) if N responses were found, (entries.len(), false) if not. /// Returns (end_index, true) if N responses were found, (entries.len(), false) if not.
fn nth_response_end(entries: &[ConversationEntry], start: usize, n: usize) -> (usize, bool) { fn nth_response_end(entries: &[ContextEntry], start: usize, n: usize) -> (usize, bool) {
let mut count = 0; let mut count = 0;
for i in start..entries.len() { for i in start..entries.len() {
if entries[i].message().role == Role::Assistant { if entries[i].entry.message().role == Role::Assistant {
count += 1; count += 1;
if count >= n { return (i + 1, true); } if count >= n { return (i + 1, true); }
} }
@ -267,16 +270,17 @@ pub async fn score_memory(
) -> anyhow::Result<f64> { ) -> anyhow::Result<f64> {
const RESPONSE_WINDOW: usize = 50; const RESPONSE_WINDOW: usize = 50;
let first_pos = match context.entries.iter().position(|e| { let entries = context.conversation.entries();
matches!(e, ConversationEntry::Memory { key: k, .. } if k == key) let first_pos = match entries.iter().position(|ce| {
matches!(&ce.entry, ConversationEntry::Memory { key: k, .. } if k == key)
}) { }) {
Some(p) => p, Some(p) => p,
None => return Ok(0.0), None => return Ok(0.0),
}; };
let (end, _) = nth_response_end(&context.entries, first_pos, RESPONSE_WINDOW); let (end, _) = nth_response_end(entries, first_pos, RESPONSE_WINDOW);
let range = first_pos..end; let range = first_pos..end;
if !context.entries[range.clone()].iter().any(|e| e.message().role == Role::Assistant) { if !entries[range.clone()].iter().any(|ce| ce.entry.message().role == Role::Assistant) {
return Ok(0.0); return Ok(0.0);
} }
@ -310,8 +314,8 @@ pub async fn score_memories_incremental(
let store = crate::hippocampus::store::Store::load().unwrap_or_default(); let store = crate::hippocampus::store::Store::load().unwrap_or_default();
for (i, entry) in context.entries.iter().enumerate() { for (i, ce) in context.conversation.entries().iter().enumerate() {
if let ConversationEntry::Memory { key, .. } = entry { if let ConversationEntry::Memory { key, .. } = &ce.entry {
if !seen.insert(key.clone()) { continue; } if !seen.insert(key.clone()) { continue; }
let last_scored = store.nodes.get(key.as_str()) let last_scored = store.nodes.get(key.as_str())
.map(|n| n.last_scored) .map(|n| n.last_scored)
@ -328,18 +332,18 @@ pub async fn score_memories_incremental(
let http = http_client(); let http = http_client();
let mut results = Vec::new(); let mut results = Vec::new();
let total_entries = context.entries.len(); let total_entries = context.conversation.entries().len();
let first_quarter = total_entries / 4; let first_quarter = total_entries / 4;
for (pos, key, _) in &candidates { for (pos, key, _) in &candidates {
let (end, full_window) = nth_response_end(&context.entries, *pos, response_window); let (end, full_window) = nth_response_end(context.conversation.entries(), *pos, response_window);
// Skip memories without a full window, unless they're in the // Skip memories without a full window, unless they're in the
// first quarter of the conversation (always score those). // first quarter of the conversation (always score those).
if !full_window && *pos >= first_quarter { if !full_window && *pos >= first_quarter {
continue; continue;
} }
let range = *pos..end; let range = *pos..end;
if !context.entries[range.clone()].iter().any(|e| e.message().role == Role::Assistant) { if !context.conversation.entries()[range.clone()].iter().any(|ce| ce.entry.message().role == Role::Assistant) {
continue; continue;
} }
@ -378,10 +382,10 @@ pub async fn score_finetune(
count: usize, count: usize,
client: &ApiClient, client: &ApiClient,
) -> anyhow::Result<Vec<(usize, f64)>> { ) -> anyhow::Result<Vec<(usize, f64)>> {
let range = context.entries.len().saturating_sub(count)..context.entries.len(); let range = context.conversation.entries().len().saturating_sub(count)..context.conversation.entries().len();
let response_positions: Vec<usize> = range.clone() let response_positions: Vec<usize> = range.clone()
.filter(|&i| context.entries[i].message().role == Role::Assistant) .filter(|&i| context.conversation.entries()[i].entry.message().role == Role::Assistant)
.collect(); .collect();
if response_positions.is_empty() { if response_positions.is_empty() {
return Ok(Vec::new()); return Ok(Vec::new());

10
src/user/chat.rs
View file

@ -376,7 +376,7 @@ pub(crate) struct InteractScreen {
call_timeout_secs: u64, call_timeout_secs: u64,
// State sync with agent — double buffer // State sync with agent — double buffer
last_generation: u64, last_generation: u64,
last_entries: Vec<crate::agent::context::ConversationEntry>, last_entries: Vec<crate::agent::context::ContextEntry>,
pending_display_count: usize, pending_display_count: usize,
/// Reference to agent for state sync /// Reference to agent for state sync
agent: std::sync::Arc<tokio::sync::Mutex<crate::agent::Agent>>, agent: std::sync::Arc<tokio::sync::Mutex<crate::agent::Agent>>,
@ -482,21 +482,21 @@ impl InteractScreen {
for i in (0..self.last_entries.len()).rev() { for i in (0..self.last_entries.len()).rev() {
// Check if this entry is out of bounds or doesn't match // Check if this entry is out of bounds or doesn't match
let matches = i < entries.len() && self.last_entries[i] == entries[i]; let matches = i < entries.len() && self.last_entries[i].entry == entries[i].entry;
if !matches { if !matches {
pop = i; pop = i;
} }
// Only stop at assistant if it matches - otherwise keep going // Only stop at assistant if it matches - otherwise keep going
if matches && self.last_entries[i].message().role == Role::Assistant { if matches && self.last_entries[i].entry.message().role == Role::Assistant {
break; break;
} }
} }
while self.last_entries.len() > pop { while self.last_entries.len() > pop {
let popped = self.last_entries.pop().unwrap(); let popped = self.last_entries.pop().unwrap();
for (target, _, _) in Self::route_entry(&popped) { for (target, _, _) in Self::route_entry(&popped.entry) {
match target { match target {
PaneTarget::Conversation | PaneTarget::ConversationAssistant PaneTarget::Conversation | PaneTarget::ConversationAssistant
=> self.conversation.pop_line(), => self.conversation.pop_line(),
@ -510,7 +510,7 @@ impl InteractScreen {
// Phase 2: push new entries // Phase 2: push new entries
let start = self.last_entries.len(); let start = self.last_entries.len();
for entry in entries.iter().skip(start) { for entry in entries.iter().skip(start) {
for (target, text, marker) in Self::route_entry(entry) { for (target, text, marker) in Self::route_entry(&entry.entry) {
match target { match target {
PaneTarget::Conversation => { PaneTarget::Conversation => {
self.conversation.current_color = Color::Cyan; self.conversation.current_color = Color::Cyan;

10
src/user/context.rs
View file

@ -23,9 +23,9 @@ impl ConsciousScreen {
Self { agent, tree: SectionTree::new() } Self { agent, tree: SectionTree::new() }
} }
fn read_context_state(&self) -> Vec<crate::agent::context::ContextSection> { fn read_context_sections(&self) -> Vec<crate::agent::context::ContextSection> {
match self.agent.try_lock() { match self.agent.try_lock() {
Ok(ag) => ag.context_state_summary(), Ok(ag) => ag.context_sections().iter().map(|s| (*s).clone()).collect(),
Err(_) => Vec::new(), Err(_) => Vec::new(),
} }
} }
@ -39,7 +39,7 @@ impl ScreenView for ConsciousScreen {
for event in events { for event in events {
if let ratatui::crossterm::event::Event::Key(key) = event { if let ratatui::crossterm::event::Event::Key(key) = event {
if key.kind != ratatui::crossterm::event::KeyEventKind::Press { continue; } if key.kind != ratatui::crossterm::event::KeyEventKind::Press { continue; }
let context_state = self.read_context_state(); let context_state = self.read_context_sections();
self.tree.handle_nav(key.code, &context_state, area.height); self.tree.handle_nav(key.code, &context_state, area.height);
} }
} }
@ -64,9 +64,9 @@ impl ScreenView for ConsciousScreen {
if !app.status.context_budget.is_empty() { if !app.status.context_budget.is_empty() {
lines.push(Line::raw(format!(" Budget: {}", app.status.context_budget))); lines.push(Line::raw(format!(" Budget: {}", app.status.context_budget)));
} }
let context_state = self.read_context_state(); let context_state = self.read_context_sections();
if !context_state.is_empty() { if !context_state.is_empty() {
let total: usize = context_state.iter().map(|s| s.tokens).sum(); let total: usize = context_state.iter().map(|s| s.tokens()).sum();
lines.push(Line::raw("")); lines.push(Line::raw(""));
lines.push(Line::styled( lines.push(Line::styled(
" (↑/↓ select, →/Enter expand, ← collapse, PgUp/PgDn scroll)", " (↑/↓ select, →/Enter expand, ← collapse, PgUp/PgDn scroll)",

26
src/user/subconscious.rs
View file

@ -16,7 +16,8 @@ use ratatui::{
use super::{App, ScreenView, screen_legend}; use super::{App, ScreenView, screen_legend};
use super::widgets::{SectionTree, pane_block_focused, render_scrollable, tree_legend, format_age, format_ts_age}; use super::widgets::{SectionTree, pane_block_focused, render_scrollable, tree_legend, format_age, format_ts_age};
use crate::agent::context::ContextSection; use crate::agent::context::{ContextSection, ContextEntry, ConversationEntry};
use crate::agent::api::Message;
#[derive(Clone, Copy, PartialEq)] #[derive(Clone, Copy, PartialEq)]
enum Pane { Agents, Outputs, History, Context } enum Pane { Agents, Outputs, History, Context }
@ -135,12 +136,13 @@ impl SubconsciousScreen {
None => return Vec::new(), None => return Vec::new(),
}; };
snap.state.iter().map(|(key, val)| { snap.state.iter().map(|(key, val)| {
ContextSection { let mut section = ContextSection::new(key.clone());
name: key.clone(), section.push(ContextEntry {
entry: ConversationEntry::Message(Message::user(val)),
tokens: 0, tokens: 0,
content: val.clone(), timestamp: None,
children: Vec::new(), });
} section
}).collect() }).collect()
} }
@ -151,7 +153,15 @@ impl SubconsciousScreen {
}; };
snap.forked_agent.as_ref() snap.forked_agent.as_ref()
.and_then(|agent| agent.try_lock().ok()) .and_then(|agent| agent.try_lock().ok())
.map(|ag| ag.conversation_sections_from(snap.fork_point)) .map(|ag| {
// Build a single section from the forked conversation entries
let entries = ag.conversation_entries_from(snap.fork_point);
let mut section = ContextSection::new("Conversation");
for e in entries {
section.push(e.clone());
}
vec![section]
})
.unwrap_or_default() .unwrap_or_default()
} }
@ -172,7 +182,7 @@ impl SubconsciousScreen {
.unwrap_or_else(|| "—".to_string()); .unwrap_or_else(|| "—".to_string());
let entries = snap.forked_agent.as_ref() let entries = snap.forked_agent.as_ref()
.and_then(|a| a.try_lock().ok()) .and_then(|a| a.try_lock().ok())
.map(|ag| ag.context.entries.len().saturating_sub(snap.fork_point)) .map(|ag| ag.context.conversation.len().saturating_sub(snap.fork_point))
.unwrap_or(0); .unwrap_or(0);
ListItem::from(Line::from(vec![ ListItem::from(Line::from(vec![
Span::styled(&snap.name, Style::default().fg(Color::Gray)), Span::styled(&snap.name, Style::default().fg(Color::Gray)),

73
src/user/widgets.rs
View file

@ -99,28 +99,25 @@ impl SectionTree {
} }
/// Total nodes in the tree (regardless of expand state). /// Total nodes in the tree (regardless of expand state).
/// Each section is 1 node, each entry within is 1 node.
fn total_nodes(&self, sections: &[ContextSection]) -> usize { fn total_nodes(&self, sections: &[ContextSection]) -> usize {
fn count_all(section: &ContextSection) -> usize { sections.iter().map(|s| 1 + s.entries().len()).sum()
1 + section.children.iter().map(|c| count_all(c)).sum::<usize>()
}
sections.iter().map(|s| count_all(s)).sum()
} }
pub fn item_count(&self, sections: &[ContextSection]) -> usize { pub fn item_count(&self, sections: &[ContextSection]) -> usize {
fn count(section: &ContextSection, expanded: &std::collections::HashSet<usize>, idx: &mut usize) -> usize { let mut idx = 0;
let my_idx = *idx; let mut total = 0;
*idx += 1; for section in sections {
let mut total = 1; let my_idx = idx;
if expanded.contains(&my_idx) { idx += 1;
for child in &section.children { total += 1;
total += count(child, expanded, idx); if self.expanded.contains(&my_idx) {
total += section.entries().len();
idx += section.entries().len();
} }
} }
total total
} }
let mut idx = 0;
sections.iter().map(|s| count(s, &self.expanded, &mut idx)).sum()
}
pub fn handle_nav(&mut self, code: KeyCode, sections: &[ContextSection], height: u16) { pub fn handle_nav(&mut self, code: KeyCode, sections: &[ContextSection], height: u16) {
let item_count = self.item_count(sections); let item_count = self.item_count(sections);
@ -193,27 +190,24 @@ impl SectionTree {
pub fn render_sections(&self, sections: &[ContextSection], lines: &mut Vec<Line>) { pub fn render_sections(&self, sections: &[ContextSection], lines: &mut Vec<Line>) {
let mut idx = 0; let mut idx = 0;
for section in sections { for section in sections {
self.render_one(section, 0, lines, &mut idx); self.render_one(section, lines, &mut idx);
} }
} }
fn render_one( fn render_one(
&self, &self,
section: &ContextSection, section: &ContextSection,
depth: usize,
lines: &mut Vec<Line>, lines: &mut Vec<Line>,
idx: &mut usize, idx: &mut usize,
) { ) {
let my_idx = *idx; let my_idx = *idx;
let selected = self.selected == Some(my_idx); let selected = self.selected == Some(my_idx);
let expanded = self.expanded.contains(&my_idx); let expanded = self.expanded.contains(&my_idx);
let has_children = !section.children.is_empty(); let expandable = !section.is_empty();
let has_content = !section.content.is_empty();
let expandable = has_children || has_content;
let indent = " ".repeat(depth + 1);
let marker = if !expandable { " " } else if expanded { "▼" } else { "▶" }; let marker = if !expandable { " " } else if expanded { "▼" } else { "▶" };
let label = format!("{}{} {:30} {:>6} tokens", indent, marker, section.name, section.tokens); let label = format!(" {} {:30} {:>6} tokens",
marker, format!("{} ({})", section.name, section.len()), section.tokens());
let style = if selected { let style = if selected {
Style::default().fg(Color::Yellow).add_modifier(Modifier::BOLD) Style::default().fg(Color::Yellow).add_modifier(Modifier::BOLD)
} else { } else {
@ -223,23 +217,41 @@ impl SectionTree {
*idx += 1; *idx += 1;
if expanded { if expanded {
if has_children { for ce in section.entries() {
for child in &section.children { let entry_selected = self.selected == Some(*idx);
self.render_one(child, depth + 1, lines, idx); let entry_expanded = self.expanded.contains(idx);
} let text = ce.entry.message().content_text();
} else if has_content { let preview: String = text.chars().take(60).collect();
let content_indent = format!("{} │ ", " ".repeat(depth + 1)); let preview = preview.replace('\n', " ");
let content_lines: Vec<&str> = section.content.lines().collect(); let label = if preview.len() < text.len() {
format!(" {}...", preview)
} else {
format!(" {}", preview)
};
let entry_marker = if text.len() > 60 {
if entry_expanded { "▼" } else { "▶" }
} else { " " };
let entry_label = format!(" {} {:>6} {}", entry_marker, ce.tokens, label);
let style = if entry_selected {
Style::default().fg(Color::Yellow).add_modifier(Modifier::BOLD)
} else {
Style::default().fg(Color::DarkGray)
};
lines.push(Line::styled(entry_label, style));
*idx += 1;
if entry_expanded {
let content_lines: Vec<&str> = text.lines().collect();
let show = content_lines.len().min(50); let show = content_lines.len().min(50);
for line in &content_lines[..show] { for line in &content_lines[..show] {
lines.push(Line::styled( lines.push(Line::styled(
format!("{}{}", content_indent, line), format!(" │ {}", line),
Style::default().fg(Color::DarkGray), Style::default().fg(Color::DarkGray),
)); ));
} }
if content_lines.len() > 50 { if content_lines.len() > 50 {
lines.push(Line::styled( lines.push(Line::styled(
format!("{}... ({} more lines)", content_indent, content_lines.len() - 50), format!(" ... ({} more lines)", content_lines.len() - 50),
Style::default().fg(Color::DarkGray), Style::default().fg(Color::DarkGray),
)); ));
} }
@ -247,3 +259,4 @@ impl SectionTree {
} }
} }
} }
}