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

ConversationEntry enum: typed memory vs conversation messages

Browse source 
Replace untyped message list with ConversationEntry enum:
- Message(Message) — regular conversation turn
- Memory { key, message } — memory content with preserved message
  for KV cache round-tripping

Budget counts memory vs conversation by matching on enum variant.
Debug screen labels memory entries with [memory: key]. No heuristic
tool-name scanning.

Custom serde: Memory serializes with a memory_key field alongside
the message fields, deserializes by checking for the field.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-02 03:26:00 -04:00
parent eb4dae04cb
commit b9e3568385
3 changed files with 153 additions and 93 deletions
Download patch file Download diff file Expand all files Collapse all files

85
src/agent/runner.rs
View file

@ -103,7 +103,7 @@ impl Agent {
journal: Vec::new(),
working_stack: Vec::new(),
loaded_nodes: Vec::new(),
messages: Vec::new(),
entries: 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);
@ -140,7 +140,7 @@ impl Agent {
if !jnl.is_empty() {
msgs.push(Message::user(jnl));
}
msgs.extend(self.context.messages.iter().cloned());
msgs.extend(self.context.entries.iter().map(|e| e.api_message().clone()));
msgs
}
@ -173,7 +173,7 @@ impl Agent {
eprintln!("warning: failed to log message: {:#}", e);
}
}
self.context.messages.push(msg);
self.context.entries.push(ConversationEntry::Message(msg));
}
/// Push a context-only message (system prompt, identity context,
@ -673,32 +673,41 @@ impl Agent {
}
// Conversation — each message as a child
let conv_messages = &self.context.messages;
let conv_messages = &self.context.entries;
let conv_children: Vec<ContextSection> = conv_messages.iter().enumerate()
.map(|(i, msg)| {
let text = msg.content.as_ref()
.map(|(i, entry)| {
let m = entry.message();
let text = m.content.as_ref()
.map(|c| c.as_text().to_string())
.unwrap_or_default();
let tool_info = msg.tool_calls.as_ref().map(|tc| {
let tool_info = m.tool_calls.as_ref().map(|tc| {
tc.iter()
.map(|c| c.function.name.clone())
.collect::<Vec<_>>()
.join(", ")
});
let label = match (&msg.role, &tool_info) {
(_, Some(tools)) => format!("[tool_call: {}]", tools),
_ => {
let preview: String = text.chars().take(60).collect();
let preview = preview.replace('\n', " ");
if text.len() > 60 { format!("{}...", preview) } else { preview }
let label = if entry.is_memory() {
if let ConversationEntry::Memory { key, .. } = entry {
format!("[memory: {}]", key)
} else { unreachable!() }
} else {
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 tokens = count(&text);
let role_name = match msg.role {
Role::Assistant => "PoC",
Role::User => "Kent",
Role::Tool => "tool",
Role::System => "system",
let role_name = if entry.is_memory() { "mem" } else {
match m.role {
Role::Assistant => "PoC",
Role::User => "Kent",
Role::Tool => "tool",
Role::System => "system",
}
};
ContextSection {
name: format!("[{}] {}: {}", i, role_name, label),
@ -846,7 +855,8 @@ 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.context.messages {
for entry in &mut self.context.entries {
let msg = entry.message_mut();
if let Some(MessageContent::Parts(parts)) = &msg.content {
let has_images = parts.iter().any(|p| matches!(p, ContentPart::ImageUrl { .. }));
if !has_images {
@ -891,7 +901,8 @@ impl Agent {
let mut strip_ids: Vec<String> = Vec::new();
let mut strip_msg_indices: Vec<usize> = Vec::new();
for (i, msg) in self.context.messages.iter().enumerate() {
for (i, entry) in self.context.entries.iter().enumerate() {
let msg = entry.message();
if msg.role != Role::Assistant {
continue;
}
@ -917,8 +928,8 @@ impl Agent {
}
// Remove in reverse order to preserve indices
self.context.messages.retain(|msg| {
// Strip the assistant messages we identified
self.context.entries.retain(|entry| {
let msg = entry.message();
if msg.role == Role::Assistant {
if let Some(calls) = &msg.tool_calls {
if calls.iter().all(|c| strip_ids.contains(&c.id)) {
@ -926,7 +937,6 @@ impl Agent {
}
}
}
// Strip matching tool results
if msg.role == Role::Tool {
if let Some(ref id) = msg.tool_call_id {
if strip_ids.contains(id) {
@ -955,7 +965,8 @@ impl Agent {
/// Internal compaction — rebuilds context window from current messages.
fn do_compact(&mut self) {
let conversation: Vec<Message> = self.context.messages.clone();
let conversation: Vec<Message> = self.context.entries.iter()
.map(|e| e.api_message().clone()).collect();
let (messages, journal) = crate::agent::context::build_context_window(
&self.context,
&conversation,
@ -963,7 +974,8 @@ impl Agent {
&self.tokenizer,
);
self.context.journal = journal::parse_journal_text(&journal);
self.context.messages = messages;
self.context.entries = messages.into_iter()
.map(ConversationEntry::Message).collect();
self.last_prompt_tokens = 0;
self.publish_context_state();
@ -1025,8 +1037,9 @@ impl Agent {
dbglog!("[restore] journal text: {} chars, {} lines",
journal.len(), journal.lines().count());
self.context.journal = journal::parse_journal_text(&journal);
self.context.messages = messages;
dbglog!("[restore] built context window: {} messages", self.context.messages.len());
self.context.entries = messages.into_iter()
.map(ConversationEntry::Message).collect();
dbglog!("[restore] built context window: {} entries", self.context.entries.len());
self.last_prompt_tokens = 0;
self.publish_context_state();
@ -1043,19 +1056,19 @@ impl Agent {
&self.client.model
}
/// Get the conversation history for persistence.
pub fn messages(&self) -> &[Message] {
&self.context.messages
/// Get the conversation entries for persistence.
pub fn entries(&self) -> &[ConversationEntry] {
&self.context.entries
}
/// Mutable access to conversation history (for /retry).
pub fn messages_mut(&mut self) -> &mut Vec<Message> {
&mut self.context.messages
/// Mutable access to conversation entries (for /retry).
pub fn entries_mut(&mut self) -> &mut Vec<ConversationEntry> {
&mut self.context.entries
}
/// Restore from a saved conversation.
pub fn restore(&mut self, messages: Vec<Message>) {
self.context.messages = messages;
/// Restore from saved conversation entries.
pub fn restore(&mut self, entries: Vec<ConversationEntry>) {
self.context.entries = entries;
}
}

127
src/agent/types.rs
View file

@ -322,19 +322,93 @@ impl ToolDef {
}
/// Mutable context state — the structured regions of the context window.
/// Conversation entry — either a regular message or memory content.
/// Memory entries preserve the original message for KV cache round-tripping.
#[derive(Debug, Clone)]
pub enum ConversationEntry {
Message(Message),
Memory { key: String, message: 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.
impl Serialize for ConversationEntry {
fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
use serde::ser::SerializeMap;
match self {
Self::Message(m) => m.serialize(s),
Self::Memory { key, message } => {
// Serialize message fields + memory_key
let json = serde_json::to_value(message).map_err(serde::ser::Error::custom)?;
let mut map = s.serialize_map(None)?;
if let serde_json::Value::Object(obj) = json {
for (k, v) in obj {
map.serialize_entry(&k, &v)?;
}
}
map.serialize_entry("memory_key", key)?;
map.end()
}
}
}
}
impl<'de> Deserialize<'de> for ConversationEntry {
fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
let mut json: serde_json::Value = serde_json::Value::deserialize(d)?;
if let Some(key) = json.as_object_mut().and_then(|o| o.remove("memory_key")) {
let key = key.as_str().unwrap_or("").to_string();
let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?;
Ok(Self::Memory { key, message })
} else {
let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?;
Ok(Self::Message(message))
}
}
}
impl ConversationEntry {
/// Get the API message for sending to the model.
pub fn api_message(&self) -> &Message {
match self {
Self::Message(m) => m,
Self::Memory { message, .. } => message,
}
}
pub fn is_memory(&self) -> bool {
matches!(self, Self::Memory { .. })
}
/// Get a reference to the inner message.
pub fn message(&self) -> &Message {
match self {
Self::Message(m) => m,
Self::Memory { message, .. } => message,
}
}
/// Get a mutable reference to the inner message.
pub fn message_mut(&mut self) -> &mut Message {
match self {
Self::Message(m) => m,
Self::Memory { message, .. } => message,
}
}
}
pub struct ContextState {
pub system_prompt: String,
pub personality: Vec<(String, String)>,
pub journal: Vec<crate::agent::journal::JournalEntry>,
pub working_stack: Vec<String>,
/// Memory nodes currently loaded in the context window.
/// Memory nodes currently loaded — for debug display and refresh.
/// Content is NOT duplicated here; the actual content is in entries
/// as ConversationEntry::Memory.
pub loaded_nodes: Vec<crate::hippocampus::memory::MemoryNode>,
/// Conversation messages (user, assistant, tool turns).
/// Does NOT include system prompt, personality, or journal —
/// those are rendered from their typed sources when assembling
/// the API call.
pub messages: Vec<Message>,
/// Conversation entries — messages and memory, interleaved in order.
/// Does NOT include system prompt, personality, or journal.
pub entries: Vec<ConversationEntry>,
}
// TODO: these should not be hardcoded absolute paths
@ -349,7 +423,12 @@ impl ContextState {
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();
let (mem, conv) = self.split_memory_conversation(count_msg);
let mut mem = 0;
let mut conv = 0;
for entry in &self.entries {
let tokens = count_msg(entry.api_message());
if entry.is_memory() { mem += tokens } else { conv += tokens }
}
ContextBudget {
identity_tokens: id,
memory_tokens: mem,
@ -359,40 +438,6 @@ impl ContextState {
}
}
/// 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<String> = 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<String> = self.personality.iter()
.map(|(name, content)| format!("## {}\n\n{}", name, content))

34
src/bin/poc-agent.rs
View file

@ -464,9 +464,9 @@ impl Session {
}
"/context" => {
if let Ok(agent) = self.agent.try_lock() {
let msgs = agent.messages();
let msgs = agent.entries();
let total_chars: usize =
msgs.iter().map(|m| m.content_text().len()).sum();
msgs.iter().map(|e| e.message().content_text().len()).sum();
let prompt_tokens = agent.last_prompt_tokens();
let threshold = compaction_threshold(agent.model(), &self.config.app);
let _ = self.ui_tx.send(UiMessage::Info(format!(
@ -587,15 +587,15 @@ impl Session {
return Command::Handled;
}
let mut agent_guard = self.agent.lock().await;
let msgs = agent_guard.messages_mut();
let entries = agent_guard.entries_mut();
let mut last_user_text = None;
while let Some(msg) = msgs.last() {
if msg.role == poc_memory::agent::types::Role::User {
while let Some(entry) = entries.last() {
if entry.message().role == poc_memory::agent::types::Role::User {
last_user_text =
Some(msgs.pop().unwrap().content_text().to_string());
Some(entries.pop().unwrap().message().content_text().to_string());
break;
}
msgs.pop();
entries.pop();
}
drop(agent_guard);
match last_user_text {
@ -936,7 +936,7 @@ async fn run(cli: cli::CliArgs) -> Result<()> {
config.context_parts.clone(),
);
if restored {
replay_session_to_ui(agent_guard.messages(), &ui_tx);
replay_session_to_ui(agent_guard.entries(), &ui_tx);
let _ = ui_tx.send(UiMessage::Info(
"--- restored from conversation log ---".into(),
));
@ -944,7 +944,7 @@ async fn run(cli: cli::CliArgs) -> Result<()> {
if let Ok(data) = std::fs::read_to_string(&session_file) {
if let Ok(messages) = serde_json::from_str(&data) {
agent_guard.restore(messages);
replay_session_to_ui(agent_guard.messages(), &ui_tx);
replay_session_to_ui(agent_guard.entries(), &ui_tx);
let _ = ui_tx.send(UiMessage::Info(
"--- restored from session file ---".into(),
));
@ -1104,7 +1104,7 @@ fn drain_ui_messages(rx: &mut ui_channel::UiReceiver, app: &mut tui::App) {
}
fn save_session(agent: &Agent, path: &PathBuf) -> Result<()> {
let data = serde_json::to_string_pretty(agent.messages())?;
let data = serde_json::to_string_pretty(agent.entries())?;
std::fs::write(path, data)?;
Ok(())
}
@ -1186,21 +1186,23 @@ async fn run_tool_tests(ui_tx: &ui_channel::UiSender, tracker: &tools::ProcessTr
/// conversation history immediately on restart. Shows user input,
/// assistant responses, and brief tool call summaries. Skips the system
/// prompt, context message, DMN plumbing, and image injection messages.
fn replay_session_to_ui(messages: &[types::Message], ui_tx: &ui_channel::UiSender) {
fn replay_session_to_ui(entries: &[types::ConversationEntry], ui_tx: &ui_channel::UiSender) {
use poc_memory::agent::ui_channel::StreamTarget;
dbglog!("[replay] replaying {} messages to UI", messages.len());
for (i, m) in messages.iter().enumerate() {
dbglog!("[replay] replaying {} entries to UI", entries.len());
for (i, e) in entries.iter().enumerate() {
let m = e.message();
let preview: String = m.content_text().chars().take(60).collect();
dbglog!("[replay] [{}] {:?} tc={} tcid={:?} {:?}",
i, m.role, m.tool_calls.as_ref().map_or(0, |t| t.len()),
dbglog!("[replay] [{}] {:?} mem={} tc={} tcid={:?} {:?}",
i, m.role, e.is_memory(), m.tool_calls.as_ref().map_or(0, |t| t.len()),
m.tool_call_id.as_deref(), preview);
}
let mut seen_first_user = false;
let mut target = StreamTarget::Conversation;
for msg in messages {
for entry in entries {
let msg = entry.message();
match msg.role {
types::Role::System => {}
types::Role::User => {