Enhance OmniSearch for LLMs with limits and include_body flags

server/src/handlers/meta.rs

@@ -339,6 +339,8 @@ impl McpTool for OmniSearchHandler {
 
     async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
         let req: OmniSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let limit = req.limit.unwrap_or(5);
+        let include_body = req.include_body.unwrap_or(false);
         let matches = {
             let idx = state.get_search_index();
             idx.search(&req.query, req.namespace.as_deref())
@@ -348,11 +350,20 @@ impl McpTool for OmniSearchHandler {
 
         let kg_json = state.read_graph(|full| {
             let mut kg_entities = std::collections::HashMap::new();
+            let mut count = 0;
             for (id, doc_type, _, _, _) in &matches {
                 if doc_type == "entity"
                     && let Some(e) = full.entities.get(id)
                 {
-                    kg_entities.insert(id, e);
+                    if count >= limit { continue; }
+                    count += 1;
+                    if !include_body {
+                        let mut summary = e.clone();
+                        summary.observations = vec![];
+                        kg_entities.insert(id.clone(), summary);
+                    } else {
+                        kg_entities.insert(id.clone(), e.clone());
+                    }
                 }
             }
             serde_json::to_value(&kg_entities).unwrap_or_default()
@@ -380,6 +391,17 @@ impl McpTool for OmniSearchHandler {
             let filtered: Vec<_> = all_tasks
                 .iter()
                 .filter(|t| matched_tasks.contains(t.id.as_str()))
+                .take(limit)
+                .map(|t| {
+                    if !include_body {
+                        let mut summary = t.clone();
+                        summary.description = "".to_string();
+                        summary.acceptance_criteria = vec![];
+                        summary
+                    } else {
+                        t.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });
@@ -388,6 +410,16 @@ impl McpTool for OmniSearchHandler {
             let filtered: Vec<_> = all_snippets
                 .iter()
                 .filter(|s| matched_snippets.contains(s.name.as_str()))
+                .take(limit)
+                .map(|s| {
+                    if !include_body {
+                        let mut summary = s.clone();
+                        summary.code = "".to_string();
+                        summary
+                    } else {
+                        s.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });
@@ -396,6 +428,18 @@ impl McpTool for OmniSearchHandler {
             let filtered: Vec<_> = all_adrs
                 .iter()
                 .filter(|a| matched_adrs.contains(a.id.as_str()))
+                .take(limit)
+                .map(|a| {
+                    if !include_body {
+                        let mut summary = a.clone();
+                        summary.context = "".to_string();
+                        summary.decision = "".to_string();
+                        summary.consequence = "".to_string();
+                        summary
+                    } else {
+                        a.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });
@@ -409,6 +453,17 @@ impl McpTool for OmniSearchHandler {
                         && (contains_ignore_ascii_case(&d.description, &q)
                             || contains_ignore_ascii_case(&d.ideal_solution, &q))
                 })
+                .take(limit)
+                .map(|d| {
+                    if !include_body {
+                        let mut summary = d.clone();
+                        summary.description = "".to_string();
+                        summary.ideal_solution = "".to_string();
+                        summary
+                    } else {
+                        d.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });
@@ -420,6 +475,16 @@ impl McpTool for OmniSearchHandler {
                     req.namespace.as_ref().is_none_or(|ns| m.namespace == *ns)
                         && contains_ignore_ascii_case(&m.content, &q)
                 })
+                .take(limit)
+                .map(|m| {
+                    if !include_body {
+                        let mut summary = m.clone();
+                        summary.content = "".to_string();
+                        summary
+                    } else {
+                        m.clone()
+                    }
+                })
                 .collect();
             serde_json::to_value(&filtered).unwrap_or_default()
         });

server/src/tools.rs

@@ -69,6 +69,10 @@ pub struct SearchNodesTool {
     pub query: String,
     /// Optional namespace to restrict the search to.
     pub namespace: Option<String>,
+    /// Limit the number of results to avoid context limit overflow. Defaults to 10.
+    pub limit: Option<usize>,
+    /// Include the full observations of the matched items. If false, returns only IDs and types (recommended for LLMs to prevent context bloat). Defaults to false.
+    pub include_body: Option<bool>,
 }
 
 /// Open and retrieve full details of specific nodes in the knowledge graph.
@@ -192,6 +196,10 @@ pub struct StoreSnippetTool {
 pub struct SearchSnippetsTool {
     /// The search query.
     pub query: String,
+    /// Limit the number of results to avoid context limit overflow. Defaults to 5.
+    pub limit: Option<usize>,
+    /// Include the full code content of the matched items. If false, returns only IDs and descriptions (recommended for LLMs to prevent context bloat). Defaults to false.
+    pub include_body: Option<bool>,
 }
 
 /// Delete a stored code snippet.
@@ -219,6 +227,10 @@ pub struct LogDecisionTool {
 pub struct QueryDecisionsTool {
     /// Optional search query.
     pub query: Option<String>,
+    /// Limit the number of results to avoid context limit overflow. Defaults to 5.
+    pub limit: Option<usize>,
+    /// Include the full body/content of the matched ADRs. If false, returns only titles and IDs (recommended for LLMs to prevent context bloat). Defaults to false.
+    pub include_body: Option<bool>,
 }
 
 /// Merge two entities in the knowledge graph into one.
@@ -265,6 +277,10 @@ pub struct LogErrorFixTool {
 pub struct SearchErrorFixesTool {
     /// The search query.
     pub query: String,
+    /// Limit the number of results to avoid context limit overflow. Defaults to 5.
+    pub limit: Option<usize>,
+    /// Include the full solution details. If false, returns only error signatures (recommended for LLMs to prevent context bloat). Defaults to false.
+    pub include_body: Option<bool>,
 }
 
 /// Pin a file to keep it explicitly in the context workspace.
@@ -523,12 +539,17 @@ pub struct ListContextWorkspacesTool {
 }
 
 /// Search across all memory stores (Graph, Tasks, Snippets, ADRs, etc.).
+/// Highly recommended for discovery. Supports Lucene query syntax.
 #[derive(Debug, Deserialize, Serialize, JsonSchema)]
 pub struct OmniSearchTool {
-    /// The search query.
+    /// The search query. Supports Lucene syntax (e.g., 'title:"auth" AND status:open').
     pub query: String,
     /// Optional namespace to restrict the search to.
     pub namespace: Option<String>,
+    /// Limit the number of results per category to avoid context limit overflow. Defaults to 5.
+    pub limit: Option<usize>,
+    /// Include the full body/content of the matched items. If false, returns only IDs and titles (recommended for LLMs to prevent context bloat). Defaults to false.
+    pub include_body: Option<bool>,
 }
 
 /// Get a health digest of the project.