1use super::registry::Tool;
20use super::schema::{ParamType, ToolParameter, ToolSchema};
21use crate::config::SearchConfig;
22use crate::error::{AacodeError, Result};
23use serde_json::{json, Value};
24use std::path::PathBuf;
25use std::sync::atomic::{AtomicBool, Ordering};
26use std::sync::{LazyLock, Mutex};
27use std::time::Duration;
28
29const CONNECT_TIMEOUT_SECS: u64 = 2;
32const PER_ENGINE_CAP_SECS: f64 = 4.0;
34const CIRCUIT_OPEN_SECS: f64 = 120.0;
36const FAILURES_TO_OPEN: u32 = 1;
38const SCRAPE_MIN_SECS: f64 = 3.0;
40
41static HTTP_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
42 reqwest::Client::builder()
43 .connect_timeout(Duration::from_secs(CONNECT_TIMEOUT_SECS))
44 .build()
45 .expect("reqwest client")
46});
47
48#[derive(Default, Clone)]
52struct EngineHealth {
53 failures: u32,
55 open_until: f64,
57 last_call: f64,
59}
60
61static ENGINE_HEALTH: std::sync::OnceLock<Mutex<std::collections::HashMap<String, EngineHealth>>> =
65 std::sync::OnceLock::new();
66
67fn engine_health() -> &'static Mutex<std::collections::HashMap<String, EngineHealth>> {
68 ENGINE_HEALTH.get_or_init(|| Mutex::new(std::collections::HashMap::new()))
69}
70
71pub struct SearchWebTool {
72 pub cfg: SearchConfig,
73 pub timeout_secs: u64,
74}
75
76enum Attempt {
78 Ok(Vec<Value>),
79 Empty,
80 TransportError(String),
81 Skipped(&'static str),
82}
83
84impl SearchWebTool {
85 pub fn new(cfg: SearchConfig, timeout_secs: u64) -> Self {
86 SearchWebTool { cfg, timeout_secs }
87 }
88
89 fn choose_best_engine(&self) -> &'static str {
91 let det = detect_engine_type(self.cfg.searxng_url.as_deref().unwrap_or(""));
92 match det {
93 "brave" if self.cfg.brave_api_key.is_some() => "brave",
94 "google_cse" if self.cfg.google_cse_key.is_some() => "google_cse",
95 "bing" if self.cfg.bing_api_key.is_some() => "bing",
96 "serpapi" if self.cfg.serpapi_key.is_some() => "serpapi",
97 _ => {
98 if self.cfg.searxng_url.is_some() {
99 "searxng"
100 } else {
101 ""
102 }
103 }
104 }
105 }
106
107 fn circuit_open(&self, engine: &str) -> bool {
109 let guard = engine_health().lock().unwrap_or_else(|e| e.into_inner());
110 guard
111 .get(engine)
112 .map(|h| current_time() < h.open_until)
113 .unwrap_or(false)
114 }
115
116 fn record_outcome(&self, engine: &str, ok: bool) {
118 let mut guard = engine_health().lock().unwrap_or_else(|e| e.into_inner());
119 let h = guard.entry(engine.to_string()).or_default();
120 if ok {
121 h.failures = 0;
122 h.open_until = 0.0;
123 } else {
124 h.failures += 1;
125 if h.failures >= FAILURES_TO_OPEN {
126 h.open_until = current_time() + CIRCUIT_OPEN_SECS;
127 }
128 }
129 }
130
131 async fn enforce_rate_limit(&self, engine: &str, remaining: f64) {
134 let do_wait = {
135 let mut guard = engine_health().lock().unwrap_or_else(|e| e.into_inner());
136 let rate = match engine {
137 "searxng" => 0.5,
138 _ => 1.0,
139 };
140 let now = current_time();
141 let h = guard.entry(engine.to_string()).or_default();
142 if h.last_call > 0.0 {
143 let wait = rate - (now - h.last_call);
144 if wait > 0.0 {
145 let wait = wait.min(remaining.max(0.0));
146 Some(wait)
147 } else {
148 h.last_call = current_time();
149 None
150 }
151 } else {
152 h.last_call = current_time();
153 None
154 }
155 };
156 if let Some(wait) = do_wait {
157 tokio::time::sleep(Duration::from_secs_f64(wait)).await;
158 let mut guard = engine_health().lock().unwrap_or_else(|e| e.into_inner());
159 guard.entry(engine.to_string()).or_default().last_call = current_time();
160 }
161 }
162
163 async fn search_with_fallback(
166 &self,
167 query: &str,
168 max_results: usize,
169 timeout: u64,
170 cancel: &AtomicBool,
171 ) -> (bool, String, Vec<Value>, Vec<String>) {
172 let deadline = current_time() + timeout as f64;
173 let mut tried: Vec<String> = Vec::new();
174
175 let mut engine = self.choose_best_engine();
176 if engine.is_empty() {
177 engine = "searxng"; }
179
180 let mut order: Vec<&str> = vec![engine];
182 for fb in ["searxng", "brave", "google_cse", "bing", "serpapi"] {
183 if fb != engine {
184 order.push(fb);
185 }
186 }
187
188 for eng in order {
189 if cancel.load(Ordering::Relaxed) {
190 tried.push(format!("{eng}:cancelled"));
191 return (false, eng.to_string(), vec![], tried);
192 }
193 let remaining = deadline - current_time();
194 if remaining <= 0.5 {
195 tried.push(format!("{eng}:skipped(budget)"));
196 continue;
197 }
198 match self.try_engine(eng, query, max_results, remaining).await {
199 Attempt::Ok(results) => {
200 tried.push(format!("{eng}:ok"));
201 self.record_outcome(eng, true);
202 return (true, eng.to_string(), results, tried);
203 }
204 Attempt::Empty => {
205 tried.push(format!("{eng}:empty"));
206 self.record_outcome(eng, true); }
208 Attempt::TransportError(e) => {
209 tried.push(format!("{eng}:error({})", brief_err(&e)));
210 self.record_outcome(eng, false);
211 }
212 Attempt::Skipped(why) => {
213 tried.push(format!("{eng}:skipped({why})"));
214 }
215 }
216 }
217
218 if !cancel.load(Ordering::Relaxed) {
221 let remaining = (deadline - current_time()).max(SCRAPE_MIN_SECS);
222 if let Some((name, results)) = fallback_scrape(query, max_results, remaining).await {
223 if !results.is_empty() {
224 tried.push(format!("{name}:ok"));
225 return (true, name, results, tried);
226 }
227 }
228 tried.push("fallback_scrape:empty".to_string());
229 }
230
231 (false, engine.to_string(), vec![], tried)
232 }
233
234 async fn try_engine(
235 &self,
236 engine: &str,
237 query: &str,
238 max_results: usize,
239 remaining: f64,
240 ) -> Attempt {
241 let configured = match engine {
243 "searxng" => self.cfg.searxng_url.is_some(),
244 "brave" => self.cfg.brave_api_key.is_some(),
245 "google_cse" => self.cfg.google_cse_key.is_some() && self.cfg.google_cse_cx.is_some(),
246 "bing" => self.cfg.bing_api_key.is_some(),
247 "serpapi" => self.cfg.serpapi_key.is_some(),
248 _ => false,
249 };
250 if !configured {
251 return Attempt::Skipped("not-configured");
252 }
253 if self.circuit_open(engine) {
254 return Attempt::Skipped("circuit-open");
255 }
256 self.enforce_rate_limit(engine, remaining).await;
257 let budget = remaining.min(PER_ENGINE_CAP_SECS);
258 let r = match engine {
259 "searxng" => searxng_search(
260 self.cfg.searxng_url.as_deref().unwrap_or("http://localhost:8080"),
261 query,
262 max_results,
263 budget,
264 ).await,
265 "brave" => brave_search(
266 self.cfg.brave_api_key.as_deref().unwrap_or(""),
267 query,
268 max_results,
269 budget,
270 ).await,
271 "google_cse" => google_cse_search(
272 self.cfg.google_cse_key.as_deref().unwrap_or(""),
273 self.cfg.google_cse_cx.as_deref().unwrap_or(""),
274 query,
275 max_results,
276 budget,
277 ).await,
278 "bing" => bing_search(
279 self.cfg.bing_api_key.as_deref().unwrap_or(""),
280 query,
281 max_results,
282 budget,
283 ).await,
284 "serpapi" => serpapi_search(
285 self.cfg.serpapi_key.as_deref().unwrap_or(""),
286 query,
287 max_results,
288 budget,
289 ).await,
290 _ => return Attempt::Skipped("unknown-engine"),
291 };
292 match r {
293 Ok(v) if !v.is_empty() => Attempt::Ok(v),
294 Ok(_) => Attempt::Empty,
295 Err(e) => Attempt::TransportError(e.to_string()),
296 }
297 }
298}
299
300fn brief_err(e: &str) -> String {
302 let first = e.lines().next().unwrap_or(e);
303 let mut s: String = first.chars().take(80).collect();
304 if first.chars().count() > 80 {
305 s.push('…');
306 }
307 s
308}
309
310#[async_trait::async_trait]
311impl Tool for SearchWebTool {
312 fn schema(&self) -> ToolSchema {
313 ToolSchema::new(
314 "search_web",
315 "Search the web. Auto-detects backend from configured URL/key. Supports SearXNG, Brave, Google CSE, Bing, SerpAPI. Falls back to HTML scraping (Bing/Sogou) when all backends fail.",
316 vec![
317 ToolParameter::new("query", ParamType::String, true, "Search keywords", &["search", "keyword", "q", "term"]),
318 ToolParameter::new("max_results", ParamType::Integer, false, "Max results (default 5)", &["limit", "count", "num", "num_results"]),
319 ToolParameter::new("timeout", ParamType::Integer, false, "Timeout seconds (default 8)", &["timeout_seconds", "time_limit"]),
320 ],
321 )
322 }
323
324 async fn call(&self, args: &Value, cancel: &AtomicBool) -> Result<String> {
325 let query = args.get("query").and_then(|v| v.as_str()).unwrap_or("");
326 let max_results = args.get("max_results").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
327 let timeout = args.get("timeout").and_then(|v| v.as_u64()).unwrap_or(self.timeout_secs);
328
329 if query.is_empty() {
330 return Ok(json!({"success": false, "error": "empty query"}).to_string());
331 }
332
333 let (success, engine, results, tried) =
334 self.search_with_fallback(query, max_results, timeout, cancel).await;
335
336 Ok(json!({
337 "success": success,
338 "query": query,
339 "engine": engine,
340 "results": results,
341 "total_results": results.len(),
342 "engines_tried": tried,
343 })
344 .to_string())
345 }
346}
347
348async fn searxng_search(base: &str, query: &str, max_results: usize, budget: f64) -> Result<Vec<Value>> {
351 let url = format!("{}/search", base.trim_end_matches('/'));
352 let resp = HTTP_CLIENT
353 .get(&url)
354 .query(&[("q", query), ("format", "json")])
355 .timeout(Duration::from_secs_f64(budget))
356 .send().await
357 .map_err(|e| AacodeError::Network(e.to_string()))?;
358 let body = resp.text().await.map_err(|e| AacodeError::Network(e.to_string()))?;
359 let v: Value = serde_json::from_str(&body)?;
360 Ok(extract_results(&v, max_results, "title", "url", "content"))
361}
362
363async fn brave_search(api_key: &str, query: &str, max_results: usize, budget: f64) -> Result<Vec<Value>> {
364 let resp = HTTP_CLIENT
365 .get("https://api.search.brave.com/res/v1/web/search")
366 .header("Accept", "application/json")
367 .header("Accept-Encoding", "gzip")
368 .header("X-Subscription-Token", api_key)
369 .query(&[("q", query), ("count", &max_results.to_string())])
370 .timeout(Duration::from_secs_f64(budget))
371 .send().await
372 .map_err(|e| AacodeError::Network(e.to_string()))?;
373 let body = resp.text().await.map_err(|e| AacodeError::Network(e.to_string()))?;
374 let v: Value = serde_json::from_str(&body)?;
375 let mut out = Vec::new();
376 if let Some(arr) = v.get("web").and_then(|r| r.get("results")).and_then(|r| r.as_array()) {
377 for item in arr.iter().take(max_results) {
378 out.push(json!({
379 "title": item.get("title").and_then(|x| x.as_str()).unwrap_or(""),
380 "url": item.get("url").and_then(|x| x.as_str()).unwrap_or(""),
381 "content": item.get("description").and_then(|x| x.as_str()).unwrap_or(""),
382 }));
383 }
384 }
385 Ok(out)
386}
387
388async fn google_cse_search(api_key: &str, cx: &str, query: &str, max_results: usize, budget: f64) -> Result<Vec<Value>> {
389 let resp = HTTP_CLIENT
390 .get("https://www.googleapis.com/customsearch/v1")
391 .query(&[("key", api_key), ("cx", cx), ("q", query), ("num", &max_results.to_string())])
392 .timeout(Duration::from_secs_f64(budget))
393 .send().await
394 .map_err(|e| AacodeError::Network(e.to_string()))?;
395 let body = resp.text().await.map_err(|e| AacodeError::Network(e.to_string()))?;
396 let v: Value = serde_json::from_str(&body)?;
397 Ok(extract_results(&v, max_results, "title", "link", "snippet"))
398}
399
400async fn bing_search(_api_key: &str, query: &str, max_results: usize, budget: f64) -> Result<Vec<Value>> {
401 let resp = HTTP_CLIENT
403 .get("https://api.bing.microsoft.com/v7.0/search")
404 .header("Ocp-Apim-Subscription-Key", _api_key)
405 .query(&[("q", query), ("count", &max_results.to_string()), ("mkt", "en-US")])
406 .timeout(Duration::from_secs_f64(budget))
407 .send().await
408 .map_err(|e| AacodeError::Network(e.to_string()))?;
409 let body = resp.text().await.map_err(|e| AacodeError::Network(e.to_string()))?;
410 let v: Value = serde_json::from_str(&body)?;
411 let mut out = Vec::new();
412 if let Some(arr) = v.get("webPages").and_then(|r| r.get("value")).and_then(|r| r.as_array()) {
413 for item in arr.iter().take(max_results) {
414 out.push(json!({
415 "title": item.get("name").and_then(|x| x.as_str()).unwrap_or(""),
416 "url": item.get("url").and_then(|x| x.as_str()).unwrap_or(""),
417 "content": item.get("snippet").and_then(|x| x.as_str()).unwrap_or(""),
418 }));
419 }
420 }
421 Ok(out)
422}
423
424async fn serpapi_search(api_key: &str, query: &str, max_results: usize, budget: f64) -> Result<Vec<Value>> {
425 let resp = HTTP_CLIENT
426 .get("https://serpapi.com/search")
427 .query(&[("api_key", api_key), ("q", query), ("engine", "google"), ("num", &max_results.to_string())])
428 .timeout(Duration::from_secs_f64(budget))
429 .send().await
430 .map_err(|e| AacodeError::Network(e.to_string()))?;
431 let body = resp.text().await.map_err(|e| AacodeError::Network(e.to_string()))?;
432 let v: Value = serde_json::from_str(&body)?;
433 Ok(extract_results(&v, max_results, "title", "link", "snippet"))
434}
435
436fn extract_results(v: &Value, max_results: usize, title_k: &str, url_k: &str, snippet_k: &str) -> Vec<Value> {
438 let arr = match v.get("results").or_else(|| v.get("items")).and_then(|r| r.as_array()) {
439 Some(a) => a,
440 None => return Vec::new(),
441 };
442 arr.iter()
443 .take(max_results)
444 .map(|item| {
445 json!({
446 "title": item.get(title_k).and_then(|x| x.as_str()).unwrap_or(""),
447 "url": item.get(url_k).and_then(|x| x.as_str()).unwrap_or(""),
448 "content": item.get(snippet_k).and_then(|x| x.as_str()).unwrap_or(""),
449 })
450 })
451 .collect()
452}
453
454async fn fallback_scrape(query: &str, max_results: usize, budget: f64) -> Option<(String, Vec<Value>)> {
457 if budget <= 0.5 {
458 return None;
459 }
460 let scrapers: [(&str, &str, &str, &str, &str); 3] = [
463 (
464 "ddg_scrape",
465 "https://html.duckduckgo.com/html/",
466 r#"(?s)<a[^>]*class="result__a"[^>]*href="([^"]+)"[^>]*>(.*?)</a>.*?class="result__snippet"[^>]*>(.*?)</"#,
467 r#"<[^>]+>"#,
468 "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
469 ),
470 (
471 "bing_scrape",
472 "https://www.bing.com/search",
473 r#"(?s)<li\s+class="b_algo"[^>]*>.*?<h2[^>]*>\s*<a[^>]*href="(https?://[^"]+)"[^>]*>(.*?)</a>.*?<p[^>]*>(.*?)</p>"#,
476 r#"<[^>]+>"#,
477 "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
478 ),
479 (
480 "sogou_scrape",
481 "https://www.sogou.com/web",
482 r#"(?s)<div[^>]*class="[^"]*vrwrap[^"]*"[^>]*>.*?<a[^>]*href="(https?://[^"]+)"[^>]*>\s*(.*?)\s*</a>.*?<p[^>]*>(.*?)</p>"#,
483 r#"<[^>]+>"#,
484 "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
485 ),
486 ];
487
488 let (tx, mut rx) = tokio::sync::mpsc::channel::<(String, Vec<Value>)>(3);
489 for (name, base_url, result_re, clean_re, ua) in &scrapers {
490 let q = query.to_string();
491 let n = name.to_string();
492 let bu = base_url.to_string();
493 let re = result_re.to_string();
494 let cr = clean_re.to_string();
495 let ua = ua.to_string();
496 let tx = tx.clone();
497 let budget = budget * 0.9; tokio::spawn(async move {
499 let scraped = scrape_one(&n, &bu, &q, max_results, budget, &re, &cr, &ua).await;
500 if let Some(v) = scraped {
501 let _ = tx.send((n, v)).await;
502 }
503 });
504 }
505 drop(tx);
506
507 tokio::time::timeout(Duration::from_secs_f64(budget), rx.recv()).await.ok().flatten()
508}
509
510async fn scrape_one(
511 name: &str,
512 base_url: &str,
513 query: &str,
514 max_results: usize,
515 budget: f64,
516 result_re: &str,
517 clean_re: &str,
518 ua: &str,
519) -> Option<Vec<Value>> {
520 let resp = HTTP_CLIENT
521 .get(base_url)
522 .header("User-Agent", ua)
523 .header("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8")
524 .query(&[("q", query)])
525 .timeout(Duration::from_secs_f64(budget))
526 .send().await;
527 let resp = match resp {
528 Ok(r) => r,
529 Err(_) => return None,
530 };
531 let html = match resp.text().await {
532 Ok(s) => s,
533 Err(_) => return None,
534 };
535 let re = regex::Regex::new(result_re).ok()?;
536 let clean = regex::Regex::new(clean_re).ok()?;
537 let mut results = Vec::new();
538 let mut seen = std::collections::HashSet::new();
539 for caps in re.captures_iter(&html) {
540 let raw_url = caps.get(1).map(|m| m.as_str().trim().to_string()).unwrap_or_default();
541 let url = normalize_result_url(&raw_url);
542 let title = html_unescape(
543 clean
544 .replace_all(caps.get(2).map(|m| m.as_str()).unwrap_or(""), "")
545 .trim(),
546 );
547 let snippet = html_unescape(
548 clean
549 .replace_all(caps.get(3).map(|m| m.as_str()).unwrap_or(""), "")
550 .trim(),
551 );
552 if !is_quality_result(&url, &title) || seen.contains(&url) {
553 continue;
554 }
555 seen.insert(url.clone());
556 results.push(json!({
557 "title": title,
558 "url": url,
559 "content": snippet,
560 "engine": name,
561 }));
562 if results.len() >= max_results {
563 break;
564 }
565 }
566 if results.is_empty() {
567 None
568 } else {
569 Some(results)
570 }
571}
572
573fn normalize_result_url(url: &str) -> String {
576 if let Some(pos) = url.find("uddg=") {
577 let rest = &url[pos + 5..];
578 let enc = rest.split('&').next().unwrap_or(rest);
579 return percent_decode(enc);
580 }
581 if let Some(stripped) = url.strip_prefix("//") {
582 return format!("https://{stripped}");
583 }
584 url.to_string()
585}
586
587fn percent_decode(s: &str) -> String {
588 let bytes = s.as_bytes();
589 let mut out = Vec::with_capacity(bytes.len());
590 let mut i = 0;
591 while i < bytes.len() {
592 if bytes[i] == b'%' && i + 2 < bytes.len() {
593 if let Ok(b) = u8::from_str_radix(&s[i + 1..i + 3], 16) {
594 out.push(b);
595 i += 3;
596 continue;
597 }
598 }
599 out.push(bytes[i]);
600 i += 1;
601 }
602 String::from_utf8_lossy(&out).to_string()
603}
604
605fn html_unescape(s: &str) -> String {
606 s.replace(" ", " ")
607 .replace("&", "&")
608 .replace("<", "<")
609 .replace(">", ">")
610 .replace(""", "\"")
611 .replace("'", "'")
612 .replace("'", "'")
613}
614
615fn is_quality_result(url: &str, title: &str) -> bool {
618 if url.is_empty() || !url.starts_with("http") {
619 return false;
620 }
621 if title.is_empty() || title.contains('›') {
622 return false;
623 }
624 if !title.contains(' ') && (title.contains("http") || title.contains(".com") || title.contains(".org")) {
626 return false;
627 }
628 if url.contains("bing.com/aclick") || url.contains("duckduckgo.com/y.js") {
630 return false;
631 }
632 true
633}
634
635fn detect_engine_type(url: &str) -> &'static str {
638 let lower = url.to_lowercase();
639 if lower.contains("brave.com") {
640 "brave"
641 } else if lower.contains("googleapis.com/customsearch") {
642 "google_cse"
643 } else if lower.contains("bing.microsoft.com") {
644 "bing"
645 } else if lower.contains("serpapi.com") {
646 "serpapi"
647 } else {
648 "searxng"
649 }
650}
651
652fn current_time() -> f64 {
653 use std::time::{SystemTime, UNIX_EPOCH};
654 SystemTime::now()
655 .duration_since(UNIX_EPOCH)
656 .map(|d| d.as_secs_f64())
657 .unwrap_or(0.0)
658}
659
660pub struct FetchUrlTool {
663 pub project_path: PathBuf,
664 pub timeout_secs: u64,
665}
666
667#[async_trait::async_trait]
668impl Tool for FetchUrlTool {
669 fn schema(&self) -> ToolSchema {
670 ToolSchema::new(
671 "fetch_url",
672 "Fetch the content of a URL (cleaned text). Also available via run_shell + curl.",
673 vec![
674 ToolParameter::new("url", ParamType::String, true, "URL to fetch", &["link", "uri", "address"]),
675 ToolParameter::new("timeout", ParamType::Integer, false, "Timeout seconds", &["time_limit", "max_time", "wait"]),
676 ToolParameter::new("max_content_length", ParamType::Integer, false, "Max cleaned chars", &["max_length", "max_chars"]),
677 ],
678 )
679 }
680
681 async fn call(&self, args: &Value, _c: &AtomicBool) -> Result<String> {
682 let url = args.get("url").and_then(|v| v.as_str()).unwrap_or("");
683 if url.is_empty() {
684 return Ok(json!({"success": false, "error": "empty url"}).to_string());
685 }
686 let timeout = args.get("timeout").and_then(|v| v.as_u64()).unwrap_or(self.timeout_secs);
687 let max_len = args.get("max_content_length").and_then(|v| v.as_u64()).unwrap_or(5000) as usize;
688
689 let resp = HTTP_CLIENT
690 .get(url)
691 .header("User-Agent", "aacode-rs/0.1")
692 .timeout(Duration::from_secs_f64(timeout as f64))
693 .send().await;
694 let (status, body) = match resp {
695 Ok(r) => {
696 let status = r.status().as_u16();
697 (status, r.text().await.unwrap_or_default())
698 }
699 Err(e) => {
700 return Ok(json!({"success": false, "url": url, "error": e.to_string()}).to_string());
701 }
702 };
703
704 let raw_len = body.len();
705 let cleaned = clean_html(&body);
706 let content: String = cleaned.chars().take(max_len).collect();
707 let saved = save_extract(&self.project_path, &cleaned);
708
709 Ok(json!({
710 "success": true,
711 "url": url,
712 "status_code": status,
713 "raw_length": raw_len,
714 "content_length": cleaned.chars().count(),
715 "content": content,
716 "extract_file": saved,
717 })
718 .to_string())
719 }
720}
721
722fn save_extract(project_path: &std::path::Path, content: &str) -> Option<String> {
723 let dir = project_path.join(".aacode").join("context");
724 std::fs::create_dir_all(&dir).ok()?;
725 let stamp = std::time::SystemTime::now()
728 .duration_since(std::time::UNIX_EPOCH)
729 .map(|d| d.as_millis())
730 .unwrap_or(0);
731 let path = dir.join(format!("web_fetch_{stamp}_{}.txt", std::process::id()));
732 std::fs::write(&path, content).ok()?;
733 if let Ok(entries) = std::fs::read_dir(&dir) {
735 let mut files: Vec<_> = entries
736 .flatten()
737 .filter(|e| e.file_name().to_string_lossy().starts_with("web_fetch_"))
738 .collect();
739 if files.len() > 20 {
740 files.sort_by_key(|e| e.metadata().and_then(|m| m.modified()).ok());
741 for old in files.iter().take(files.len() - 20) {
742 let _ = std::fs::remove_file(old.path());
743 }
744 }
745 }
746 Some(path.to_string_lossy().to_string())
747}
748
749pub fn clean_html(html: &str) -> String {
753 let stripped = remove_blocks(html);
754 let mut out = String::with_capacity(stripped.len() / 2);
755 let mut in_tag = false;
756 for c in stripped.chars() {
757 match c {
758 '<' => in_tag = true,
759 '>' => { in_tag = false; out.push(' '); }
760 _ if !in_tag => out.push(c),
761 _ => {}
762 }
763 }
764 let decoded = out
765 .replace(" ", " ")
766 .replace("&", "&")
767 .replace("<", "<")
768 .replace(">", ">")
769 .replace(""", "\"");
770 decoded.split_whitespace().collect::<Vec<_>>().join(" ")
771}
772
773fn remove_blocks(html: &str) -> String {
774 let lower = html.to_lowercase();
775 let chars: Vec<char> = html.chars().collect();
776 let low: Vec<char> = lower.chars().collect();
777 let mut out = String::with_capacity(html.len());
778 let mut i = 0;
779 let n = chars.len();
780 while i < n {
781 let rest: String = low[i..].iter().take(8).collect();
782 if rest.starts_with("<script") || rest.starts_with("<style") {
783 let close = if rest.starts_with("<script") { "</script>" } else { "</style>" };
784 let low_rest: String = low[i..].iter().collect();
785 if let Some(pos) = low_rest.find(close) {
786 i += pos + close.len();
787 continue;
788 } else {
789 break;
790 }
791 }
792 out.push(chars[i]);
793 i += 1;
794 }
795 out
796}
797
798pub struct SearchCodeTool {
801 pub cfg: SearchConfig,
802 pub timeout_secs: u64,
803}
804
805#[async_trait::async_trait]
806impl Tool for SearchCodeTool {
807 fn schema(&self) -> ToolSchema {
808 ToolSchema::new(
809 "search_code",
810 "Search code examples. If SearXNG is configured, uses it with categories=it to find code. Falls back to GitHub repository search.",
811 vec![
812 ToolParameter::new("query", ParamType::String, true, "Search keywords", &["q", "keyword", "search"]),
813 ToolParameter::new("max_results", ParamType::Integer, false, "Max results", &["limit", "count"]),
814 ],
815 )
816 }
817 async fn call(&self, args: &Value, _c: &AtomicBool) -> Result<String> {
818 let query = args.get("query").and_then(|v| v.as_str()).unwrap_or("");
819 if query.is_empty() {
820 return Ok(json!({"success": false, "error": "empty query"}).to_string());
821 }
822 let max = args.get("max_results").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
823
824 if let Some(base) = &self.cfg.searxng_url {
827 let url = format!("{}/search", base.trim_end_matches('/'));
828 match HTTP_CLIENT
829 .get(&url)
830 .query(&[("q", query), ("format", "json"), ("categories", "it")])
831 .timeout(Duration::from_secs_f64(PER_ENGINE_CAP_SECS))
832 .send().await
833 {
834 Ok(r) => {
835 let body = r.text().await.unwrap_or_default();
836 if let Ok(v) = serde_json::from_str::<Value>(&body) {
837 let results = extract_results(&v, max, "title", "url", "content");
838 if !results.is_empty() {
839 return Ok(json!({"success": true, "query": query, "results": results}).to_string());
840 }
841 }
842 }
843 Err(_) => {}
844 }
845 }
846
847 let url = format!(
849 "https://api.github.com/search/repositories?q={}&per_page={}",
850 urlencode(query),
851 max
852 );
853 let resp = HTTP_CLIENT
854 .get(&url)
855 .header("User-Agent", "aacode-rs")
856 .header("Accept", "application/vnd.github+json")
857 .timeout(Duration::from_secs_f64(self.timeout_secs as f64))
858 .send().await;
859 match resp {
860 Ok(r) => {
861 let body = r.text().await.unwrap_or_default();
862 let v: Value = serde_json::from_str(&body).unwrap_or(json!({}));
863 let mut items = Vec::new();
864 if let Some(arr) = v.get("items").and_then(|x| x.as_array()) {
865 for it in arr.iter().take(max) {
866 items.push(json!({
867 "name": it.get("full_name").and_then(|x| x.as_str()).unwrap_or(""),
868 "url": it.get("html_url").and_then(|x| x.as_str()).unwrap_or(""),
869 "description": it.get("description").and_then(|x| x.as_str()).unwrap_or(""),
870 "stars": it.get("stargazers_count").and_then(|x| x.as_u64()).unwrap_or(0),
871 }));
872 }
873 }
874 Ok(json!({"success": true, "query": query, "results": items}).to_string())
875 }
876 Err(e) => Ok(json!({"success": false, "error": e.to_string()}).to_string()),
877 }
878 }
879}
880
881fn urlencode(s: &str) -> String {
882 let mut out = String::new();
883 for b in s.bytes() {
884 match b {
885 b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => out.push(b as char),
886 b' ' => out.push_str("%20"),
887 _ => out.push_str(&format!("%{b:02X}")),
888 }
889 }
890 out
891}
892
893#[cfg(test)]
894mod tests {
895 use super::*;
896
897 #[test]
898 fn clean_html_strips_tags_and_scripts() {
899 let html = "<html><head><style>.a{color:red}</style><script>var x=1;</script></head><body><p>Hello <b>World</b></p></body></html>";
900 let cleaned = clean_html(html);
901 assert!(cleaned.contains("Hello"));
902 assert!(cleaned.contains("World"));
903 assert!(!cleaned.contains("color:red"));
904 assert!(!cleaned.contains("var x"));
905 }
906
907 #[test]
908 fn urlencode_works() {
909 assert_eq!(urlencode("a b"), "a%20b");
910 assert_eq!(urlencode("rust lang"), "rust%20lang");
911 }
912
913 #[test]
914 fn detect_engine() {
915 assert_eq!(detect_engine_type("https://api.search.brave.com"), "brave");
916 assert_eq!(detect_engine_type("https://api.bing.microsoft.com"), "bing");
917 assert_eq!(detect_engine_type("https://myserver.com"), "searxng");
918 }
919
920 #[tokio::test]
921 async fn search_web_no_backend() {
922 let t = SearchWebTool::new(SearchConfig::default(), 5);
923 let cancel = AtomicBool::new(false);
924 let out = t.call(&json!({"query": "rust"}), &cancel).await.unwrap();
925 let v: Value = serde_json::from_str(&out).unwrap();
926 let _ = v["success"].as_bool();
928 }
929
930 #[tokio::test]
931 async fn fetch_url_empty() {
932 let t = FetchUrlTool { project_path: std::env::temp_dir(), timeout_secs: 5 };
933 let cancel = AtomicBool::new(false);
934 let out = t.call(&json!({"url": ""}), &cancel).await.unwrap();
935 assert_eq!(serde_json::from_str::<Value>(&out).unwrap()["success"], false);
936 }
937
938 #[test]
939 fn schema_defined() {
940 let s = SearchWebTool::new(SearchConfig::default(), 5).schema();
941 assert_eq!(s.name, "search_web");
942 }
943
944 #[test]
945 fn circuit_breaker_opens_after_failure() {
946 let t = SearchWebTool::new(SearchConfig::default(), 5);
948 assert!(!t.circuit_open("test_engine_cb"));
949 t.record_outcome("test_engine_cb", false);
950 assert!(t.circuit_open("test_engine_cb"));
951 t.record_outcome("test_engine_cb", true);
953 assert!(!t.circuit_open("test_engine_cb"));
954 }
955
956 #[test]
957 fn circuit_breaker_shared_across_instances() {
958 let a = SearchWebTool::new(SearchConfig::default(), 5);
960 let b = SearchWebTool::new(SearchConfig::default(), 5);
961 a.record_outcome("test_engine_shared", false);
962 assert!(b.circuit_open("test_engine_shared"), "breaker must be shared");
963 b.record_outcome("test_engine_shared", true);
964 assert!(!a.circuit_open("test_engine_shared"));
965 }
966
967 #[tokio::test]
968 async fn unconfigured_engines_are_skipped_fast() {
969 let t = SearchWebTool::new(SearchConfig::default(), 5);
972 let start = std::time::Instant::now();
973 let a = t.try_engine("brave", "q", 3, 5.0).await;
974 assert!(matches!(a, Attempt::Skipped("not-configured")));
975 let a = t.try_engine("google_cse", "q", 3, 5.0).await;
976 assert!(matches!(a, Attempt::Skipped("not-configured")));
977 let a = t.try_engine("serpapi", "q", 3, 5.0).await;
978 assert!(matches!(a, Attempt::Skipped("not-configured")));
979 let a = t.try_engine("searxng", "q", 3, 5.0).await;
980 assert!(matches!(a, Attempt::Skipped("not-configured")));
981 assert!(start.elapsed().as_millis() < 200, "skips must not hit the network");
982 }
983
984 static SEARXNG_TEST_LOCK: Mutex<()> = Mutex::new(());
986
987 #[tokio::test]
988 async fn refused_searxng_fails_fast_and_opens_circuit() {
989 let _l = SEARXNG_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
990 let cfg = SearchConfig {
993 searxng_url: Some("http://127.0.0.1:59999".into()),
994 ..Default::default()
995 };
996 let t = SearchWebTool::new(cfg, 5);
997 t.record_outcome("searxng", true); let a = t.try_engine("searxng", "q", 3, 5.0).await;
999 assert!(matches!(a, Attempt::TransportError(_)));
1000 t.record_outcome("searxng", false);
1001 let a = t.try_engine("searxng", "q", 3, 5.0).await;
1002 assert!(matches!(a, Attempt::Skipped("circuit-open")));
1003 t.record_outcome("searxng", true); }
1005
1006 #[tokio::test]
1007 async fn search_reports_engines_tried() {
1008 let _l = SEARXNG_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
1009 let cfg = SearchConfig {
1010 searxng_url: Some("http://127.0.0.1:59999".into()),
1011 ..Default::default()
1012 };
1013 let t = SearchWebTool::new(cfg, 1);
1014 let cancel = AtomicBool::new(false);
1015 let out = t.call(&json!({"query": "rust", "timeout": 1}), &cancel).await.unwrap();
1016 let v: Value = serde_json::from_str(&out).unwrap();
1017 let tried = v["engines_tried"].as_array().unwrap();
1018 assert!(!tried.is_empty());
1019 let first = tried[0].as_str().unwrap();
1020 assert!(first.starts_with("searxng:"), "first tried should be searxng, got {first}");
1021 t.record_outcome("searxng", true); }
1023
1024 #[tokio::test]
1025 async fn cancelled_search_returns_immediately() {
1026 let cfg = SearchConfig {
1027 searxng_url: Some("http://127.0.0.1:59999".into()),
1028 ..Default::default()
1029 };
1030 let t = SearchWebTool::new(cfg, 8);
1031 let cancel = AtomicBool::new(true);
1032 let start = std::time::Instant::now();
1033 let out = t.call(&json!({"query": "rust"}), &cancel).await.unwrap();
1034 let v: Value = serde_json::from_str(&out).unwrap();
1035 assert_eq!(v["success"], false);
1036 assert!(start.elapsed().as_secs() < 2);
1037 }
1038
1039 #[test]
1040 fn brief_err_truncates() {
1041 let long = "x".repeat(300);
1042 assert!(brief_err(&long).chars().count() <= 81);
1043 assert_eq!(brief_err("short"), "short");
1044 }
1045
1046 #[test]
1047 fn ddg_redirect_url_unwrapped() {
1048 let u = "//duckduckgo.com/l/?uddg=https%3A%2F%2Fblog.rust%2Dlang.org%2F2024%2F07%2F25%2FRust%2D1.80.0.html&rut=abc";
1049 assert_eq!(
1050 normalize_result_url(u),
1051 "https://blog.rust-lang.org/2024/07/25/Rust-1.80.0.html"
1052 );
1053 assert_eq!(normalize_result_url("https://a.com/x"), "https://a.com/x");
1055 assert_eq!(normalize_result_url("//a.com/x"), "https://a.com/x");
1057 }
1058
1059 #[test]
1060 fn percent_decode_works() {
1061 assert_eq!(percent_decode("a%20b"), "a b");
1062 assert_eq!(percent_decode("100%"), "100%");
1063 assert_eq!(percent_decode("%2Fpath"), "/path");
1064 }
1065
1066 #[test]
1067 fn quality_filter_drops_junk() {
1068 assert!(!is_quality_result("https://a.com", "rust-lang.orghttps://rust-lang.org › zh-CN"));
1070 assert!(!is_quality_result("https://a.com", "rust-lang.org"));
1072 assert!(!is_quality_result("https://www.bing.com/aclick?x=1", "Real Title"));
1074 assert!(is_quality_result(
1076 "https://blog.rust-lang.org/2024/07/25/Rust-1.80.0.html",
1077 "Announcing Rust 1.80.0"
1078 ));
1079 }
1080
1081 #[test]
1082 fn html_unescape_entities() {
1083 assert_eq!(html_unescape("a & b's <tag>"), "a & b's <tag>");
1084 }
1085}