mcp_servers/coding/tools/web_search/
mod.rs

1pub mod search_client;
2
3pub use search_client::{BraveSearchClient, RawSearchResult, SearchClient, SearchParams};
4
5#[cfg(test)]
6pub use search_client::FakeSearchClient;
7
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10use std::collections::HashSet;
11
12use crate::coding::error::WebSearchError;
13use mcp_utils::display_meta::{ToolDisplayMeta, ToolResultMeta};
14
15const DEFAULT_COUNT: u32 = 10;
16const MAX_COUNT: u32 = 20;
17
18/// Input parameters for `web_search` tool
19#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
20#[serde(rename_all = "camelCase")]
21pub struct WebSearchInput {
22    /// The search query
23    pub query: String,
24
25    /// Maximum number of results (default: 10, max: 20)
26    pub count: Option<u32>,
27
28    /// Only include results from these domains
29    #[serde(alias = "allowed_domains")]
30    pub allowed_domains: Option<Vec<String>>,
31
32    /// Exclude results from these domains
33    #[serde(alias = "blocked_domains")]
34    pub blocked_domains: Option<Vec<String>>,
35}
36
37/// Output from `web_search` tool
38#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
39#[serde(rename_all = "camelCase")]
40pub struct WebSearchOutput {
41    /// Search results
42    pub results: Vec<SearchResult>,
43
44    /// The original query
45    pub query: String,
46
47    /// Display metadata for human-friendly rendering
48    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
49    #[schemars(skip)]
50    pub meta: Option<ToolResultMeta>,
51}
52
53/// Individual search result
54#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
55#[serde(rename_all = "camelCase")]
56pub struct SearchResult {
57    /// Result title
58    pub title: String,
59
60    /// Result URL
61    pub url: String,
62
63    /// Result snippet/description
64    pub snippet: String,
65}
66
67/// Web searcher that performs searches and filters results
68#[derive(Debug, Clone)]
69pub struct WebSearcher<C: SearchClient> {
70    client: C,
71}
72
73impl<C: SearchClient> WebSearcher<C> {
74    /// Creates a new `WebSearcher` with the given client
75    pub fn with_client(client: C) -> Self {
76        Self { client }
77    }
78
79    /// Performs a web search with the given parameters
80    pub async fn search(&self, args: WebSearchInput) -> Result<WebSearchOutput, WebSearchError> {
81        let query = args.query.trim();
82
83        if query.is_empty() {
84            return Err(WebSearchError::InvalidQuery("Search query cannot be empty".to_string()));
85        }
86
87        let count = args.count.unwrap_or(DEFAULT_COUNT).min(MAX_COUNT);
88
89        let params = SearchParams { query: query.to_string(), count };
90
91        let mut results = self.client.search(params).await?;
92
93        // Ensure we don't return more than requested
94        if results.len() > count as usize {
95            results.truncate(count as usize);
96        }
97
98        // Apply domain filters if provided
99        if let Some(ref allowed) = args.allowed_domains {
100            results = filter_allowed_domains(results, allowed);
101        }
102
103        if let Some(ref blocked) = args.blocked_domains {
104            results = filter_blocked_domains(results, blocked);
105        }
106
107        // Convert to output format
108        let results: Vec<SearchResult> = results
109            .into_iter()
110            .map(|r| SearchResult { title: r.title, url: r.url.clone(), snippet: r.description })
111            .collect();
112
113        let display_meta = ToolDisplayMeta::new("Web search", format!("'{}' ({} results)", query, results.len()));
114
115        Ok(WebSearchOutput { results, query: query.to_string(), meta: Some(display_meta.into()) })
116    }
117}
118
119impl WebSearcher<BraveSearchClient> {
120    /// Creates a new `WebSearcher` with Brave Search API
121    pub fn try_new() -> Result<Self, WebSearchError> {
122        let client = BraveSearchClient::new()?;
123        Ok(Self::with_client(client))
124    }
125
126    /// Creates a new `WebSearcher` with a custom Brave API key
127    pub fn with_api_key(api_key: String) -> Self {
128        let client = BraveSearchClient::with_api_key(api_key);
129        Self::with_client(client)
130    }
131}
132
133impl Default for WebSearcher<BraveSearchClient> {
134    fn default() -> Self {
135        Self::with_api_key(std::env::var("BRAVE_SEARCH_API_KEY").expect("BRAVE_SEARCH_API_KEY must be set"))
136    }
137}
138
139fn filter_allowed_domains(results: Vec<RawSearchResult>, allowed_domains: &[String]) -> Vec<RawSearchResult> {
140    let allowed_set: HashSet<String> = allowed_domains.iter().cloned().collect();
141
142    results
143        .into_iter()
144        .filter(|r| {
145            if let Some(host) = extract_domain(&r.url) {
146                return allowed_set.iter().any(|domain| host.ends_with(domain));
147            }
148            false
149        })
150        .collect()
151}
152
153fn filter_blocked_domains(results: Vec<RawSearchResult>, blocked_domains: &[String]) -> Vec<RawSearchResult> {
154    let blocked_set: HashSet<String> = blocked_domains.iter().cloned().collect();
155
156    results
157        .into_iter()
158        .filter(|r| {
159            if let Some(host) = extract_domain(&r.url) {
160                return !blocked_set.iter().any(|domain| host.ends_with(domain));
161            }
162            true
163        })
164        .collect()
165}
166
167/// Extract domain/hostname from a URL
168fn extract_domain(url: &str) -> Option<String> {
169    // Remove protocol
170    let url = url.strip_prefix("http://").or_else(|| url.strip_prefix("https://")).unwrap_or(url);
171
172    // Remove path and query string
173    let domain = url.split('/').next()?;
174
175    // Remove port if present
176    let domain = domain.split(':').next()?;
177
178    Some(domain.to_string())
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    fn make_result(url: &str) -> RawSearchResult {
186        RawSearchResult { title: "Test".to_string(), url: url.to_string(), description: "Test description".to_string() }
187    }
188
189    fn input(query: &str) -> WebSearchInput {
190        WebSearchInput { query: query.to_string(), count: None, allowed_domains: None, blocked_domains: None }
191    }
192
193    fn strs(vals: &[&str]) -> Vec<String> {
194        vals.iter().map(std::string::ToString::to_string).collect()
195    }
196
197    fn searcher_with(query: &str, results: Vec<RawSearchResult>) -> WebSearcher<FakeSearchClient> {
198        WebSearcher::with_client(FakeSearchClient::new().with_results(query, results))
199    }
200
201    #[test]
202    fn test_domain_filtering() {
203        let urls = ["https://example.com/page", "https://test.org/page", "https://sub.example.com/page"];
204        let results: Vec<_> = urls.iter().map(|u| make_result(u)).collect();
205
206        // Allowed domains: keeps example.com and sub.example.com
207        let filtered = filter_allowed_domains(results.clone(), &strs(&["example.com"]));
208        assert_eq!(filtered.len(), 2);
209        assert!(filtered.iter().all(|r| r.url.contains("example.com")));
210
211        // Blocked domains: removes example.com
212        let results2 = vec![
213            make_result("https://example.com/page"),
214            make_result("https://test.org/page"),
215            make_result("https://allowed.com/page"),
216        ];
217        let filtered = filter_blocked_domains(results2, &strs(&["example.com"]));
218        assert_eq!(filtered.len(), 2);
219        assert!(!filtered.iter().any(|r| r.url.contains("example.com")));
220
221        // Combined: allow example.com + test.org, then block test.org => only example.com
222        let results3: Vec<_> = urls.iter().map(|u| make_result(u)).collect();
223        let filtered = filter_allowed_domains(results3, &strs(&["example.com", "test.org"]));
224        let filtered = filter_blocked_domains(filtered, &strs(&["test.org"]));
225        assert_eq!(filtered.len(), 2);
226        assert!(filtered.iter().all(|r| r.url.contains("example.com")));
227    }
228
229    #[tokio::test]
230    async fn test_empty_query_returns_error() {
231        let searcher = WebSearcher::with_client(FakeSearchClient::new());
232        let result = searcher.search(input("   ")).await;
233        assert!(matches!(result, Err(WebSearchError::InvalidQuery(_))));
234    }
235
236    #[tokio::test]
237    async fn test_count_limiting() {
238        let many_results: Vec<_> = (0..30)
239            .map(|i| RawSearchResult {
240                title: format!("Result {i}"),
241                url: format!("https://example.com/{i}"),
242                description: format!("Description {i}"),
243            })
244            .collect();
245        let searcher = searcher_with("test query", many_results);
246
247        let mut q = input("test query");
248        q.count = Some(15);
249        let output = searcher.search(q).await.unwrap();
250        assert_eq!(output.results.len(), 15);
251    }
252
253    #[tokio::test]
254    async fn test_domain_filtering_in_web_searcher() {
255        let searcher = searcher_with(
256            "test query",
257            vec![
258                make_result("https://example.com/page1"),
259                make_result("https://blocked.com/page2"),
260                make_result("https://example.com/page3"),
261            ],
262        );
263
264        let mut q = input("test query");
265        q.blocked_domains = Some(strs(&["blocked.com"]));
266        let output = searcher.search(q).await.unwrap();
267        assert_eq!(output.results.len(), 2);
268        assert!(!output.results.iter().any(|r| r.url.contains("blocked.com")));
269    }
270
271    #[tokio::test]
272    async fn test_searcher_tracks_results() {
273        let searcher = searcher_with(
274            "rust programming",
275            vec![RawSearchResult {
276                title: "Rust Language".to_string(),
277                url: "https://rust-lang.org".to_string(),
278                description: "A systems language".to_string(),
279            }],
280        );
281        let output = searcher.search(input("rust programming")).await.unwrap();
282        assert_eq!(output.results.len(), 1);
283        assert_eq!(output.results[0].title, "Rust Language");
284        assert_eq!(output.results[0].snippet, "A systems language");
285    }
286
287    #[test]
288    fn web_search_input_accepts_snake_case_domains() {
289        let args: WebSearchInput = serde_json::from_value(serde_json::json!({
290            "query": "rust",
291            "count": 5,
292            "allowed_domains": ["docs.rs"],
293            "blocked_domains": ["example.com"]
294        }))
295        .unwrap();
296
297        assert_eq!(args.query, "rust");
298        assert_eq!(args.count, Some(5));
299        assert_eq!(args.allowed_domains, Some(vec!["docs.rs".to_string()]));
300        assert_eq!(args.blocked_domains, Some(vec!["example.com".to_string()]));
301    }
302}
mcp_servers/coding/tools/web_search/mod.rs

mcp_servers/coding/tools/web_search/
mod.rs
