Address code review feedback - refactor AI query handling and improve MarkdownRenderer

Co-authored-by: 0xrinegade <[email protected]>

Author: Copilot

src/main.rs

@@ -2,14 +2,18 @@
 #![allow(unused)]
 
 use {
-    crate::config::Config, // Added
+    crate::config::Config,
     crate::utils::diagnostics::DiagnosticCoordinator,
     crate::utils::{dashboard, ebpf_deploy, examples, nodes, ssh_deploy, svm_info},
     crate::utils::markdown_renderer::MarkdownRenderer,
     clparse::parse_command_line,
     solana_client::rpc_client::RpcClient,
-    solana_sdk::{native_token::Sol, pubkey::Pubkey, signature::Signer}, // Modified (removed CommitmentConfig) - bad formatting
-    std::{process::exit, str::FromStr},                                 // Modified
+    solana_sdk::{
+        native_token::Sol, 
+        pubkey::Pubkey, 
+        signature::Signer
+    },
+    std::{process::exit, str::FromStr},
 };
 
 // Helper function to handle the type mismatch between clap v2 and v4
@@ -57,22 +61,30 @@ async fn handle_ai_query(
     sub_matches: &clap::ArgMatches,
     app_matches: &clap::ArgMatches,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    // For external subcommands, clap provides the additional arguments differently
-    // We need to collect them from the raw args since clap doesn't know about them
-    let args: Vec<String> = std::env::args().collect();
-
-    // Find where our subcommand starts and collect everything after osvm
-    let mut query_parts = Vec::new();
-    let mut found_osvm = false;
-
-    for arg in args.iter().skip(1) {
-        // Skip the binary name
-        if !found_osvm {
-            found_osvm = true;
-        }
-        // Skip global flags like --help, --version, etc.
-        if !arg.starts_with('-') {
-            query_parts.push(arg.clone());
+    // For external subcommands, clap collects additional arguments in subcommand_value
+    // This is the proper way to handle external subcommands with clap
+    let mut query_parts = vec![sub_command.to_string()];
+    
+    // Get additional arguments from clap's external subcommand handling
+    if let Some(external_args) = sub_matches.get_many::<String>("") {
+        query_parts.extend(external_args.cloned());
+    }
+    
+    // If clap doesn't provide args (fallback), parse from environment
+    // This maintains compatibility while documenting the limitation
+    if query_parts.len() == 1 {
+        let args: Vec<String> = std::env::args().collect();
+        
+        // Collect non-flag arguments starting from the subcommand
+        let mut found_subcommand = false;
+        for arg in args.iter().skip(1) {
+            if found_subcommand {
+                if !arg.starts_with('-') {
+                    query_parts.push(arg.clone());
+                }
+            } else if arg == sub_command {
+                found_subcommand = true;
+            }
         }
     }
 
@@ -94,11 +106,7 @@ async fn handle_ai_query(
             }
             // Render the response as markdown for better formatting
             let renderer = MarkdownRenderer::new();
-            if let Err(e) = renderer.render(&response) {
-                // Fallback to plain text if markdown rendering fails
-                eprintln!("⚠️  Markdown rendering failed: {}", e);
-                println!("{}", response);
-            }
+            renderer.render(&response);
         }
         Err(e) => {
             eprintln!("❌ AI query failed: {}", e);

src/utils/markdown_renderer.rs

@@ -1,6 +1,7 @@
 use termimad::{MadSkin};
 
 /// A markdown renderer for CLI output
+#[derive(Debug, Clone)]
 pub struct MarkdownRenderer {
     skin: MadSkin,
 }
@@ -12,16 +13,24 @@ impl MarkdownRenderer {
         Self { skin }
     }
 
+    /// Create a new markdown renderer with custom theme configuration
+    pub fn with_theme() -> Self {
+        // For now, use default theme since termimad color API is complex
+        // TODO: Implement proper color configuration once termimad usage is clarified
+        let skin = MadSkin::default();
+        Self { skin }
+    }
+
     /// Render markdown text to the terminal
-    pub fn render(&self, markdown: &str) -> Result<(), Box<dyn std::error::Error>> {
+    /// Note: termimad's print_text never actually fails, so we return unit type
+    pub fn render(&self, markdown: &str) {
         self.skin.print_text(markdown);
-        Ok(())
     }
 
     /// Render markdown text and return as string (for testing)
-    pub fn render_to_string(&self, markdown: &str) -> Result<String, Box<dyn std::error::Error>> {
+    pub fn render_to_string(&self, markdown: &str) -> String {
         let fmt_text = self.skin.text(markdown, None);
-        Ok(fmt_text.to_string())
+        fmt_text.to_string()
     }
 }
 
@@ -38,16 +47,22 @@ mod tests {
     #[test]
     fn test_markdown_renderer_creation() {
         let renderer = MarkdownRenderer::new();
-        assert!(renderer.render_to_string("# Test").is_ok());
+        let result = renderer.render_to_string("# Test");
+        assert!(result.contains("Test"));
     }
 
     #[test]
     fn test_basic_markdown_rendering() {
         let renderer = MarkdownRenderer::new();
         let markdown = "# Header\nSome **bold** text\n```rust\nlet x = 5;\n```";
-        let result = renderer.render_to_string(markdown);
-        assert!(result.is_ok());
-        let output = result.unwrap();
+        let output = renderer.render_to_string(markdown);
         assert!(output.contains("Header"));
     }
+
+    #[test]
+    fn test_themed_renderer() {
+        let renderer = MarkdownRenderer::with_theme();
+        let result = renderer.render_to_string("# Themed Test");
+        assert!(result.contains("Themed Test"));
+    }
 }