WIP: ContextEntry/ContextSection data structures for incremental token counting

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

View file

@ -99,27 +99,24 @@ impl SectionTree {
}
/// 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 count_all(section: &ContextSection) -> usize {
1 + section.children.iter().map(|c| count_all(c)).sum::<usize>()
}
sections.iter().map(|s| count_all(s)).sum()
sections.iter().map(|s| 1 + s.entries().len()).sum()
}
pub fn item_count(&self, sections: &[ContextSection]) -> usize {
fn count(section: &ContextSection, expanded: &std::collections::HashSet<usize>, idx: &mut usize) -> usize {
let my_idx = *idx;
*idx += 1;
let mut total = 1;
if expanded.contains(&my_idx) {
for child in &section.children {
total += count(child, expanded, idx);
}
}
total
}
let mut idx = 0;
sections.iter().map(|s| count(s, &self.expanded, &mut idx)).sum()
let mut total = 0;
for section in sections {
let my_idx = idx;
idx += 1;
total += 1;
if self.expanded.contains(&my_idx) {
total += section.entries().len();
idx += section.entries().len();
}
}
total
}
pub fn handle_nav(&mut self, code: KeyCode, sections: &[ContextSection], height: u16) {
@ -193,27 +190,24 @@ impl SectionTree {
pub fn render_sections(&self, sections: &[ContextSection], lines: &mut Vec<Line>) {
let mut idx = 0;
for section in sections {
self.render_one(section, 0, lines, &mut idx);
self.render_one(section, lines, &mut idx);
}
}
fn render_one(
&self,
section: &ContextSection,
depth: usize,
lines: &mut Vec<Line>,
idx: &mut usize,
) {
let my_idx = *idx;
let selected = self.selected == Some(my_idx);
let expanded = self.expanded.contains(&my_idx);
let has_children = !section.children.is_empty();
let has_content = !section.content.is_empty();
let expandable = has_children || has_content;
let expandable = !section.is_empty();
let indent = " ".repeat(depth + 1);
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 {
Style::default().fg(Color::Yellow).add_modifier(Modifier::BOLD)
} else {
@ -223,25 +217,44 @@ impl SectionTree {
*idx += 1;
if expanded {
if has_children {
for child in &section.children {
self.render_one(child, depth + 1, lines, idx);
}
} else if has_content {
let content_indent = format!("{}", " ".repeat(depth + 1));
let content_lines: Vec<&str> = section.content.lines().collect();
let show = content_lines.len().min(50);
for line in &content_lines[..show] {
lines.push(Line::styled(
format!("{}{}", content_indent, line),
Style::default().fg(Color::DarkGray),
));
}
if content_lines.len() > 50 {
lines.push(Line::styled(
format!("{}... ({} more lines)", content_indent, content_lines.len() - 50),
Style::default().fg(Color::DarkGray),
));
for ce in section.entries() {
let entry_selected = self.selected == Some(*idx);
let entry_expanded = self.expanded.contains(idx);
let text = ce.entry.message().content_text();
let preview: String = text.chars().take(60).collect();
let preview = preview.replace('\n', " ");
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);
for line in &content_lines[..show] {
lines.push(Line::styled(
format!("{}", line),
Style::default().fg(Color::DarkGray),
));
}
if content_lines.len() > 50 {
lines.push(Line::styled(
format!(" ... ({} more lines)", content_lines.len() - 50),
Style::default().fg(Color::DarkGray),
));
}
}
}
}