harness/tools/
web_search.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use async_trait::async_trait;
5use futures::StreamExt;
6use serde::{Deserialize, Serialize};
7use serde_json::{json, Value};
8use tokio_util::sync::CancellationToken;
9
10use crate::tools::{
11    invalid_input_failure, ToolFailure, ToolFailureKind, ToolInvocation, ToolOutcome, ToolRuntime,
12    ToolRuntimeError, ToolSpec,
13};
14
15const DEFAULT_RESULT_COUNT: usize = 5;
16const MAX_RESULT_COUNT: usize = 10;
17const MAX_TITLE_CHARS: usize = 500;
18const MAX_SNIPPET_CHARS: usize = 2_000;
19const MAX_PROVIDER_RESPONSE_BYTES: usize = 2 * 1024 * 1024;
20const DEFAULT_TIMEOUT: Duration = Duration::from_secs(20);
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct WebSearchRequest {
24    pub query: String,
25    pub count: usize,
26}
27
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29pub struct WebSearchResult {
30    pub title: String,
31    pub url: String,
32    pub snippet: String,
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub published_at: Option<String>,
35}
36
37#[derive(Debug, thiserror::Error)]
38pub enum WebSearchProviderError {
39    #[error("authentication failed: {0}")]
40    Auth(String),
41    #[error("request timed out: {0}")]
42    Timeout(String),
43    #[error("provider request failed: {0}")]
44    Request(String),
45    #[error("provider returned an invalid response: {0}")]
46    InvalidResponse(String),
47    #[error("cancelled")]
48    Cancelled,
49}
50
51#[async_trait]
52pub trait WebSearchProvider: Send + Sync {
53    fn id(&self) -> &str;
54
55    async fn search(
56        &self,
57        request: WebSearchRequest,
58        cancel: Option<&CancellationToken>,
59    ) -> Result<Vec<WebSearchResult>, WebSearchProviderError>;
60}
61
62/// A standalone managed `web_search` tool. Compose it with filesystem or MCP
63/// runtimes through `CompositeToolRuntime`; it is intentionally not part of
64/// `builtin_tool_specs()` because search credentials are optional.
65#[derive(Clone)]
66pub struct WebSearchToolRuntime {
67    provider: Arc<dyn WebSearchProvider>,
68}
69
70impl WebSearchToolRuntime {
71    pub fn new(provider: Arc<dyn WebSearchProvider>) -> Self {
72        Self { provider }
73    }
74
75    pub fn from_provider(provider: impl WebSearchProvider + 'static) -> Self {
76        Self::new(Arc::new(provider))
77    }
78}
79
80#[async_trait]
81impl ToolRuntime for WebSearchToolRuntime {
82    fn specs(&self) -> Vec<ToolSpec> {
83        vec![web_search_spec()]
84    }
85
86    async fn invoke(&self, invocation: ToolInvocation) -> Result<ToolOutcome, ToolRuntimeError> {
87        self.invoke_cancellable(invocation, None).await
88    }
89
90    async fn invoke_cancellable(
91        &self,
92        invocation: ToolInvocation,
93        cancel: Option<&CancellationToken>,
94    ) -> Result<ToolOutcome, ToolRuntimeError> {
95        if invocation.name != "web_search" {
96            return Err(ToolRuntimeError::UnknownTool(invocation.name));
97        }
98
99        let request = match parse_request(&invocation) {
100            Ok(request) => request,
101            Err(failure) => {
102                return Ok(ToolOutcome {
103                    output: Err(failure),
104                    attachments: vec![],
105                });
106            }
107        };
108        let query = request.query.clone();
109        let requested_count = request.count;
110        let provider = self.provider.id().to_string();
111        let outcome = self.provider.search(request, cancel).await;
112        let output = match outcome {
113            Ok(results) => {
114                let mut truncated = results.len() > requested_count;
115                let results = results
116                    .into_iter()
117                    .take(requested_count)
118                    .filter_map(|result| normalize_result(result, &mut truncated))
119                    .collect::<Vec<_>>();
120                Ok(json!({
121                    "query": query,
122                    "provider": provider,
123                    "results": results,
124                    "count": results.len(),
125                    "truncated": truncated,
126                    "external_content": {
127                        "untrusted": true,
128                        "source": "web_search"
129                    }
130                }))
131            }
132            Err(error) => Err(provider_failure(error)),
133        };
134        Ok(ToolOutcome {
135            output,
136            attachments: vec![],
137        })
138    }
139}
140
141pub fn web_search_spec() -> ToolSpec {
142    ToolSpec {
143        name: "web_search".into(),
144        description: "Search the public web and return current titles, URLs, and snippets. Use web_fetch to read a selected result in full.".into(),
145        input_schema: json!({
146            "type": "object",
147            "properties": {
148                "query": {
149                    "type": "string",
150                    "description": "Search query."
151                },
152                "count": {
153                    "type": "integer",
154                    "description": "Number of results to return (default 5, maximum 10).",
155                    "minimum": 1,
156                    "maximum": MAX_RESULT_COUNT
157                }
158            },
159            "required": ["query"],
160            "additionalProperties": false
161        }),
162    }
163}
164
165fn parse_request(invocation: &ToolInvocation) -> Result<WebSearchRequest, ToolFailure> {
166    let query = invocation
167        .input
168        .get("query")
169        .and_then(Value::as_str)
170        .map(str::trim)
171        .filter(|query| !query.is_empty())
172        .ok_or_else(|| {
173            invalid(
174                invocation,
175                "missing required non-empty string field `query`",
176            )
177        })?;
178    let count = match invocation.input.get("count") {
179        Some(value) => value
180            .as_u64()
181            .and_then(|count| usize::try_from(count).ok())
182            .ok_or_else(|| invalid(invocation, "count must be an integer from 1 to 10"))?,
183        None => DEFAULT_RESULT_COUNT,
184    };
185    if !(1..=MAX_RESULT_COUNT).contains(&count) {
186        return Err(invalid(invocation, "count must be an integer from 1 to 10"));
187    }
188    Ok(WebSearchRequest {
189        query: query.to_string(),
190        count,
191    })
192}
193
194fn invalid(invocation: &ToolInvocation, message: &str) -> ToolFailure {
195    ToolFailure::new(
196        ToolFailureKind::InvalidInput,
197        invalid_input_failure("web_search", message, &invocation.input, None).message,
198    )
199}
200
201fn provider_failure(error: WebSearchProviderError) -> ToolFailure {
202    let kind = match error {
203        WebSearchProviderError::Timeout(_) => ToolFailureKind::Timeout,
204        WebSearchProviderError::Cancelled => ToolFailureKind::Runtime,
205        WebSearchProviderError::Auth(_)
206        | WebSearchProviderError::Request(_)
207        | WebSearchProviderError::InvalidResponse(_) => ToolFailureKind::Runtime,
208    };
209    ToolFailure::new(kind, error.to_string())
210}
211
212fn normalize_result(mut result: WebSearchResult, truncated: &mut bool) -> Option<WebSearchResult> {
213    let parsed = match reqwest::Url::parse(result.url.trim()) {
214        Ok(parsed) => parsed,
215        Err(_) => {
216            *truncated = true;
217            return None;
218        }
219    };
220    if !matches!(parsed.scheme(), "http" | "https") {
221        *truncated = true;
222        return None;
223    }
224    result.url = parsed.to_string();
225    result.title = wrap_untrusted(&truncate_chars(&result.title, MAX_TITLE_CHARS, truncated));
226    result.snippet = wrap_untrusted(&truncate_chars(
227        &result.snippet,
228        MAX_SNIPPET_CHARS,
229        truncated,
230    ));
231    result.published_at = result
232        .published_at
233        .as_deref()
234        .map(|value| wrap_untrusted(&truncate_chars(value, 100, truncated)));
235    if result.title.is_empty() && result.snippet.is_empty() {
236        return None;
237    }
238    Some(result)
239}
240
241fn wrap_untrusted(value: &str) -> String {
242    if value.is_empty() {
243        return String::new();
244    }
245    // Prevent provider text from forging our own boundary markers.
246    let escaped = value.replace("<<<", "< < <").replace(">>>", "> > >");
247    format!(
248        "<<<EXTERNAL_UNTRUSTED_CONTENT source=\"web_search\">>>\n{escaped}\n<<<END_EXTERNAL_UNTRUSTED_CONTENT>>>"
249    )
250}
251
252fn truncate_chars(value: &str, max: usize, truncated: &mut bool) -> String {
253    if value.chars().count() <= max {
254        return value.to_string();
255    }
256    *truncated = true;
257    value.chars().take(max).collect()
258}
259
260#[derive(Clone)]
261pub struct BraveSearchConfig {
262    pub api_key: String,
263    pub timeout: Duration,
264}
265
266impl BraveSearchConfig {
267    pub fn new(api_key: impl Into<String>) -> Self {
268        Self {
269            api_key: api_key.into(),
270            timeout: DEFAULT_TIMEOUT,
271        }
272    }
273}
274
275#[derive(Clone)]
276pub struct BraveSearchProvider {
277    http: reqwest::Client,
278    config: BraveSearchConfig,
279}
280
281impl BraveSearchProvider {
282    pub fn new(config: BraveSearchConfig) -> Self {
283        let http = reqwest::Client::builder()
284            .connect_timeout(Duration::from_secs(10))
285            .build()
286            .unwrap_or_else(|_| reqwest::Client::new());
287        Self { http, config }
288    }
289}
290
291#[derive(Deserialize)]
292struct BraveResponse {
293    #[serde(default)]
294    web: Option<BraveWeb>,
295}
296
297#[derive(Deserialize)]
298struct BraveWeb {
299    #[serde(default)]
300    results: Vec<BraveResult>,
301}
302
303#[derive(Deserialize)]
304struct BraveResult {
305    #[serde(default)]
306    title: String,
307    #[serde(default)]
308    url: String,
309    #[serde(default)]
310    description: String,
311    age: Option<String>,
312}
313
314#[async_trait]
315impl WebSearchProvider for BraveSearchProvider {
316    fn id(&self) -> &str {
317        "brave"
318    }
319
320    async fn search(
321        &self,
322        request: WebSearchRequest,
323        cancel: Option<&CancellationToken>,
324    ) -> Result<Vec<WebSearchResult>, WebSearchProviderError> {
325        if self.config.api_key.trim().is_empty() {
326            return Err(WebSearchProviderError::Auth(
327                "Brave Search API key is empty".into(),
328            ));
329        }
330        let send = self
331            .http
332            .get("https://api.search.brave.com/res/v1/web/search")
333            .header("Accept", "application/json")
334            .header("X-Subscription-Token", &self.config.api_key)
335            .query(&[("q", request.query), ("count", request.count.to_string())])
336            .timeout(self.config.timeout)
337            .send();
338        let response = if let Some(cancel) = cancel {
339            tokio::select! {
340                biased;
341                _ = cancel.cancelled() => return Err(WebSearchProviderError::Cancelled),
342                response = send => response,
343            }
344        } else {
345            send.await
346        }
347        .map_err(|error| {
348            if error.is_timeout() {
349                WebSearchProviderError::Timeout(error.to_string())
350            } else {
351                WebSearchProviderError::Request(error.to_string())
352            }
353        })?;
354        let status = response.status();
355        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
356            return Err(WebSearchProviderError::Auth(format!(
357                "Brave Search returned HTTP {}",
358                status.as_u16()
359            )));
360        }
361        if !status.is_success() {
362            return Err(WebSearchProviderError::Request(format!(
363                "Brave Search returned HTTP {}",
364                status.as_u16()
365            )));
366        }
367        let mut stream = response.bytes_stream();
368        let mut body = Vec::new();
369        loop {
370            let next = if let Some(cancel) = cancel {
371                tokio::select! {
372                    biased;
373                    _ = cancel.cancelled() => return Err(WebSearchProviderError::Cancelled),
374                    next = stream.next() => next,
375                }
376            } else {
377                stream.next().await
378            };
379            let Some(chunk) = next else {
380                break;
381            };
382            let chunk =
383                chunk.map_err(|error| WebSearchProviderError::Request(error.to_string()))?;
384            if body.len() + chunk.len() > MAX_PROVIDER_RESPONSE_BYTES {
385                return Err(WebSearchProviderError::InvalidResponse(format!(
386                    "response exceeded {MAX_PROVIDER_RESPONSE_BYTES} bytes"
387                )));
388            }
389            body.extend_from_slice(&chunk);
390        }
391        let payload = serde_json::from_slice::<BraveResponse>(&body)
392            .map_err(|error| WebSearchProviderError::InvalidResponse(error.to_string()))?;
393        Ok(payload
394            .web
395            .map(|web| web.results)
396            .unwrap_or_default()
397            .into_iter()
398            .map(|result| WebSearchResult {
399                title: result.title,
400                url: result.url,
401                snippet: result.description,
402                published_at: result.age,
403            })
404            .collect())
405    }
406}
407
408#[derive(Clone)]
409pub struct ExaSearchConfig {
410    pub api_key: String,
411    /// API prefix, e.g. `https://api.exa.ai`. The `/search` route is appended.
412    pub base_url: String,
413    pub timeout: Duration,
414}
415
416impl ExaSearchConfig {
417    pub const DEFAULT_BASE_URL: &'static str = "https://api.exa.ai";
418
419    pub fn new(api_key: impl Into<String>) -> Self {
420        Self {
421            api_key: api_key.into(),
422            base_url: Self::DEFAULT_BASE_URL.into(),
423            timeout: DEFAULT_TIMEOUT,
424        }
425    }
426}
427
428/// Exa (`https://exa.ai`) managed-search adapter. Exa returns extract
429/// `highlights` per result, which make better snippets than the raw page
430/// `text`, so we request those and fall back to truncated text.
431#[derive(Clone)]
432pub struct ExaSearchProvider {
433    http: reqwest::Client,
434    config: ExaSearchConfig,
435}
436
437impl ExaSearchProvider {
438    pub fn new(config: ExaSearchConfig) -> Self {
439        let http = reqwest::Client::builder()
440            .connect_timeout(Duration::from_secs(10))
441            .build()
442            .unwrap_or_else(|_| reqwest::Client::new());
443        Self { http, config }
444    }
445}
446
447#[derive(Deserialize)]
448struct ExaResponse {
449    #[serde(default)]
450    results: Vec<ExaResult>,
451}
452
453#[derive(Deserialize)]
454struct ExaResult {
455    #[serde(default)]
456    title: String,
457    #[serde(default)]
458    url: String,
459    #[serde(default, rename = "publishedDate")]
460    published_date: Option<String>,
461    #[serde(default)]
462    highlights: Vec<String>,
463    #[serde(default)]
464    text: Option<String>,
465}
466
467/// Snippet source preference: joined highlights, else truncated page text.
468/// The runtime still applies `MAX_SNIPPET_CHARS` on top, so the fallback cap
469/// here only needs to keep pathological pages from dominating the body cap.
470const MAX_EXA_TEXT_FALLBACK_CHARS: usize = 1_000;
471
472fn parse_exa_response(body: &[u8]) -> Result<Vec<WebSearchResult>, WebSearchProviderError> {
473    let payload = serde_json::from_slice::<ExaResponse>(body)
474        .map_err(|error| WebSearchProviderError::InvalidResponse(error.to_string()))?;
475    Ok(payload
476        .results
477        .into_iter()
478        .map(|result| {
479            let snippet = if result.highlights.is_empty() {
480                result
481                    .text
482                    .map(|text| text.chars().take(MAX_EXA_TEXT_FALLBACK_CHARS).collect())
483                    .unwrap_or_default()
484            } else {
485                result.highlights.join("\n")
486            };
487            WebSearchResult {
488                title: result.title,
489                url: result.url,
490                snippet,
491                published_at: result.published_date,
492            }
493        })
494        .collect())
495}
496
497#[async_trait]
498impl WebSearchProvider for ExaSearchProvider {
499    fn id(&self) -> &str {
500        "exa"
501    }
502
503    async fn search(
504        &self,
505        request: WebSearchRequest,
506        cancel: Option<&CancellationToken>,
507    ) -> Result<Vec<WebSearchResult>, WebSearchProviderError> {
508        if self.config.api_key.trim().is_empty() {
509            return Err(WebSearchProviderError::Auth(
510                "Exa API key is empty".into(),
511            ));
512        }
513        let send = self
514            .http
515            .post(format!("{}/search", self.config.base_url))
516            .header("x-api-key", &self.config.api_key)
517            .json(&json!({
518                "query": request.query,
519                "numResults": request.count,
520                "contents": {
521                    "highlights": { "numSentences": 2, "highlightsPerUrl": 1 }
522                }
523            }))
524            .timeout(self.config.timeout)
525            .send();
526        let response = if let Some(cancel) = cancel {
527            tokio::select! {
528                biased;
529                _ = cancel.cancelled() => return Err(WebSearchProviderError::Cancelled),
530                response = send => response,
531            }
532        } else {
533            send.await
534        }
535        .map_err(|error| {
536            if error.is_timeout() {
537                WebSearchProviderError::Timeout(error.to_string())
538            } else {
539                WebSearchProviderError::Request(error.to_string())
540            }
541        })?;
542        let status = response.status();
543        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
544            return Err(WebSearchProviderError::Auth(format!(
545                "Exa returned HTTP {}",
546                status.as_u16()
547            )));
548        }
549        // Quota exhaustion (402/429) and every other non-2xx degrade to a
550        // tool-visible failure — the model is told search is unavailable and
551        // the turn keeps going.
552        if !status.is_success() {
553            return Err(WebSearchProviderError::Request(format!(
554                "Exa returned HTTP {}",
555                status.as_u16()
556            )));
557        }
558        let mut stream = response.bytes_stream();
559        let mut body = Vec::new();
560        loop {
561            let next = if let Some(cancel) = cancel {
562                tokio::select! {
563                    biased;
564                    _ = cancel.cancelled() => return Err(WebSearchProviderError::Cancelled),
565                    next = stream.next() => next,
566                }
567            } else {
568                stream.next().await
569            };
570            let Some(chunk) = next else {
571                break;
572            };
573            let chunk =
574                chunk.map_err(|error| WebSearchProviderError::Request(error.to_string()))?;
575            if body.len() + chunk.len() > MAX_PROVIDER_RESPONSE_BYTES {
576                return Err(WebSearchProviderError::InvalidResponse(format!(
577                    "response exceeded {MAX_PROVIDER_RESPONSE_BYTES} bytes"
578                )));
579            }
580            body.extend_from_slice(&chunk);
581        }
582        parse_exa_response(&body)
583    }
584}
585
586#[cfg(test)]
587mod tests {
588    use super::*;
589
590    #[derive(Clone)]
591    struct FakeProvider;
592
593    #[async_trait]
594    impl WebSearchProvider for FakeProvider {
595        fn id(&self) -> &str {
596            "fake"
597        }
598
599        async fn search(
600            &self,
601            _request: WebSearchRequest,
602            _cancel: Option<&CancellationToken>,
603        ) -> Result<Vec<WebSearchResult>, WebSearchProviderError> {
604            Ok(vec![
605                WebSearchResult {
606                    title: "Valid".into(),
607                    url: "https://example.com/result".into(),
608                    snippet: "Current information".into(),
609                    published_at: None,
610                },
611                WebSearchResult {
612                    title: "Unsafe".into(),
613                    url: "file:///etc/passwd".into(),
614                    snippet: "discard me".into(),
615                    published_at: None,
616                },
617            ])
618        }
619    }
620
621    fn invocation(input: Value) -> ToolInvocation {
622        ToolInvocation {
623            id: "search-1".into(),
624            name: "web_search".into(),
625            input,
626            raw_emitted_args: None,
627        }
628    }
629
630    #[tokio::test]
631    async fn runtime_normalizes_results_and_rejects_non_http_urls() {
632        let runtime = WebSearchToolRuntime::from_provider(FakeProvider);
633        let output = runtime
634            .invoke(invocation(json!({"query": "rust", "count": 5})))
635            .await
636            .unwrap()
637            .output
638            .unwrap();
639        assert_eq!(output["provider"], "fake");
640        assert_eq!(output["results"].as_array().unwrap().len(), 1);
641        assert_eq!(output["external_content"]["untrusted"], true);
642        assert!(output["results"][0]["title"]
643            .as_str()
644            .unwrap()
645            .contains("EXTERNAL_UNTRUSTED_CONTENT"));
646    }
647
648    #[tokio::test]
649    async fn runtime_rejects_invalid_count() {
650        let runtime = WebSearchToolRuntime::from_provider(FakeProvider);
651        for count in [json!(11), json!("5")] {
652            let failure = runtime
653                .invoke(invocation(json!({"query": "rust", "count": count})))
654                .await
655                .unwrap()
656                .output
657                .unwrap_err();
658            assert_eq!(failure.kind, ToolFailureKind::InvalidInput);
659        }
660    }
661
662    #[test]
663    fn exa_parse_prefers_highlights_over_text() {
664        let body = json!({
665            "results": [{
666                "title": "Rust 1.90",
667                "url": "https://blog.rust-lang.org/1.90",
668                "publishedDate": "2025-09-18T00:00:00.000Z",
669                "highlights": ["first highlight", "second highlight"],
670                "text": "full page text that should not be used"
671            }]
672        })
673        .to_string();
674        let results = parse_exa_response(body.as_bytes()).unwrap();
675        assert_eq!(results.len(), 1);
676        assert_eq!(results[0].title, "Rust 1.90");
677        assert_eq!(results[0].snippet, "first highlight\nsecond highlight");
678        assert_eq!(
679            results[0].published_at.as_deref(),
680            Some("2025-09-18T00:00:00.000Z")
681        );
682    }
683
684    #[test]
685    fn exa_parse_falls_back_to_truncated_text() {
686        let long_text = "x".repeat(MAX_EXA_TEXT_FALLBACK_CHARS + 100);
687        let body = json!({
688            "results": [{
689                "title": "No highlights",
690                "url": "https://example.com/a",
691                "text": long_text
692            }, {
693                "url": "https://example.com/b"
694            }]
695        })
696        .to_string();
697        let results = parse_exa_response(body.as_bytes()).unwrap();
698        assert_eq!(results.len(), 2);
699        assert_eq!(
700            results[0].snippet.chars().count(),
701            MAX_EXA_TEXT_FALLBACK_CHARS
702        );
703        // Missing title/highlights/text all default to empty — the runtime's
704        // normalize_result drops results with neither title nor snippet.
705        assert_eq!(results[1].title, "");
706        assert_eq!(results[1].snippet, "");
707        assert_eq!(results[1].published_at, None);
708    }
709
710    #[test]
711    fn exa_parse_handles_empty_results_and_rejects_garbage() {
712        let results = parse_exa_response(br#"{"results": []}"#).unwrap();
713        assert!(results.is_empty());
714        // Missing `results` key defaults to empty rather than erroring.
715        let results = parse_exa_response(br#"{"requestId": "abc"}"#).unwrap();
716        assert!(results.is_empty());
717        let error = parse_exa_response(b"not json").unwrap_err();
718        assert!(matches!(error, WebSearchProviderError::InvalidResponse(_)));
719    }
720
721    #[tokio::test]
722    async fn exa_empty_api_key_fails_closed_without_request() {
723        let provider = ExaSearchProvider::new(ExaSearchConfig::new("  "));
724        let error = provider
725            .search(
726                WebSearchRequest {
727                    query: "rust".into(),
728                    count: 5,
729                },
730                None,
731            )
732            .await
733            .unwrap_err();
734        assert!(matches!(error, WebSearchProviderError::Auth(_)));
735    }
736}
harness/tools/web_search.rs

harness/tools/
web_search.rs
