added tui improvements

Ammar-Alnagar · Ammar-Alnagar · commit 3e8b31e63dc8 · 2025-09-07T18:57:20.000+03:00
diff --git a/rust_tui_coder/src/agent.rs b/rust_tui_coder/src/agent.rs
@@ -1066,7 +1066,7 @@ TOOL: {"name": "RUN_COMMAND", "parameters": {"command": "cargo build --release"}
         None
     }
 
-    pub async fn run(&mut self, config: &LlmConfig, user_prompt: String) -> Result<(String, Vec<String>), Box<dyn std::error::Error>> {
+    pub async fn run(&mut self, config: &LlmConfig, user_prompt: String, app: &mut crate::app::App) -> Result<(String, Vec<String>), Box<dyn std::error::Error>> {
         // Add system message if this is the first interaction
         if self.messages.is_empty() {
             self.messages.push(Message {
@@ -1116,7 +1116,9 @@ TOOL: {"name": "RUN_COMMAND", "parameters": {"command": "cargo build --release"}
             attempts += 1;
             
             // Get response from LLM
-            let response = llm::ask_llm_with_messages(config, &self.messages).await?;
+            let (response, tokens_used) = llm::ask_llm_with_messages(config, &self.messages).await?;
+            app.increment_tokens(tokens_used);
+            app.increment_requests();
 
             // Check if response contains a tool call
             if let Some(tool) = self.parse_tool_call(&response) {
@@ -1152,6 +1154,7 @@ TOOL: {"name": "RUN_COMMAND", "parameters": {"command": "cargo build --release"}
                 // Execute the tool
                 let tool_result = match tool.execute() {
                     Ok(result) => {
+                        app.increment_tools_executed();
                         tool_logs.push(format!("✅ Success: {}", result));
                         result
                     }
@@ -1178,12 +1181,15 @@ TOOL: {"name": "RUN_COMMAND", "parameters": {"command": "cargo build --release"}
                 // Check if we should continue or if the task is complete
                 if attempts >= MAX_ATTEMPTS {
                     // Get final response after max attempts
-                    let final_response = llm::ask_llm_with_messages(config, &self.messages).await?;
+                    let (final_response, final_tokens) = llm::ask_llm_with_messages(config, &self.messages).await?;
+                    app.increment_tokens(final_tokens);
+                    app.increment_requests();
+
                     self.messages.push(Message {
                         role: "assistant".to_string(),
                         content: final_response.clone(),
                     });
-                    
+
                     all_tool_logs.push(format!("⚠️  Reached maximum attempts ({})", MAX_ATTEMPTS));
                     return Ok((final_response, all_tool_logs));
                 }
diff --git a/rust_tui_coder/src/app.rs b/rust_tui_coder/src/app.rs
@@ -5,21 +5,75 @@ pub struct App {
     pub tool_logs: Vec<String>,
     pub is_executing_tool: bool,
     pub current_tool: String,
+    // Usage tracking
+    pub session_start_time: std::time::Instant,
+    pub tokens_used: u64,
+    pub total_requests: u64,
+    pub total_tools_executed: u64,
 }
 
 impl App {
     pub fn new() -> Self {
         Self {
             user_input: String::new(),
             conversation: Vec::new(),
-            status_message: "Type in /quit to exit".to_string(),
+            status_message: "Commands: /quit (exit), /stats (usage stats)".to_string(),
             tool_logs: Vec::new(),
             is_executing_tool: false,
             current_tool: String::new(),
+            // Initialize usage tracking
+            session_start_time: std::time::Instant::now(),
+            tokens_used: 0,
+            total_requests: 0,
+            total_tools_executed: 0,
         }
     }
 
     pub fn add_tool_log(&mut self, log: String) {
         self.tool_logs.push(log);
     }
+
+    // Usage tracking methods
+    pub fn increment_tokens(&mut self, tokens: u64) {
+        self.tokens_used += tokens;
+    }
+
+    pub fn increment_requests(&mut self) {
+        self.total_requests += 1;
+    }
+
+    pub fn increment_tools_executed(&mut self) {
+        self.total_tools_executed += 1;
+    }
+
+    pub fn get_session_duration(&self) -> std::time::Duration {
+        self.session_start_time.elapsed()
+    }
+
+    pub fn get_usage_summary(&self) -> String {
+        let duration = self.get_session_duration();
+        let hours = duration.as_secs() / 3600;
+        let minutes = (duration.as_secs() % 3600) / 60;
+        let seconds = duration.as_secs() % 60;
+
+        format!(
+            "Session Summary:\n\
+             • Duration: {:02}:{:02}:{:02}\n\
+             • Tokens Used: {}\n\
+             • LLM Requests: {}\n\
+             • Tools Executed: {}\n\
+             • Average Tokens/Request: {}",
+            hours,
+            minutes,
+            seconds,
+            self.tokens_used,
+            self.total_requests,
+            self.total_tools_executed,
+            if self.total_requests > 0 {
+                self.tokens_used / self.total_requests
+            } else {
+                0
+            }
+        )
+    }
 }
diff --git a/rust_tui_coder/src/llm.rs b/rust_tui_coder/src/llm.rs
@@ -4,6 +4,21 @@ use crate::config::LlmConfig;
 use std::error::Error;
 use std::fmt;
 
+// Token estimation function (rough approximation based on GPT tokenization)
+pub fn estimate_token_count(text: &str) -> u64 {
+    // Rough approximation: ~4 characters per token for English text
+    // This is not exact but provides a reasonable estimate
+    let char_count = text.chars().count() as f64;
+    let estimated_tokens = (char_count / 4.0).ceil() as u64;
+
+    // Ensure minimum of 1 token for non-empty text
+    if estimated_tokens == 0 && !text.trim().is_empty() {
+        1
+    } else {
+        estimated_tokens
+    }
+}
+
 #[derive(Debug)]
 pub enum LlmError {
     RequestFailed(reqwest::Error),
@@ -61,10 +76,17 @@ struct Choice {
     message: Message,
 }
 
-pub async fn ask_llm_with_messages(config: &LlmConfig, messages: &[Message]) -> Result<String, LlmError> {
+pub async fn ask_llm_with_messages(config: &LlmConfig, messages: &[Message]) -> Result<(String, u64), LlmError> {
     let client = Client::new();
     let _provider = config.provider.as_deref().unwrap_or("openai");
 
+    // Calculate input tokens
+    let mut input_tokens = 0u64;
+    for message in messages {
+        input_tokens += estimate_token_count(&message.role);
+        input_tokens += estimate_token_count(&message.content);
+    }
+
     // Determine model to use (support AUTODETECT from /models endpoint for OpenAI-compatible APIs)
     let model_to_use = if config.model_name.eq_ignore_ascii_case("AUTODETECT") || config.model_name.trim().is_empty() {
         // Fetch available models and use the first one
@@ -133,9 +155,15 @@ pub async fn ask_llm_with_messages(config: &LlmConfig, messages: &[Message]) ->
     match serde_json::from_str::<ChatCompletionResponse>(&response_text) {
         Ok(parsed_response) => {
             if let Some(choice) = parsed_response.choices.into_iter().next() {
-                Ok(choice.message.content)
+                let response_content = choice.message.content;
+                let output_tokens = estimate_token_count(&response_content);
+                let total_tokens = input_tokens + output_tokens;
+                Ok((response_content, total_tokens))
             } else {
-                Ok("No response content available.".to_string())
+                let response_content = "No response content available.".to_string();
+                let output_tokens = estimate_token_count(&response_content);
+                let total_tokens = input_tokens + output_tokens;
+                Ok((response_content, total_tokens))
             }
         }
         Err(e) => {
diff --git a/rust_tui_coder/src/main.rs b/rust_tui_coder/src/main.rs
@@ -73,15 +73,29 @@ async fn run_app<B: Backend>(
                     
                     // Check for quit command
                     if user_input.trim() == "/quit" {
+                        // Display usage summary before quitting
+                        let summary = app.get_usage_summary();
+                        app.conversation.push(format!("System: {}", summary));
+                        terminal.draw(|f| ui::ui(f, &app))?;
+
+                        // Give user a moment to see the summary
+                        std::thread::sleep(std::time::Duration::from_secs(2));
                         return Ok(());
                     }
+
+                    // Check for stats command
+                    if user_input.trim() == "/stats" {
+                        let summary = app.get_usage_summary();
+                        app.conversation.push(format!("System: {}", summary));
+                        continue;
+                    }
                     
                     app.conversation.push(format!("User: {}", user_input));
 
                     app.status_message = "Thinking...".to_string();
                     terminal.draw(|f| ui::ui(f, &app))?;
 
-                    match agent.run(&config.llm, user_input).await {
+                    match agent.run(&config.llm, user_input, &mut app).await {
                         Ok((response, tool_logs)) => {
                             // Add tool logs to the app
                             for log in tool_logs {
diff --git a/rust_tui_coder/src/ui.rs b/rust_tui_coder/src/ui.rs
@@ -7,34 +7,86 @@ use ratatui::{
 };
 use crate::app::App;
 
-// Helper function to wrap text to fit within a given width
+// Enhanced helper function to wrap text to fit within a given width
 fn wrap_text(text: &str, max_width: usize) -> Vec<String> {
-    let words: Vec<&str> = text.split_whitespace().collect();
+    if max_width == 0 {
+        return vec![text.to_string()];
+    }
+
     let mut lines = Vec::new();
     let mut current_line = String::new();
-    
-    for word in words {
-        if current_line.len() + word.len() + 1 <= max_width {
-            if !current_line.is_empty() {
-                current_line.push(' ');
-            }
-            current_line.push_str(word);
-        } else {
+
+    // Handle empty text
+    if text.trim().is_empty() {
+        return vec![text.to_string()];
+    }
+
+    for line in text.lines() {
+        if line.trim().is_empty() {
+            // Handle empty lines by preserving them
             if !current_line.is_empty() {
                 lines.push(current_line.clone());
+                current_line.clear();
+            }
+            lines.push(String::new());
+            continue;
+        }
+
+        let words: Vec<&str> = line.split_whitespace().collect();
+
+        for word in words {
+            // Handle words longer than max_width by breaking them
+            if word.len() > max_width {
+                if !current_line.is_empty() {
+                    lines.push(current_line.clone());
+                    current_line.clear();
+                }
+                // Break long word into chunks
+                let mut remaining_word = word;
+                while !remaining_word.is_empty() {
+                    let chunk_size = std::cmp::min(max_width, remaining_word.len());
+                    let chunk = &remaining_word[..chunk_size];
+                    lines.push(chunk.to_string());
+                    remaining_word = &remaining_word[chunk_size..];
+                }
+                continue;
+            }
+
+            // Check if word fits on current line
+            let space_needed = if current_line.is_empty() { word.len() } else { current_line.len() + 1 + word.len() };
+
+            if space_needed <= max_width {
+                if !current_line.is_empty() {
+                    current_line.push(' ');
+                }
+                current_line.push_str(word);
+            } else {
+                // Word doesn't fit, start new line
+                if !current_line.is_empty() {
+                    lines.push(current_line.clone());
+                    current_line.clear();
+                }
+                current_line.push_str(word);
             }
-            current_line = word.to_string();
+        }
+
+        // Add the current line if it's not empty
+        if !current_line.is_empty() {
+            lines.push(current_line.clone());
+            current_line.clear();
         }
     }
-    
+
+    // Handle case where we have pending content
     if !current_line.is_empty() {
         lines.push(current_line);
     }
-    
+
+    // Ensure we always return at least the original text if no wrapping occurred
     if lines.is_empty() {
         lines.push(text.to_string());
     }
-    
+
     lines
 }
 
@@ -152,8 +204,19 @@ pub fn ui(f: &mut Frame, app: &App) {
         .position(tool_scroll_position);
     f.render_stateful_widget(tool_scrollbar, chunks[1], &mut tool_scrollbar_state);
 
+    // Enhanced input field with text wrapping
     let input_block = Block::default().title("Input").borders(Borders::ALL);
-    let input = Paragraph::new(app.user_input.as_str())
+    let input_max_width = chunks[2].width.saturating_sub(4) as usize; // Account for borders and padding
+
+    // Wrap the input text for better display
+    let wrapped_input = if app.user_input.is_empty() {
+        vec![String::new()]
+    } else {
+        wrap_text(&app.user_input, input_max_width)
+    };
+
+    let input_text = wrapped_input.join("\n");
+    let input = Paragraph::new(input_text)
         .block(input_block);
     f.render_widget(input, chunks[2]);
 
