Move conversation messages into ContextState

ContextState now owns everything in the context window:
system_prompt, personality, journal, working_stack, loaded_nodes,
and conversation messages. No duplication — each piece exists once
in its typed form.

assemble_api_messages() renders the full message list on the fly
from typed sources. measure_budget() counts each bucket from its
source directly. push_context() removed — identity/journal are
never pushed as messages.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-02 02:47:32 -04:00
parent 4580f5dade
commit a0aacfc552
2 changed files with 44 additions and 45 deletions

View file

@ -53,7 +53,6 @@ struct DispatchState {
pub struct Agent {
client: ApiClient,
messages: Vec<Message>,
tool_defs: Vec<ToolDef>,
/// Last known prompt token count from the API (tracks context size).
last_prompt_tokens: u32,
@ -79,6 +78,7 @@ pub struct Agent {
}
fn render_journal(entries: &[journal::JournalEntry]) -> String {
if entries.is_empty() { return String::new(); }
let mut text = String::from("[Earlier — from your journal]\n\n");
for entry in entries {
use std::fmt::Write;
@ -105,12 +105,12 @@ impl Agent {
journal: Vec::new(),
working_stack: Vec::new(),
loaded_nodes: Vec::new(),
messages: Vec::new(),
};
let session_id = format!("poc-agent-{}", chrono::Utc::now().format("%Y%m%d-%H%M%S"));
let agent_cycles = crate::subconscious::subconscious::AgentCycleState::new(&session_id);
let mut agent = Self {
client,
messages: Vec::new(),
tool_defs,
last_prompt_tokens: 0,
process_tracker: ProcessTracker::new(),
@ -124,23 +124,30 @@ impl Agent {
agent_cycles,
};
// Load recent journal entries at startup for orientation
agent.load_startup_journal();
agent.load_working_stack();
agent.push_context(Message::system(system_prompt));
let rendered = agent.context.render_context_message();
if !rendered.is_empty() {
agent.push_context(Message::user(rendered));
}
if !agent.context.journal.is_empty() {
agent.push_context(Message::user(render_journal(&agent.context.journal)));
}
agent.measure_budget();
agent.publish_context_state();
agent
}
/// Assemble the full message list for the API call from typed sources.
/// System prompt + personality context + journal + conversation messages.
fn assemble_api_messages(&self) -> Vec<Message> {
let mut msgs = Vec::new();
msgs.push(Message::system(&self.context.system_prompt));
let ctx = self.context.render_context_message();
if !ctx.is_empty() {
msgs.push(Message::user(ctx));
}
let jnl = render_journal(&self.context.journal);
if !jnl.is_empty() {
msgs.push(Message::user(jnl));
}
msgs.extend(self.context.messages.iter().cloned());
msgs
}
/// Run agent orchestration cycle and return formatted output to inject.
fn run_agent_cycle(&mut self) -> Option<String> {
let transcript_path = self.conversation_log.as_ref()
@ -170,16 +177,12 @@ impl Agent {
eprintln!("warning: failed to log message: {:#}", e);
}
}
self.messages.push(msg);
self.context.messages.push(msg);
}
/// Push a context-only message (system prompt, identity context,
/// journal summaries). Not logged — these are reconstructed on
/// every startup/compaction.
fn push_context(&mut self, msg: Message) {
self.messages.push(msg);
}
/// Measure context window usage by category. Uses the BPE tokenizer
/// for direct token counting (no chars/4 approximation).
fn measure_budget(&mut self) {
@ -192,9 +195,8 @@ impl Agent {
.map(|e| count(&e.content)).sum();
let mem_tokens: usize = self.context.loaded_nodes.iter()
.map(|node| count(&node.render())).sum();
let total: usize = self.messages.iter()
let conv_tokens: usize = self.context.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,
@ -239,8 +241,9 @@ impl Agent {
// Stream events from the API — we route each event to the
// appropriate UI pane rather than letting the API layer do it.
let api_messages = self.assemble_api_messages();
let mut rx = self.client.start_stream(
&self.messages,
&api_messages,
Some(&self.tool_defs),
ui_tx,
&self.reasoning_effort,
@ -691,10 +694,10 @@ impl Agent {
}
// Conversation — each message as a child
let conv_start = self.messages.iter()
let conv_start = self.context.messages.iter()
.position(|m| m.role == Role::Assistant || m.role == Role::Tool)
.unwrap_or(self.messages.len());
let conv_messages = &self.messages[conv_start..];
.unwrap_or(self.context.messages.len());
let conv_messages = &self.context.messages[conv_start..];
let conv_children: Vec<ContextSection> = conv_messages.iter().enumerate()
.map(|(i, msg)| {
let text = msg.content.as_ref()
@ -824,13 +827,13 @@ impl Agent {
dbg_log!("[journal] context.journal now has {} entries", self.context.journal.len());
}
/// Re-render the context message in self.messages from live ContextState.
/// Re-render the context message in self.context.messages from live ContextState.
/// Called after any change to context state (working stack, etc).
fn refresh_context_message(&mut self) {
let rendered = self.context.render_context_message();
// The context message is the first user message (index 1, after system prompt)
if self.messages.len() >= 2 && self.messages[1].role == Role::User {
self.messages[1] = Message::user(rendered);
if self.context.messages.len() >= 2 && self.context.messages[1].role == Role::User {
self.context.messages[1] = Message::user(rendered);
}
self.publish_context_state();
self.save_working_stack();
@ -864,7 +867,7 @@ impl Agent {
/// all previous ones. The tool result message (right before each image
/// message) already records what was loaded, so no info is lost.
fn age_out_images(&mut self) {
for msg in &mut self.messages {
for msg in &mut self.context.messages {
if let Some(MessageContent::Parts(parts)) = &msg.content {
let has_images = parts.iter().any(|p| matches!(p, ContentPart::ImageUrl { .. }));
if !has_images {
@ -909,7 +912,7 @@ impl Agent {
let mut strip_ids: Vec<String> = Vec::new();
let mut strip_msg_indices: Vec<usize> = Vec::new();
for (i, msg) in self.messages.iter().enumerate() {
for (i, msg) in self.context.messages.iter().enumerate() {
if msg.role != Role::Assistant {
continue;
}
@ -935,7 +938,7 @@ impl Agent {
}
// Remove in reverse order to preserve indices
self.messages.retain(|msg| {
self.context.messages.retain(|msg| {
// Strip the assistant messages we identified
if msg.role == Role::Assistant {
if let Some(calls) = &msg.tool_calls {
@ -973,14 +976,7 @@ impl Agent {
/// Internal compaction — rebuilds context window from current messages.
fn do_compact(&mut self) {
// Find where actual conversation starts (after system + context)
let conv_start = self
.messages
.iter()
.position(|m| m.role == Role::Assistant || m.role == Role::Tool)
.unwrap_or(self.messages.len());
let conversation: Vec<Message> = self.messages[conv_start..].to_vec();
let conversation: Vec<Message> = self.context.messages.clone();
let (messages, journal) = crate::agent::context::build_context_window(
&self.context,
&conversation,
@ -988,7 +984,7 @@ impl Agent {
&self.tokenizer,
);
self.context.journal = journal::parse_journal_text(&journal);
self.messages = messages;
self.context.messages = messages;
self.last_prompt_tokens = 0;
self.measure_budget();
self.publish_context_state();
@ -1050,8 +1046,8 @@ impl Agent {
dbglog!("[restore] journal text: {} chars, {} lines",
journal.len(), journal.lines().count());
self.context.journal = journal::parse_journal_text(&journal);
self.messages = messages;
dbglog!("[restore] built context window: {} messages", self.messages.len());
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();
@ -1070,17 +1066,17 @@ impl Agent {
/// Get the conversation history for persistence.
pub fn messages(&self) -> &[Message] {
&self.messages
&self.context.messages
}
/// Mutable access to conversation history (for /retry).
pub fn messages_mut(&mut self) -> &mut Vec<Message> {
&mut self.messages
&mut self.context.messages
}
/// Restore from a saved conversation.
pub fn restore(&mut self, messages: Vec<Message>) {
self.messages = messages;
self.context.messages = messages;
}
}