feat(mcp): limit max tokens for docs search and logs (#2114)

* Upgrade rmcp to v0.6.1 and refactor MCP server with new macros

Co-authored-by: johan <[email protected]>

* Refactor mcp.00 code style: improve formatting and line breaks

Co-authored-by: johan <[email protected]>

* Specify max lines for mcp logs

* clippy fix

* logs limit fix

* max tokens for docs search

---------

Co-authored-by: Cursor Agent <[email protected]>
Co-authored-by: johan <[email protected]>
Co-authored-by: jonaro00 <[email protected]>

Author: 4 people

Cargo.toml

@@ -82,9 +82,7 @@ syn = "2"
 tempfile = "3"
 thiserror = "2"
 tokio = "1.40.0"
-tokio-tungstenite = { version = "0.27", features = [
-  "rustls-tls-webpki-roots",
-] }
+tokio-tungstenite = { version = "0.27", features = ["rustls-tls-webpki-roots"] }
 toml = "0.9"
 toml_edit = "0.23"
 tracing = { version = "0.1.37", default-features = false }

mcp/src/mcp.rs

@@ -8,12 +8,6 @@ use tracing::instrument;
 use crate::tools::{deployment::*, docs::*, logs::*, project::*};
 use crate::utils::run_tool;
 
-#[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
-struct SearchDocsArgs {
-    #[schemars(description = "Search query for documentation")]
-    query: String,
-}
-
 #[derive(Clone)]
 pub struct ShuttleMcpServer {
     pub(crate) tool_router: ToolRouter<Self>,
@@ -80,7 +74,7 @@ impl ShuttleMcpServer {
         &self,
         Parameters(args): Parameters<SearchDocsArgs>,
     ) -> Result<String, String> {
-        run_tool(search_docs(args.query)).await
+        run_tool(search_docs(args)).await
     }
 }
 

mcp/src/tools/docs.rs

@@ -1,9 +1,18 @@
 use crate::{constants::SHUTTLE_DOCS_SEARCH_BASE_URL, utils::build_client};
 
-pub async fn search_docs(query: String) -> Result<String, String> {
+#[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
+pub struct SearchDocsArgs {
+    #[schemars(description = "Search query for documentation")]
+    query: String,
+    #[schemars(description = "Maximum number of tokens to retrieve (default: 4000)")]
+    max_tokens: Option<u32>,
+}
+
+pub async fn search_docs(params: SearchDocsArgs) -> Result<String, String> {
     let url = format!(
-        "{SHUTTLE_DOCS_SEARCH_BASE_URL}/search?q={}",
-        urlencoding::encode(&query)
+        "{SHUTTLE_DOCS_SEARCH_BASE_URL}/search?q={}&maxTokens={}",
+        urlencoding::encode(&params.query),
+        params.max_tokens.unwrap_or(4000)
     );
 
     let client = build_client()?;
@@ -12,13 +21,13 @@ pub async fn search_docs(query: String) -> Result<String, String> {
         .get(&url)
         .send()
         .await
-        .map_err(|e| format!("Request failed: {}", e))?;
+        .map_err(|e| format!("Request failed: {e}"))?;
 
     if response.status().is_success() {
         response
             .text()
             .await
-            .map_err(|e| format!("Failed to read response: {}", e))
+            .map_err(|e| format!("Failed to read response: {e}"))
     } else {
         Err(format!("Request failed with status: {}", response.status()))
     }

mcp/src/tools/logs.rs

@@ -1,5 +1,16 @@
 use crate::utils::execute_command;
 
+fn limit_to_last_n_lines(text: &str, max_lines: u32) -> String {
+    let lines: Vec<&str> = text.lines().collect();
+
+    if lines.len() <= max_lines as usize {
+        return text.to_string();
+    }
+
+    let start_index = lines.len() - max_lines as usize;
+    lines[start_index..].join("\n")
+}
+
 #[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
 pub struct LogsArgs {
     #[schemars(description = "Specify the working directory")]
@@ -12,6 +23,8 @@ pub struct LogsArgs {
     name: Option<String>,
     #[schemars(description = "Specify the id of the project")]
     project_id: Option<String>,
+    #[schemars(description = "Maximum number of lines to return")]
+    lines: Option<u32>,
 }
 
 pub async fn logs(params: LogsArgs) -> Result<String, String> {
@@ -35,5 +48,9 @@ pub async fn logs(params: LogsArgs) -> Result<String, String> {
         args.push(id);
     }
 
-    execute_command("shuttle", args, &params.cwd).await
+    let output = execute_command("shuttle", args, &params.cwd).await?;
+
+    // Limit the output to the last N lines (default 50)
+    let max_lines = params.lines.unwrap_or(50);
+    Ok(limit_to_last_n_lines(&output, max_lines))
 }

mcp/src/utils.rs

@@ -46,5 +46,5 @@ pub fn build_client() -> Result<reqwest::Client, String> {
     reqwest::Client::builder()
         .default_headers(headers)
         .build()
-        .map_err(|e| format!("Failed to build client: {}", e))
+        .map_err(|e| format!("Failed to build client: {e}"))
 }