mcp_servers/coding/tools/web_search/
mod.rs

1pub mod search_client;
2
3pub use search_client::{BraveSearchClient, SearchClient, SearchParams};
4
5#[cfg(test)]
6pub use search_client::FakeSearchClient;
7
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10use std::future::Future;
11use std::time::Duration;
12use url::Url;
13
14use crate::coding::error::WebSearchError;
15
16const DEFAULT_COUNT: u32 = 10;
17const MAX_COUNT: u32 = 20;
18const MAX_RATE_LIMIT_RETRIES: u32 = 3;
19const MAX_RETRY_AFTER: Duration = Duration::from_mins(1);
20const DEFAULT_RETRY_DELAY: Duration = Duration::from_secs(2);
21
22/// Input parameters for `web_search` tool
23#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
24#[serde(rename_all = "camelCase")]
25pub struct WebSearchInput {
26    /// The search query
27    pub query: String,
28
29    /// Maximum number of results (default: 10, max: 20)
30    pub count: Option<u32>,
31
32    /// Only include results from these domains
33    #[serde(alias = "allowed_domains")]
34    pub allowed_domains: Option<Vec<String>>,
35
36    /// Exclude results from these domains
37    #[serde(alias = "blocked_domains")]
38    pub blocked_domains: Option<Vec<String>>,
39}
40
41/// Output from `web_search` tool
42#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
43#[serde(rename_all = "camelCase")]
44pub struct WebSearchOutput {
45    /// Search results
46    pub results: Vec<SearchResult>,
47
48    /// The original query
49    pub query: String,
50}
51
52/// Individual search result
53#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
54#[serde(rename_all = "camelCase")]
55pub struct SearchResult {
56    /// Result title
57    pub title: String,
58
59    /// Result URL
60    pub url: String,
61
62    /// Result snippet/description
63    #[serde(alias = "description")]
64    pub snippet: String,
65}
66
67/// Web searcher that performs searches and filters results
68pub struct WebSearcher {
69    client: Box<dyn SearchClient>,
70}
71
72impl WebSearcher {
73    pub fn with_client(client: impl SearchClient + 'static) -> Self {
74        Self { client: Box::new(client) }
75    }
76
77    pub fn try_new() -> Result<Self, WebSearchError> {
78        Ok(Self::with_client(BraveSearchClient::new()?))
79    }
80
81    /// Performs a web search with the given parameters
82    pub async fn search<T, U>(&self, args: WebSearchInput, mut on_retry: T) -> Result<WebSearchOutput, WebSearchError>
83    where
84        T: FnMut(Duration) -> U,
85        U: Future<Output = ()>,
86    {
87        let query = args.query.trim();
88
89        if query.is_empty() {
90            return Err(WebSearchError::InvalidQuery("Search query cannot be empty".to_string()));
91        }
92
93        let count = args.count.unwrap_or(DEFAULT_COUNT).min(MAX_COUNT);
94
95        let query_owned = query.to_string();
96        let mut retries = 0;
97        let mut results = loop {
98            let params = SearchParams { query: query_owned.clone(), count };
99
100            match self.client.search(params).await {
101                Ok(results) => break results,
102                Err(WebSearchError::RateLimited { retry_after, .. }) if retries < MAX_RATE_LIMIT_RETRIES => {
103                    let delay = retry_after.unwrap_or(DEFAULT_RETRY_DELAY).min(MAX_RETRY_AFTER);
104                    retries += 1;
105                    on_retry(delay).await;
106                    tokio::time::sleep(delay).await;
107                }
108                Err(error) => return Err(error),
109            }
110        };
111
112        results.truncate(count as usize);
113        results.retain(|result| {
114            let url = Url::parse(&result.url).ok();
115            let host = url.as_ref().and_then(Url::host_str);
116            let in_any =
117                |domains: &[String]| host.is_some_and(|host| domains.iter().any(|d| host.ends_with(d.as_str())));
118            args.allowed_domains.as_deref().is_none_or(in_any) && !args.blocked_domains.as_deref().is_some_and(in_any)
119        });
120
121        Ok(WebSearchOutput { results, query: query.to_string() })
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use std::sync::{Arc, Mutex};
128    use std::time::Duration;
129
130    use super::*;
131
132    fn make_result(url: &str) -> SearchResult {
133        SearchResult { title: "Test".to_string(), url: url.to_string(), snippet: "Test description".to_string() }
134    }
135
136    fn input(query: &str) -> WebSearchInput {
137        WebSearchInput { query: query.to_string(), count: None, allowed_domains: None, blocked_domains: None }
138    }
139
140    fn strs(vals: &[&str]) -> Vec<String> {
141        vals.iter().map(std::string::ToString::to_string).collect()
142    }
143
144    fn searcher_with(query: &str, results: Vec<SearchResult>) -> WebSearcher {
145        WebSearcher::with_client(FakeSearchClient::default().with_results(query, results))
146    }
147
148    fn rate_limited(retry_after: Option<Duration>) -> WebSearchError {
149        WebSearchError::RateLimited { message: "too many requests".to_string(), retry_after }
150    }
151
152    #[tokio::test]
153    async fn test_domain_filtering() {
154        let urls = ["https://example.com/page", "https://test.org/page", "https://sub.example.com:8080/page"];
155        let cases = [
156            (Some(strs(&["example.com"])), None, vec![urls[0], urls[2]]),
157            (None, Some(strs(&["example.com"])), vec![urls[1]]),
158            (Some(strs(&["example.com", "test.org"])), Some(strs(&["test.org"])), vec![urls[0], urls[2]]),
159        ];
160
161        for (allowed_domains, blocked_domains, expected) in cases {
162            let searcher = searcher_with("rust", urls.iter().map(|url| make_result(url)).collect());
163            let query = WebSearchInput { allowed_domains, blocked_domains, ..input("rust") };
164            let output = searcher.search(query, |_| async {}).await.unwrap();
165            let urls: Vec<_> = output.results.iter().map(|r| r.url.as_str()).collect();
166            assert_eq!(urls, expected);
167        }
168    }
169
170    #[tokio::test]
171    async fn test_empty_query_returns_error() {
172        let searcher = WebSearcher::with_client(FakeSearchClient::default());
173        let result = searcher.search(input("   "), |_| async {}).await;
174        assert!(matches!(result, Err(WebSearchError::InvalidQuery(_))));
175    }
176
177    #[tokio::test(start_paused = true)]
178    async fn retries_rate_limited_search_after_retry_after_period() {
179        let client = FakeSearchClient::default()
180            .with_results("rust", vec![make_result("https://example.com")])
181            .with_failures("rust", [rate_limited(Some(Duration::ZERO))]);
182        let searcher = WebSearcher::with_client(client);
183
184        let retry_delays = Arc::new(Mutex::new(Vec::new()));
185        let notified_delays = retry_delays.clone();
186        let output = searcher
187            .search(input("rust"), move |delay| {
188                let notified_delays = notified_delays.clone();
189                async move {
190                    notified_delays.lock().unwrap().push(delay);
191                }
192            })
193            .await
194            .unwrap();
195
196        assert_eq!(output.results.len(), 1);
197        assert_eq!(*retry_delays.lock().unwrap(), vec![Duration::ZERO]);
198    }
199
200    #[tokio::test(start_paused = true)]
201    async fn retries_rate_limited_search_when_retry_after_header_is_missing() {
202        let client = FakeSearchClient::default()
203            .with_results("rust", vec![make_result("https://example.com")])
204            .with_failures("rust", [rate_limited(None)]);
205        let searcher = WebSearcher::with_client(client);
206
207        let retry_delays = Arc::new(Mutex::new(Vec::new()));
208        let notified_delays = retry_delays.clone();
209        let output = searcher
210            .search(input("rust"), move |delay| {
211                let notified_delays = notified_delays.clone();
212                async move {
213                    notified_delays.lock().unwrap().push(delay);
214                }
215            })
216            .await
217            .unwrap();
218
219        assert_eq!(output.results.len(), 1);
220        assert_eq!(retry_delays.lock().unwrap().len(), 1, "should have retried once");
221    }
222
223    #[tokio::test(start_paused = true)]
224    async fn returns_rate_limited_error_after_retry_budget_is_exhausted() {
225        let mut failures = vec![rate_limited(Some(Duration::ZERO)); 4];
226        failures.push(WebSearchError::ApiError("retry budget was exceeded".to_string()));
227        let searcher = WebSearcher::with_client(FakeSearchClient::default().with_failures("rust", failures));
228        let result = searcher.search(input("rust"), |_| async {}).await;
229
230        assert!(matches!(result, Err(WebSearchError::RateLimited { .. })));
231    }
232
233    #[tokio::test]
234    async fn test_count_limiting() {
235        let searcher =
236            searcher_with("test query", (0..30).map(|i| make_result(&format!("https://example.com/{i}"))).collect());
237
238        let output = searcher.search(WebSearchInput { count: Some(15), ..input("test query") }, |_| async {}).await;
239        assert_eq!(output.unwrap().results.len(), 15);
240    }
241
242    #[test]
243    fn search_result_deserializes_brave_description_as_snippet() {
244        let result: SearchResult = serde_json::from_value(serde_json::json!({
245            "title": "Rust Language",
246            "url": "https://rust-lang.org",
247            "description": "A systems language"
248        }))
249        .unwrap();
250
251        assert_eq!(result.snippet, "A systems language");
252    }
253
254    #[test]
255    fn web_search_input_accepts_snake_case_domains() {
256        let args: WebSearchInput = serde_json::from_value(serde_json::json!({
257            "query": "rust",
258            "count": 5,
259            "allowed_domains": ["docs.rs"],
260            "blocked_domains": ["example.com"]
261        }))
262        .unwrap();
263
264        assert_eq!(args.query, "rust");
265        assert_eq!(args.count, Some(5));
266        assert_eq!(args.allowed_domains, Some(vec!["docs.rs".to_string()]));
267        assert_eq!(args.blocked_domains, Some(vec!["example.com".to_string()]));
268    }
269}
mcp_servers/coding/tools/web_search/mod.rs

mcp_servers/coding/tools/web_search/
mod.rs
