Skip to main content

bamboo_tools/tools/
web_search.rs

1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolClass, ToolCtx, ToolError, ToolOutcome, ToolResult};
3use parking_lot::RwLock;
4use regex::Regex;
5use serde::Deserialize;
6use serde_json::json;
7use std::collections::{HashMap, HashSet};
8use std::sync::OnceLock;
9use std::time::{Duration, Instant};
10use url::Url;
11
12const CACHE_TTL: Duration = Duration::from_secs(15 * 60);
13const DEFAULT_MAX_RESULTS: usize = 10;
14const ABSOLUTE_MAX_RESULTS: usize = 20;
15const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36";
16const WEB_SEARCH_ENDPOINTS_ENV: &str = "BAMBOO_WEB_SEARCH_ENDPOINTS";
17const DEFAULT_WEB_SEARCH_ENDPOINTS: [&str; 2] = [
18    "https://html.duckduckgo.com/html/",
19    "https://lite.duckduckgo.com/lite/",
20];
21
22#[derive(Debug, Deserialize)]
23struct WebSearchArgs {
24    query: String,
25    #[serde(default)]
26    allowed_domains: Option<Vec<String>>,
27    #[serde(default)]
28    blocked_domains: Option<Vec<String>>,
29    #[serde(default)]
30    max_results: Option<usize>,
31}
32
33struct CachedSearch {
34    results: serde_json::Value,
35    expires_at: Instant,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq)]
39struct SearchResult {
40    title: String,
41    url: String,
42    domain: String,
43    snippet: Option<String>,
44}
45
46impl SearchResult {
47    fn into_json(self) -> serde_json::Value {
48        let mut value = json!({
49            "title": self.title,
50            "url": self.url,
51            "domain": self.domain,
52        });
53        if let Some(snippet) = self.snippet {
54            value["snippet"] = json!(snippet);
55        }
56        value
57    }
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
61enum ParsedSearchPage {
62    Results(Vec<SearchResult>),
63    AntiBot,
64    Unrecognized,
65}
66
67static SEARCH_CACHE: OnceLock<RwLock<HashMap<String, CachedSearch>>> = OnceLock::new();
68
69fn search_cache() -> &'static RwLock<HashMap<String, CachedSearch>> {
70    SEARCH_CACHE.get_or_init(|| RwLock::new(HashMap::new()))
71}
72
73// Static, compile-time-constant patterns: compile each exactly once and reuse.
74// `expect` is safe here because the patterns are hardcoded and verified valid.
75static ANCHOR_RE: OnceLock<Regex> = OnceLock::new();
76static TAG_RE: OnceLock<Regex> = OnceLock::new();
77static HREF_RE: OnceLock<Regex> = OnceLock::new();
78static CLASS_RE: OnceLock<Regex> = OnceLock::new();
79static LITE_SNIPPET_RE: OnceLock<Regex> = OnceLock::new();
80
81pub struct WebSearchTool;
82
83impl WebSearchTool {
84    pub fn new() -> Self {
85        Self
86    }
87
88    fn cache_key(
89        query: &str,
90        allowed: &Option<Vec<String>>,
91        blocked: &Option<Vec<String>>,
92        max_results: usize,
93        endpoints: &[Url],
94    ) -> String {
95        let endpoints = endpoints.iter().map(Url::as_str).collect::<Vec<_>>();
96        serde_json::to_string(&(query, allowed, blocked, max_results, endpoints))
97            .expect("web search cache key components are serializable")
98    }
99
100    fn try_cache(key: &str) -> Option<serde_json::Value> {
101        let cache = search_cache().read();
102        let entry = cache.get(key)?;
103        if entry.expires_at > Instant::now() {
104            Some(entry.results.clone())
105        } else {
106            None
107        }
108    }
109
110    fn put_cache(key: String, results: serde_json::Value) {
111        let mut cache = search_cache().write();
112        cache.insert(
113            key,
114            CachedSearch {
115                results,
116                expires_at: Instant::now() + CACHE_TTL,
117            },
118        );
119    }
120
121    fn anchor_re() -> &'static Regex {
122        ANCHOR_RE
123            .get_or_init(|| Regex::new(r"(?is)<a\b([^>]*)>(.*?)</a>").expect("valid static regex"))
124    }
125
126    fn tag_re() -> &'static Regex {
127        TAG_RE.get_or_init(|| Regex::new(r"(?is)<[^>]+>").expect("valid static regex"))
128    }
129
130    fn href_re() -> &'static Regex {
131        HREF_RE.get_or_init(|| {
132            Regex::new(r#"(?i)\bhref\s*=\s*[\"']([^\"']+)[\"']"#).expect("valid static regex")
133        })
134    }
135
136    fn class_re() -> &'static Regex {
137        CLASS_RE.get_or_init(|| {
138            Regex::new(r#"(?i)\bclass\s*=\s*[\"']([^\"']+)[\"']"#).expect("valid static regex")
139        })
140    }
141
142    fn lite_snippet_re() -> &'static Regex {
143        LITE_SNIPPET_RE.get_or_init(|| {
144            Regex::new(
145                r#"(?is)<td\b[^>]*class\s*=\s*[\"'][^\"']*\bresult-snippet\b[^\"']*[\"'][^>]*>(.*?)</td>"#,
146            )
147            .expect("valid static regex")
148        })
149    }
150
151    fn attr_value<'a>(attrs: &'a str, pattern: &Regex) -> Option<&'a str> {
152        pattern
153            .captures(attrs)
154            .and_then(|capture| capture.get(1))
155            .map(|value| value.as_str())
156    }
157
158    fn has_class(attrs: &str, expected: &str) -> bool {
159        Self::attr_value(attrs, Self::class_re()).is_some_and(|classes| {
160            classes
161                .split_ascii_whitespace()
162                .any(|class| class == expected)
163        })
164    }
165
166    fn decode_html_entities(value: &str) -> String {
167        value
168            .replace("&quot;", "\"")
169            .replace("&#x27;", "'")
170            .replace("&#39;", "'")
171            .replace("&lt;", "<")
172            .replace("&gt;", ">")
173            .replace("&nbsp;", " ")
174            .replace("&amp;", "&")
175    }
176
177    fn clean_html_text(value: &str) -> String {
178        let without_tags = Self::tag_re().replace_all(value, " ");
179        Self::decode_html_entities(&without_tags)
180            .split_whitespace()
181            .collect::<Vec<_>>()
182            .join(" ")
183    }
184
185    fn make_absolute_url(raw: &str) -> Option<String> {
186        let raw = Self::decode_html_entities(raw.trim());
187        if raw.is_empty() {
188            return None;
189        }
190        if raw.starts_with("//") {
191            return Some(format!("https:{raw}"));
192        }
193        if raw.starts_with('/') {
194            return Url::parse("https://duckduckgo.com")
195                .ok()?
196                .join(&raw)
197                .ok()
198                .map(|url| url.to_string());
199        }
200        Some(raw)
201    }
202
203    fn decode_duckduckgo_url(raw: &str) -> Option<String> {
204        let absolute = Self::make_absolute_url(raw)?;
205        let parsed = Url::parse(&absolute).ok()?;
206        if !matches!(parsed.scheme(), "http" | "https") || parsed.host_str().is_none() {
207            return None;
208        }
209
210        let host = parsed.host_str()?.to_ascii_lowercase();
211        if Self::domain_matches(&host, "duckduckgo.com") && parsed.path() == "/l/" {
212            if let Some(value) = parsed
213                .query_pairs()
214                .find(|(key, _)| key == "uddg")
215                .map(|(_, value)| value.to_string())
216            {
217                let target = Self::make_absolute_url(&value)?;
218                let target = Url::parse(&target).ok()?;
219                if !matches!(target.scheme(), "http" | "https") || target.host_str().is_none() {
220                    return None;
221                }
222                return Some(target.to_string());
223            }
224        }
225
226        Some(parsed.to_string())
227    }
228
229    fn host_of(url: &str) -> Option<String> {
230        url::Url::parse(url)
231            .ok()
232            .and_then(|parsed| parsed.host_str().map(|host| host.to_ascii_lowercase()))
233    }
234
235    fn domain_matches(host: &str, domain: &str) -> bool {
236        host == domain || host.ends_with(&format!(".{}", domain))
237    }
238
239    fn result_from_parts(
240        raw_url: &str,
241        title_html: &str,
242        snippet_html: Option<&str>,
243        allowed: &Option<HashSet<String>>,
244        blocked: &HashSet<String>,
245    ) -> Option<SearchResult> {
246        let url = Self::decode_duckduckgo_url(raw_url)?;
247        let host = Self::host_of(&url)?;
248
249        if blocked
250            .iter()
251            .any(|blocked_domain| Self::domain_matches(&host, blocked_domain))
252        {
253            return None;
254        }
255        if let Some(allowed_set) = allowed {
256            if !allowed_set
257                .iter()
258                .any(|allowed_domain| Self::domain_matches(&host, allowed_domain))
259            {
260                return None;
261            }
262        }
263
264        let title = Self::clean_html_text(title_html);
265        let snippet = snippet_html
266            .map(Self::clean_html_text)
267            .filter(|value| !value.is_empty());
268        Some(SearchResult {
269            title: if title.is_empty() { url.clone() } else { title },
270            url,
271            domain: host,
272            snippet,
273        })
274    }
275
276    fn recognized_empty_page(html: &str, layout_path: &str) -> bool {
277        let lower = html.to_ascii_lowercase();
278        lower.contains(layout_path)
279            && (lower.contains("no results")
280                || lower.contains("no-results")
281                || lower.contains("result--no-result"))
282    }
283
284    fn parse_html_layout(
285        html: &str,
286        allowed: &Option<HashSet<String>>,
287        blocked: &HashSet<String>,
288        max_results: usize,
289    ) -> Option<Vec<SearchResult>> {
290        let mut snippets = HashMap::new();
291        let mut links = Vec::new();
292
293        for capture in Self::anchor_re().captures_iter(html) {
294            let attrs = capture.get(1)?.as_str();
295            let content = capture.get(2)?.as_str();
296            let Some(href) = Self::attr_value(attrs, Self::href_re()) else {
297                continue;
298            };
299            if Self::has_class(attrs, "result__snippet") {
300                if let Some(url) = Self::decode_duckduckgo_url(href) {
301                    snippets.insert(url, content.to_string());
302                }
303            } else if Self::has_class(attrs, "result__a") {
304                links.push((href.to_string(), content.to_string()));
305            }
306        }
307
308        if links.is_empty() {
309            return Self::recognized_empty_page(html, "/html/").then(Vec::new);
310        }
311        if !links
312            .iter()
313            .any(|(raw_url, _)| Self::decode_duckduckgo_url(raw_url).is_some())
314        {
315            return None;
316        }
317
318        let mut results = Vec::new();
319        for (raw_url, title) in links {
320            let decoded = Self::decode_duckduckgo_url(&raw_url);
321            let snippet = decoded.as_ref().and_then(|url| snippets.get(url));
322            if let Some(result) = Self::result_from_parts(
323                &raw_url,
324                &title,
325                snippet.map(String::as_str),
326                allowed,
327                blocked,
328            ) {
329                results.push(result);
330                if results.len() >= max_results {
331                    break;
332                }
333            }
334        }
335        Some(results)
336    }
337
338    fn parse_lite_layout(
339        html: &str,
340        allowed: &Option<HashSet<String>>,
341        blocked: &HashSet<String>,
342        max_results: usize,
343    ) -> Option<Vec<SearchResult>> {
344        let mut links = Vec::new();
345        for capture in Self::anchor_re().captures_iter(html) {
346            let full = capture.get(0)?;
347            let attrs = capture.get(1)?.as_str();
348            if !Self::has_class(attrs, "result-link") {
349                continue;
350            }
351            let Some(href) = Self::attr_value(attrs, Self::href_re()) else {
352                continue;
353            };
354            links.push((
355                full.start(),
356                full.end(),
357                href.to_string(),
358                capture.get(2)?.as_str().to_string(),
359            ));
360        }
361
362        if links.is_empty() {
363            return Self::recognized_empty_page(html, "/lite/").then(Vec::new);
364        }
365        if !links
366            .iter()
367            .any(|(_, _, raw_url, _)| Self::decode_duckduckgo_url(raw_url).is_some())
368        {
369            return None;
370        }
371
372        let mut results = Vec::new();
373        for (index, (_, end, raw_url, title)) in links.iter().enumerate() {
374            let next_start = links
375                .get(index + 1)
376                .map(|(start, _, _, _)| *start)
377                .unwrap_or(html.len());
378            let snippet = Self::lite_snippet_re()
379                .captures(&html[*end..next_start])
380                .and_then(|capture| capture.get(1))
381                .map(|value| value.as_str());
382            if let Some(result) = Self::result_from_parts(raw_url, title, snippet, allowed, blocked)
383            {
384                results.push(result);
385                if results.len() >= max_results {
386                    break;
387                }
388            }
389        }
390        Some(results)
391    }
392
393    fn is_anti_bot_page(html: &str) -> bool {
394        let lower = html.to_ascii_lowercase();
395        lower.contains("unfortunately, bots use duckduckgo too")
396            || lower.contains("anomaly-modal")
397            || lower.contains("challenge-form")
398    }
399
400    fn parse_search_page(
401        html: &str,
402        allowed: &Option<HashSet<String>>,
403        blocked: &HashSet<String>,
404        max_results: usize,
405    ) -> ParsedSearchPage {
406        if Self::is_anti_bot_page(html) {
407            return ParsedSearchPage::AntiBot;
408        }
409        if let Some(results) = Self::parse_html_layout(html, allowed, blocked, max_results) {
410            return ParsedSearchPage::Results(results);
411        }
412        if let Some(results) = Self::parse_lite_layout(html, allowed, blocked, max_results) {
413            return ParsedSearchPage::Results(results);
414        }
415        ParsedSearchPage::Unrecognized
416    }
417
418    fn parse_endpoint_list(override_value: Option<&str>) -> Result<Vec<Url>, String> {
419        let raw_endpoints: Vec<&str> = match override_value {
420            Some(value) => value
421                .split(',')
422                .map(str::trim)
423                .filter(|value| !value.is_empty())
424                .collect(),
425            None => DEFAULT_WEB_SEARCH_ENDPOINTS.to_vec(),
426        };
427        if raw_endpoints.is_empty() {
428            return Err(format!(
429                "{WEB_SEARCH_ENDPOINTS_ENV} must contain at least one endpoint"
430            ));
431        }
432
433        raw_endpoints
434            .into_iter()
435            .enumerate()
436            .map(|(index, raw)| {
437                let url = Url::parse(raw).map_err(|_| {
438                    format!(
439                        "{WEB_SEARCH_ENDPOINTS_ENV} entry {} must be an absolute HTTP(S) URL",
440                        index + 1
441                    )
442                })?;
443                if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
444                    return Err(format!(
445                        "{WEB_SEARCH_ENDPOINTS_ENV} entry {} must be an absolute HTTP(S) URL",
446                        index + 1
447                    ));
448                }
449                Ok(url)
450            })
451            .collect()
452    }
453
454    fn configured_endpoints() -> Result<Vec<Url>, String> {
455        match std::env::var(WEB_SEARCH_ENDPOINTS_ENV) {
456            Ok(value) => Self::parse_endpoint_list(Some(&value)),
457            Err(std::env::VarError::NotPresent) => Self::parse_endpoint_list(None),
458            Err(std::env::VarError::NotUnicode(_)) => Err(format!(
459                "{WEB_SEARCH_ENDPOINTS_ENV} must contain valid UTF-8"
460            )),
461        }
462    }
463
464    fn endpoint_label(endpoint: &Url) -> String {
465        let host = endpoint.host_str().unwrap_or("unknown-host");
466        let port = endpoint
467            .port()
468            .map(|port| format!(":{port}"))
469            .unwrap_or_default();
470        format!("{}://{host}{port}{}", endpoint.scheme(), endpoint.path())
471    }
472
473    fn build_http_client() -> Result<reqwest::Client, String> {
474        reqwest::Client::builder()
475            .timeout(Duration::from_secs(30))
476            // Configured endpoints are the complete trust boundary. Do not let a
477            // remote response replay the query POST to an unvalidated redirect.
478            .redirect(reqwest::redirect::Policy::none())
479            .build()
480            .map_err(|error| format!("Failed to build HTTP client: {error}"))
481    }
482
483    async fn search_endpoint_chain(
484        client: &reqwest::Client,
485        endpoints: &[Url],
486        query: &str,
487        allowed: &Option<HashSet<String>>,
488        blocked: &HashSet<String>,
489        max_results: usize,
490    ) -> Result<Vec<SearchResult>, String> {
491        let mut failures = Vec::with_capacity(endpoints.len());
492        for endpoint in endpoints {
493            let label = Self::endpoint_label(endpoint);
494            let response = match client
495                .post(endpoint.clone())
496                .header("User-Agent", USER_AGENT)
497                .form(&[("q", query)])
498                .send()
499                .await
500            {
501                Ok(response) => response,
502                Err(error) => {
503                    failures.push(format!("{label}: request failed ({})", error.without_url()));
504                    continue;
505                }
506            };
507            if !response.status().is_success() {
508                failures.push(format!("{label}: HTTP {}", response.status()));
509                continue;
510            }
511            let html = match response.text().await {
512                Ok(html) => html,
513                Err(error) => {
514                    failures.push(format!(
515                        "{label}: response decode failed ({})",
516                        error.without_url()
517                    ));
518                    continue;
519                }
520            };
521            match Self::parse_search_page(&html, allowed, blocked, max_results) {
522                ParsedSearchPage::Results(results) => return Ok(results),
523                ParsedSearchPage::AntiBot => {
524                    failures.push(format!("{label}: blocked by anti-bot protection"));
525                }
526                ParsedSearchPage::Unrecognized => {
527                    failures.push(format!("{label}: unrecognized search response"));
528                }
529            }
530        }
531
532        Err(format!(
533            "all configured web search endpoints failed: {}",
534            failures.join("; ")
535        ))
536    }
537}
538
539impl Default for WebSearchTool {
540    fn default() -> Self {
541        Self::new()
542    }
543}
544
545#[async_trait]
546impl Tool for WebSearchTool {
547    fn name(&self) -> &str {
548        "WebSearch"
549    }
550
551    fn description(&self) -> &str {
552        "Search DuckDuckGo and return up to 10 filtered results (title, url, domain, snippet) with optional allow/block domain filters."
553    }
554
555    fn classify(&self, _args: &serde_json::Value) -> ToolClass {
556        ToolClass::READONLY_PARALLEL.promotable()
557    }
558
559    fn parameters_schema(&self) -> serde_json::Value {
560        json!({
561            "type": "object",
562            "properties": {
563                "query": {
564                    "type": "string",
565                    "minLength": 2,
566                    "description": "The search query to use"
567                },
568                "allowed_domains": {
569                    "type": "array",
570                    "items": { "type": "string" },
571                    "description": "Only include results from these domains"
572                },
573                "blocked_domains": {
574                    "type": "array",
575                    "items": { "type": "string" },
576                    "description": "Never include results from these domains"
577                },
578                "max_results": {
579                    "type": "number",
580                    "description": "Maximum results to return (default 10, max 20)"
581                }
582            },
583            "required": ["query"],
584            "additionalProperties": false
585        })
586    }
587
588    async fn invoke(
589        &self,
590        args: serde_json::Value,
591        ctx: ToolCtx,
592    ) -> Result<ToolOutcome, ToolError> {
593        let parsed: WebSearchArgs = serde_json::from_value(args)
594            .map_err(|e| ToolError::InvalidArguments(format!("Invalid WebSearch args: {}", e)))?;
595
596        let query = parsed.query.trim();
597        if query.len() < 2 {
598            return Err(ToolError::InvalidArguments(
599                "query must be at least 2 characters".to_string(),
600            ));
601        }
602
603        let allowed_domains = parsed.allowed_domains.filter(|v| !v.is_empty());
604        let blocked_domains = parsed.blocked_domains.filter(|v| !v.is_empty());
605
606        // Mutual-exclusion validation
607        if allowed_domains.is_some() && blocked_domains.is_some() {
608            return Err(ToolError::InvalidArguments(
609                "Cannot specify both allowed_domains and blocked_domains in the same request"
610                    .to_string(),
611            ));
612        }
613
614        let max_results = parsed
615            .max_results
616            .unwrap_or(DEFAULT_MAX_RESULTS)
617            .min(ABSOLUTE_MAX_RESULTS);
618        let endpoints = Self::configured_endpoints().map_err(ToolError::Execution)?;
619
620        // Check cache
621        let cache_key = Self::cache_key(
622            query,
623            &allowed_domains,
624            &blocked_domains,
625            max_results,
626            &endpoints,
627        );
628        if let Some(cached) = Self::try_cache(&cache_key) {
629            ctx.emit_tool_token("Using cached search results\n").await;
630            return Ok(ToolOutcome::Completed(ToolResult {
631                success: true,
632                result: cached.to_string(),
633                display_preference: Some("Collapsible".to_string()),
634                images: Vec::new(),
635            }));
636        }
637
638        ctx.emit_tool_token(format!("Searching: {}\n", query)).await;
639
640        let client = Self::build_http_client().map_err(ToolError::Execution)?;
641
642        let allowed: Option<HashSet<String>> = allowed_domains.map(|domains| {
643            domains
644                .into_iter()
645                .map(|value| value.to_ascii_lowercase())
646                .collect()
647        });
648        let blocked: HashSet<String> = blocked_domains
649            .unwrap_or_default()
650            .into_iter()
651            .map(|value| value.to_ascii_lowercase())
652            .collect();
653        let results = Self::search_endpoint_chain(
654            &client,
655            &endpoints,
656            query,
657            &allowed,
658            &blocked,
659            max_results,
660        )
661        .await
662        .map_err(ToolError::Execution)?;
663
664        ctx.emit_tool_token(format!(
665            "Found {} results for \"{}\"\n",
666            results.len(),
667            query
668        ))
669        .await;
670
671        let results: Vec<serde_json::Value> =
672            results.into_iter().map(SearchResult::into_json).collect();
673
674        let result_value = if results.is_empty() {
675            json!({
676                "query": parsed.query,
677                "results": [],
678                "note": "No results found for this query.",
679            })
680        } else {
681            json!({
682                "query": parsed.query,
683                "results": results,
684            })
685        };
686
687        // Store in cache
688        Self::put_cache(cache_key, result_value.clone());
689
690        let mut result_string = result_value.to_string();
691        result_string.push_str("\n\nREMINDER: You MUST include a Sources section at the end of your response, listing all relevant URLs as markdown hyperlinks: [Title](URL)");
692
693        Ok(ToolOutcome::Completed(ToolResult {
694            success: true,
695            result: result_string,
696            display_preference: Some("Collapsible".to_string()),
697            images: Vec::new(),
698        }))
699    }
700}
701
702#[cfg(test)]
703mod tests {
704    use super::*;
705    use wiremock::matchers::{method, path};
706    use wiremock::{Mock, MockServer, ResponseTemplate};
707
708    const HTML_FIXTURE: &str = include_str!("fixtures/web_search_html.html");
709    const LITE_FIXTURE: &str = include_str!("fixtures/web_search_lite.html");
710    const ANTI_BOT_FIXTURE: &str = include_str!("fixtures/web_search_antibot.html");
711
712    fn empty_filters() -> (Option<HashSet<String>>, HashSet<String>) {
713        (None, HashSet::new())
714    }
715
716    #[test]
717    fn domain_matches_supports_subdomains() {
718        assert!(WebSearchTool::domain_matches("example.com", "example.com"));
719        assert!(WebSearchTool::domain_matches(
720            "docs.example.com",
721            "example.com"
722        ));
723        assert!(!WebSearchTool::domain_matches(
724            "notexample.com",
725            "example.com"
726        ));
727        assert!(!WebSearchTool::domain_matches(
728            "evil-example.com",
729            "example.com"
730        ));
731    }
732
733    #[test]
734    fn host_of_normalizes_case() {
735        let host = WebSearchTool::host_of("https://Docs.Example.Com/path").unwrap();
736        assert_eq!(host, "docs.example.com");
737    }
738
739    #[test]
740    fn decode_duckduckgo_url_extracts_uddg_param() {
741        let raw = "https://duckduckgo.com/l/?uddg=https%3A%2F%2Fexample.com%2Fpage&rut=whatever";
742        let decoded = WebSearchTool::decode_duckduckgo_url(raw).unwrap();
743        assert_eq!(decoded, "https://example.com/page");
744    }
745
746    #[test]
747    fn decode_duckduckgo_url_handles_root_and_protocol_relative_links() {
748        let root_relative = "/l/?uddg=https%3A%2F%2Fexample.com%2Froot&amp;rut=irrelevant";
749        assert_eq!(
750            WebSearchTool::decode_duckduckgo_url(root_relative).as_deref(),
751            Some("https://example.com/root")
752        );
753
754        let protocol_relative = "//duckduckgo.com/l/?uddg=https%3A%2F%2Fexample.com%2Fprotocol";
755        assert_eq!(
756            WebSearchTool::decode_duckduckgo_url(protocol_relative).as_deref(),
757            Some("https://example.com/protocol")
758        );
759    }
760
761    #[test]
762    fn decode_duckduckgo_url_rejects_non_http_schemes() {
763        for raw in [
764            "ftp://example.com/archive",
765            "file://example.com/private",
766            "javascript://example.com/alert",
767            "/l/?uddg=ftp%3A%2F%2Fexample.com%2Farchive",
768        ] {
769            assert_eq!(WebSearchTool::decode_duckduckgo_url(raw), None, "{raw}");
770        }
771    }
772
773    #[test]
774    fn html_fixture_parses_results_snippets_and_filters() {
775        let (allowed, blocked) = empty_filters();
776        let ParsedSearchPage::Results(results) =
777            WebSearchTool::parse_search_page(HTML_FIXTURE, &allowed, &blocked, 10)
778        else {
779            panic!("HTML fixture should be recognized");
780        };
781        assert_eq!(results.len(), 2);
782        assert_eq!(results[0].title, "Example Guide");
783        assert_eq!(results[0].url, "https://example.com/guide");
784        assert_eq!(
785            results[0].snippet.as_deref(),
786            Some("A useful guide & reference.")
787        );
788
789        let allowed = Some(HashSet::from(["example.com".to_string()]));
790        let ParsedSearchPage::Results(filtered) =
791            WebSearchTool::parse_search_page(HTML_FIXTURE, &allowed, &blocked, 10)
792        else {
793            panic!("HTML fixture should be recognized after filtering");
794        };
795        assert_eq!(filtered.len(), 1);
796        assert_eq!(filtered[0].domain, "example.com");
797    }
798
799    #[test]
800    fn lite_fixture_uses_lite_parser_and_block_filter() {
801        let allowed = None;
802        let blocked = HashSet::from(["example.org".to_string()]);
803        let ParsedSearchPage::Results(results) =
804            WebSearchTool::parse_search_page(LITE_FIXTURE, &allowed, &blocked, 10)
805        else {
806            panic!("Lite fixture should be recognized by cross-layout parsing");
807        };
808        assert_eq!(results.len(), 1);
809        assert_eq!(results[0].url, "https://example.net/lite");
810        assert_eq!(
811            results[0].snippet.as_deref(),
812            Some("A compact result & summary.")
813        );
814    }
815
816    #[test]
817    fn anti_bot_and_unknown_pages_are_hard_endpoint_failures() {
818        let (allowed, blocked) = empty_filters();
819        assert_eq!(
820            WebSearchTool::parse_search_page(ANTI_BOT_FIXTURE, &allowed, &blocked, 10),
821            ParsedSearchPage::AntiBot
822        );
823        assert_eq!(
824            WebSearchTool::parse_search_page(
825                "<html><body>proxy login page</body></html>",
826                &allowed,
827                &blocked,
828                10,
829            ),
830            ParsedSearchPage::Unrecognized
831        );
832    }
833
834    #[test]
835    fn recognized_empty_page_is_a_successful_empty_result() {
836        let (allowed, blocked) = empty_filters();
837        let empty = r#"<html><body class="body--html"><form action="/html/"></form><div class="no-results">No results found</div></body></html>"#;
838        assert_eq!(
839            WebSearchTool::parse_search_page(empty, &allowed, &blocked, 10),
840            ParsedSearchPage::Results(Vec::new())
841        );
842    }
843
844    #[test]
845    fn malformed_result_links_are_unrecognized_but_filtered_http_links_are_valid() {
846        let (allowed, blocked) = empty_filters();
847        let malformed = r#"<html><body><form action="/html/"></form><a class="result__a" href="javascript:void(0)">Broken</a></body></html>"#;
848        assert_eq!(
849            WebSearchTool::parse_search_page(malformed, &allowed, &blocked, 10),
850            ParsedSearchPage::Unrecognized
851        );
852
853        let allowed = Some(HashSet::from(["allowed.example".to_string()]));
854        let filtered = r#"<html><body><form action="/html/"></form><a class="result__a" href="https://blocked.example/result">Filtered</a></body></html>"#;
855        assert_eq!(
856            WebSearchTool::parse_search_page(filtered, &allowed, &blocked, 10),
857            ParsedSearchPage::Results(Vec::new())
858        );
859    }
860
861    #[test]
862    fn endpoint_override_preserves_order_and_rejects_invalid_entries() {
863        let endpoints = WebSearchTool::parse_endpoint_list(Some(
864            " https://mirror.example/search , http://127.0.0.1:8080/lite ",
865        ))
866        .unwrap();
867        assert_eq!(endpoints.len(), 2);
868        assert_eq!(endpoints[0].as_str(), "https://mirror.example/search");
869        assert_eq!(endpoints[1].as_str(), "http://127.0.0.1:8080/lite");
870
871        assert!(WebSearchTool::parse_endpoint_list(Some("  , ")).is_err());
872        assert!(WebSearchTool::parse_endpoint_list(Some("file:///tmp/results.html")).is_err());
873        assert!(WebSearchTool::parse_endpoint_list(Some("not a url")).is_err());
874    }
875
876    #[test]
877    fn default_endpoint_chain_uses_html_then_lite() {
878        let endpoints = WebSearchTool::parse_endpoint_list(None).unwrap();
879        assert_eq!(
880            endpoints.iter().map(Url::as_str).collect::<Vec<_>>(),
881            DEFAULT_WEB_SEARCH_ENDPOINTS
882        );
883    }
884
885    #[tokio::test]
886    async fn endpoint_chain_falls_back_after_an_outage() {
887        let server = MockServer::start().await;
888        Mock::given(method("POST"))
889            .and(path("/unavailable"))
890            .respond_with(ResponseTemplate::new(503))
891            .expect(1)
892            .mount(&server)
893            .await;
894        Mock::given(method("POST"))
895            .and(path("/lite"))
896            .respond_with(ResponseTemplate::new(200).set_body_string(LITE_FIXTURE))
897            .expect(1)
898            .mount(&server)
899            .await;
900
901        let endpoints = vec![
902            Url::parse(&format!("{}/unavailable", server.uri())).unwrap(),
903            Url::parse(&format!("{}/lite", server.uri())).unwrap(),
904        ];
905        let client = WebSearchTool::build_http_client().unwrap();
906        let (allowed, blocked) = empty_filters();
907        let results = WebSearchTool::search_endpoint_chain(
908            &client, &endpoints, "rust", &allowed, &blocked, 10,
909        )
910        .await
911        .unwrap();
912
913        assert_eq!(results.len(), 2);
914        let requests = server.received_requests().await.unwrap();
915        assert_eq!(requests.len(), 2);
916        for request in requests {
917            assert_eq!(request.method.as_str(), "POST");
918            assert_eq!(String::from_utf8(request.body).unwrap(), "q=rust");
919        }
920    }
921
922    #[tokio::test]
923    async fn endpoint_chain_falls_back_after_anti_bot_response() {
924        let server = MockServer::start().await;
925        Mock::given(method("POST"))
926            .and(path("/blocked"))
927            .respond_with(ResponseTemplate::new(200).set_body_string(ANTI_BOT_FIXTURE))
928            .expect(1)
929            .mount(&server)
930            .await;
931        Mock::given(method("POST"))
932            .and(path("/html"))
933            .respond_with(ResponseTemplate::new(200).set_body_string(HTML_FIXTURE))
934            .expect(1)
935            .mount(&server)
936            .await;
937
938        let endpoints = vec![
939            Url::parse(&format!("{}/blocked", server.uri())).unwrap(),
940            Url::parse(&format!("{}/html", server.uri())).unwrap(),
941        ];
942        let client = WebSearchTool::build_http_client().unwrap();
943        let (allowed, blocked) = empty_filters();
944        let results = WebSearchTool::search_endpoint_chain(
945            &client, &endpoints, "rust", &allowed, &blocked, 10,
946        )
947        .await
948        .unwrap();
949        assert_eq!(results.len(), 2);
950    }
951
952    #[tokio::test]
953    async fn endpoint_chain_falls_back_after_malformed_result_links() {
954        let server = MockServer::start().await;
955        let malformed = r#"<html><body><form action="/html/"></form><a class="result__a" href="javascript:void(0)">Broken</a></body></html>"#;
956        Mock::given(method("POST"))
957            .and(path("/malformed"))
958            .respond_with(ResponseTemplate::new(200).set_body_string(malformed))
959            .expect(1)
960            .mount(&server)
961            .await;
962        Mock::given(method("POST"))
963            .and(path("/lite"))
964            .respond_with(ResponseTemplate::new(200).set_body_string(LITE_FIXTURE))
965            .expect(1)
966            .mount(&server)
967            .await;
968
969        let endpoints = vec![
970            Url::parse(&format!("{}/malformed", server.uri())).unwrap(),
971            Url::parse(&format!("{}/lite", server.uri())).unwrap(),
972        ];
973        let client = WebSearchTool::build_http_client().unwrap();
974        let (allowed, blocked) = empty_filters();
975        let results = WebSearchTool::search_endpoint_chain(
976            &client, &endpoints, "rust", &allowed, &blocked, 10,
977        )
978        .await
979        .unwrap();
980
981        assert_eq!(results[0].url, "https://example.net/lite");
982    }
983
984    #[tokio::test]
985    async fn endpoint_chain_reports_all_hard_failures() {
986        let server = MockServer::start().await;
987        Mock::given(method("POST"))
988            .and(path("/unavailable"))
989            .respond_with(ResponseTemplate::new(503))
990            .mount(&server)
991            .await;
992        Mock::given(method("POST"))
993            .and(path("/unknown"))
994            .respond_with(ResponseTemplate::new(200).set_body_string("not a search page"))
995            .mount(&server)
996            .await;
997
998        let endpoints = vec![
999            Url::parse(&format!("{}/unavailable", server.uri())).unwrap(),
1000            Url::parse(&format!("{}/unknown", server.uri())).unwrap(),
1001        ];
1002        let client = WebSearchTool::build_http_client().unwrap();
1003        let (allowed, blocked) = empty_filters();
1004        let error = WebSearchTool::search_endpoint_chain(
1005            &client, &endpoints, "rust", &allowed, &blocked, 10,
1006        )
1007        .await
1008        .unwrap_err();
1009
1010        assert!(error.contains("all configured web search endpoints failed"));
1011        assert!(error.contains("HTTP 503"));
1012        assert!(error.contains("unrecognized search response"));
1013    }
1014
1015    #[tokio::test]
1016    async fn endpoint_chain_does_not_follow_redirects_before_fallback() {
1017        let server = MockServer::start().await;
1018        Mock::given(method("POST"))
1019            .and(path("/redirect"))
1020            .respond_with(
1021                ResponseTemplate::new(307)
1022                    .insert_header("Location", format!("{}/trap", server.uri())),
1023            )
1024            .expect(1)
1025            .mount(&server)
1026            .await;
1027        Mock::given(method("POST"))
1028            .and(path("/lite"))
1029            .respond_with(ResponseTemplate::new(200).set_body_string(LITE_FIXTURE))
1030            .expect(1)
1031            .mount(&server)
1032            .await;
1033
1034        let endpoints = vec![
1035            Url::parse(&format!("{}/redirect", server.uri())).unwrap(),
1036            Url::parse(&format!("{}/lite", server.uri())).unwrap(),
1037        ];
1038        let client = WebSearchTool::build_http_client().unwrap();
1039        let (allowed, blocked) = empty_filters();
1040        let results = WebSearchTool::search_endpoint_chain(
1041            &client, &endpoints, "rust", &allowed, &blocked, 10,
1042        )
1043        .await
1044        .unwrap();
1045
1046        assert_eq!(results[0].url, "https://example.net/lite");
1047        let paths = server
1048            .received_requests()
1049            .await
1050            .unwrap()
1051            .into_iter()
1052            .map(|request| request.url.path().to_string())
1053            .collect::<Vec<_>>();
1054        assert_eq!(paths, ["/redirect", "/lite"]);
1055    }
1056
1057    #[test]
1058    fn cache_key_is_stable_and_isolates_result_limits_and_endpoints() {
1059        let endpoints = WebSearchTool::parse_endpoint_list(None).unwrap();
1060        let allowed = Some(vec!["doc.rust-lang.org".to_string()]);
1061        let k1 = WebSearchTool::cache_key("rust", &allowed, &None, 10, &endpoints);
1062        let k2 = WebSearchTool::cache_key("rust", &allowed, &None, 10, &endpoints);
1063        assert_eq!(k1, k2);
1064
1065        let k3 = WebSearchTool::cache_key(
1066            "rust",
1067            &None,
1068            &Some(vec!["bad.com".to_string()]),
1069            10,
1070            &endpoints,
1071        );
1072        assert_ne!(k1, k3);
1073
1074        let fewer_results = WebSearchTool::cache_key("rust", &allowed, &None, 1, &endpoints);
1075        assert_ne!(k1, fewer_results);
1076
1077        let mirror =
1078            WebSearchTool::parse_endpoint_list(Some("https://mirror.example/search")).unwrap();
1079        let other_endpoint = WebSearchTool::cache_key("rust", &allowed, &None, 10, &mirror);
1080        assert_ne!(k1, other_endpoint);
1081    }
1082}