added --max_tokens ### to stdio mode

Author: elasticdotventures

src/bin/cratedocs.rs

@@ -32,6 +32,9 @@ enum Commands {
         /// Summarize output by stripping LICENSE and VERSION sections (TL;DR mode)
         #[arg(long)]
         tldr: bool,
+        /// Maximum number of tokens for output (token-aware truncation)
+        #[arg(long)]
+        max_tokens: Option<usize>,
     },
     /// Run the server with HTTP/SSE interface
     Http {
@@ -112,7 +115,7 @@ async fn main() -> Result<()> {
             println!("{}", env!("CARGO_PKG_VERSION"));
             Ok(())
         },
-        Commands::Stdio { debug, tldr } => run_stdio_server(debug, tldr).await,
+        Commands::Stdio { debug, tldr, max_tokens } => run_stdio_server(debug, tldr, max_tokens).await,
         Commands::Http { address, debug } => run_http_server(address, debug).await,
         Commands::Test {
             tool,
@@ -148,7 +151,7 @@ async fn main() -> Result<()> {
     }
 }
 
-async fn run_stdio_server(debug: bool, tldr: bool) -> Result<()> {
+async fn run_stdio_server(debug: bool, tldr: bool, max_tokens: Option<usize>) -> Result<()> {
     // Set up file appender for logging
     let file_appender = RollingFileAppender::new(Rotation::DAILY, "logs", "stdio-server.log");
 
@@ -168,7 +171,7 @@ async fn run_stdio_server(debug: bool, tldr: bool) -> Result<()> {
 
     // Create an instance of our documentation router
     // If tldr is needed globally, you may want to pass it to DocRouter or handle it in tool output
-    let router = RouterService(DocRouter::new_with_tldr(tldr));
+    let router = RouterService(DocRouter::new_with_tldr_and_max_tokens(tldr, max_tokens));
 
     // Create and run the server
     let server = Server::new(router);

src/tools/docs/docs.rs

@@ -49,6 +49,7 @@ pub struct DocRouter {
     pub client: Client,
     pub cache: DocCache,
     pub tldr: bool,
+    pub max_tokens: Option<usize>,
 }
 
 impl Default for DocRouter {
@@ -58,15 +59,19 @@ impl Default for DocRouter {
 }
 
 impl DocRouter {
-    pub fn new_with_tldr(tldr: bool) -> Self {
+    pub fn new_with_tldr_and_max_tokens(tldr: bool, max_tokens: Option<usize>) -> Self {
         Self {
             client: Client::new(),
             cache: DocCache::new(),
             tldr,
+            max_tokens,
         }
     }
+    pub fn new_with_tldr(tldr: bool) -> Self {
+        Self::new_with_tldr_and_max_tokens(tldr, None)
+    }
     pub fn new() -> Self {
-        Self::new_with_tldr(false)
+        Self::new_with_tldr_and_max_tokens(false, None)
     }
 
     // Fetch crate documentation from docs.rs
@@ -370,6 +375,7 @@ impl mcp_server::Router for DocRouter {
         let tool_name = tool_name.to_string();
         let arguments = arguments.clone();
         let tldr = self.tldr;
+        let max_tokens = self.max_tokens;
 
         Box::pin(async move {
             let mut result = match tool_name.as_str() {
@@ -473,6 +479,27 @@ impl mcp_server::Router for DocRouter {
                 }
             }
 
+            // Apply max_tokens truncation if enabled
+            if let Some(max_tokens) = max_tokens {
+                for content in &mut result {
+                    if let Content::Text(text) = content {
+                        if let Ok(token_count) = crate::tools::count_tokens(&text.text) {
+                            if token_count > max_tokens {
+                                let mut truncated = text.text.clone();
+                                while crate::tools::count_tokens(&truncated).map_or(0, |c| c) > max_tokens && !truncated.is_empty() {
+                                    truncated.pop();
+                                }
+                                if let Some(last_space) = truncated.rfind(' ') {
+                                    truncated.truncate(last_space);
+                                }
+                                truncated.push_str(" 内容被截断");
+                                text.text = truncated;
+                            }
+                        }
+                    }
+                }
+            }
+
             Ok(result)
         })
     }