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

codetether_agent/tui/
message_formatter.rs
