1use super::traits::{CacheableTool, ModeTool, Tool};
4use crate::config::constants::tools;
5use crate::tools::grep_search::{GrepSearchInput, GrepSearchManager};
6use anyhow::{Result, anyhow};
7use async_trait::async_trait;
8use serde_json::{Value, json};
9use std::path::PathBuf;
10use std::sync::Arc;
11
12#[derive(Clone)]
14pub struct SearchTool {
15 workspace_root: PathBuf,
16 grep_search: Arc<GrepSearchManager>,
17}
18
19impl SearchTool {
20 pub fn new(workspace_root: PathBuf, grep_search: Arc<GrepSearchManager>) -> Self {
21 Self {
22 workspace_root,
23 grep_search,
24 }
25 }
26
27 async fn execute_exact(&self, args: Value) -> Result<Value> {
29 let pattern = args
30 .get("pattern")
31 .and_then(|p| p.as_str())
32 .ok_or_else(|| anyhow!("Error: Missing 'pattern'. Example: grep_search({{\"pattern\": \"TODO|FIXME\", \"path\": \"src\"}})"))?;
33
34 let input = GrepSearchInput {
35 pattern: pattern.to_string(),
36 path: args
37 .get("path")
38 .and_then(|p| p.as_str())
39 .unwrap_or(".")
40 .to_string(),
41 max_results: Some(
42 args.get("max_results")
43 .and_then(|m| m.as_u64())
44 .unwrap_or(100) as usize,
45 ),
46 case_sensitive: Some(
47 args.get("case_sensitive")
48 .and_then(|c| c.as_bool())
49 .unwrap_or(true),
50 ),
51 literal: Some(false),
52 glob_pattern: None,
53 context_lines: Some(0),
54 include_hidden: Some(false),
55 };
56
57 let result = self.grep_search.perform_search(input.clone()).await?;
58
59 let concise = args
61 .get("response_format")
62 .and_then(|v| v.as_str())
63 .map(|s| s.eq_ignore_ascii_case("concise"))
64 .unwrap_or(true);
65
66 let mut body = if concise {
67 let concise_matches = transform_matches_to_concise(&result.matches);
68 json!({
69 "success": true,
70 "matches": concise_matches,
71 "mode": "exact",
72 "response_format": "concise"
73 })
74 } else {
75 json!({
76 "success": true,
77 "matches": result.matches,
78 "mode": "exact",
79 "response_format": "detailed"
80 })
81 };
82
83 if let Some(max) = input.max_results {
84 if let Some(arr) = body.get("matches").and_then(|m| m.as_array())
86 && arr.len() >= max
87 {
88 body["message"] = json!(format!(
89 "Showing {} results (limit). Narrow your query or use more specific patterns to reduce tokens.",
90 max
91 ));
92 }
93 }
94 Ok(body)
95 }
96
97 async fn execute_fuzzy(&self, args: Value) -> Result<Value> {
99 let pattern = args
100 .get("pattern")
101 .and_then(|p| p.as_str())
102 .ok_or_else(|| anyhow!("Error: Missing 'pattern'. Example: grep_search({{\"mode\": \"fuzzy\", \"pattern\": \"todo\", \"path\": \"src\"}})"))?;
103
104 let input = GrepSearchInput {
105 pattern: pattern.to_string(),
106 path: args
107 .get("path")
108 .and_then(|p| p.as_str())
109 .unwrap_or(".")
110 .to_string(),
111 max_results: Some(
112 args.get("max_results")
113 .and_then(|m| m.as_u64())
114 .unwrap_or(100) as usize,
115 ),
116 case_sensitive: Some(
117 args.get("case_sensitive")
118 .and_then(|c| c.as_bool())
119 .unwrap_or(false), ),
121 literal: Some(false),
122 glob_pattern: None,
123 context_lines: Some(0),
124 include_hidden: Some(false),
125 };
126
127 let result = self.grep_search.perform_search(input.clone()).await?;
128
129 let concise = args
131 .get("response_format")
132 .and_then(|v| v.as_str())
133 .map(|s| s.eq_ignore_ascii_case("concise"))
134 .unwrap_or(true);
135
136 let mut body = if concise {
137 let concise_matches = transform_matches_to_concise(&result.matches);
138 json!({
139 "success": true,
140 "matches": concise_matches,
141 "mode": "fuzzy",
142 "case_sensitive": false,
143 "response_format": "concise"
144 })
145 } else {
146 json!({
147 "success": true,
148 "matches": result.matches,
149 "mode": "fuzzy",
150 "case_sensitive": false,
151 "response_format": "detailed"
152 })
153 };
154
155 if let Some(max) = input.max_results {
156 if let Some(arr) = body.get("matches").and_then(|m| m.as_array())
158 && arr.len() >= max
159 {
160 body["message"] = json!(format!(
161 "Showing {} results (limit). Narrow your query or use more specific patterns to reduce tokens.",
162 max
163 ));
164 }
165 }
166 Ok(body)
167 }
168
169 async fn execute_multi(&self, args: Value) -> Result<Value> {
171 let args_obj = args
172 .as_object()
173 .ok_or_else(|| anyhow!("Error: Invalid 'multi' arguments. Required: {{ patterns: string[] }}. Optional: {{ logic: 'AND'|'OR' }}. Example: grep_search({{\"mode\": \"multi\", \"patterns\": [\"fn \\w+\", \"use \\w+\"], \"logic\": \"AND\"}})"))?;
174
175 let patterns = args_obj
176 .get("patterns")
177 .and_then(|p| p.as_array())
178 .ok_or_else(|| anyhow!("Missing patterns array for multi mode"))?;
179
180 let logic = args_obj
181 .get("logic")
182 .and_then(|l| l.as_str())
183 .unwrap_or("AND");
184
185 let mut all_results = Vec::new();
186
187 for pattern in patterns {
189 if let Some(pattern_str) = pattern.as_str() {
190 let mut pattern_args = args.clone();
191 if let Some(obj) = pattern_args.as_object_mut() {
192 obj.insert("pattern".to_string(), json!(pattern_str));
193 }
194
195 match self.execute_exact(pattern_args).await {
196 Ok(result) => {
197 if let Some(matches) = result.get("matches").and_then(|m| m.as_array()) {
198 all_results.extend(matches.clone());
199 }
200 }
201 Err(_) => continue, }
203 }
204 }
205
206 let final_results = if logic == "AND" {
208 self.apply_and_logic(all_results, patterns.len())
209 } else {
210 self.apply_or_logic(all_results)
211 };
212
213 Ok(json!({
214 "success": true,
215 "matches": final_results,
216 "mode": "multi",
217 "logic": logic,
218 "pattern_count": patterns.len()
219 }))
220 }
221
222 async fn execute_similarity(&self, args: Value) -> Result<Value> {
224 let args_obj = args
225 .as_object()
226 .ok_or_else(|| anyhow!("Error: Invalid 'similarity' arguments. Required: {{ reference_file: string }}. Optional: {{ content_type: 'structure'|'imports'|'functions'|'all' }}. Example: grep_search({{\"mode\": \"similarity\", \"reference_file\": \"src/lib.rs\", \"content_type\": \"functions\"}})"))?;
227
228 let reference_file = args_obj
229 .get("reference_file")
230 .and_then(|f| f.as_str())
231 .ok_or_else(|| anyhow!("Error: Missing 'reference_file'. Example: grep_search({{\"mode\": \"similarity\", \"reference_file\": \"src/main.rs\"}})"))?;
232
233 let content_type = args_obj
234 .get("content_type")
235 .and_then(|c| c.as_str())
236 .unwrap_or("all");
237
238 let ref_path = self.workspace_root.join(reference_file);
240 let ref_content = tokio::fs::read_to_string(&ref_path).await.map_err(|e| {
241 anyhow!(
242 "Error: Failed to read reference file '{}': {}",
243 reference_file,
244 e
245 )
246 })?;
247
248 let patterns = self.extract_similarity_patterns(&ref_content, content_type)?;
250
251 let mut search_args = args.clone();
253 if let Some(obj) = search_args.as_object_mut() {
254 obj.insert("patterns".to_string(), json!(patterns));
255 obj.insert("logic".to_string(), json!("OR"));
256 }
257
258 self.execute_multi(search_args).await
259 }
260
261 fn apply_and_logic(&self, results: Vec<Value>, pattern_count: usize) -> Vec<Value> {
263 use std::collections::HashMap;
264
265 let mut file_matches: HashMap<String, Vec<Value>> = HashMap::new();
266
267 for result in results {
269 if let Some(path) = result.get("path").and_then(|p| p.as_str()) {
270 file_matches
271 .entry(path.to_string())
272 .or_default()
273 .push(result);
274 }
275 }
276
277 file_matches
279 .into_iter()
280 .filter(|(_, matches)| matches.len() >= pattern_count)
281 .flat_map(|(_, matches)| matches)
282 .collect()
283 }
284
285 fn apply_or_logic(&self, results: Vec<Value>) -> Vec<Value> {
287 use std::collections::HashSet;
288
289 let mut seen = HashSet::new();
290 let mut unique_results = Vec::new();
291
292 for result in results {
293 let key = format!(
294 "{}:{}:{}",
295 result.get("path").and_then(|p| p.as_str()).unwrap_or(""),
296 result
297 .get("line_number")
298 .and_then(|l| l.as_u64())
299 .unwrap_or(0),
300 result.get("column").and_then(|c| c.as_u64()).unwrap_or(0)
301 );
302
303 if seen.insert(key) {
304 unique_results.push(result);
305 }
306 }
307
308 unique_results
309 }
310
311 fn extract_similarity_patterns(
313 &self,
314 content: &str,
315 content_type: &str,
316 ) -> Result<Vec<String>> {
317 let mut patterns = Vec::new();
318
319 match content_type {
320 "functions" => {
321 for line in content.lines() {
323 if (line.trim_start().starts_with("fn ")
324 || line.trim_start().starts_with("pub fn "))
325 && let Some(name) = self.extract_function_name(line)
326 {
327 patterns.push(format!("fn {}", name));
328 }
329 }
330 }
331 "imports" => {
332 for line in content.lines() {
334 if line.trim_start().starts_with("use ") {
335 patterns.push(line.trim().to_string());
336 }
337 }
338 }
339 "structure" => {
340 for line in content.lines() {
342 let trimmed = line.trim_start();
343 if trimmed.starts_with("struct ") || trimmed.starts_with("enum ") {
344 patterns.push(
345 trimmed
346 .split_whitespace()
347 .take(2)
348 .collect::<Vec<_>>()
349 .join(" "),
350 );
351 }
352 }
353 }
354 _ => {
355 patterns.extend(self.extract_keywords(content));
357 }
358 }
359
360 if patterns.is_empty() {
361 return Err(anyhow!(
362 "No patterns extracted from reference file. Try content_type='all' or provide a different reference_file."
363 ));
364 }
365
366 Ok(patterns)
367 }
368
369 fn extract_function_name(&self, line: &str) -> Option<String> {
371 let parts: Vec<&str> = line.split_whitespace().collect();
372 for (i, part) in parts.iter().enumerate() {
373 if *part == "fn" && i + 1 < parts.len() {
374 let name = parts[i + 1];
375 if let Some(paren_pos) = name.find('(') {
376 return Some(name[..paren_pos].to_string());
377 }
378 return Some(name.to_string());
379 }
380 }
381 None
382 }
383
384 fn extract_keywords(&self, content: &str) -> Vec<String> {
386 let keywords = ["fn ", "struct ", "enum ", "impl ", "trait ", "use ", "mod "];
387 let mut patterns = Vec::new();
388
389 for line in content.lines() {
390 for keyword in &keywords {
391 if line.contains(keyword) {
392 patterns.push(keyword.trim().to_string());
393 }
394 }
395 }
396
397 patterns.sort();
398 patterns.dedup();
399 patterns
400 }
401}
402
403#[async_trait]
404impl Tool for SearchTool {
405 async fn execute(&self, args: Value) -> Result<Value> {
406 let args_clone = args.clone();
407 let mode = args_clone
408 .get("mode")
409 .and_then(|m| m.as_str())
410 .unwrap_or("exact");
411
412 self.execute_mode(mode, args).await
413 }
414
415 fn name(&self) -> &'static str {
416 tools::GREP_SEARCH
417 }
418
419 fn description(&self) -> &'static str {
420 "Enhanced unified search tool with multiple modes: exact (default), fuzzy, multi-pattern, and similarity search"
421 }
422}
423
424#[async_trait]
425impl ModeTool for SearchTool {
426 fn supported_modes(&self) -> Vec<&'static str> {
427 vec!["exact", "fuzzy", "multi", "similarity"]
428 }
429
430 async fn execute_mode(&self, mode: &str, args: Value) -> Result<Value> {
431 match mode {
432 "exact" => self.execute_exact(args).await,
433 "fuzzy" => self.execute_fuzzy(args).await,
434 "multi" => self.execute_multi(args).await,
435 "similarity" => self.execute_similarity(args).await,
436 _ => Err(anyhow!("Unsupported search mode: {}", mode)),
437 }
438 }
439}
440
441#[async_trait]
442impl CacheableTool for SearchTool {
443 fn cache_key(&self, args: &Value) -> String {
444 format!(
445 "search:{}:{}",
446 args.get("pattern").and_then(|p| p.as_str()).unwrap_or(""),
447 args.get("mode").and_then(|m| m.as_str()).unwrap_or("exact")
448 )
449 }
450
451 fn should_cache(&self, args: &Value) -> bool {
452 let mode = args.get("mode").and_then(|m| m.as_str()).unwrap_or("exact");
454 matches!(mode, "exact" | "fuzzy")
455 }
456}
457
458pub(crate) fn transform_matches_to_concise(events: &[Value]) -> Vec<Value> {
461 let mut out = Vec::new();
462 for ev in events {
463 if ev.get("type").and_then(|t| t.as_str()) != Some("match") {
464 continue;
465 }
466 if let Some(data) = ev.get("data") {
467 let path = data
468 .get("path")
469 .and_then(|p| p.get("text"))
470 .and_then(|t| t.as_str())
471 .unwrap_or("");
472 let line = data
473 .get("line_number")
474 .and_then(|n| n.as_u64())
475 .unwrap_or(0);
476 let preview = data
477 .get("lines")
478 .and_then(|l| l.get("text"))
479 .and_then(|t| t.as_str())
480 .unwrap_or("")
481 .trim_end_matches(['\r', '\n']);
482
483 out.push(json!({
484 "path": path,
485 "line_number": line,
486 "text": preview,
487 }));
488 }
489 }
490 out
491}
492
493#[cfg(test)]
494mod tests {
495 use super::*;
496
497 #[test]
498 fn test_transform_matches_to_concise() {
499 let raw = vec![
500 json!({
501 "type": "match",
502 "data": {
503 "path": {"text": "src/main.rs"},
504 "line_number": 10,
505 "lines": {"text": "fn main() {}\n"}
506 }
507 }),
508 json!({"type": "begin"}),
509 ];
510 let concise = transform_matches_to_concise(&raw);
511 assert_eq!(concise.len(), 1);
512 assert_eq!(concise[0]["path"], "src/main.rs");
513 assert_eq!(concise[0]["line_number"], 10);
514 assert_eq!(concise[0]["text"], "fn main() {}");
515 }
516}