codetether_agent/tui/
message_formatter.rs

1use ratatui::{
2    style::{Color, Modifier, Style},
3    text::{Line, Span},
4};
5use std::sync::LazyLock;
6use syntect::{
7    easy::HighlightLines,
8    highlighting::{Style as SyntectStyle, ThemeSet},
9    parsing::SyntaxSet,
10    util::LinesWithEndings,
11};
12
13/// Global syntax set and theme set for syntax highlighting
14static SYNTAX_SET: LazyLock<SyntaxSet> = LazyLock::new(SyntaxSet::load_defaults_newlines);
15static THEME_SET: LazyLock<ThemeSet> = LazyLock::new(ThemeSet::load_defaults);
16
17/// Enhanced message formatter with syntax highlighting and improved styling
18pub struct MessageFormatter {
19    max_width: usize,
20}
21
22impl MessageFormatter {
23    pub fn new(max_width: usize) -> Self {
24        Self { max_width }
25    }
26
27    /// Format message content with enhanced features
28    pub fn format_content(&self, content: &str, role: &str) -> Vec<Line<'static>> {
29        let mut lines = Vec::new();
30        let mut in_code_block = false;
31        let mut code_block_start = false;
32        let mut code_block_language = String::new();
33        let mut code_block_lines = Vec::new();
34
35        for line in content.lines() {
36            // Detect code blocks
37            if line.trim().starts_with("```") {
38                if in_code_block {
39                    // End of code block - render with syntax highlighting
40                    if !code_block_lines.is_empty() {
41                        lines.extend(
42                            self.render_code_block(&code_block_lines, &code_block_language),
43                        );
44                        code_block_lines.clear();
45                        code_block_language.clear();
46                    }
47                    in_code_block = false;
48                    code_block_start = false;
49                } else {
50                    // Start of code block - extract language
51                    in_code_block = true;
52                    code_block_start = true;
53                    let lang = line.trim().trim_start_matches('`').trim();
54                    code_block_language = lang.to_string();
55                }
56                continue;
57            }
58
59            if in_code_block {
60                if code_block_start {
61                    // First line after opening ``` might be language specifier
62                    code_block_start = false;
63                    if !line.trim().is_empty() && code_block_language.is_empty() {
64                        code_block_language = line.trim().to_string();
65                    } else {
66                        code_block_lines.push(line.to_string());
67                    }
68                } else {
69                    code_block_lines.push(line.to_string());
70                }
71                continue;
72            }
73
74            // Handle regular text with enhanced formatting
75            if line.trim().is_empty() {
76                lines.push(Line::from(""));
77                continue;
78            }
79
80            // Handle markdown-like formatting
81            let formatted_line = self.format_inline_text(line, role);
82            lines.extend(self.wrap_line(formatted_line, self.max_width.saturating_sub(4)));
83        }
84
85        // Handle unclosed code blocks
86        if !code_block_lines.is_empty() {
87            lines.extend(self.render_code_block(&code_block_lines, &code_block_language));
88        }
89
90        if lines.is_empty() {
91            lines.push(Line::from(""));
92        }
93
94        lines
95    }
96
97    /// Render a code block with syntax highlighting and styling
98    fn render_code_block(&self, lines: &[String], language: &str) -> Vec<Line<'static>> {
99        let mut result = Vec::new();
100        let block_width = self.max_width.saturating_sub(4);
101
102        // Header with language indicator
103        let header = if language.is_empty() {
104            "┌─ Code ─".to_string() + &"─".repeat(block_width.saturating_sub(9))
105        } else {
106            let lang_header = format!("┌─ {} Code ─", language);
107            let header_len = lang_header.len();
108            lang_header + &"─".repeat(block_width.saturating_sub(header_len))
109        };
110
111        result.push(Line::from(Span::styled(
112            header,
113            Style::default()
114                .fg(Color::DarkGray)
115                .add_modifier(Modifier::BOLD),
116        )));
117
118        // Use syntect for syntax highlighting
119        let highlighted_lines = self.highlight_code_block_syntect(lines, language);
120
121        for line in highlighted_lines {
122            let formatted_line = if line.trim().is_empty() {
123                "│".to_string()
124            } else {
125                format!("│ {}", line)
126            };
127
128            result.push(Line::from(Span::styled(
129                formatted_line,
130                Style::default().fg(Color::DarkGray),
131            )));
132        }
133
134        result.push(Line::from(Span::styled(
135            "└".to_string() + &"─".repeat(block_width.saturating_sub(1)),
136            Style::default().fg(Color::DarkGray),
137        )));
138
139        result
140    }
141
142    /// Advanced syntax highlighting using syntect
143    fn highlight_code_block_syntect(&self, lines: &[String], language: &str) -> Vec<String> {
144        let syntax_set = &*SYNTAX_SET;
145        let theme_set = &*THEME_SET;
146
147        // Use a dark theme suitable for terminal
148        let theme = &theme_set.themes["base16-ocean.dark"];
149
150        // Find the appropriate syntax
151        let syntax = if language.is_empty() {
152            syntax_set.find_syntax_plain_text()
153        } else {
154            syntax_set
155                .find_syntax_by_token(language)
156                .unwrap_or_else(|| syntax_set.find_syntax_plain_text())
157        };
158
159        let mut highlighter = HighlightLines::new(syntax, theme);
160
161        let mut highlighted_lines = Vec::new();
162        let code = lines.join("\n");
163
164        for line in LinesWithEndings::from(&code) {
165            let ranges = match highlighter.highlight_line(line, syntax_set) {
166                Ok(r) => r,
167                Err(_) => {
168                    // On error, just return plain text
169                    highlighted_lines.push(line.trim_end().to_string());
170                    continue;
171                }
172            };
173            let mut line_result = String::new();
174
175            for (style, text) in ranges {
176                let fg_color = style.foreground;
177                let _color = Color::Rgb(fg_color.r, fg_color.g, fg_color.b);
178
179                // Convert the styled text to a ratatui-compatible string
180                // Since we can't use actual terminal colors in ratatui spans easily,
181                // we'll use the closest ratatui colors
182                let _ratatui_color = self.map_syntect_color_to_ratatui(&fg_color);
183
184                // For now, we'll just return the text without ANSI codes
185                // since ratatui handles styling through its own Span system
186                line_result.push_str(text);
187            }
188
189            highlighted_lines.push(line_result.trim_end().to_string());
190        }
191
192        highlighted_lines
193    }
194
195    /// Map syntect colors to ratatui colors (simplified for now)
196    fn map_syntect_color_to_ratatui(&self, color: &syntect::highlighting::Color) -> Color {
197        // Simple mapping - we'll use the closest ratatui color
198        Color::Rgb(color.r, color.g, color.b)
199    }
200
201    /// Format inline text with basic markdown-like formatting
202    fn format_inline_text(&self, line: &str, role: &str) -> Vec<Span<'static>> {
203        let mut spans = Vec::new();
204        let mut current = String::new();
205        let mut in_bold = false;
206        let mut in_italic = false;
207        let mut in_code = false;
208
209        let role_color = match role {
210            "user" => Color::White,
211            "assistant" => Color::Cyan,
212            "system" => Color::Yellow,
213            "tool" => Color::Green,
214            _ => Color::White,
215        };
216
217        let mut chars = line.chars().peekable();
218
219        while let Some(c) = chars.next() {
220            match c {
221                '*' => {
222                    if chars.peek() == Some(&'*') {
223                        // Bold
224                        if !current.is_empty() {
225                            spans.push(Span::styled(
226                                current.clone(),
227                                Style::default().fg(role_color).add_modifier(if in_bold {
228                                    Modifier::BOLD
229                                } else {
230                                    Modifier::empty()
231                                }),
232                            ));
233                            current.clear();
234                        }
235                        chars.next(); // consume second '*'
236                        in_bold = !in_bold;
237                    } else {
238                        // Italic
239                        if !current.is_empty() {
240                            spans.push(Span::styled(
241                                current.clone(),
242                                Style::default().fg(role_color).add_modifier(if in_italic {
243                                    Modifier::ITALIC
244                                } else {
245                                    Modifier::empty()
246                                }),
247                            ));
248                            current.clear();
249                        }
250                        in_italic = !in_italic;
251                    }
252                }
253                '`' => {
254                    if !current.is_empty() {
255                        spans.push(Span::styled(
256                            current.clone(),
257                            Style::default().fg(role_color),
258                        ));
259                        current.clear();
260                    }
261                    in_code = !in_code;
262                }
263                _ => {
264                    current.push(c);
265                }
266            }
267        }
268
269        if !current.is_empty() {
270            spans.push(Span::styled(current, Style::default().fg(role_color)));
271        }
272
273        if spans.is_empty() {
274            spans.push(Span::styled(
275                line.to_string(),
276                Style::default().fg(role_color),
277            ));
278        }
279
280        spans
281    }
282
283    /// Wrap text to fit within width
284    fn wrap_line(&self, spans: Vec<Span<'static>>, width: usize) -> Vec<Line<'static>> {
285        if spans.is_empty() {
286            return vec![Line::from("")];
287        }
288
289        // Simple wrapping - for now, just return as single line
290        vec![Line::from(spans)]
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    #[test]
299    fn test_code_block_detection() {
300        let formatter = MessageFormatter::new(80);
301        let content = "```rust\nfn main() {\n    println!(\"Hello, world!\");\n}\n```";
302        let lines = formatter.format_content(content, "assistant");
303        assert!(!lines.is_empty());
304    }
305
306    #[test]
307    fn test_syntax_highlighting() {
308        let formatter = MessageFormatter::new(80);
309        let lines = vec![
310            "fn main() {".to_string(),
311            "    println!(\"Hello!\");".to_string(),
312            "}".to_string(),
313        ];
314        let highlighted = formatter.highlight_code_block_syntect(&lines, "rust");
315        assert_eq!(highlighted.len(), 3);
316    }
317}
codetether_agent/tui/message_formatter.rs

codetether_agent/tui/
message_formatter.rs
