ConversationEntry enum: typed memory vs conversation messages

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

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 => {