add syntax highlighting

Author: alexylon

src/conversation.rs

@@ -105,6 +105,7 @@ Your goal is to help users with coding tasks efficiently and accurately."#,
 
     pub fn restore_messages(&mut self, messages: Vec<Message>) {
         self.messages = messages;
+        self.trim_if_needed();
     }
 
     pub fn _len(&self) -> usize {
@@ -131,11 +132,11 @@ mod tests {
     fn test_message_limit_trimming() {
         let mut history = ConversationHistory::new();
 
-        for i in 0..60 {
+        for i in 0..510 {
             history.add_user_message(format!("Message {}", i));
         }
 
-        assert_eq!(history.messages().len(), 50);
+        assert_eq!(history.messages().len(), 500);
 
         if let crate::api::MessageContent::Text { content } = &history.messages()[0].content {
             assert_eq!(content, "Message 10");
@@ -146,14 +147,14 @@ mod tests {
     fn test_message_limit_with_blocks() {
         let mut history = ConversationHistory::new();
 
-        for i in 0..30 {
+        for i in 0..260 {
             history.add_user_message(format!("User {}", i));
             history.add_assistant_with_blocks(vec![MessageContentBlock::Text {
                 text: format!("Assistant {}", i),
             }]);
         }
 
-        assert_eq!(history.messages().len(), 50);
+        assert_eq!(history.messages().len(), 500);
     }
 
     #[test]

src/main.rs

@@ -5,6 +5,7 @@ mod error;
 mod history;
 mod repl;
 mod session_selector;
+mod syntax;
 mod tools;
 
 use api::MorphClient;

src/repl.rs

@@ -2,6 +2,7 @@ use crate::api::{AnthropicClient, ContentBlock, CreateMessageRequest, MorphClien
 use crate::conversation::ConversationHistory;
 use crate::error::{Result, SofosError};
 use crate::history::HistoryManager;
+use crate::syntax::SyntaxHighlighter;
 use crate::tools::{add_code_search_tool, get_tools, get_tools_with_morph, ToolExecutor};
 use colored::Colorize;
 use rustyline::error::ReadlineError;
@@ -18,6 +19,7 @@ pub struct Repl {
     conversation: ConversationHistory,
     tool_executor: ToolExecutor,
     history_manager: HistoryManager,
+    highlighter: SyntaxHighlighter,
     editor: DefaultEditor,
     model: String,
     max_tokens: u32,
@@ -47,11 +49,14 @@ impl Repl {
 
         let session_id = HistoryManager::generate_session_id();
 
+        let highlighter = SyntaxHighlighter::new();
+
         Ok(Self {
             client,
             conversation,
             tool_executor,
             history_manager,
+            highlighter,
             editor,
             model,
             max_tokens,
@@ -221,7 +226,8 @@ impl Repl {
         if !text_output.is_empty() {
             println!("{}", "Assistant:".bright_blue().bold());
             for text in &text_output {
-                println!("{}", text);
+                let highlighted = self.highlighter.highlight_text(text);
+                println!("{}", highlighted);
             }
             println!();
         }

src/syntax.rs

@@ -0,0 +1,115 @@
+use colored::Colorize;
+use syntect::easy::HighlightLines;
+use syntect::highlighting::{Style, ThemeSet};
+use syntect::parsing::SyntaxSet;
+use syntect::util::as_24_bit_terminal_escaped;
+
+pub struct SyntaxHighlighter {
+    syntax_set: SyntaxSet,
+    theme_set: ThemeSet,
+}
+
+impl SyntaxHighlighter {
+    pub fn new() -> Self {
+        Self {
+            syntax_set: SyntaxSet::load_defaults_newlines(),
+            theme_set: ThemeSet::load_defaults(),
+        }
+    }
+
+    pub fn highlight_text(&self, text: &str) -> String {
+        let mut result = String::new();
+        let mut in_code_block = false;
+        let mut code_block = String::new();
+        let mut language = String::new();
+
+        for line in text.lines() {
+            if line.starts_with("```") {
+                if in_code_block {
+                    result.push_str(&self.highlight_code(&code_block, &language));
+                    result.push('\n');
+                    code_block.clear();
+                    language.clear();
+                    in_code_block = false;
+                } else {
+                    language = line.trim_start_matches('`').trim().to_string();
+                    in_code_block = true;
+                }
+            } else if in_code_block {
+                code_block.push_str(line);
+                code_block.push('\n');
+            } else {
+                result.push_str(line);
+                result.push('\n');
+            }
+        }
+
+        if in_code_block {
+            result.push_str(&self.highlight_code(&code_block, &language));
+        }
+
+        result.trim_end().to_string()
+    }
+
+    fn highlight_code(&self, code: &str, language: &str) -> String {
+        let syntax = if language.is_empty() {
+            self.syntax_set.find_syntax_plain_text()
+        } else {
+            self.syntax_set
+                .find_syntax_by_token(language)
+                .or_else(|| self.syntax_set.find_syntax_by_extension(language))
+                .unwrap_or_else(|| self.syntax_set.find_syntax_plain_text())
+        };
+
+        let theme = &self.theme_set.themes["base16-ocean.dark"];
+        let mut highlighter = HighlightLines::new(syntax, theme);
+
+        let mut result = String::new();
+        result.push_str(&format!("{}\n", "┌─────".dimmed()));
+
+        for line in code.lines() {
+            let ranges: Vec<(Style, &str)> = highlighter
+                .highlight_line(line, &self.syntax_set)
+                .unwrap_or_default();
+            let escaped = as_24_bit_terminal_escaped(&ranges[..], false);
+            result.push_str(&format!("{}  {}\n", "│".dimmed(), escaped));
+        }
+
+        result.push_str(&format!("{}", "└─────".dimmed()));
+        result
+    }
+}
+
+impl Default for SyntaxHighlighter {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_highlighter_creation() {
+        let highlighter = SyntaxHighlighter::new();
+        assert!(!highlighter.syntax_set.syntaxes().is_empty());
+    }
+
+    #[test]
+    fn test_plain_text() {
+        let highlighter = SyntaxHighlighter::new();
+        let text = "Hello, world!";
+        let result = highlighter.highlight_text(text);
+        assert!(result.contains("Hello, world!"));
+    }
+
+    #[test]
+    fn test_code_block_detection() {
+        let highlighter = SyntaxHighlighter::new();
+        let text = "Here is some code:\n```rust\nfn main() {}\n```";
+        let result = highlighter.highlight_text(text);
+        assert!(result.contains("Here is some code:"));
+        assert!(result.contains("main"));
+    }
+}