From b9a95161282a7f0031d2eb2cd2bf2cb69b7714db Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 09:16:35 -0300
Subject: [PATCH 01/10] refactor: types

---
 src/lib.rs                     | 104 +--------------------------------
 src/render_template.rs         |   3 +-
 src/schema.rs                  |   2 +-
 src/types.rs                   | 103 ++++++++++++++++++++++++++++++++
 tests/common/common.rs         |   5 +-
 tests/render_template_tests.rs |   3 +-
 tests/test_std_output.rs       |   2 +-
 7 files changed, 116 insertions(+), 106 deletions(-)
 create mode 100644 src/types.rs

diff --git a/src/lib.rs b/src/lib.rs
index 95f7ff6..a3a8535 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,6 +5,7 @@ pub mod render_template;
 pub mod schema;
 pub mod std_output;
 pub mod templates;
+pub mod types;
 pub mod utils;
 
 pub use cli::{CommandArguments, DiscoveryCommand, LogLevel, PrintOptions, Template, WriteOptions};
@@ -12,9 +13,10 @@ use colored::Colorize;
 use error::{DiscoveryError, DiscoveryResult};
 use render_template::{detect_render_markers, render_template};
 use schema::tool_params;
-use std::{fmt::Display, io::stdout};
+use std::io::stdout;
 use std_output::{print_header, print_list, print_summary};
 pub use templates::OutputTemplate;
+use types::{McpCapabilities, McpServerInfo, McpToolMeta};
 
 use std::sync::Arc;
 
@@ -30,106 +32,6 @@ use rust_mcp_sdk::{
     McpClient, StdioTransport, TransportOptions,
 };
 
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub struct McpCapabilities {
-    pub tools: bool,
-    pub prompts: bool,
-    pub resources: bool,
-    pub logging: bool,
-    pub experimental: bool,
-}
-
-impl Display for McpCapabilities {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "tools:{}, prompts:{}, resources:{}, logging:{}, experimental:{}",
-            self.tools, self.prompts, self.resources, self.logging, self.experimental
-        )
-    }
-}
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub enum ParamTypes {
-    Primitive(String),
-    Object(Vec<McpToolSParams>),
-    Array(Vec<ParamTypes>),
-}
-
-impl Display for ParamTypes {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let type_name = match self {
-            ParamTypes::Primitive(type_name) => type_name.to_owned(),
-            ParamTypes::Object(items) => {
-                format!(
-                    "{{{}}}",
-                    items
-                        .iter()
-                        .map(|t| format!("{} : {}", t.param_name, t.param_type))
-                        .collect::<Vec<String>>()
-                        .join(", ")
-                )
-            }
-            ParamTypes::Array(items) => format!("{} [ ]", items[0]),
-        };
-        write!(f, "{}", type_name)
-    }
-}
-
-// impl Serialize for ParamTypes {
-//     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-//     where
-//         S: Serializer,
-//     {
-//         match self {
-//             ParamTypes::Primitive(s) => {
-//                 let mut map = serializer.serialize_map(Some(1))?;
-//                 map.serialize_entry("Primitive", s)?;
-//                 map.end()
-//             }
-//
-//             ParamTypes::Object(params) => params.serialize(serializer), // Inline as array
-//             ParamTypes::Array(arr) => {
-//                 let mut map = serializer.serialize_map(Some(1))?;
-//                 map.serialize_entry("Array", arr)?;
-//                 map.end()
-//             }
-//         }
-//     }
-// }
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub struct McpToolSParams {
-    pub param_name: String,
-    pub param_type: ParamTypes,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub param_description: Option<String>,
-    pub required: bool,
-}
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub struct McpToolMeta {
-    pub name: String,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub description: Option<String>,
-    pub params: Vec<McpToolSParams>,
-}
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub struct McpServerInfo {
-    pub name: String,
-    pub version: String,
-    pub capabilities: McpCapabilities,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub tools: Option<Vec<McpToolMeta>>,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub prompts: Option<Vec<Prompt>>,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub resources: Option<Vec<Resource>>,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub resource_templates: Option<Vec<ResourceTemplate>>,
-}
-
 pub struct McpDiscovery {
     options: DiscoveryCommand,
     pub server_info: Option<McpServerInfo>,
diff --git a/src/render_template.rs b/src/render_template.rs
index 04ece57..80d38ad 100644
--- a/src/render_template.rs
+++ b/src/render_template.rs
@@ -15,10 +15,11 @@ use unicode_width::UnicodeWidthStr;
 use crate::{
     error::{DiscoveryError, DiscoveryResult},
     templates::{InlineTemplateInfo, PARTIALS},
+    types::ParamTypes,
     utils::{
         boolean_indicator, line_ending, match_template, RenderTemplateInfo, UpdateTemplateInfo,
     },
-    McpServerInfo, OutputTemplate, ParamTypes, Template, WriteOptions,
+    McpServerInfo, OutputTemplate, Template, WriteOptions,
 };
 
 /// Properties parsed from the `MCP_DISCOVERY_TEMPLATE_START` marker line in template files.
diff --git a/src/schema.rs b/src/schema.rs
index bfa7363..a004c1b 100644
--- a/src/schema.rs
+++ b/src/schema.rs
@@ -4,7 +4,7 @@ use serde_json::{Map, Value};
 
 use crate::{
     error::{DiscoveryError, DiscoveryResult},
-    McpToolSParams, ParamTypes,
+    types::{McpToolSParams, ParamTypes},
 };
 
 /// Parses an object schema into a vector of `McpToolSParams`.
diff --git a/src/types.rs b/src/types.rs
new file mode 100644
index 0000000..8ca2db7
--- /dev/null
+++ b/src/types.rs
@@ -0,0 +1,103 @@
+use std::fmt::Display;
+
+use rust_mcp_schema::{Prompt, Resource, ResourceTemplate};
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub struct McpCapabilities {
+    pub tools: bool,
+    pub prompts: bool,
+    pub resources: bool,
+    pub logging: bool,
+    pub experimental: bool,
+}
+
+impl Display for McpCapabilities {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "tools:{}, prompts:{}, resources:{}, logging:{}, experimental:{}",
+            self.tools, self.prompts, self.resources, self.logging, self.experimental
+        )
+    }
+}
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub enum ParamTypes {
+    Primitive(String),
+    Object(Vec<McpToolSParams>),
+    Array(Vec<ParamTypes>),
+}
+
+impl Display for ParamTypes {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let type_name = match self {
+            ParamTypes::Primitive(type_name) => type_name.to_owned(),
+            ParamTypes::Object(items) => {
+                format!(
+                    "{{{}}}",
+                    items
+                        .iter()
+                        .map(|t| format!("{} : {}", t.param_name, t.param_type))
+                        .collect::<Vec<String>>()
+                        .join(", ")
+                )
+            }
+            ParamTypes::Array(items) => format!("{} [ ]", items[0]),
+        };
+        write!(f, "{}", type_name)
+    }
+}
+
+// impl Serialize for ParamTypes {
+//     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+//     where
+//         S: Serializer,
+//     {
+//         match self {
+//             ParamTypes::Primitive(s) => {
+//                 let mut map = serializer.serialize_map(Some(1))?;
+//                 map.serialize_entry("Primitive", s)?;
+//                 map.end()
+//             }
+//
+//             ParamTypes::Object(params) => params.serialize(serializer), // Inline as array
+//             ParamTypes::Array(arr) => {
+//                 let mut map = serializer.serialize_map(Some(1))?;
+//                 map.serialize_entry("Array", arr)?;
+//                 map.end()
+//             }
+//         }
+//     }
+// }
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub struct McpToolSParams {
+    pub param_name: String,
+    pub param_type: ParamTypes,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub param_description: Option<String>,
+    pub required: bool,
+}
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub struct McpToolMeta {
+    pub name: String,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub description: Option<String>,
+    pub params: Vec<McpToolSParams>,
+}
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub struct McpServerInfo {
+    pub name: String,
+    pub version: String,
+    pub capabilities: McpCapabilities,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub tools: Option<Vec<McpToolMeta>>,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub prompts: Option<Vec<Prompt>>,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub resources: Option<Vec<Resource>>,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub resource_templates: Option<Vec<ResourceTemplate>>,
+}
diff --git a/tests/common/common.rs b/tests/common/common.rs
index c03d380..382ae17 100644
--- a/tests/common/common.rs
+++ b/tests/common/common.rs
@@ -1,5 +1,8 @@
 use clap::Parser;
-use mcp_discovery::{CommandArguments, McpCapabilities, McpServerInfo};
+use mcp_discovery::{
+    types::{McpCapabilities, McpServerInfo},
+    CommandArguments,
+};
 use rust_mcp_sdk::macros::JsonSchema;
 
 // Helper function to parse arguments from a vector of strings
diff --git a/tests/render_template_tests.rs b/tests/render_template_tests.rs
index a2867e6..e879a77 100644
--- a/tests/render_template_tests.rs
+++ b/tests/render_template_tests.rs
@@ -5,7 +5,8 @@ use common::default_mcp_server_info;
 use handlebars::Handlebars;
 use mcp_discovery::{
     render_template::{detect_render_markers, extract_template_file, register_helpers},
-    OutputTemplate, ParamTypes, WriteOptions,
+    types::ParamTypes,
+    OutputTemplate, WriteOptions,
 };
 use serde_json::json;
 use std::fs::write;
diff --git a/tests/test_std_output.rs b/tests/test_std_output.rs
index 6bffcae..d64d14e 100644
--- a/tests/test_std_output.rs
+++ b/tests/test_std_output.rs
@@ -2,7 +2,7 @@ use mcp_discovery::{
     std_output::{
         print_header, print_json, print_list, print_summary, table_bottom, table_content, table_top,
     },
-    McpCapabilities, McpServerInfo,
+    types::{McpCapabilities, McpServerInfo},
 };
 use std::io::{self};
 use unicode_width::UnicodeWidthStr;

From 46229b62e2c176659b737d7f6605e1f7472d5847 Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 09:25:07 -0300
Subject: [PATCH 02/10] refactor: sub type modules

---
 src/types.rs              | 106 ++------------------------------------
 src/types/capabilities.rs | 102 ++++++++++++++++++++++++++++++++++++
 src/types/commands.rs     |   1 +
 3 files changed, 107 insertions(+), 102 deletions(-)
 create mode 100644 src/types/capabilities.rs
 create mode 100644 src/types/commands.rs

diff --git a/src/types.rs b/src/types.rs
index 8ca2db7..7333909 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -1,103 +1,5 @@
-use std::fmt::Display;
+mod capabilities;
+mod commands;
 
-use rust_mcp_schema::{Prompt, Resource, ResourceTemplate};
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub struct McpCapabilities {
-    pub tools: bool,
-    pub prompts: bool,
-    pub resources: bool,
-    pub logging: bool,
-    pub experimental: bool,
-}
-
-impl Display for McpCapabilities {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "tools:{}, prompts:{}, resources:{}, logging:{}, experimental:{}",
-            self.tools, self.prompts, self.resources, self.logging, self.experimental
-        )
-    }
-}
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub enum ParamTypes {
-    Primitive(String),
-    Object(Vec<McpToolSParams>),
-    Array(Vec<ParamTypes>),
-}
-
-impl Display for ParamTypes {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let type_name = match self {
-            ParamTypes::Primitive(type_name) => type_name.to_owned(),
-            ParamTypes::Object(items) => {
-                format!(
-                    "{{{}}}",
-                    items
-                        .iter()
-                        .map(|t| format!("{} : {}", t.param_name, t.param_type))
-                        .collect::<Vec<String>>()
-                        .join(", ")
-                )
-            }
-            ParamTypes::Array(items) => format!("{} [ ]", items[0]),
-        };
-        write!(f, "{}", type_name)
-    }
-}
-
-// impl Serialize for ParamTypes {
-//     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-//     where
-//         S: Serializer,
-//     {
-//         match self {
-//             ParamTypes::Primitive(s) => {
-//                 let mut map = serializer.serialize_map(Some(1))?;
-//                 map.serialize_entry("Primitive", s)?;
-//                 map.end()
-//             }
-//
-//             ParamTypes::Object(params) => params.serialize(serializer), // Inline as array
-//             ParamTypes::Array(arr) => {
-//                 let mut map = serializer.serialize_map(Some(1))?;
-//                 map.serialize_entry("Array", arr)?;
-//                 map.end()
-//             }
-//         }
-//     }
-// }
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub struct McpToolSParams {
-    pub param_name: String,
-    pub param_type: ParamTypes,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub param_description: Option<String>,
-    pub required: bool,
-}
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub struct McpToolMeta {
-    pub name: String,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub description: Option<String>,
-    pub params: Vec<McpToolSParams>,
-}
-
-#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
-pub struct McpServerInfo {
-    pub name: String,
-    pub version: String,
-    pub capabilities: McpCapabilities,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub tools: Option<Vec<McpToolMeta>>,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub prompts: Option<Vec<Prompt>>,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub resources: Option<Vec<Resource>>,
-    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
-    pub resource_templates: Option<Vec<ResourceTemplate>>,
-}
+pub use capabilities::*;
+pub use commands::*;
diff --git a/src/types/capabilities.rs b/src/types/capabilities.rs
new file mode 100644
index 0000000..0cbac6c
--- /dev/null
+++ b/src/types/capabilities.rs
@@ -0,0 +1,102 @@
+use rust_mcp_schema::{Prompt, Resource, ResourceTemplate};
+use std::fmt::Display;
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub struct McpCapabilities {
+    pub tools: bool,
+    pub prompts: bool,
+    pub resources: bool,
+    pub logging: bool,
+    pub experimental: bool,
+}
+
+impl Display for McpCapabilities {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "tools:{}, prompts:{}, resources:{}, logging:{}, experimental:{}",
+            self.tools, self.prompts, self.resources, self.logging, self.experimental
+        )
+    }
+}
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub enum ParamTypes {
+    Primitive(String),
+    Object(Vec<McpToolSParams>),
+    Array(Vec<ParamTypes>),
+}
+
+impl Display for ParamTypes {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let type_name = match self {
+            ParamTypes::Primitive(type_name) => type_name.to_owned(),
+            ParamTypes::Object(items) => {
+                format!(
+                    "{{{}}}",
+                    items
+                        .iter()
+                        .map(|t| format!("{} : {}", t.param_name, t.param_type))
+                        .collect::<Vec<String>>()
+                        .join(", ")
+                )
+            }
+            ParamTypes::Array(items) => format!("{} [ ]", items[0]),
+        };
+        write!(f, "{}", type_name)
+    }
+}
+
+// impl Serialize for ParamTypes {
+//     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+//     where
+//         S: Serializer,
+//     {
+//         match self {
+//             ParamTypes::Primitive(s) => {
+//                 let mut map = serializer.serialize_map(Some(1))?;
+//                 map.serialize_entry("Primitive", s)?;
+//                 map.end()
+//             }
+//
+//             ParamTypes::Object(params) => params.serialize(serializer), // Inline as array
+//             ParamTypes::Array(arr) => {
+//                 let mut map = serializer.serialize_map(Some(1))?;
+//                 map.serialize_entry("Array", arr)?;
+//                 map.end()
+//             }
+//         }
+//     }
+// }
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub struct McpToolSParams {
+    pub param_name: String,
+    pub param_type: ParamTypes,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub param_description: Option<String>,
+    pub required: bool,
+}
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub struct McpToolMeta {
+    pub name: String,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub description: Option<String>,
+    pub params: Vec<McpToolSParams>,
+}
+
+#[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
+pub struct McpServerInfo {
+    pub name: String,
+    pub version: String,
+    pub capabilities: McpCapabilities,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub tools: Option<Vec<McpToolMeta>>,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub prompts: Option<Vec<Prompt>>,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub resources: Option<Vec<Resource>>,
+    #[serde(default, skip_serializing_if = "::std::option::Option::is_none")]
+    pub resource_templates: Option<Vec<ResourceTemplate>>,
+}
diff --git a/src/types/commands.rs b/src/types/commands.rs
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/src/types/commands.rs
@@ -0,0 +1 @@
+

From 4885d3f6fea2e9d2e2cd2c895a3e1cfcfb2b3e39 Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 11:43:31 -0300
Subject: [PATCH 03/10] chore: update tests

---
 src/cli.rs                     | 316 +++++++++++++++++++++++++--------
 src/error.rs                   |   2 +
 src/lib.rs                     |   6 +-
 src/main.rs                    |  15 +-
 src/render_template.rs         |   9 +-
 src/templates.rs               |   2 +-
 src/types/commands.rs          | 151 ++++++++++++++++
 src/utils.rs                   |   2 +-
 tests/common/common.rs         |  11 +-
 tests/render_template_tests.rs |   4 +-
 tests/test_cli.rs              | 180 -------------------
 11 files changed, 421 insertions(+), 277 deletions(-)
 delete mode 100644 tests/test_cli.rs

diff --git a/src/cli.rs b/src/cli.rs
index fd14aad..8a70af7 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -1,24 +1,30 @@
-use std::{
-    fmt::Display,
-    io::{self, ErrorKind},
-    path::PathBuf,
-};
+use std::path::PathBuf;
 
 use clap::{arg, command, Parser, Subcommand, ValueEnum};
-
-use crate::{error::DiscoveryResult, utils::match_template, OutputTemplate};
+use mcp_discovery::types::{DiscoveryCommand, LogLevel, PrintOptions, Template, WriteOptions};
 
 #[derive(Debug, Clone, ValueEnum, PartialEq)]
-pub enum Template {
+pub enum CliTemplate {
     Md,
     MdPlain,
     Html,
     Txt,
 }
 
+impl Into<Template> for CliTemplate {
+    fn into(self) -> Template {
+        match self {
+            CliTemplate::Md => Template::Md,
+            CliTemplate::MdPlain => Template::MdPlain,
+            CliTemplate::Html => Template::Html,
+            CliTemplate::Txt => Template::Txt,
+        }
+    }
+}
+
 #[derive(Debug, Clone, ValueEnum, PartialEq)]
 #[allow(non_camel_case_types)]
-pub enum LogLevel {
+pub enum CliLogLevel {
     error,
     warn,
     info,
@@ -26,36 +32,36 @@ pub enum LogLevel {
     trace,
 }
 
-impl Display for LogLevel {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+impl Into<LogLevel> for CliLogLevel {
+    fn into(self) -> LogLevel {
         match self {
-            LogLevel::error => write!(f, "error"),
-            LogLevel::warn => write!(f, "warn"),
-            LogLevel::info => write!(f, "info"),
-            LogLevel::debug => write!(f, "debug"),
-            LogLevel::trace => write!(f, "trace"),
+            CliLogLevel::error => LogLevel::error,
+            CliLogLevel::warn => LogLevel::warn,
+            CliLogLevel::info => LogLevel::info,
+            CliLogLevel::debug => LogLevel::debug,
+            CliLogLevel::trace => LogLevel::trace,
         }
     }
 }
 
 #[derive(Subcommand, Debug)]
-pub enum DiscoveryCommand {
+pub enum CliDiscoveryCommand {
     /// Displays MCP server capability details in the terminal.
-    Print(PrintOptions),
+    Print(CliPrintOptions),
     /// Creates a file with MCP server capability details.
-    Create(WriteOptions),
+    Create(CliWriteOptions),
     /// Updates a file by adding MCP server capability information between specified markers.
-    Update(WriteOptions),
+    Update(CliWriteOptions),
 }
 
 #[derive(Parser, Debug)]
-pub struct WriteOptions {
+pub struct CliWriteOptions {
     #[arg(short, long)]
     pub filename: PathBuf,
 
     /// Select an output template from the built-in options.
     #[arg(short, long, value_enum, conflicts_with_all = ["template_file", "template_string"])]
-    pub template: Option<Template>,
+    pub template: Option<CliTemplate>,
 
     /// Path to a custom template file written in the Handlebars format.
     #[arg(long, short = 'p',
@@ -72,7 +78,7 @@ pub struct WriteOptions {
 
     /// Specifies the logging level for the application (default: info)
     #[arg(long, short)]
-    pub log_level: Option<LogLevel>,
+    pub log_level: Option<CliLogLevel>,
     /// Command and arguments to launch the MCP server.
     #[arg(
         value_name = "MCP Launch Command",
@@ -83,33 +89,24 @@ pub struct WriteOptions {
     pub mcp_server_cmd: Vec<String>,
 }
 
-impl WriteOptions {
-    pub fn match_template(&self) -> DiscoveryResult<OutputTemplate> {
-        match_template(
-            Some(&self.filename),
-            &self.template,
-            &self.template_file,
-            &self.template_string,
-        )
-    }
-
-    pub fn validate(&self) -> DiscoveryResult<()> {
-        if !self.filename.exists() {
-            return Err(io::Error::new(
-                ErrorKind::NotFound,
-                format!("File '{}' not found", self.filename.to_string_lossy()),
-            )
-            .into());
+impl Into<WriteOptions> for CliWriteOptions {
+    fn into(self) -> WriteOptions {
+        WriteOptions {
+            filename: self.filename,
+            template: self.template.map(|t| t.into()),
+            template_file: self.template_file,
+            template_string: self.template_string,
+            log_level: self.log_level.map(|l| l.into()),
+            mcp_server_cmd: self.mcp_server_cmd,
         }
-        Ok(())
     }
 }
 
 #[derive(Parser, Debug)]
-pub struct PrintOptions {
+pub struct CliPrintOptions {
     /// Select an output template from the built-in options.
     #[arg(short, long, value_enum, conflicts_with_all = ["template_file", "template_string"])]
-    pub template: Option<Template>,
+    pub template: Option<CliTemplate>,
 
     /// Path to a custom template file written in the Handlebars format.
     #[arg(long, short = 'p',
@@ -126,7 +123,7 @@ conflicts_with_all = ["template", "template_string"])]
 
     /// Specifies the logging level for the application (default: info)
     #[arg(long, short)]
-    pub log_level: Option<LogLevel>,
+    pub log_level: Option<CliLogLevel>,
 
     /// Command and arguments to launch the MCP server.
     #[arg(
@@ -138,31 +135,30 @@ conflicts_with_all = ["template", "template_string"])]
     pub mcp_server_cmd: Vec<String>,
 }
 
-impl PrintOptions {
-    pub fn match_template(&self) -> DiscoveryResult<OutputTemplate> {
-        match_template(
-            None,
-            &self.template,
-            &self.template_file,
-            &self.template_string,
-        )
-    }
-}
-
-impl DiscoveryCommand {
-    pub fn mcp_launch_command(&self) -> &Vec<String> {
-        match self {
-            DiscoveryCommand::Create(create_options) => &create_options.mcp_server_cmd,
-            DiscoveryCommand::Update(update_options) => &update_options.mcp_server_cmd,
-            DiscoveryCommand::Print(print_args) => &print_args.mcp_server_cmd,
+impl Into<PrintOptions> for CliPrintOptions {
+    fn into(self) -> PrintOptions {
+        PrintOptions {
+            template: self.template.map(|t| t.into()),
+            template_file: self.template_file,
+            template_string: self.template_string,
+            log_level: self.log_level.map(|l| l.into()),
+            mcp_server_cmd: self.mcp_server_cmd,
         }
     }
+}
 
-    pub fn log_level(&self) -> &Option<LogLevel> {
+impl Into<DiscoveryCommand> for CliDiscoveryCommand {
+    fn into(self) -> DiscoveryCommand {
         match self {
-            DiscoveryCommand::Create(create_options) => &create_options.log_level,
-            DiscoveryCommand::Update(update_options) => &update_options.log_level,
-            DiscoveryCommand::Print(print_args) => &print_args.log_level,
+            CliDiscoveryCommand::Print(cli_print_options) => {
+                DiscoveryCommand::Print(cli_print_options.into())
+            }
+            CliDiscoveryCommand::Create(cli_write_options) => {
+                DiscoveryCommand::Create(cli_write_options.into())
+            }
+            CliDiscoveryCommand::Update(cli_write_options) => {
+                DiscoveryCommand::Update(cli_write_options.into())
+            }
         }
     }
 }
@@ -175,11 +171,11 @@ long_about = None)]
 #[clap(args_conflicts_with_subcommands = true)]
 pub struct CommandArguments {
     #[command(subcommand)]
-    pub command: Option<DiscoveryCommand>,
+    pub command: Option<CliDiscoveryCommand>,
 
     /// Select an output template from the built-in options.
     #[arg(short, long, value_enum, conflicts_with_all = ["template_file", "template_string"])]
-    pub template: Option<Template>,
+    pub template: Option<CliTemplate>,
 
     /// Path to a custom template file written in the Handlebars format.
     #[arg(long, short = 'p',
@@ -196,7 +192,7 @@ pub struct CommandArguments {
 
     /// Specifies the logging level for the application (default: info)
     #[arg(long, short)]
-    pub log_level: Option<LogLevel>,
+    pub log_level: Option<CliLogLevel>,
 
     /// Command and arguments to launch the MCP server.
     #[arg(
@@ -208,8 +204,186 @@ pub struct CommandArguments {
     pub mcp_server_cmd: Vec<String>,
 }
 
-impl CommandArguments {
-    pub fn mcp_launch_command(&self) -> &Vec<String> {
-        self.command.as_ref().unwrap().mcp_launch_command()
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use mcp_discovery::types::DiscoveryCommand;
+
+    // Helper function to parse arguments from a vector of strings
+    pub fn parse_args(args: Vec<&str>) -> CommandArguments {
+        CommandArguments::parse_from(args)
+    }
+
+    #[test]
+    fn test_version_flag() {
+        let args = vec!["mcp-tool", "--version"];
+        let result = CommandArguments::try_parse_from(args);
+        assert!(result.is_err(), "Expected clap to handle --version flag");
+        // Note: clap automatically handles --version and exits, so this test verifies it doesn't parse normally
+    }
+
+    #[test]
+    fn test_conflicting_template_and_template_file() {
+        let args = vec![
+            "mcp-tool",
+            "create",
+            "--filename",
+            "output.md",
+            "--template",
+            "md",
+            "--template-file",
+            "custom.hbs",
+            "--",
+            "mcp-server",
+        ];
+        let result = CommandArguments::try_parse_from(args);
+        assert!(
+            result.is_err(),
+            "Expected error due to conflicting template options"
+        );
+    }
+
+    #[test]
+    fn test_missing_required_mcp_server_cmd() {
+        let args = vec!["mcp-tool", "print"];
+        let result = CommandArguments::try_parse_from(args);
+        assert!(
+            result.is_err(),
+            "Expected error due to missing mcp_server_cmd"
+        );
+    }
+
+    #[test]
+    fn test_file_options_match_template_custom() {
+        let file_options = WriteOptions {
+            filename: PathBuf::from("output.html"),
+            template: None,
+            template_file: Some(PathBuf::from("templates/markdown/markdown_template.md")),
+            mcp_server_cmd: vec!["mcp-server".to_string()],
+            template_string: None,
+            log_level: None,
+        };
+
+        let result = file_options.match_template();
+
+        assert!(
+            result.is_ok(),
+            "Expected successful template matching with custom file"
+        );
+    }
+
+    #[test]
+    fn test_create_command_parsing() {
+        let args = vec![
+            "mcp-tool",
+            "create",
+            "--filename",
+            "output.md",
+            "--template",
+            "md",
+            "--",
+            "mcp-server",
+            "--some-params",
+            "some-values",
+        ];
+        let parsed = parse_args(args);
+
+        match parsed.command {
+            Some(CliDiscoveryCommand::Create(file_options)) => {
+                assert_eq!(file_options.filename, PathBuf::from("output.md"));
+                assert_eq!(file_options.template, Some(CliTemplate::Md));
+                assert_eq!(file_options.template_file, None);
+                assert_eq!(
+                    file_options.mcp_server_cmd,
+                    vec!["mcp-server", "--some-params", "some-values"]
+                );
+            }
+            _ => panic!("Expected Create command"),
+        }
+    }
+
+    #[test]
+    fn test_update_command_with_template_file() {
+        let args = vec![
+            "mcp-tool",
+            "update",
+            "--filename",
+            "output.html",
+            "--template-file",
+            "custom.hbs",
+            "--",
+            "mcp-server",
+            "--param",
+            "90",
+        ];
+        let parsed = parse_args(args);
+
+        match parsed.command {
+            Some(CliDiscoveryCommand::Update(file_options)) => {
+                assert_eq!(file_options.filename, PathBuf::from("output.html"));
+                assert_eq!(file_options.template, None);
+                assert_eq!(
+                    file_options.template_file,
+                    Some(PathBuf::from("custom.hbs"))
+                );
+                assert_eq!(
+                    file_options.mcp_server_cmd,
+                    vec!["mcp-server", "--param", "90"]
+                );
+            }
+            _ => panic!("Expected Update command"),
+        }
+    }
+
+    #[test]
+    fn test_print_command_with_json() {
+        let args = vec!["mcp-tool", "print", "--", "mcp-server", "--verbose"];
+        let parsed = parse_args(args);
+
+        match parsed.command {
+            Some(CliDiscoveryCommand::Print(print_options)) => {
+                assert_eq!(
+                    print_options.mcp_server_cmd,
+                    vec!["mcp-server", "--verbose"]
+                );
+            }
+            _ => panic!("Expected Print command"),
+        }
+    }
+
+    #[test]
+    fn test_mcp_launch_command_retrieval() {
+        let args = vec![
+            "mcp-tool",
+            "create",
+            "--filename",
+            "output.txt",
+            "--template",
+            "txt",
+            "--",
+            "mcp-server",
+            "--port",
+            "9090",
+        ];
+        let command: DiscoveryCommand = parse_args(args).command.unwrap().into();
+
+        let launch_cmd = command.mcp_launch_command();
+        assert_eq!(launch_cmd, &vec!["mcp-server", "--port", "9090"]);
+    }
+
+    #[test]
+    fn test_file_options_match_template_builtin() {
+        let file_options = WriteOptions {
+            filename: PathBuf::from("output.md"),
+            template: Some(Template::Md),
+            template_file: None,
+            mcp_server_cmd: vec!["mcp-server".to_string()],
+            template_string: None,
+            log_level: None,
+        };
+
+        let result = file_options.match_template();
+        assert!(result.is_ok(), "Expected successful template matching");
+        // Note: Cannot assert specific OutputTemplate without knowing its structure
     }
 }
diff --git a/src/error.rs b/src/error.rs
index 5cb60f1..23009c8 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -7,6 +7,8 @@ pub type DiscoveryResult<T> = core::result::Result<T, DiscoveryError>;
 pub enum DiscoveryError {
     #[error("The MCP Server failed to initialize successfully.")]
     ServerNotInitialized,
+    #[error("'{0}' is not a valid template!")]
+    InvalidTemplate(String),
     #[error("{0}")]
     InvalidSchema(String),
     #[error("{0}")]
diff --git a/src/lib.rs b/src/lib.rs
index a3a8535..546c67d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,3 @@
-mod cli;
 pub mod error;
 mod handler;
 pub mod render_template;
@@ -8,7 +7,6 @@ pub mod templates;
 pub mod types;
 pub mod utils;
 
-pub use cli::{CommandArguments, DiscoveryCommand, LogLevel, PrintOptions, Template, WriteOptions};
 use colored::Colorize;
 use error::{DiscoveryError, DiscoveryResult};
 use render_template::{detect_render_markers, render_template};
@@ -16,7 +14,9 @@ use schema::tool_params;
 use std::io::stdout;
 use std_output::{print_header, print_list, print_summary};
 pub use templates::OutputTemplate;
-use types::{McpCapabilities, McpServerInfo, McpToolMeta};
+use types::{
+    DiscoveryCommand, McpCapabilities, McpServerInfo, McpToolMeta, PrintOptions, WriteOptions,
+};
 
 use std::sync::Arc;
 
diff --git a/src/main.rs b/src/main.rs
index 68066c2..137a980 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,21 +1,28 @@
+mod cli;
 use clap::Parser;
+use cli::{CliDiscoveryCommand, CliPrintOptions, CommandArguments};
 use colored::Colorize;
-use mcp_discovery::{CommandArguments, DiscoveryCommand, LogLevel, McpDiscovery, PrintOptions};
+
+use mcp_discovery::{
+    types::{DiscoveryCommand, LogLevel},
+    McpDiscovery,
+};
 use tracing_subscriber::{self, EnvFilter};
 
 #[tokio::main]
 async fn main() {
     let args = CommandArguments::parse();
 
-    let command = args
+    let command: DiscoveryCommand = args
         .command
-        .unwrap_or(DiscoveryCommand::Print(PrintOptions {
+        .unwrap_or(CliDiscoveryCommand::Print(CliPrintOptions {
             mcp_server_cmd: args.mcp_server_cmd,
             template: args.template,
             template_file: args.template_file,
             template_string: args.template_string,
             log_level: args.log_level,
-        }));
+        }))
+        .into();
 
     let filter = format!(
         "{}={}",
diff --git a/src/render_template.rs b/src/render_template.rs
index 80d38ad..3d8f6f9 100644
--- a/src/render_template.rs
+++ b/src/render_template.rs
@@ -1,8 +1,7 @@
 //! Module for rendering templates using Handlebars and handling MCP server template markers.
 
-use std::path::PathBuf;
+use std::{path::PathBuf, str::FromStr};
 
-use clap::ValueEnum;
 use handlebars::{
     handlebars_helper, Context, Handlebars, Helper, HelperDef, HelperResult, Output, RenderContext,
     RenderError,
@@ -15,11 +14,11 @@ use unicode_width::UnicodeWidthStr;
 use crate::{
     error::{DiscoveryError, DiscoveryResult},
     templates::{InlineTemplateInfo, PARTIALS},
-    types::ParamTypes,
+    types::{ParamTypes, Template, WriteOptions},
     utils::{
         boolean_indicator, line_ending, match_template, RenderTemplateInfo, UpdateTemplateInfo,
     },
-    McpServerInfo, OutputTemplate, Template, WriteOptions,
+    McpServerInfo, OutputTemplate,
 };
 
 /// Properties parsed from the `MCP_DISCOVERY_TEMPLATE_START` marker line in template files.
@@ -482,7 +481,7 @@ pub fn extract_template_prop(line: &str) -> Option<Template> {
         .captures(line)
         .and_then(|cap| cap.get(2).map(|m| m.as_str().trim().to_string()));
 
-    let template = Template::from_str(&prop_value.unwrap_or_default(), true);
+    let template = Template::from_str(&prop_value.unwrap_or_default());
     template.ok()
 }
 
diff --git a/src/templates.rs b/src/templates.rs
index bab39e8..9f496d6 100644
--- a/src/templates.rs
+++ b/src/templates.rs
@@ -1,7 +1,7 @@
 use crate::{
-    cli::Template,
     error::DiscoveryResult,
     render_template,
+    types::Template,
     utils::{find_template_file, line_ending},
     McpServerInfo,
 };
diff --git a/src/types/commands.rs b/src/types/commands.rs
index 8b13789..2e62197 100644
--- a/src/types/commands.rs
+++ b/src/types/commands.rs
@@ -1 +1,152 @@
+use std::{
+    fmt::Display,
+    io::{self, ErrorKind},
+    path::PathBuf,
+    str::FromStr,
+};
 
+use crate::{
+    error::{DiscoveryError, DiscoveryResult},
+    utils::match_template,
+    OutputTemplate,
+};
+
+#[derive(Debug)]
+pub enum DiscoveryCommand {
+    /// Displays MCP server capability details in the terminal.
+    Print(PrintOptions),
+    /// Creates a file with MCP server capability details.
+    Create(WriteOptions),
+    /// Updates a file by adding MCP server capability information between specified markers.
+    Update(WriteOptions),
+}
+#[derive(Debug, Clone, PartialEq)]
+pub enum Template {
+    Md,
+    MdPlain,
+    Html,
+    Txt,
+}
+
+impl FromStr for Template {
+    type Err = DiscoveryError;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "md" => Ok(Template::Md),
+            "md-plain" => Ok(Template::MdPlain),
+            "html" => Ok(Template::Html),
+            "txt" => Ok(Template::Txt),
+            _ => Err(DiscoveryError::InvalidTemplate(s.to_string())),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+#[allow(non_camel_case_types)]
+pub enum LogLevel {
+    error,
+    warn,
+    info,
+    debug,
+    trace,
+}
+
+impl Display for LogLevel {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            LogLevel::error => write!(f, "error"),
+            LogLevel::warn => write!(f, "warn"),
+            LogLevel::info => write!(f, "info"),
+            LogLevel::debug => write!(f, "debug"),
+            LogLevel::trace => write!(f, "trace"),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct PrintOptions {
+    /// Select an output template from the built-in options.
+    pub template: Option<Template>,
+
+    /// Path to a custom template file written in the Handlebars format.
+    pub template_file: Option<PathBuf>,
+
+    /// Template content provided as a string.
+    pub template_string: Option<String>,
+
+    /// Specifies the logging level for the application (default: info)
+    pub log_level: Option<LogLevel>,
+
+    /// Command and arguments to launch the MCP server.
+    pub mcp_server_cmd: Vec<String>,
+}
+
+impl PrintOptions {
+    pub fn match_template(&self) -> DiscoveryResult<OutputTemplate> {
+        match_template(
+            None,
+            &self.template,
+            &self.template_file,
+            &self.template_string,
+        )
+    }
+}
+
+#[derive(Debug)]
+pub struct WriteOptions {
+    pub filename: PathBuf,
+
+    /// Select an output template from the built-in options.
+    pub template: Option<Template>,
+
+    /// Path to a custom template file written in the Handlebars format.
+    pub template_file: Option<PathBuf>,
+
+    /// Template content provided as a string.
+    pub template_string: Option<String>,
+
+    /// Specifies the logging level for the application (default: info)
+    pub log_level: Option<LogLevel>,
+    /// Command and arguments to launch the MCP server.
+    pub mcp_server_cmd: Vec<String>,
+}
+
+impl WriteOptions {
+    pub fn match_template(&self) -> DiscoveryResult<OutputTemplate> {
+        match_template(
+            Some(&self.filename),
+            &self.template,
+            &self.template_file,
+            &self.template_string,
+        )
+    }
+
+    pub fn validate(&self) -> DiscoveryResult<()> {
+        if !self.filename.exists() {
+            return Err(io::Error::new(
+                ErrorKind::NotFound,
+                format!("File '{}' not found", self.filename.to_string_lossy()),
+            )
+            .into());
+        }
+        Ok(())
+    }
+}
+
+impl DiscoveryCommand {
+    pub fn mcp_launch_command(&self) -> &Vec<String> {
+        match self {
+            DiscoveryCommand::Create(create_options) => &create_options.mcp_server_cmd,
+            DiscoveryCommand::Update(update_options) => &update_options.mcp_server_cmd,
+            DiscoveryCommand::Print(print_args) => &print_args.mcp_server_cmd,
+        }
+    }
+
+    pub fn log_level(&self) -> &Option<LogLevel> {
+        match self {
+            DiscoveryCommand::Create(create_options) => &create_options.log_level,
+            DiscoveryCommand::Update(update_options) => &update_options.log_level,
+            DiscoveryCommand::Print(print_args) => &print_args.log_level,
+        }
+    }
+}
diff --git a/src/utils.rs b/src/utils.rs
index 847ff75..25360c4 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,6 +1,6 @@
 use path_clean::PathClean;
 
-use crate::{cli::Template, error::DiscoveryResult, OutputTemplate};
+use crate::{error::DiscoveryResult, types::Template, OutputTemplate};
 use std::{
     io::{self, ErrorKind},
     path::{Path, PathBuf},
diff --git a/tests/common/common.rs b/tests/common/common.rs
index 382ae17..cfb0402 100644
--- a/tests/common/common.rs
+++ b/tests/common/common.rs
@@ -1,15 +1,6 @@
-use clap::Parser;
-use mcp_discovery::{
-    types::{McpCapabilities, McpServerInfo},
-    CommandArguments,
-};
+use mcp_discovery::types::{McpCapabilities, McpServerInfo};
 use rust_mcp_sdk::macros::JsonSchema;
 
-// Helper function to parse arguments from a vector of strings
-pub fn parse_args(args: Vec<&str>) -> CommandArguments {
-    CommandArguments::parse_from(args)
-}
-
 pub fn default_mcp_server_info() -> McpServerInfo {
     McpServerInfo {
         name: Default::default(),
diff --git a/tests/render_template_tests.rs b/tests/render_template_tests.rs
index e879a77..9727dcd 100644
--- a/tests/render_template_tests.rs
+++ b/tests/render_template_tests.rs
@@ -5,8 +5,8 @@ use common::default_mcp_server_info;
 use handlebars::Handlebars;
 use mcp_discovery::{
     render_template::{detect_render_markers, extract_template_file, register_helpers},
-    types::ParamTypes,
-    OutputTemplate, WriteOptions,
+    types::{ParamTypes, WriteOptions},
+    OutputTemplate,
 };
 use serde_json::json;
 use std::fs::write;
diff --git a/tests/test_cli.rs b/tests/test_cli.rs
deleted file mode 100644
index 86571bc..0000000
--- a/tests/test_cli.rs
+++ /dev/null
@@ -1,180 +0,0 @@
-#[path = "common/common.rs"]
-pub mod common;
-
-use clap::Parser;
-use common::parse_args;
-use mcp_discovery::{CommandArguments, DiscoveryCommand, Template, WriteOptions};
-use std::path::PathBuf;
-
-#[test]
-fn test_create_command_parsing() {
-    let args = vec![
-        "mcp-tool",
-        "create",
-        "--filename",
-        "output.md",
-        "--template",
-        "md",
-        "--",
-        "mcp-server",
-        "--some-params",
-        "some-values",
-    ];
-    let parsed = parse_args(args);
-
-    match parsed.command {
-        Some(DiscoveryCommand::Create(file_options)) => {
-            assert_eq!(file_options.filename, PathBuf::from("output.md"));
-            assert_eq!(file_options.template, Some(Template::Md));
-            assert_eq!(file_options.template_file, None);
-            assert_eq!(
-                file_options.mcp_server_cmd,
-                vec!["mcp-server", "--some-params", "some-values"]
-            );
-        }
-        _ => panic!("Expected Create command"),
-    }
-}
-
-#[test]
-fn test_update_command_with_template_file() {
-    let args = vec![
-        "mcp-tool",
-        "update",
-        "--filename",
-        "output.html",
-        "--template-file",
-        "custom.hbs",
-        "--",
-        "mcp-server",
-        "--param",
-        "90",
-    ];
-    let parsed = parse_args(args);
-
-    match parsed.command {
-        Some(DiscoveryCommand::Update(file_options)) => {
-            assert_eq!(file_options.filename, PathBuf::from("output.html"));
-            assert_eq!(file_options.template, None);
-            assert_eq!(
-                file_options.template_file,
-                Some(PathBuf::from("custom.hbs"))
-            );
-            assert_eq!(
-                file_options.mcp_server_cmd,
-                vec!["mcp-server", "--param", "90"]
-            );
-        }
-        _ => panic!("Expected Update command"),
-    }
-}
-
-#[test]
-fn test_print_command_with_json() {
-    let args = vec!["mcp-tool", "print", "--", "mcp-server", "--verbose"];
-    let parsed = parse_args(args);
-
-    match parsed.command {
-        Some(DiscoveryCommand::Print(print_options)) => {
-            assert_eq!(
-                print_options.mcp_server_cmd,
-                vec!["mcp-server", "--verbose"]
-            );
-        }
-        _ => panic!("Expected Print command"),
-    }
-}
-
-#[test]
-fn test_mcp_launch_command_retrieval() {
-    let args = vec![
-        "mcp-tool",
-        "create",
-        "--filename",
-        "output.txt",
-        "--template",
-        "txt",
-        "--",
-        "mcp-server",
-        "--port",
-        "9090",
-    ];
-    let parsed = parse_args(args);
-
-    let launch_cmd = parsed.mcp_launch_command();
-    assert_eq!(launch_cmd, &vec!["mcp-server", "--port", "9090"]);
-}
-
-#[test]
-fn test_file_options_match_template_builtin() {
-    let file_options = WriteOptions {
-        filename: PathBuf::from("output.md"),
-        template: Some(Template::Md),
-        template_file: None,
-        mcp_server_cmd: vec!["mcp-server".to_string()],
-        template_string: None,
-        log_level: None,
-    };
-
-    let result = file_options.match_template();
-    assert!(result.is_ok(), "Expected successful template matching");
-    // Note: Cannot assert specific OutputTemplate without knowing its structure
-}
-
-#[test]
-fn test_file_options_match_template_custom() {
-    let file_options = WriteOptions {
-        filename: PathBuf::from("output.html"),
-        template: None,
-        template_file: Some(PathBuf::from("templates/markdown/markdown_template.md")),
-        mcp_server_cmd: vec!["mcp-server".to_string()],
-        template_string: None,
-        log_level: None,
-    };
-
-    let result = file_options.match_template();
-
-    assert!(
-        result.is_ok(),
-        "Expected successful template matching with custom file"
-    );
-}
-
-#[test]
-fn test_missing_required_mcp_server_cmd() {
-    let args = vec!["mcp-tool", "print"];
-    let result = CommandArguments::try_parse_from(args);
-    assert!(
-        result.is_err(),
-        "Expected error due to missing mcp_server_cmd"
-    );
-}
-
-#[test]
-fn test_conflicting_template_and_template_file() {
-    let args = vec![
-        "mcp-tool",
-        "create",
-        "--filename",
-        "output.md",
-        "--template",
-        "md",
-        "--template-file",
-        "custom.hbs",
-        "--",
-        "mcp-server",
-    ];
-    let result = CommandArguments::try_parse_from(args);
-    assert!(
-        result.is_err(),
-        "Expected error due to conflicting template options"
-    );
-}
-
-#[test]
-fn test_version_flag() {
-    let args = vec!["mcp-tool", "--version"];
-    let result = CommandArguments::try_parse_from(args);
-    assert!(result.is_err(), "Expected clap to handle --version flag");
-    // Note: clap automatically handles --version and exits, so this test verifies it doesn't parse normally
-}

From 5aa11a5f3a14fa3b2a1fc3f8564e8dfce2213007 Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 13:31:35 -0300
Subject: [PATCH 04/10] refactor: getters

---
 src/cli.rs    |  3 +--
 src/lib.rs    | 37 ++++++++++++++++++-------------------
 src/main.rs   |  6 +-----
 src/schema.rs | 12 ++++++------
 src/utils.rs  |  4 ++--
 5 files changed, 28 insertions(+), 34 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index 8a70af7..614fe13 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -1,7 +1,7 @@
 use std::path::PathBuf;
 
 use clap::{arg, command, Parser, Subcommand, ValueEnum};
-use mcp_discovery::types::{DiscoveryCommand, LogLevel, PrintOptions, Template, WriteOptions};
+use mcp_discovery::{DiscoveryCommand, LogLevel, PrintOptions, Template, WriteOptions};
 
 #[derive(Debug, Clone, ValueEnum, PartialEq)]
 pub enum CliTemplate {
@@ -207,7 +207,6 @@ pub struct CommandArguments {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use mcp_discovery::types::DiscoveryCommand;
 
     // Helper function to parse arguments from a vector of strings
     pub fn parse_args(args: Vec<&str>) -> CommandArguments {
diff --git a/src/lib.rs b/src/lib.rs
index 546c67d..1178403 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,11 +1,17 @@
 pub mod error;
 mod handler;
-pub mod render_template;
-pub mod schema;
-pub mod std_output;
-pub mod templates;
-pub mod types;
-pub mod utils;
+mod render_template;
+mod schema;
+mod std_output;
+mod templates;
+mod types;
+mod utils;
+
+pub use templates::OutputTemplate;
+pub use types::{
+    DiscoveryCommand, LogLevel, McpCapabilities, McpServerInfo, McpToolMeta, ParamTypes,
+    PrintOptions, Template, WriteOptions,
+};
 
 use colored::Colorize;
 use error::{DiscoveryError, DiscoveryResult};
@@ -13,10 +19,6 @@ use render_template::{detect_render_markers, render_template};
 use schema::tool_params;
 use std::io::stdout;
 use std_output::{print_header, print_list, print_summary};
-pub use templates::OutputTemplate;
-use types::{
-    DiscoveryCommand, McpCapabilities, McpServerInfo, McpToolMeta, PrintOptions, WriteOptions,
-};
 
 use std::sync::Arc;
 
@@ -326,10 +328,7 @@ impl McpDiscovery {
         Ok(Some(tools))
     }
 
-    async fn get_prompts(
-        &self,
-        client: Arc<ClientRuntime>,
-    ) -> DiscoveryResult<Option<Vec<Prompt>>> {
+    async fn prompts(&self, client: Arc<ClientRuntime>) -> DiscoveryResult<Option<Vec<Prompt>>> {
         if !client.server_has_prompts().unwrap_or(false) {
             return Ok(None);
         }
@@ -343,7 +342,7 @@ impl McpDiscovery {
         Ok(Some(prompts))
     }
 
-    async fn get_resources(
+    async fn resources(
         &self,
         client: Arc<ClientRuntime>,
     ) -> DiscoveryResult<Option<Vec<Resource>>> {
@@ -361,7 +360,7 @@ impl McpDiscovery {
         Ok(Some(resources))
     }
 
-    async fn get_resource_templates(
+    async fn resource_templates(
         &self,
         client: Arc<ClientRuntime>,
     ) -> DiscoveryResult<Option<Vec<ResourceTemplate>>> {
@@ -417,9 +416,9 @@ impl McpDiscovery {
         tracing::trace!("Capabilities: {}", capabilities);
 
         let tools = self.tools(Arc::clone(&client)).await?;
-        let prompts = self.get_prompts(Arc::clone(&client)).await?;
-        let resources = self.get_resources(Arc::clone(&client)).await?;
-        let resource_templates = self.get_resource_templates(Arc::clone(&client)).await?;
+        let prompts = self.prompts(Arc::clone(&client)).await?;
+        let resources = self.resources(Arc::clone(&client)).await?;
+        let resource_templates = self.resource_templates(Arc::clone(&client)).await?;
 
         let server_info = McpServerInfo {
             name: server_version.name,
diff --git a/src/main.rs b/src/main.rs
index 137a980..0b3a3dd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,11 +2,7 @@ mod cli;
 use clap::Parser;
 use cli::{CliDiscoveryCommand, CliPrintOptions, CommandArguments};
 use colored::Colorize;
-
-use mcp_discovery::{
-    types::{DiscoveryCommand, LogLevel},
-    McpDiscovery,
-};
+use mcp_discovery::{DiscoveryCommand, LogLevel, McpDiscovery};
 use tracing_subscriber::{self, EnvFilter};
 
 #[tokio::main]
diff --git a/src/schema.rs b/src/schema.rs
index a004c1b..d05e80e 100644
--- a/src/schema.rs
+++ b/src/schema.rs
@@ -8,7 +8,7 @@ use crate::{
 };
 
 /// Parses an object schema into a vector of `McpToolSParams`.
-pub fn get_param_object(object_map: &Map<String, Value>) -> DiscoveryResult<Vec<McpToolSParams>> {
+pub fn param_object(object_map: &Map<String, Value>) -> DiscoveryResult<Vec<McpToolSParams>> {
     let properties = object_map
         .get("properties")
         .and_then(|v| v.as_object())
@@ -31,7 +31,7 @@ pub fn get_param_object(object_map: &Map<String, Value>) -> DiscoveryResult<Vec<
                     "Property '{}' is not an object",
                     param_name
                 )))?;
-            let param_type = get_param_type(param_value)?;
+            let param_type = param_type(param_value)?;
             let param_description = object_map
                 .get("description")
                 .and_then(|v| v.as_str())
@@ -50,7 +50,7 @@ pub fn get_param_object(object_map: &Map<String, Value>) -> DiscoveryResult<Vec<
 }
 
 /// Determines the parameter type from a schema definition.
-pub fn get_param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTypes> {
+pub fn param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTypes> {
     let type_name =
         type_info
             .get("type")
@@ -66,9 +66,9 @@ pub fn get_param_type(type_info: &Map<String, Value>) -> DiscoveryResult<ParamTy
                     "Missing or invalid 'items' field in array type".to_string(),
                 ),
             )?;
-            Ok(ParamTypes::Array(vec![get_param_type(items_map)?]))
+            Ok(ParamTypes::Array(vec![param_type(items_map)?]))
         }
-        "object" => Ok(ParamTypes::Object(get_param_object(type_info)?)),
+        "object" => Ok(ParamTypes::Object(param_object(type_info)?)),
         _ => Ok(ParamTypes::Primitive(type_name.to_string())),
     }
 }
@@ -81,7 +81,7 @@ pub fn tool_params(
             .iter()
             .map(|(prop_name, prop_map)| {
                 let param_name = prop_name.to_owned();
-                let prop_type = get_param_type(prop_map).unwrap();
+                let prop_type = param_type(prop_map).unwrap();
                 let prop_description = prop_map
                     .get("description")
                     .and_then(|v| v.as_str())
diff --git a/src/utils.rs b/src/utils.rs
index 25360c4..7b3062e 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -76,7 +76,7 @@ pub fn boolean_indicator(boolean: bool) -> char {
     }
 }
 
-pub fn get_absolute_from_relative(
+pub fn absolute_from_relative(
     base_file_path: &Path,
     relative_file_path: &Path,
 ) -> io::Result<PathBuf> {
@@ -98,7 +98,7 @@ pub fn find_template_file(
     template_file: &Path,
     base_file: Option<&PathBuf>,
 ) -> DiscoveryResult<PathBuf> {
-    let relative_path = base_file.and_then(|f| get_absolute_from_relative(f, template_file).ok());
+    let relative_path = base_file.and_then(|f| absolute_from_relative(f, template_file).ok());
 
     let search_paths: Vec<_> = relative_path
         .into_iter()

From ba92bda018cdd2c9df88059c6995f096231ec1c6 Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 13:52:16 -0300
Subject: [PATCH 05/10] refactor: tests

---
 src/render_template.rs         | 257 +++++++++++++++++++++++++++++++++
 src/std_output.rs              | 124 +++++++++++++++-
 tests/common/common.rs         |   2 +-
 tests/render_template_tests.rs | 242 -------------------------------
 tests/test_std_output.rs       | 143 ------------------
 5 files changed, 374 insertions(+), 394 deletions(-)
 delete mode 100644 tests/render_template_tests.rs
 delete mode 100644 tests/test_std_output.rs

diff --git a/src/render_template.rs b/src/render_template.rs
index 3d8f6f9..0513874 100644
--- a/src/render_template.rs
+++ b/src/render_template.rs
@@ -491,3 +491,260 @@ pub fn extract_render_props(line: &str) -> RenderTemplateProps {
         template_file: extract_template_file(line).map(PathBuf::from),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+    use std::fs::write;
+    use tempfile::NamedTempFile;
+
+    use super::*;
+    use crate::*;
+
+    pub fn default_mcp_server_info() -> McpServerInfo {
+        McpServerInfo {
+            name: Default::default(),
+            version: Default::default(),
+            capabilities: McpCapabilities {
+                tools: false,
+                prompts: false,
+                resources: false,
+                logging: false,
+                experimental: false,
+            },
+            tools: Default::default(),
+            prompts: Default::default(),
+            resources: Default::default(),
+            resource_templates: Default::default(),
+        }
+    }
+
+    #[test]
+    fn test_register_helpers() {
+        let mut handlebar = Handlebars::new();
+        register_helpers(&mut handlebar);
+
+        // Test plus_one helper
+        let result = handlebar
+            .render_template("{{plus_one 5}}", &json!({}))
+            .expect("Failed to render plus_one");
+        assert_eq!(result, "6");
+
+        // Test capability_tag helper (supported)
+        let result = handlebar
+            .render_template("{{capability_tag \"Feature\" true 42}}", &json!({}))
+            .expect("Failed to render capability_tag");
+        assert_eq!(result, "🟢 Feature (42)");
+
+        // Test capability_tag helper (not supported)
+        let result = handlebar
+            .render_template("{{{capability_tag \"Feature\" false 0}}}", &json!({}))
+            .expect("Failed to render capability_tag");
+        assert_eq!(result, r#"<span style="opacity:0.6">🔴 Feature</span>"#);
+
+        // Test capability helper
+        let result = handlebar
+            .render_template("{{{capability \"Feature\" true 42}}}", &json!({}))
+            .expect("Failed to render capability");
+        assert_eq!(result, "🟢 Feature (42)");
+
+        // Test underline helper
+        let result = handlebar
+            .render_template("{{underline \"Title\"}}", &json!({}))
+            .expect("Failed to render underline");
+        assert_eq!(result, "Title\n─────");
+
+        // Test capability_title helper
+        let result = handlebar
+            .render_template("{{capability_title \"Title\" 10 true}}", &json!({}))
+            .expect("Failed to render capability_title");
+        assert_eq!(result, "Title(10)\n─────────");
+
+        // Test replace_regex helper
+        let result = handlebar
+            .render_template(
+                "{{replace_regex \"Hello World\" \"World\" \"Rust\"}}",
+                &json!({}),
+            )
+            .expect("Failed to render replace_regex");
+        assert_eq!(result, "Hello Rust");
+
+        // Test json helper (pretty)
+        let result = handlebar
+            .render_template("{{json true}}", &json!({"key": "value"}))
+            .expect("Failed to render json");
+        assert_eq!(result, "{\n  \"key\": \"value\"\n}");
+
+        // Test tool_param_type helper
+        let result = handlebar
+            .render_template(
+                "{{tool_param_type param}}",
+                &json!({"param": ParamTypes::Primitive("string".to_string())}),
+            )
+            .expect("Failed to render tool_param_type");
+        assert_eq!(result, "string");
+    }
+
+    #[test]
+    fn test_render_template() {
+        let template = OutputTemplate::TemplateString("Hello, {{name}}!".to_string());
+        let data = json!({"name": "World"});
+        let result =
+            render_template::render_template(&template, &data).expect("Failed to render template");
+        assert_eq!(result, "Hello, World!");
+
+        // Test with helper
+        let template = OutputTemplate::TemplateString("{{plus_one 5}}".to_string());
+        let result = render_template::render_template(&template, &json!({}))
+            .expect("Failed to render template");
+        assert_eq!(result, "6");
+
+        // Test invalid template
+        let template = OutputTemplate::TemplateString("{{#invalid}}".to_string());
+        let result = render_template::render_template(&template, &json!({}));
+        assert!(result.is_err(), "Expected error for invalid template");
+    }
+
+    #[test]
+    fn test_detect_render_markers_valid() {
+        let file = NamedTempFile::new().unwrap();
+        let content = r#"mcp-discovery-render
+        mcp-discovery-template
+        Template content
+        mcp-discovery-template-end
+        mcp-discovery-render-end"#;
+        write(&file, content).unwrap();
+        let options = WriteOptions {
+            filename: file.path().to_path_buf(),
+            template: None,
+            template_file: None,
+            mcp_server_cmd: vec!["mcp-server".to_string()],
+            template_string: None,
+            log_level: None,
+        };
+        let server_info = default_mcp_server_info();
+        let result = detect_render_markers(&options, &server_info);
+        assert!(result.is_ok(), "Expected valid marker detection");
+        let update_info = result.unwrap();
+        assert_eq!(update_info.render_locations.len(), 1);
+        assert_eq!(update_info.line_ending, "\n");
+        assert_eq!(update_info.render_locations[0].render_location, (1, 5));
+    }
+
+    #[test]
+    fn test_detect_render_markers_duplicate_template_start() {
+        let file = NamedTempFile::new().unwrap();
+        let content = r#"mcp-discovery-render
+        mcp-discovery-template
+        mcp-discovery-template
+        Template content
+        mcp-discovery-template-end
+        mcp-discovery-render-end"#;
+        write(&file, content).unwrap();
+        let options = WriteOptions {
+            filename: file.path().to_path_buf(),
+            template: None,
+            template_file: None,
+            mcp_server_cmd: vec!["mcp-server".to_string()],
+            template_string: None,
+            log_level: None,
+        };
+        let server_info = default_mcp_server_info();
+        let result = detect_render_markers(&options, &server_info);
+        assert!(result.is_err());
+        assert!(result
+            .unwrap_err()
+            .to_string()
+            .contains("Duplicate template start marker"));
+    }
+
+    #[test]
+    fn test_detect_render_markers_template_outside_render() {
+        let file = NamedTempFile::new().unwrap();
+        let content = r#"
+        mcp-discovery-template
+        Template content
+        mcp-discovery-template-end
+        "#;
+        write(&file, content).unwrap();
+        let options = WriteOptions {
+            filename: file.path().to_path_buf(),
+            template: None,
+            template_file: None,
+            mcp_server_cmd: vec!["mcp-server".to_string()],
+            template_string: None,
+            log_level: None,
+        };
+        let server_info = default_mcp_server_info();
+        let result = detect_render_markers(&options, &server_info);
+        assert!(result.is_err());
+        assert!(result
+            .unwrap_err()
+            .to_string()
+            .contains("outside a render section"));
+    }
+
+    #[test]
+    fn test_detect_render_markers_unmatched_template_end() {
+        let file = NamedTempFile::new().unwrap();
+        let content = r#"
+        mcp-discovery-render
+        mcp-discovery-template-end
+        mcp-discovery-render-end
+        "#;
+        write(&file, content).unwrap();
+        let options = WriteOptions {
+            filename: file.path().to_path_buf(),
+            template: None,
+            template_file: None,
+            mcp_server_cmd: vec!["mcp-server".to_string()],
+            template_string: None,
+            log_level: None,
+        };
+        let server_info = default_mcp_server_info();
+        let result = detect_render_markers(&options, &server_info);
+        assert!(result.is_err());
+        assert!(result
+            .unwrap_err()
+            .to_string()
+            .contains("no matching start marker"));
+    }
+
+    #[test]
+    fn test_detect_render_markers_conflicting_templates() {
+        let file = NamedTempFile::new().unwrap();
+        let content = r#"
+        mcp-discovery-render template-file=./template.hbs
+        mcp-discovery-template
+        Template content
+        mcp-discovery-template-end
+        mcp-discovery-render-end"#;
+        write(&file, content).unwrap();
+        let options = WriteOptions {
+            filename: file.path().to_path_buf(),
+            template: None,
+            template_file: None,
+            mcp_server_cmd: vec!["mcp-server".to_string()],
+            template_string: None,
+            log_level: None,
+        };
+        let server_info = default_mcp_server_info();
+        let result = detect_render_markers(&options, &server_info);
+        assert!(result.is_err());
+        assert!(result
+            .unwrap_err()
+            .to_string()
+            .contains("both a 'template-file' and an inline template"));
+    }
+
+    #[test]
+    fn test_extract_template_file() {
+        let line = "mcp-discovery-render template-file=./template.hbs";
+        let result = extract_template_file(line);
+        assert_eq!(result, Some("./template.hbs".to_string()));
+
+        let line = "mcp-discovery-render";
+        let result = extract_template_file(line);
+        assert_eq!(result, None);
+    }
+}
diff --git a/src/std_output.rs b/src/std_output.rs
index 2688a64..4505e79 100644
--- a/src/std_output.rs
+++ b/src/std_output.rs
@@ -2,7 +2,7 @@ use colored::Colorize;
 use std::io::{self, Write};
 use unicode_width::UnicodeWidthStr;
 
-use crate::{error::DiscoveryResult, utils::boolean_indicator, McpServerInfo};
+use crate::{utils::boolean_indicator, McpServerInfo};
 
 const SUMMARY_HEADER_SIZE: usize = 44;
 
@@ -50,13 +50,6 @@ pub fn table_content(width: usize, content: &str) -> String {
     format!("│{}{}{}│", " ".repeat(l_pad), content, " ".repeat(r_pad))
 }
 
-/// Function to print the server information as JSON.
-pub fn print_json<W: Write>(mut w: W, server_info: &McpServerInfo) -> DiscoveryResult<()> {
-    let json_content = serde_json::to_string(server_info)?;
-    writeln!(w, "{}", json_content)?;
-    Ok(())
-}
-
 /// Function to print a formatted summary of the server information.
 pub fn print_summary<W: Write>(mut w: W, server_info: &McpServerInfo) -> io::Result<usize> {
     let server_name = format!("{} {}", server_info.name, server_info.version);
@@ -97,3 +90,118 @@ pub fn print_summary<W: Write>(mut w: W, server_info: &McpServerInfo) -> io::Res
     writeln!(w, "{}", table_bottom(table_size))?;
     Ok(table_size)
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::McpCapabilities;
+
+    use super::*;
+
+    #[test]
+    fn test_table_top_bottom() {
+        assert_eq!(table_top(5), "┌─────┐");
+        assert_eq!(table_bottom(5), "└─────┘");
+    }
+
+    #[test]
+    fn test_table_content_centering() {
+        let result = table_content(10, "Hi");
+        assert!(result.starts_with("│"));
+        assert!(result.ends_with("│"));
+        assert_eq!(strip_ansi_escapes::strip_str(&result).width(), 12); // includes borders
+        assert!(result.contains("Hi"));
+    }
+
+    #[test]
+    fn test_print_list() {
+        let items = vec![
+            ("Item1".to_string(), "Value1".to_string()),
+            ("Item2".to_string(), "Value2".to_string()),
+        ];
+
+        let mut buffer = Vec::new();
+        print_list(&mut buffer, items).unwrap();
+
+        let output = String::from_utf8(buffer).unwrap();
+        let output = strip_ansi_escapes::strip_str(output);
+
+        assert!(output.contains("1. Item1: Value1"));
+        assert!(output.contains("2. Item2: Value2"));
+    }
+
+    #[test]
+    fn test_print_header_structure() {
+        let mut buffer = Vec::new();
+        print_header(&mut buffer, "My Header", 20).unwrap();
+
+        // Create the String first
+        let output = String::from_utf8(buffer).unwrap();
+
+        // Now borrow the String with `lines()`
+        let lines: Vec<&str> = output.lines().collect();
+
+        assert_eq!(lines.len(), 3);
+        assert!(lines[0].starts_with("┌") && lines[0].ends_with("┐"));
+        assert!(lines[2].starts_with("└") && lines[2].ends_with("┘"));
+        assert!(lines[1].contains("My Header"));
+    }
+
+    #[test]
+    fn test_print_summary_layout() {
+        let info = McpServerInfo {
+            name: "MyApp".to_string(),
+            version: "0.9".to_string(),
+            capabilities: McpCapabilities {
+                tools: true,
+                prompts: false,
+                resources: false,
+                logging: true,
+                experimental: false,
+            },
+            tools: None,
+            prompts: None,
+            resources: None,
+            resource_templates: None,
+        };
+
+        let mut buffer = Vec::new();
+        let size = print_summary(&mut buffer, &info).unwrap();
+
+        let output = String::from_utf8(buffer).unwrap();
+        assert!(output.contains("MyApp 0.9"));
+        assert!(size >= 44);
+    }
+
+    #[test]
+    fn test_print_to_stdout() {
+        let items = vec![
+            ("Tool1".to_string(), "Enabled".to_string()),
+            ("Tool2".to_string(), "Disabled".to_string()),
+        ];
+
+        // Print directly to stdout
+        print_list(io::stdout(), items).unwrap();
+    }
+
+    #[test]
+    fn test_print_summary_to_stdout() {
+        let info = McpServerInfo {
+            name: "SampleServer".to_string(),
+            version: "1.1".to_string(),
+            capabilities: McpCapabilities {
+                tools: true,
+                prompts: true,
+                resources: true,
+                logging: false,
+                experimental: false,
+            },
+            tools: None,
+            prompts: None,
+            resources: None,
+            resource_templates: None,
+        };
+
+        // Print directly to stdout
+        print_summary(io::stdout(), &info).unwrap();
+    }
+}
diff --git a/tests/common/common.rs b/tests/common/common.rs
index cfb0402..a955b37 100644
--- a/tests/common/common.rs
+++ b/tests/common/common.rs
@@ -1,4 +1,4 @@
-use mcp_discovery::types::{McpCapabilities, McpServerInfo};
+use mcp_discovery::{McpCapabilities, McpServerInfo};
 use rust_mcp_sdk::macros::JsonSchema;
 
 pub fn default_mcp_server_info() -> McpServerInfo {
diff --git a/tests/render_template_tests.rs b/tests/render_template_tests.rs
deleted file mode 100644
index 9727dcd..0000000
--- a/tests/render_template_tests.rs
+++ /dev/null
@@ -1,242 +0,0 @@
-#[path = "common/common.rs"]
-pub mod common;
-
-use common::default_mcp_server_info;
-use handlebars::Handlebars;
-use mcp_discovery::{
-    render_template::{detect_render_markers, extract_template_file, register_helpers},
-    types::{ParamTypes, WriteOptions},
-    OutputTemplate,
-};
-use serde_json::json;
-use std::fs::write;
-use tempfile::NamedTempFile;
-
-#[test]
-fn test_register_helpers() {
-    let mut handlebar = Handlebars::new();
-    register_helpers(&mut handlebar);
-
-    // Test plus_one helper
-    let result = handlebar
-        .render_template("{{plus_one 5}}", &json!({}))
-        .expect("Failed to render plus_one");
-    assert_eq!(result, "6");
-
-    // Test capability_tag helper (supported)
-    let result = handlebar
-        .render_template("{{capability_tag \"Feature\" true 42}}", &json!({}))
-        .expect("Failed to render capability_tag");
-    assert_eq!(result, "🟢 Feature (42)");
-
-    // Test capability_tag helper (not supported)
-    let result = handlebar
-        .render_template("{{{capability_tag \"Feature\" false 0}}}", &json!({}))
-        .expect("Failed to render capability_tag");
-    assert_eq!(result, r#"<span style="opacity:0.6">🔴 Feature</span>"#);
-
-    // Test capability helper
-    let result = handlebar
-        .render_template("{{{capability \"Feature\" true 42}}}", &json!({}))
-        .expect("Failed to render capability");
-    assert_eq!(result, "🟢 Feature (42)");
-
-    // Test underline helper
-    let result = handlebar
-        .render_template("{{underline \"Title\"}}", &json!({}))
-        .expect("Failed to render underline");
-    assert_eq!(result, "Title\n─────");
-
-    // Test capability_title helper
-    let result = handlebar
-        .render_template("{{capability_title \"Title\" 10 true}}", &json!({}))
-        .expect("Failed to render capability_title");
-    assert_eq!(result, "Title(10)\n─────────");
-
-    // Test replace_regex helper
-    let result = handlebar
-        .render_template(
-            "{{replace_regex \"Hello World\" \"World\" \"Rust\"}}",
-            &json!({}),
-        )
-        .expect("Failed to render replace_regex");
-    assert_eq!(result, "Hello Rust");
-
-    // Test json helper (pretty)
-    let result = handlebar
-        .render_template("{{json true}}", &json!({"key": "value"}))
-        .expect("Failed to render json");
-    assert_eq!(result, "{\n  \"key\": \"value\"\n}");
-
-    // Test tool_param_type helper
-    let result = handlebar
-        .render_template(
-            "{{tool_param_type param}}",
-            &json!({"param": ParamTypes::Primitive("string".to_string())}),
-        )
-        .expect("Failed to render tool_param_type");
-    assert_eq!(result, "string");
-}
-
-#[test]
-fn test_render_template() {
-    let template = OutputTemplate::TemplateString("Hello, {{name}}!".to_string());
-    let data = json!({"name": "World"});
-    let result = mcp_discovery::render_template::render_template(&template, &data)
-        .expect("Failed to render template");
-    assert_eq!(result, "Hello, World!");
-
-    // Test with helper
-    let template = OutputTemplate::TemplateString("{{plus_one 5}}".to_string());
-    let result = mcp_discovery::render_template::render_template(&template, &json!({}))
-        .expect("Failed to render template");
-    assert_eq!(result, "6");
-
-    // Test invalid template
-    let template = OutputTemplate::TemplateString("{{#invalid}}".to_string());
-    let result = mcp_discovery::render_template::render_template(&template, &json!({}));
-    assert!(result.is_err(), "Expected error for invalid template");
-}
-
-#[test]
-fn test_detect_render_markers_valid() {
-    let file = NamedTempFile::new().unwrap();
-    let content = r#"mcp-discovery-render
-    mcp-discovery-template
-    Template content
-    mcp-discovery-template-end
-    mcp-discovery-render-end"#;
-    write(&file, content).unwrap();
-    let options = WriteOptions {
-        filename: file.path().to_path_buf(),
-        template: None,
-        template_file: None,
-        mcp_server_cmd: vec!["mcp-server".to_string()],
-        template_string: None,
-        log_level: None,
-    };
-    let server_info = default_mcp_server_info();
-    let result = detect_render_markers(&options, &server_info);
-    assert!(result.is_ok(), "Expected valid marker detection");
-    let update_info = result.unwrap();
-    assert_eq!(update_info.render_locations.len(), 1);
-    assert_eq!(update_info.line_ending, "\n");
-    assert_eq!(update_info.render_locations[0].render_location, (1, 5));
-}
-
-#[test]
-fn test_detect_render_markers_duplicate_template_start() {
-    let file = NamedTempFile::new().unwrap();
-    let content = r#"mcp-discovery-render
-    mcp-discovery-template
-    mcp-discovery-template
-    Template content
-    mcp-discovery-template-end
-    mcp-discovery-render-end"#;
-    write(&file, content).unwrap();
-    let options = WriteOptions {
-        filename: file.path().to_path_buf(),
-        template: None,
-        template_file: None,
-        mcp_server_cmd: vec!["mcp-server".to_string()],
-        template_string: None,
-        log_level: None,
-    };
-    let server_info = default_mcp_server_info();
-    let result = detect_render_markers(&options, &server_info);
-    assert!(result.is_err());
-    assert!(result
-        .unwrap_err()
-        .to_string()
-        .contains("Duplicate template start marker"));
-}
-
-#[test]
-fn test_detect_render_markers_template_outside_render() {
-    let file = NamedTempFile::new().unwrap();
-    let content = r#"
-    mcp-discovery-template
-    Template content
-    mcp-discovery-template-end
-    "#;
-    write(&file, content).unwrap();
-    let options = WriteOptions {
-        filename: file.path().to_path_buf(),
-        template: None,
-        template_file: None,
-        mcp_server_cmd: vec!["mcp-server".to_string()],
-        template_string: None,
-        log_level: None,
-    };
-    let server_info = default_mcp_server_info();
-    let result = detect_render_markers(&options, &server_info);
-    assert!(result.is_err());
-    assert!(result
-        .unwrap_err()
-        .to_string()
-        .contains("outside a render section"));
-}
-
-#[test]
-fn test_detect_render_markers_unmatched_template_end() {
-    let file = NamedTempFile::new().unwrap();
-    let content = r#"
-    mcp-discovery-render
-    mcp-discovery-template-end
-    mcp-discovery-render-end
-    "#;
-    write(&file, content).unwrap();
-    let options = WriteOptions {
-        filename: file.path().to_path_buf(),
-        template: None,
-        template_file: None,
-        mcp_server_cmd: vec!["mcp-server".to_string()],
-        template_string: None,
-        log_level: None,
-    };
-    let server_info = default_mcp_server_info();
-    let result = detect_render_markers(&options, &server_info);
-    assert!(result.is_err());
-    assert!(result
-        .unwrap_err()
-        .to_string()
-        .contains("no matching start marker"));
-}
-
-#[test]
-fn test_detect_render_markers_conflicting_templates() {
-    let file = NamedTempFile::new().unwrap();
-    let content = r#"
-    mcp-discovery-render template-file=./template.hbs
-    mcp-discovery-template
-    Template content
-    mcp-discovery-template-end
-    mcp-discovery-render-end"#;
-    write(&file, content).unwrap();
-    let options = WriteOptions {
-        filename: file.path().to_path_buf(),
-        template: None,
-        template_file: None,
-        mcp_server_cmd: vec!["mcp-server".to_string()],
-        template_string: None,
-        log_level: None,
-    };
-    let server_info = default_mcp_server_info();
-    let result = detect_render_markers(&options, &server_info);
-    assert!(result.is_err());
-    assert!(result
-        .unwrap_err()
-        .to_string()
-        .contains("both a 'template-file' and an inline template"));
-}
-
-#[test]
-fn test_extract_template_file() {
-    let line = "mcp-discovery-render template-file=./template.hbs";
-    let result = extract_template_file(line);
-    assert_eq!(result, Some("./template.hbs".to_string()));
-
-    let line = "mcp-discovery-render";
-    let result = extract_template_file(line);
-    assert_eq!(result, None);
-}
diff --git a/tests/test_std_output.rs b/tests/test_std_output.rs
deleted file mode 100644
index d64d14e..0000000
--- a/tests/test_std_output.rs
+++ /dev/null
@@ -1,143 +0,0 @@
-use mcp_discovery::{
-    std_output::{
-        print_header, print_json, print_list, print_summary, table_bottom, table_content, table_top,
-    },
-    types::{McpCapabilities, McpServerInfo},
-};
-use std::io::{self};
-use unicode_width::UnicodeWidthStr;
-
-#[test]
-fn test_table_top_bottom() {
-    assert_eq!(table_top(5), "┌─────┐");
-    assert_eq!(table_bottom(5), "└─────┘");
-}
-
-#[test]
-fn test_table_content_centering() {
-    let result = table_content(10, "Hi");
-    assert!(result.starts_with("│"));
-    assert!(result.ends_with("│"));
-    assert_eq!(strip_ansi_escapes::strip_str(&result).width(), 12); // includes borders
-    assert!(result.contains("Hi"));
-}
-
-#[test]
-fn test_print_list() {
-    let items = vec![
-        ("Item1".to_string(), "Value1".to_string()),
-        ("Item2".to_string(), "Value2".to_string()),
-    ];
-
-    let mut buffer = Vec::new();
-    print_list(&mut buffer, items).unwrap();
-
-    let output = String::from_utf8(buffer).unwrap();
-    let output = strip_ansi_escapes::strip_str(output);
-
-    assert!(output.contains("1. Item1: Value1"));
-    assert!(output.contains("2. Item2: Value2"));
-}
-
-#[test]
-fn test_print_header_structure() {
-    let mut buffer = Vec::new();
-    print_header(&mut buffer, "My Header", 20).unwrap();
-
-    // Create the String first
-    let output = String::from_utf8(buffer).unwrap();
-
-    // Now borrow the String with `lines()`
-    let lines: Vec<&str> = output.lines().collect();
-
-    assert_eq!(lines.len(), 3);
-    assert!(lines[0].starts_with("┌") && lines[0].ends_with("┐"));
-    assert!(lines[2].starts_with("└") && lines[2].ends_with("┘"));
-    assert!(lines[1].contains("My Header"));
-}
-
-#[test]
-fn test_print_json_output() {
-    let info = McpServerInfo {
-        name: "Test".to_string(),
-        version: "1.0".to_string(),
-        capabilities: McpCapabilities {
-            tools: true,
-            prompts: false,
-            resources: true,
-            logging: false,
-            experimental: true,
-        },
-        tools: None,
-        prompts: None,
-        resources: None,
-        resource_templates: None,
-    };
-
-    let mut buffer = Vec::new();
-    print_json(&mut buffer, &info).unwrap();
-
-    let output = String::from_utf8(buffer).unwrap();
-    let json_value: serde_json::Value = serde_json::from_str(output.trim()).unwrap();
-    assert_eq!(json_value["name"], "Test");
-    assert_eq!(json_value["version"], "1.0");
-}
-
-#[test]
-fn test_print_summary_layout() {
-    let info = McpServerInfo {
-        name: "MyApp".to_string(),
-        version: "0.9".to_string(),
-        capabilities: McpCapabilities {
-            tools: true,
-            prompts: false,
-            resources: false,
-            logging: true,
-            experimental: false,
-        },
-        tools: None,
-        prompts: None,
-        resources: None,
-        resource_templates: None,
-    };
-
-    let mut buffer = Vec::new();
-    let size = print_summary(&mut buffer, &info).unwrap();
-
-    let output = String::from_utf8(buffer).unwrap();
-    assert!(output.contains("MyApp 0.9"));
-    assert!(size >= 44);
-}
-
-#[test]
-fn test_print_to_stdout() {
-    let items = vec![
-        ("Tool1".to_string(), "Enabled".to_string()),
-        ("Tool2".to_string(), "Disabled".to_string()),
-    ];
-
-    // Print directly to stdout
-    print_list(io::stdout(), items).unwrap();
-}
-
-#[test]
-fn test_print_summary_to_stdout() {
-    let info = McpServerInfo {
-        name: "SampleServer".to_string(),
-        version: "1.1".to_string(),
-        capabilities: McpCapabilities {
-            tools: true,
-            prompts: true,
-            resources: true,
-            logging: false,
-            experimental: false,
-        },
-        tools: None,
-        prompts: None,
-        resources: None,
-        resource_templates: None,
-    };
-
-    // Print directly to stdout
-    print_summary(io::stdout(), &info).unwrap();
-}

From 9db6ca5cf62dce6b95804c46d68d835b51e843f8 Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 18:25:23 -0300
Subject: [PATCH 06/10] chore: add code comments

---
 src/lib.rs                | 36 +++++++++++++++++++-----------------
 src/main.rs               |  2 +-
 src/templates.rs          | 16 ++++++++++++----
 src/types/capabilities.rs |  5 +++++
 src/types/commands.rs     | 10 ++++++++++
 5 files changed, 47 insertions(+), 22 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 1178403..0ca6c87 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -34,8 +34,11 @@ use rust_mcp_sdk::{
     McpClient, StdioTransport, TransportOptions,
 };
 
+/// Core struct representing the discovery mechanism for the MCP server.
 pub struct McpDiscovery {
+    /// Discovery action and its options
     options: DiscoveryCommand,
+    /// Collected server capabilities and metadata
     pub server_info: Option<McpServerInfo>,
 }
 
@@ -47,6 +50,7 @@ impl McpDiscovery {
         }
     }
 
+    /// Entry point to execute the discovery workflow based on the command.
     pub async fn start(&mut self) -> DiscoveryResult<()> {
         // launch mcp server and discover capabilities
 
@@ -60,13 +64,14 @@ impl McpDiscovery {
                 self.update_document(update_options).await?;
             }
             DiscoveryCommand::Print(print_options) => {
-                self.print_mcp_capabilities(print_options).await?;
+                self.print_server_capabilities(print_options).await?;
             }
         };
         Ok(())
     }
 
-    pub async fn print_mcp_capabilities(
+    /// Prints MCP server capabilities using a specific template or default view.
+    pub async fn print_server_capabilities(
         &self,
         print_options: &PrintOptions,
     ) -> DiscoveryResult<()> {
@@ -90,6 +95,7 @@ impl McpDiscovery {
         Ok(())
     }
 
+    /// Creates a new file using a specified template and discovered server info.
     pub async fn create_document(&self, create_options: &WriteOptions) -> DiscoveryResult<()> {
         tracing::trace!("Creating '{}' ", create_options.filename.to_string_lossy());
 
@@ -120,6 +126,7 @@ impl McpDiscovery {
         Ok(())
     }
 
+    /// Updates an existing file by replacing only templated sections.
     pub async fn update_document(&self, update_options: &WriteOptions) -> DiscoveryResult<()> {
         tracing::trace!("Updating '{}' ", update_options.filename.to_string_lossy());
 
@@ -160,7 +167,8 @@ impl McpDiscovery {
         Ok(())
     }
 
-    pub fn print_summary(&self) -> DiscoveryResult<usize> {
+    /// Print a brief summary of the discovered server information.
+    fn print_summary(&self) -> DiscoveryResult<usize> {
         let server_info = self
             .server_info
             .as_ref()
@@ -168,19 +176,8 @@ impl McpDiscovery {
         Ok(print_summary(stdout(), server_info)?)
     }
 
-    pub fn render_with_template(&self, template: OutputTemplate) -> DiscoveryResult<()> {
-        let server_info = self
-            .server_info
-            .as_ref()
-            .ok_or(DiscoveryError::NotDiscovered)?;
-
-        let content = template.render_template(server_info)?;
-
-        println!("{}", content);
-        Ok(())
-    }
-
-    pub fn print_server_details(&self) -> DiscoveryResult<()> {
+    /// Prints summary and then detailed info about tools, prompts, resources, and templates from server.
+    fn print_server_details(&self) -> DiscoveryResult<()> {
         let table_size = self.print_summary()?;
 
         let server_info = self
@@ -296,6 +293,7 @@ impl McpDiscovery {
         Ok(())
     }
 
+    /// Retrieves tools metadata from the MCP server.
     pub async fn tools(
         &self,
         client: Arc<ClientRuntime>,
@@ -342,6 +340,7 @@ impl McpDiscovery {
         Ok(Some(prompts))
     }
 
+    /// Retrieves resources from the server.
     async fn resources(
         &self,
         client: Arc<ClientRuntime>,
@@ -360,6 +359,7 @@ impl McpDiscovery {
         Ok(Some(resources))
     }
 
+    /// Retrieves resource templates from the server.
     async fn resource_templates(
         &self,
         client: Arc<ClientRuntime>,
@@ -382,6 +382,7 @@ impl McpDiscovery {
         }
     }
 
+    /// Discovers all MCP server capabilities and stores them internally.
     pub async fn discover(&mut self) -> DiscoveryResult<&McpServerInfo> {
         let client = self.launch_mcp_server().await?;
 
@@ -435,7 +436,8 @@ impl McpDiscovery {
         Ok(self.server_info.as_ref().unwrap())
     }
 
-    pub async fn launch_mcp_server(&self) -> SdkResult<Arc<ClientRuntime>> {
+    /// Launches the MCP server as a subprocess and initializes the client.
+    async fn launch_mcp_server(&self) -> SdkResult<Arc<ClientRuntime>> {
         let client_details: InitializeRequestParams = InitializeRequestParams {
             capabilities: ClientCapabilities::default(),
             client_info: Implementation {
diff --git a/src/main.rs b/src/main.rs
index 0b3a3dd..57c068e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,7 @@ mod cli;
 use clap::Parser;
 use cli::{CliDiscoveryCommand, CliPrintOptions, CommandArguments};
 use colored::Colorize;
-use mcp_discovery::{DiscoveryCommand, LogLevel, McpDiscovery};
+use mcp_discovery::{DiscoveryCommand, LogLevel, McpDiscovery, OutputTemplate};
 use tracing_subscriber::{self, EnvFilter};
 
 #[tokio::main]
diff --git a/src/templates.rs b/src/templates.rs
index 9f496d6..487ab33 100644
--- a/src/templates.rs
+++ b/src/templates.rs
@@ -74,7 +74,8 @@ pub static PARTIALS: [(&str, &str); 21] = [
     ("txt-summary", TEXT_SUMMARY),
 ];
 
-// Struct to hold information about inline templates
+/// Struct to hold information about inline templates
+/// Used for templates embedded within other content with specific markers
 #[derive(Debug)]
 pub struct InlineTemplateInfo {
     pub template: String,
@@ -82,7 +83,8 @@ pub struct InlineTemplateInfo {
     pub marker_end: String,
 }
 
-// Enum representing different types of output templates
+/// Enum representing different types of output templates
+/// Used to specify the type of template to render
 #[derive(Debug)]
 pub enum OutputTemplate {
     /// Markdown template (generates tables)
@@ -104,12 +106,15 @@ pub enum OutputTemplate {
 }
 
 impl OutputTemplate {
+    /// Creates an OutputTemplate from a file path
+    /// Resolves the template file path relative to a base file if provided
+    /// Returns a DiscoveryResult containing the CustomTemplate variant
     pub fn from_file(template_file: &Path, base_file: Option<&PathBuf>) -> DiscoveryResult<Self> {
         let actual_template_file = find_template_file(template_file, base_file)?;
         Ok(OutputTemplate::CustomTemplate(actual_template_file))
     }
 
-    // Returns the content of the template as a Cow string
+    /// Returns the content of the template as a Cow string
     pub fn content(&self) -> Cow<'_, str> {
         match &self {
             // Return borrowed references to static templates
@@ -133,7 +138,8 @@ impl OutputTemplate {
         }
     }
 
-    // Generates formatted inline template, including the markers and proper line endings
+    /// Generates formatted inline template, including markers and proper line endings
+    /// Used for InlineTemplate variant to format the output with start/end markers
     fn inline_template(&self, inline_template_info: &InlineTemplateInfo) -> String {
         let line_ending = line_ending(&inline_template_info.template, None);
         format!(
@@ -147,6 +153,8 @@ impl OutputTemplate {
         )
     }
 
+    /// Renders the template with provided server information
+    /// Returns the rendered output as a `DiscoveryResult<String>`
     pub fn render_template(&self, server_info: &McpServerInfo) -> DiscoveryResult<String> {
         let rendered = render_template(self, server_info)?;
         match self {
diff --git a/src/types/capabilities.rs b/src/types/capabilities.rs
index 0cbac6c..08c8d56 100644
--- a/src/types/capabilities.rs
+++ b/src/types/capabilities.rs
@@ -1,6 +1,7 @@
 use rust_mcp_schema::{Prompt, Resource, ResourceTemplate};
 use std::fmt::Display;
 
+/// Represents the capabilities of an MCP server, indicating which features are supported.
 #[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
 pub struct McpCapabilities {
     pub tools: bool,
@@ -20,6 +21,7 @@ impl Display for McpCapabilities {
     }
 }
 
+/// Defines the possible parameter types for tools, which can be primitive (e.g., string, int), objects, or arrays.
 #[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
 pub enum ParamTypes {
     Primitive(String),
@@ -69,6 +71,7 @@ impl Display for ParamTypes {
 //     }
 // }
 
+/// Represents a tool parameter, including its name, type, description, and whether it's required.
 #[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
 pub struct McpToolSParams {
     pub param_name: String,
@@ -78,6 +81,7 @@ pub struct McpToolSParams {
     pub required: bool,
 }
 
+/// Represents metadata for an MCP tool, including its name, description, and parameters.
 #[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
 pub struct McpToolMeta {
     pub name: String,
@@ -86,6 +90,7 @@ pub struct McpToolMeta {
     pub params: Vec<McpToolSParams>,
 }
 
+/// Represents the MCP server's information, including its name, version, capabilities, and supported features.
 #[derive(::serde::Deserialize, ::serde::Serialize, Clone, Debug)]
 pub struct McpServerInfo {
     pub name: String,
diff --git a/src/types/commands.rs b/src/types/commands.rs
index 2e62197..c126433 100644
--- a/src/types/commands.rs
+++ b/src/types/commands.rs
@@ -11,6 +11,7 @@ use crate::{
     OutputTemplate,
 };
 
+/// Enum representing the main actions that can be performed for MCP discovery.
 #[derive(Debug)]
 pub enum DiscoveryCommand {
     /// Displays MCP server capability details in the terminal.
@@ -20,6 +21,8 @@ pub enum DiscoveryCommand {
     /// Updates a file by adding MCP server capability information between specified markers.
     Update(WriteOptions),
 }
+
+/// Enum defining the types of built-in templates supported for output formatting.
 #[derive(Debug, Clone, PartialEq)]
 pub enum Template {
     Md,
@@ -41,6 +44,7 @@ impl FromStr for Template {
     }
 }
 
+/// Enum representing supported log levels for controlling output verbosity.
 #[derive(Debug, Clone, PartialEq)]
 #[allow(non_camel_case_types)]
 pub enum LogLevel {
@@ -63,6 +67,7 @@ impl Display for LogLevel {
     }
 }
 
+/// Options used when running the `Print` variant of `DiscoveryCommand`.
 #[derive(Debug)]
 pub struct PrintOptions {
     /// Select an output template from the built-in options.
@@ -82,6 +87,7 @@ pub struct PrintOptions {
 }
 
 impl PrintOptions {
+    /// Resolves the output template (built-in, file, or string) based on user input.
     pub fn match_template(&self) -> DiscoveryResult<OutputTemplate> {
         match_template(
             None,
@@ -92,6 +98,7 @@ impl PrintOptions {
     }
 }
 
+/// Options used when running the `Create` or `Update` variants of `DiscoveryCommand`.
 #[derive(Debug)]
 pub struct WriteOptions {
     pub filename: PathBuf,
@@ -112,6 +119,7 @@ pub struct WriteOptions {
 }
 
 impl WriteOptions {
+    /// Resolves the output template (built-in, file, or string) based on user input.
     pub fn match_template(&self) -> DiscoveryResult<OutputTemplate> {
         match_template(
             Some(&self.filename),
@@ -134,6 +142,7 @@ impl WriteOptions {
 }
 
 impl DiscoveryCommand {
+    /// Retrieves the MCP server launch command for the current variant.
     pub fn mcp_launch_command(&self) -> &Vec<String> {
         match self {
             DiscoveryCommand::Create(create_options) => &create_options.mcp_server_cmd,
@@ -142,6 +151,7 @@ impl DiscoveryCommand {
         }
     }
 
+    /// Retrieves the configured log level for the current variant.
     pub fn log_level(&self) -> &Option<LogLevel> {
         match self {
             DiscoveryCommand::Create(create_options) => &create_options.log_level,

From e66e6c2546439bfc8217cd6a64ea08e4826618e0 Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 18:48:39 -0300
Subject: [PATCH 07/10] chore: clippy

---
 src/cli.rs  | 78 ++++++++++++++++++++++++++---------------------------
 src/main.rs |  2 +-
 2 files changed, 39 insertions(+), 41 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index 614fe13..5529c4e 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -11,13 +11,13 @@ pub enum CliTemplate {
     Txt,
 }
 
-impl Into<Template> for CliTemplate {
-    fn into(self) -> Template {
-        match self {
-            CliTemplate::Md => Template::Md,
-            CliTemplate::MdPlain => Template::MdPlain,
-            CliTemplate::Html => Template::Html,
-            CliTemplate::Txt => Template::Txt,
+impl From<CliTemplate> for Template {
+    fn from(value: CliTemplate) -> Self {
+        match value {
+            CliTemplate::Md => Self::Md,
+            CliTemplate::MdPlain => Self::MdPlain,
+            CliTemplate::Html => Self::Html,
+            CliTemplate::Txt => Self::Txt,
         }
     }
 }
@@ -32,14 +32,14 @@ pub enum CliLogLevel {
     trace,
 }
 
-impl Into<LogLevel> for CliLogLevel {
-    fn into(self) -> LogLevel {
-        match self {
-            CliLogLevel::error => LogLevel::error,
-            CliLogLevel::warn => LogLevel::warn,
-            CliLogLevel::info => LogLevel::info,
-            CliLogLevel::debug => LogLevel::debug,
-            CliLogLevel::trace => LogLevel::trace,
+impl From<CliLogLevel> for LogLevel {
+    fn from(value: CliLogLevel) -> Self {
+        match value {
+            CliLogLevel::error => Self::error,
+            CliLogLevel::warn => Self::warn,
+            CliLogLevel::info => Self::info,
+            CliLogLevel::debug => Self::debug,
+            CliLogLevel::trace => Self::trace,
         }
     }
 }
@@ -89,15 +89,15 @@ pub struct CliWriteOptions {
     pub mcp_server_cmd: Vec<String>,
 }
 
-impl Into<WriteOptions> for CliWriteOptions {
-    fn into(self) -> WriteOptions {
-        WriteOptions {
-            filename: self.filename,
-            template: self.template.map(|t| t.into()),
-            template_file: self.template_file,
-            template_string: self.template_string,
-            log_level: self.log_level.map(|l| l.into()),
-            mcp_server_cmd: self.mcp_server_cmd,
+impl From<CliWriteOptions> for WriteOptions {
+    fn from(value: CliWriteOptions) -> Self {
+        Self {
+            filename: value.filename,
+            template: value.template.map(|t| t.into()),
+            template_file: value.template_file,
+            template_string: value.template_string,
+            log_level: value.log_level.map(|l| l.into()),
+            mcp_server_cmd: value.mcp_server_cmd,
         }
     }
 }
@@ -135,29 +135,27 @@ conflicts_with_all = ["template", "template_string"])]
     pub mcp_server_cmd: Vec<String>,
 }
 
-impl Into<PrintOptions> for CliPrintOptions {
-    fn into(self) -> PrintOptions {
-        PrintOptions {
-            template: self.template.map(|t| t.into()),
-            template_file: self.template_file,
-            template_string: self.template_string,
-            log_level: self.log_level.map(|l| l.into()),
-            mcp_server_cmd: self.mcp_server_cmd,
+impl From<CliPrintOptions> for PrintOptions {
+    fn from(value: CliPrintOptions) -> Self {
+        Self {
+            template: value.template.map(|t| t.into()),
+            template_file: value.template_file,
+            template_string: value.template_string,
+            log_level: value.log_level.map(|l| l.into()),
+            mcp_server_cmd: value.mcp_server_cmd,
         }
     }
 }
 
-impl Into<DiscoveryCommand> for CliDiscoveryCommand {
-    fn into(self) -> DiscoveryCommand {
-        match self {
-            CliDiscoveryCommand::Print(cli_print_options) => {
-                DiscoveryCommand::Print(cli_print_options.into())
-            }
+impl From<CliDiscoveryCommand> for DiscoveryCommand {
+    fn from(value: CliDiscoveryCommand) -> Self {
+        match value {
+            CliDiscoveryCommand::Print(cli_print_options) => Self::Print(cli_print_options.into()),
             CliDiscoveryCommand::Create(cli_write_options) => {
-                DiscoveryCommand::Create(cli_write_options.into())
+                Self::Create(cli_write_options.into())
             }
             CliDiscoveryCommand::Update(cli_write_options) => {
-                DiscoveryCommand::Update(cli_write_options.into())
+                Self::Update(cli_write_options.into())
             }
         }
     }
diff --git a/src/main.rs b/src/main.rs
index 57c068e..0b3a3dd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,7 @@ mod cli;
 use clap::Parser;
 use cli::{CliDiscoveryCommand, CliPrintOptions, CommandArguments};
 use colored::Colorize;
-use mcp_discovery::{DiscoveryCommand, LogLevel, McpDiscovery, OutputTemplate};
+use mcp_discovery::{DiscoveryCommand, LogLevel, McpDiscovery};
 use tracing_subscriber::{self, EnvFilter};
 
 #[tokio::main]

From 954cb010c46f27ad5373b10cbf6ba9378a2f7e8b Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 19:02:32 -0300
Subject: [PATCH 08/10] chore: update cargo

---
 .github/workflows/publish.yml | 32 ++++++++++++++++++++++++++++++++
 Cargo.lock                    | 16 ++++++++--------
 Cargo.toml                    |  2 +-
 3 files changed, 41 insertions(+), 9 deletions(-)
 create mode 100644 .github/workflows/publish.yml

diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
new file mode 100644
index 0000000..4ad9d1b
--- /dev/null
+++ b/.github/workflows/publish.yml
@@ -0,0 +1,32 @@
+name: Publish to Crates.io
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v4
+
+      - name: Cache Rust
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.rustup/toolchains
+            ~/.cargo/registry
+            ~/.cargo/git
+            target
+          key: ${{ runner.os }}-rust-${{ steps.toolchain.outputs.cachekey }}
+          restore-keys: ${{ runner.os }}-rust-
+
+      - name: Install Rust Toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Publish to Crates.io
+        env:
+          CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
+        run: cargo publish --token $CARGO_REGISTRY_TOKEN
diff --git a/Cargo.lock b/Cargo.lock
index 5f05d0b..2be7db5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -777,9 +777,9 @@ checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c"
 
 [[package]]
 name = "rust-mcp-macros"
-version = "0.2.0"
+version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8cd50f8b98c96ea529cf1a00fc4e418ac19e8f42d6760edb08d2440f504ce706"
+checksum = "46f7a75e2d2a87e62815e6c8f2b0657c8b7a89d059263e44a3ef90e1065793f5"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -800,9 +800,9 @@ dependencies = [
 
 [[package]]
 name = "rust-mcp-sdk"
-version = "0.2.2"
+version = "0.2.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "40cafcf7e092641b82937c16ca47dd4ca5723ac95fd61b6427d8bcc484752994"
+checksum = "d0948d498301f9b26e20b4b0a3a82493e18edc98270d05e47cfa0b33cbd1ed1f"
 dependencies = [
  "async-trait",
  "futures",
@@ -838,9 +838,9 @@ checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f"
 
 [[package]]
 name = "rustix"
-version = "1.0.5"
+version = "1.0.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d97817398dd4bb2e6da002002db259209759911da105da92bec29ccb12cf58bf"
+checksum = "c71e83d6afe7ff64890ec6b71d6a69bb8a610ab78ce364b3352876bb4c801266"
 dependencies = [
  "bitflags",
  "errno",
@@ -895,9 +895,9 @@ dependencies = [
 
 [[package]]
 name = "sha2"
-version = "0.10.8"
+version = "0.10.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8"
+checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
 dependencies = [
  "cfg-if",
  "cpufeatures",
diff --git a/Cargo.toml b/Cargo.toml
index 119df52..a4e152f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,7 +5,7 @@ edition = "2021"
 repository = "https://github.com/rust-mcp-stack/mcp-discovery"
 authors = ["Ali Hashemi"]
 description = "A command-line tool written in Rust for discovering and documenting MCP Server capabilities."
-homepage = "https://github.com/rust-mcp-stack/mcp-discovery"
+homepage = "https://rust-mcp-stack.github.io/mcp-discovery"
 
 [package.metadata.wix]
 upgrade-guid = "F8C005B2-0006-40E0-93CF-01A2DA37CAB7"

From 47d6a1621c4ac63f4ec65802be285ddc8660366b Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 19:12:34 -0300
Subject: [PATCH 09/10] chore: update docs

---
 docs/_coverpage.md | 1 +
 docs/quickstart.md | 6 ++++++
 2 files changed, 7 insertions(+)

diff --git a/docs/_coverpage.md b/docs/_coverpage.md
index 6c33074..f3e8829 100644
--- a/docs/_coverpage.md
+++ b/docs/_coverpage.md
@@ -16,6 +16,7 @@
 - ⚙️ CI/CD-friendly
 
 [GitHub](https://github.com/rust-mcp-stack/mcp-discovery)
+[Install](https://rust-mcp-stack.github.io/mcp-discovery#/quickstart)
 [Get Started](#mcp-discovery)
 
 <!-- background color -->
diff --git a/docs/quickstart.md b/docs/quickstart.md
index e97d86b..e6b0882 100644
--- a/docs/quickstart.md
+++ b/docs/quickstart.md
@@ -26,6 +26,12 @@ powershell -ExecutionPolicy Bypass -c "irm https://github.com/rust-mcp-stack/mcp
 brew install rust-mcp-stack/tap/mcp-discovery
 ```
 
+#### **Cargo**
+
+```sh
+cargo install mcp-discovery --locked
+```
+
 #### **Download Binaries**
 
 <table>

From 8ac2b15a0b4140d7e8abb84abef359de3195abe6 Mon Sep 17 00:00:00 2001
From: Ali Hashemi <programmer.ice@gmail.com>
Date: Thu, 1 May 2025 19:14:42 -0300
Subject: [PATCH 10/10] chore: update website link

---
 wix/main.wxs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wix/main.wxs b/wix/main.wxs
index f085a45..33ebe08 100644
--- a/wix/main.wxs
+++ b/wix/main.wxs
@@ -173,7 +173,7 @@
         <!--<Icon Id='ProductICO' SourceFile='wix\Product.ico'/>-->
         <!--<Property Id='ARPPRODUCTICON' Value='ProductICO' />-->
 
-        <Property Id='ARPHELPLINK' Value='https://github.com/rust-mcp-stack/mcp-discovery'/>
+        <Property Id='ARPHELPLINK' Value='https://rust-mcp-stack.github.io/mcp-discovery'/>
         
         <UI>
             <UIRef Id='WixUI_FeatureTree'/>