baml_agent_tui/
chat.rs

1use ratatui::{
2    prelude::*,
3    widgets::{Block, Borders, List, ListItem, ListState},
4};
5
6/// Maximum chars for tool output preview before collapsing.
7const TOOL_OUTPUT_PREVIEW_LEN: usize = 120;
8
9/// Shared chat panel state — messages list + scroll.
10///
11/// This is the core of any agent TUI: a scrollable list of messages
12/// with role-based coloring and prefix detection.
13pub struct ChatState {
14    pub messages: Vec<ChatMessage>,
15    pub list_state: ListState,
16}
17
18/// A single chat message with metadata.
19pub struct ChatMessage {
20    pub text: String,
21    pub timestamp: String,
22    /// For tool output: whether the full text is expanded.
23    pub expanded: bool,
24}
25
26impl ChatMessage {
27    fn new(text: String) -> Self {
28        let timestamp = chrono_now();
29        Self {
30            text,
31            timestamp,
32            expanded: false,
33        }
34    }
35
36    /// Is this a long tool output that can be collapsed?
37    fn is_collapsible(&self) -> bool {
38        self.text.starts_with("  = ") && self.text.len() > TOOL_OUTPUT_PREVIEW_LEN
39    }
40
41    /// Display text, respecting collapsed state.
42    fn display_text(&self) -> String {
43        if self.is_collapsible() && !self.expanded {
44            let lines: Vec<&str> = self.text.lines().collect();
45            let first_line = lines.first().map_or("", |l| l);
46            let remaining = lines.len().saturating_sub(1);
47            if remaining > 0 {
48                format!("{} [+{} lines]", first_line, remaining)
49            } else {
50                let preview = &self.text[..TOOL_OUTPUT_PREVIEW_LEN.min(self.text.len())];
51                format!("{}... [+more]", preview)
52            }
53        } else {
54            self.text.clone()
55        }
56    }
57}
58
59fn chrono_now() -> String {
60    let now = std::time::SystemTime::now()
61        .duration_since(std::time::UNIX_EPOCH)
62        .unwrap_or_default()
63        .as_secs();
64    let hours = (now % 86400) / 3600;
65    let minutes = (now % 3600) / 60;
66    format!("{:02}:{:02}", hours, minutes)
67}
68
69impl ChatState {
70    pub fn new() -> Self {
71        let mut list_state = ListState::default();
72        list_state.select(Some(0));
73        Self {
74            messages: Vec::new(),
75            list_state,
76        }
77    }
78
79    pub fn push(&mut self, msg: String) {
80        self.messages.push(ChatMessage::new(msg));
81        self.scroll_to_bottom();
82    }
83
84    pub fn clear(&mut self) {
85        self.messages.clear();
86        self.list_state.select(Some(0));
87    }
88
89    pub fn scroll_to_bottom(&mut self) {
90        if !self.messages.is_empty() {
91            self.list_state.select(Some(self.messages.len() - 1));
92        }
93    }
94
95    /// Scroll up by one page (10 messages).
96    pub fn page_up(&mut self) {
97        let cur = self.list_state.selected().unwrap_or(0);
98        let next = cur.saturating_sub(10);
99        self.list_state.select(Some(next));
100    }
101
102    /// Scroll down by one page (10 messages).
103    pub fn page_down(&mut self) {
104        let cur = self.list_state.selected().unwrap_or(0);
105        let max = self.messages.len().saturating_sub(1);
106        let next = (cur + 10).min(max);
107        self.list_state.select(Some(next));
108    }
109
110    /// Toggle expanded state of currently selected message.
111    pub fn toggle_expand(&mut self) {
112        if let Some(idx) = self.list_state.selected() {
113            if let Some(msg) = self.messages.get_mut(idx) {
114                if msg.is_collapsible() {
115                    msg.expanded = !msg.expanded;
116                }
117            }
118        }
119    }
120
121    /// Replace the last message if it starts with `prefix`, otherwise push new.
122    pub fn replace_or_push(&mut self, prefix: &str, msg: String) {
123        if let Some(last) = self.messages.last_mut() {
124            if last.text.starts_with(prefix) {
125                last.text = msg;
126                return;
127            }
128        }
129        self.push(msg);
130    }
131
132    /// Append text to the last message that starts with `prefix`.
133    /// If no such message, push a new one with `prefix + chunk`.
134    pub fn append_stream_chunk(&mut self, prefix: &str, chunk: &str) {
135        if let Some(last) = self.messages.last_mut() {
136            if last.text.starts_with(prefix) {
137                last.text.push_str(chunk);
138                self.scroll_to_bottom();
139                return;
140            }
141        }
142        self.push(format!("{}{}", prefix, chunk));
143    }
144
145    /// Render messages as a List widget with role-based styling and timestamps.
146    pub fn render(&mut self, area: Rect, buf: &mut Buffer, title: &str) {
147        let inner_width = area.width.saturating_sub(2) as usize; // borders
148
149        let items: Vec<ListItem> = self
150            .messages
151            .iter()
152            .map(|m| {
153                let style = Self::message_style(&m.text);
154                let display = m.display_text();
155
156                // Timestamp prefix for non-system messages.
157                let line_text = if m.text.starts_with('=') || m.text.starts_with("Type ") {
158                    display
159                } else {
160                    format!("[{}] {}", m.timestamp, display)
161                };
162
163                // Word-wrap long lines.
164                let wrapped = wrap_text(&line_text, inner_width);
165                let lines: Vec<Line> = wrapped
166                    .into_iter()
167                    .map(|l| Line::from(l).style(style))
168                    .collect();
169
170                ListItem::new(Text::from(lines))
171            })
172            .collect();
173
174        let list = List::new(items)
175            .block(Block::default().borders(Borders::ALL).title(title))
176            .highlight_style(Style::default());
177
178        ratatui::widgets::StatefulWidget::render(list, area, buf, &mut self.list_state);
179    }
180
181    /// Classify message and return appropriate style.
182    fn message_style(msg: &str) -> Style {
183        if msg.starts_with('>') {
184            Style::default().fg(Color::Cyan)
185        } else if msg.starts_with("[THINK]") || msg.starts_with("[STREAM]") {
186            Style::default()
187                .fg(Color::DarkGray)
188                .add_modifier(Modifier::ITALIC)
189        } else if msg.starts_with("[TOOL]") {
190            Style::default().fg(Color::Green)
191        } else if msg.starts_with("[ERR]") {
192            Style::default().fg(Color::Red)
193        } else if msg.starts_with("[WARN]") {
194            Style::default().fg(Color::Yellow)
195        } else if msg.starts_with("[NOTE]") {
196            Style::default().fg(Color::Magenta)
197        } else if msg.starts_with("[DONE]") {
198            Style::default()
199                .fg(Color::Green)
200                .add_modifier(Modifier::BOLD)
201        } else if msg.starts_with("Analysis:") {
202            Style::default().fg(Color::White)
203        } else {
204            Style::default()
205        }
206    }
207}
208
209impl Default for ChatState {
210    fn default() -> Self {
211        Self::new()
212    }
213}
214
215/// Simple word-wrap: split into lines that fit within `max_width`.
216fn wrap_text(text: &str, max_width: usize) -> Vec<String> {
217    if max_width == 0 {
218        return vec![text.to_string()];
219    }
220
221    let mut lines = Vec::new();
222    for input_line in text.lines() {
223        if input_line.len() <= max_width {
224            lines.push(input_line.to_string());
225        } else {
226            let mut remaining = input_line;
227            while remaining.len() > max_width {
228                // Try to break at a space.
229                let break_at = remaining[..max_width]
230                    .rfind(' ')
231                    .map_or(max_width, |pos| pos + 1);
232                lines.push(remaining[..break_at].trim_end().to_string());
233                remaining = &remaining[break_at..];
234            }
235            if !remaining.is_empty() {
236                lines.push(remaining.to_string());
237            }
238        }
239    }
240
241    if lines.is_empty() {
242        lines.push(String::new());
243    }
244    lines
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn wrap_short_text() {
253        let result = wrap_text("hello world", 80);
254        assert_eq!(result, vec!["hello world"]);
255    }
256
257    #[test]
258    fn wrap_long_text() {
259        let result = wrap_text("the quick brown fox jumps over the lazy dog", 20);
260        assert!(result.len() > 1);
261        for line in &result {
262            assert!(line.len() <= 20);
263        }
264    }
265
266    #[test]
267    fn collapsible_tool_output() {
268        let long_output = format!("  = {}", "x".repeat(200));
269        let msg = ChatMessage::new(long_output.clone());
270        assert!(msg.is_collapsible());
271        assert!(msg.display_text().contains("[+more]"));
272    }
273
274    #[test]
275    fn page_up_down() {
276        let mut chat = ChatState::new();
277        for i in 0..30 {
278            chat.push(format!("msg {}", i));
279        }
280        assert_eq!(chat.list_state.selected(), Some(29));
281        chat.page_up();
282        assert_eq!(chat.list_state.selected(), Some(19));
283        chat.page_down();
284        assert_eq!(chat.list_state.selected(), Some(29));
285    }
286}
baml_agent_tui/chat.rs

baml_agent_tui/
chat.rs
