Skip to main content

mermaid_cli/providers/tool/
web.rs

1//! Web tools: `web_search` and `web_fetch`.
2//!
3//! Each tool holds a pluggable backend (`web_client::SearchProvider` /
4//! `FetchProvider`) selected from `[web]` config: `web_fetch` defaults to a
5//! native in-process fetch (no key), `web_search` to Ollama Cloud or a
6//! self-hosted SearXNG. This tool layer owns cancellation plumbing, the SSRF
7//! guard, and multi-query fan-out; the backend owns the transport.
8
9use std::sync::Arc;
10
11use async_trait::async_trait;
12
13use crate::app::{FetchBackend, SearchBackend, WebConfig};
14use crate::domain::{ToolDefinition, ToolMetadata, ToolOutcome, ToolRunMetadata};
15
16use super::super::ctx::{ExecContext, ProgressEvent};
17use super::ToolExecutor;
18use super::web_client::{
19    FetchProvider, ManagedSearxngBackend, NativeFetchClient, OllamaWebClient, SearchProvider,
20    SearxngClient, WebFetchResult, format_results,
21};
22
23/// Build the `web_fetch` tool for the configured backend. `native` always
24/// yields a tool; `ollama` yields one only when `OLLAMA_API_KEY` resolves
25/// (otherwise the tool would 401 on every call, so we don't register it).
26pub fn web_fetch_tool(web: &WebConfig) -> Option<WebFetchTool> {
27    match web.fetch_backend {
28        FetchBackend::Native => Some(WebFetchTool::native()),
29        FetchBackend::Ollama => {
30            crate::utils::resolve_provider_key("ollama", "OLLAMA_API_KEY", None)
31                .map(WebFetchTool::ollama)
32        },
33    }
34}
35
36/// Build the `web_search` tool for the configured backend. `auto` (the default)
37/// and `searxng` always yield a tool; `ollama` yields one only when
38/// `OLLAMA_API_KEY` resolves.
39pub fn web_search_tool(web: &WebConfig) -> Option<WebSearchTool> {
40    match web.search_backend {
41        // Zero-config default: Ollama Cloud when a key is present, otherwise an
42        // auto-managed local SearXNG (started lazily on the first search).
43        SearchBackend::Auto => Some(
44            crate::utils::resolve_provider_key("ollama", "OLLAMA_API_KEY", None)
45                .map(WebSearchTool::ollama)
46                .unwrap_or_else(WebSearchTool::managed_searxng),
47        ),
48        SearchBackend::Ollama => {
49            crate::utils::resolve_provider_key("ollama", "OLLAMA_API_KEY", None)
50                .map(WebSearchTool::ollama)
51        },
52        SearchBackend::Searxng => Some(WebSearchTool::searxng(web.searxng_url.clone())),
53    }
54}
55
56/// `web_search` — query the configured search backend. Accepts a single
57/// `{query, max_results}` OR a list of `{queries: [{query, max_results}]}` for
58/// parallel fan-out.
59pub struct WebSearchTool {
60    backend: Arc<dyn SearchProvider>,
61}
62
63impl WebSearchTool {
64    /// Search via Ollama Cloud (bearer `OLLAMA_API_KEY`).
65    pub fn ollama(api_key: String) -> Self {
66        Self {
67            backend: Arc::new(OllamaWebClient::new(api_key)),
68        }
69    }
70
71    /// Search via a self-hosted SearXNG instance at `base_url`.
72    pub fn searxng(base_url: String) -> Self {
73        Self {
74            backend: Arc::new(SearxngClient::new(base_url)),
75        }
76    }
77
78    /// Search via an auto-managed local SearXNG container (zero-config default).
79    pub fn managed_searxng() -> Self {
80        Self {
81            backend: Arc::new(ManagedSearxngBackend),
82        }
83    }
84}
85
86#[async_trait]
87impl ToolExecutor for WebSearchTool {
88    fn name(&self) -> &'static str {
89        "web_search"
90    }
91
92    fn schema(&self) -> ToolDefinition {
93        ToolDefinition {
94            name: "web_search".to_string(),
95            description:
96                "Search the web. Takes either a single `query` + `max_results`, or an array of `queries` for parallel fan-out."
97                    .to_string(),
98            input_schema: serde_json::json!({
99                "type": "object",
100                "properties": {
101                    "query": { "type": "string" },
102                    "max_results": { "type": "integer", "minimum": 1, "maximum": 10, "default": 5 },
103                    "queries": {
104                        "type": "array",
105                        "items": {
106                            "type": "object",
107                            "properties": {
108                                "query": { "type": "string" },
109                                "max_results": { "type": "integer", "minimum": 1, "maximum": 10 }
110                            },
111                            "required": ["query"]
112                        }
113                    }
114                }
115            }),
116        }
117    }
118
119    async fn execute(&self, args: serde_json::Value, ctx: ExecContext) -> ToolOutcome {
120        let queries = match parse_queries(&args) {
121            Ok(q) => q,
122            Err(e) => return ToolOutcome::error(e, 0.0),
123        };
124        if queries.is_empty() {
125            return ToolOutcome::error("web_search requires at least one query", 0.0);
126        }
127        if let Some(blocked) = super::policy_gate::gate_external(
128            &ctx,
129            "web_search",
130            crate::runtime::ToolCategory::Web,
131            format!("web_search ({} queries)", queries.len()),
132            &args,
133        )
134        .await
135        {
136            return blocked;
137        }
138
139        let start = std::time::Instant::now();
140        let mut combined = String::new();
141        let mut result_count = 0usize;
142        let mut sources = Vec::new();
143        let mut errors: Vec<String> = Vec::new();
144        for (idx, (query, count)) in queries.iter().enumerate() {
145            let _ = ctx
146                .progress
147                .send(ProgressEvent::Status(format!(
148                    "searching {}/{}: {}",
149                    idx + 1,
150                    queries.len(),
151                    query
152                )))
153                .await;
154
155            let search = self.backend.search(query, *count);
156            let result = tokio::select! {
157                biased;
158                _ = ctx.token.cancelled() => return ToolOutcome::cancelled(),
159                result = search => result,
160            };
161            // A single query returning nothing or erroring does NOT abort the
162            // batch — record it and carry on so the other queries' results
163            // survive (a partial answer beats none).
164            let section = match result {
165                Ok(results) => {
166                    result_count += results.len();
167                    sources.extend(results.iter().map(|result| result.url.clone()));
168                    if results.is_empty() {
169                        "[SEARCH_RESULTS]\n(no results found)\n[/SEARCH_RESULTS]\n".to_string()
170                    } else {
171                        format_results(&results)
172                    }
173                },
174                Err(e) => {
175                    errors.push(format!("{query}: {e}"));
176                    format!("(search failed: {e})\n")
177                },
178            };
179            if queries.len() > 1 {
180                combined.push_str(&format!("=== query: {query} ===\n{section}\n\n"));
181            } else {
182                combined = section;
183            }
184        }
185
186        // Only a total failure — every query hit a backend error — is a tool
187        // error. An empty-but-reachable search, or a partial success, returns
188        // normally so the model sees what did come back.
189        if errors.len() == queries.len() {
190            return ToolOutcome::error(
191                format!("web_search failed: {}", errors.join("; ")),
192                start.elapsed().as_secs_f64(),
193            );
194        }
195
196        // Cap the aggregate output. Per-result content is already truncated to
197        // WEB_CONTENT_MAX_CHARS, but many results across many queries can still
198        // bloat context (and memory) past what any single result's cap bounds (#28).
199        let combined = crate::utils::truncate_middle(
200            &combined,
201            crate::constants::WEB_SEARCH_AGGREGATE_MAX_CHARS,
202        );
203
204        let duration_secs = start.elapsed().as_secs_f64();
205        let requested_count = queries.iter().map(|(_, count)| *count).sum();
206        let query_texts = queries.iter().map(|(query, _)| query.clone()).collect();
207        ToolOutcome::success(
208            combined,
209            format!(
210                "{} {} returned",
211                result_count,
212                if result_count == 1 {
213                    "result"
214                } else {
215                    "results"
216                }
217            ),
218            duration_secs,
219        )
220        .with_metadata(ToolRunMetadata {
221            detail: ToolMetadata::WebSearch {
222                queries: query_texts,
223                requested_count,
224                result_count,
225                sources,
226            },
227            result_count: Some(result_count),
228            ..ToolRunMetadata::default()
229        })
230    }
231}
232
233/// `web_fetch` — retrieve a URL's readable content as markdown. Single URL,
234/// single response. Native by default (fetches + converts in-process, no key);
235/// can be backed by Ollama Cloud instead.
236pub struct WebFetchTool {
237    backend: Arc<dyn FetchProvider>,
238}
239
240impl WebFetchTool {
241    /// Fetch the URL in-process and convert its HTML to markdown (no API key).
242    pub fn native() -> Self {
243        Self {
244            backend: Arc::new(NativeFetchClient::new()),
245        }
246    }
247
248    /// Fetch via Ollama Cloud's server-side `/api/web_fetch`.
249    pub fn ollama(api_key: String) -> Self {
250        Self {
251            backend: Arc::new(OllamaWebClient::new(api_key)),
252        }
253    }
254}
255
256#[async_trait]
257impl ToolExecutor for WebFetchTool {
258    fn name(&self) -> &'static str {
259        "web_fetch"
260    }
261
262    fn schema(&self) -> ToolDefinition {
263        ToolDefinition {
264            name: "web_fetch".to_string(),
265            description: "Retrieve a single URL's main content as markdown. Optional 'pattern' \
266                          finds matches in the page instead of returning the whole body: plain \
267                          case-insensitive substring matching, applied per line (a pattern \
268                          containing a newline never matches), returning each match with \
269                          surrounding context lines."
270                .to_string(),
271            input_schema: serde_json::json!({
272                "type": "object",
273                "properties": {
274                    "url": { "type": "string" },
275                    "pattern": {
276                        "type": "string",
277                        "description": "Case-insensitive substring to find in the page (not a regex)"
278                    },
279                    "context_lines": {
280                        "type": "integer",
281                        "description": "Context lines around each match (default 2, max 10)"
282                    }
283                },
284                "required": ["url"]
285            }),
286        }
287    }
288
289    async fn execute(&self, args: serde_json::Value, ctx: ExecContext) -> ToolOutcome {
290        let Some(url) = args.get("url").and_then(|v| v.as_str()) else {
291            return ToolOutcome::error("web_fetch requires 'url' (string)", 0.0);
292        };
293        let pattern = args
294            .get("pattern")
295            .and_then(|v| v.as_str())
296            .map(str::trim)
297            .filter(|p| !p.is_empty());
298        let context_lines = args
299            .get("context_lines")
300            .and_then(|v| v.as_u64())
301            .unwrap_or(2)
302            .min(10) as usize;
303        if let Err(reason) = validate_fetch_url(url) {
304            return ToolOutcome::error(format!("web_fetch: {reason}"), 0.0);
305        }
306        if let Some(blocked) = super::policy_gate::gate_external(
307            &ctx,
308            "web_fetch",
309            crate::runtime::ToolCategory::Web,
310            format!("web_fetch {}", url),
311            &args,
312        )
313        .await
314        {
315            return blocked;
316        }
317        let start = std::time::Instant::now();
318        let fetch = self.backend.fetch(url);
319
320        tokio::select! {
321            biased;
322            _ = ctx.token.cancelled() => ToolOutcome::cancelled(),
323            result = fetch => match result {
324                Ok(page) => {
325                    let output = format_fetch(url, &page, pattern, context_lines);
326                    let duration_secs = start.elapsed().as_secs_f64();
327                    let line_count = output.lines().count();
328                    let byte_count = output.len();
329                    let title = if page.title.is_empty() {
330                        None
331                    } else {
332                        Some(page.title)
333                    };
334                    ToolOutcome::success(
335                        output,
336                        format!("{} {} fetched", line_count, if line_count == 1 { "line" } else { "lines" }),
337                        duration_secs,
338                    )
339                    .with_metadata(ToolRunMetadata {
340                        detail: ToolMetadata::WebFetch {
341                            url: url.to_string(),
342                            title,
343                            line_count,
344                            byte_count,
345                        },
346                        line_count: Some(line_count),
347                        byte_count: Some(byte_count),
348                        ..ToolRunMetadata::default()
349                    })
350                },
351                Err(e) => ToolOutcome::error(
352                    format!("web_fetch({}): {}", url, e),
353                    start.elapsed().as_secs_f64(),
354                ),
355            },
356        }
357    }
358}
359
360/// Cap on a single `web_fetch` body (#F46). The raw fetch is bounded only by the
361/// 16 MB HTTP body limit, so without this one URL could dump megabytes into model
362/// context. A full page warrants more room than a `web_search` snippet
363/// (`WEB_CONTENT_MAX_CHARS`), so this mirrors web_search's per-call aggregate
364/// budget. Applied as a byte cap; truncation is char-boundary safe.
365const WEB_FETCH_MAX_CHARS: usize = crate::constants::WEB_SEARCH_AGGREGATE_MAX_CHARS;
366
367/// Truncate a fetched page body to `WEB_FETCH_MAX_CHARS` bytes, char-boundary
368/// safe, appending a marker — consistent with how `web_search` bounds the
369/// content it returns (#F46). Borrows when no truncation is needed.
370fn cap_fetch_content(content: &str) -> std::borrow::Cow<'_, str> {
371    if content.len() <= WEB_FETCH_MAX_CHARS {
372        return std::borrow::Cow::Borrowed(content);
373    }
374    let cut = content.floor_char_boundary(WEB_FETCH_MAX_CHARS);
375    std::borrow::Cow::Owned(format!("{}\n\n...[content truncated]", &content[..cut]))
376}
377
378/// Cap on find-in-page match BLOCKS per call (merged context windows).
379/// Matched lines beyond the included blocks are summarized as a
380/// `(+N more matches)` tail so the model knows the page has more.
381const MAX_PATTERN_MATCHES: usize = 20;
382
383fn format_fetch(
384    url: &str,
385    page: &WebFetchResult,
386    pattern: Option<&str>,
387    ctx_lines: usize,
388) -> String {
389    let title = if page.title.is_empty() {
390        "(no title)"
391    } else {
392        page.title.as_str()
393    };
394    // Find-in-page runs on the FULL readable markdown, then the report goes
395    // through the cap — capping first would hide matches in the tail.
396    if let Some(pattern) = pattern {
397        let body = match extract_matches(&page.content, pattern, ctx_lines, MAX_PATTERN_MATCHES) {
398            Some(report) => report,
399            // No match: say so, then the capped head as usual so the model
400            // keeps its orientation on the page.
401            None => format!(
402                "No matches for \"{}\".\n\n{}",
403                pattern,
404                cap_fetch_content(&page.content)
405            ),
406        };
407        let report = format!("# {}\n\nURL: {}\n\n{}", title, url, body);
408        return cap_fetch_content(&report).into_owned();
409    }
410    let content = cap_fetch_content(&page.content);
411    format!("# {}\n\nURL: {}\n\n{}", title, url, content)
412}
413
414/// Find-in-page core: plain case-insensitive SUBSTRING match per line (not a
415/// regex — model-supplied metacharacters must mean themselves), each match
416/// reported as a `L<n>:`-prefixed context block. Overlapping or adjacent
417/// windows merge into one block; blocks are separated by `---` lines and
418/// capped at `max_blocks`, with a `(+N more matches)` tail counting the
419/// matched lines that didn't fit. Returns `None` when nothing matches.
420fn extract_matches(
421    content: &str,
422    pattern: &str,
423    context_lines: usize,
424    max_blocks: usize,
425) -> Option<String> {
426    let needle = pattern.to_lowercase();
427    let lines: Vec<&str> = content.lines().collect();
428    let matched: Vec<usize> = lines
429        .iter()
430        .enumerate()
431        .filter(|(_, l)| l.to_lowercase().contains(&needle))
432        .map(|(i, _)| i)
433        .collect();
434    if matched.is_empty() {
435        return None;
436    }
437
438    // Merge each match's [i-ctx, i+ctx] window with overlapping/adjacent
439    // neighbors; count how many matched lines the included blocks cover.
440    let mut blocks: Vec<(usize, usize)> = Vec::new();
441    for &i in &matched {
442        let start = i.saturating_sub(context_lines);
443        let end = (i + context_lines).min(lines.len() - 1);
444        match blocks.last_mut() {
445            Some((_, last_end)) if start <= *last_end + 1 => *last_end = (*last_end).max(end),
446            _ => blocks.push((start, end)),
447        }
448    }
449    let included = &blocks[..blocks.len().min(max_blocks)];
450    let cutoff = included.last().map(|&(_, end)| end).unwrap_or(0);
451    let dropped = matched.iter().filter(|&&i| i > cutoff).count();
452
453    let mut out = format!(
454        "{} match{} for \"{}\":\n",
455        matched.len(),
456        if matched.len() == 1 { "" } else { "es" },
457        pattern
458    );
459    for (bi, &(start, end)) in included.iter().enumerate() {
460        if bi > 0 {
461            out.push_str("---\n");
462        }
463        for (offset, line) in lines[start..=end].iter().enumerate() {
464            // 1-based line numbers, matching how editors and grep report.
465            out.push_str(&format!("L{}: {}\n", start + offset + 1, line));
466        }
467    }
468    if dropped > 0 {
469        out.push_str(&format!(
470            "(+{dropped} more match{})\n",
471            if dropped == 1 { "" } else { "es" }
472        ));
473    }
474    Some(out)
475}
476
477fn parse_queries(args: &serde_json::Value) -> Result<Vec<(String, usize)>, String> {
478    if let Some(arr) = args.get("queries").and_then(|v| v.as_array()) {
479        if arr.len() > crate::constants::MAX_BATCH_TOOL_ITEMS {
480            return Err(format!(
481                "web_search: too many queries ({}); cap is {} per call — split the request",
482                arr.len(),
483                crate::constants::MAX_BATCH_TOOL_ITEMS
484            ));
485        }
486        let mut out = Vec::with_capacity(arr.len());
487        for v in arr {
488            let Some(obj) = v.as_object() else {
489                return Err(
490                    "web_search: 'queries' must be an array of {query, max_results}".to_string(),
491                );
492            };
493            let Some(query) = obj.get("query").and_then(|x| x.as_str()) else {
494                return Err("web_search: each query entry needs 'query' (string)".to_string());
495            };
496            let count = obj
497                .get("max_results")
498                .or_else(|| obj.get("result_count"))
499                .and_then(|x| x.as_u64())
500                .unwrap_or(5)
501                .clamp(1, 10) as usize;
502            out.push((query.to_string(), count));
503        }
504        return Ok(out);
505    }
506    if let Some(query) = args.get("query").and_then(|v| v.as_str()) {
507        let count = args
508            .get("max_results")
509            .or_else(|| args.get("result_count"))
510            .and_then(|v| v.as_u64())
511            .unwrap_or(5)
512            .clamp(1, 10) as usize;
513        return Ok(vec![(query.to_string(), count)]);
514    }
515    Err("web_search requires 'query' (string) or 'queries' (array)".to_string())
516}
517
518/// Reject obviously-unsafe fetch URLs before the backend runs: only
519/// `http`/`https`, and no loopback / link-local / private / metadata hosts.
520/// For the native backend this is the primary SSRF boundary (the request
521/// leaves from this process, so `web_client::guard_resolved_ips` also checks
522/// the resolved addresses); for the Ollama backend it's defense-in-depth ahead
523/// of Ollama's own server-side fetch. Guards against model-supplied URLs.
524/// Reject anything that isn't a plain `http(s)` URL. A `file:`, `javascript:`,
525/// `data:`, or otherwise exotic scheme has no business reaching an HTTP fetch or
526/// an OS browser launcher. Returns the parsed URL so callers can inspect the
527/// host without re-parsing. Note: this deliberately does NOT block loopback —
528/// `open_url` legitimately opens a just-started local dev server.
529pub(crate) fn require_http_scheme(url: &str) -> Result<reqwest::Url, String> {
530    let parsed = reqwest::Url::parse(url).map_err(|e| format!("invalid URL: {e}"))?;
531    match parsed.scheme() {
532        "http" | "https" => Ok(parsed),
533        other => Err(format!(
534            "unsupported URL scheme '{other}' (only http/https allowed)"
535        )),
536    }
537}
538
539fn validate_fetch_url(url: &str) -> Result<(), String> {
540    let parsed = require_http_scheme(url)?;
541    let host = parsed
542        .host_str()
543        .ok_or_else(|| "URL has no host".to_string())?;
544    if is_blocked_host(host) {
545        return Err(format!("refusing to fetch internal/loopback host '{host}'"));
546    }
547    Ok(())
548}
549
550/// Cloud-metadata DNS hostnames that resolve (inside the relevant cloud) to a
551/// link-local metadata IP — `169.254.169.254` and friends — but are LEXICALLY
552/// public, so the IP-only `classify_host` waves them through as
553/// [`crate::utils::HostClass::Public`]. We block them by name as well. Matched
554/// after lowercasing + trimming any surrounding `[]` and trailing FQDN dot.
555const METADATA_HOSTNAMES: &[&str] = &[
556    "metadata.google.internal",   // GCP (canonical)
557    "metadata.goog",              // GCP (alternate)
558    "metadata",                   // GCP/Azure short name (http://metadata/ responds)
559    "instance-data",              // AWS (cloud-init alias)
560    "instance-data.ec2.internal", // AWS
561];
562
563/// F57: client-side SSRF denylist for `web_fetch`.
564///
565/// For the NATIVE backend the request originates from this process, so this
566/// lexical check plus `web_client::guard_resolved_ips` (which classifies the
567/// resolved addresses) is the authoritative boundary — modulo the
568/// DNS-rebinding TOCTOU that no pre-connect check can fully close.
569///
570/// For the OLLAMA backend the URL is POSTed to Ollama's server-side
571/// `/api/web_fetch` and the in-process client only ever connects to
572/// `ollama.com`; there the authoritative boundary is server-side (Ollama) and
573/// this check is defense-in-depth so a model can't trivially aim the server at
574/// an obvious internal target.
575///
576/// Either way we reject:
577/// - every non-public IP form via the shared `classify_host` (loopback,
578///   RFC-1918/ULA, link-local incl. `169.254.169.254`, CGNAT, unspecified
579///   `0.0.0.0`/`::`, plus the IPv4-mapped-IPv6 / ULA / link-local-IPv6 / `[::1]`
580///   forms a hand-rolled IPv4 check would miss); and
581/// - the well-known cloud-metadata HOSTNAMES (`metadata.google.internal`, …)
582///   that are lexically public but front a metadata service.
583fn is_blocked_host(host: &str) -> bool {
584    let normalized = host
585        .trim_start_matches('[')
586        .trim_end_matches(']')
587        .trim_end_matches('.')
588        .to_ascii_lowercase();
589    if METADATA_HOSTNAMES.contains(&normalized.as_str()) {
590        return true;
591    }
592    crate::utils::classify_host(host).is_internal()
593}
594
595#[cfg(test)]
596mod tests {
597    use super::*;
598
599    #[test]
600    fn require_http_scheme_accepts_http_rejects_exotic() {
601        // http/https pass — including loopback, since `open_url` legitimately
602        // opens a just-started local dev server (so this must NOT block localhost).
603        for good in [
604            "http://example.com",
605            "https://example.com/path?a=1&b=2",
606            "http://localhost:3000",
607            "http://127.0.0.1:8080",
608        ] {
609            assert!(require_http_scheme(good).is_ok(), "{good} should pass");
610        }
611        // Non-http(s) schemes and unparseable input are rejected.
612        for bad in [
613            "file:///etc/passwd",
614            "javascript:alert(1)",
615            "data:text/html,<script>",
616            "ftp://example.com",
617            "not a url",
618        ] {
619            assert!(
620                require_http_scheme(bad).is_err(),
621                "{bad} should be rejected"
622            );
623        }
624    }
625
626    #[test]
627    fn validate_fetch_url_blocks_unsafe_targets() {
628        // #9: scheme + internal-host guards.
629        for bad in [
630            "file:///etc/passwd",
631            "ftp://example.com/x",
632            "http://localhost/admin",
633            "http://127.0.0.1:8080",
634            "http://169.254.169.254/latest/meta-data/",
635            "http://10.0.0.5/",
636            "http://192.168.1.1/",
637            "http://[::1]/",
638            // #27/#80: IPv6/CGNAT bypasses the old IPv4-centric blocklist missed.
639            "http://[::ffff:169.254.169.254]/latest/meta-data/",
640            "http://[fc00::1]/",
641            "http://[fe80::1]/",
642            "http://100.100.100.200/",
643            // F57: cloud-metadata hostnames are lexically public (IP-only
644            // classify_host waves them through) but front a metadata service.
645            "http://metadata.google.internal/computeMetadata/v1/",
646            "http://metadata.goog/",
647            "http://metadata/",
648            "http://instance-data/latest/meta-data/",
649            "https://METADATA.GOOGLE.INTERNAL./",
650            "not a url",
651        ] {
652            assert!(
653                validate_fetch_url(bad).is_err(),
654                "expected reject for {bad:?}",
655            );
656        }
657        for good in [
658            "https://example.com",
659            "http://example.com/page?x=1",
660            "https://docs.rs/serde",
661        ] {
662            assert!(
663                validate_fetch_url(good).is_ok(),
664                "expected accept for {good:?}",
665            );
666        }
667    }
668
669    #[test]
670    fn is_blocked_host_covers_metadata_names_and_ip_forms() {
671        // F57: cloud-metadata HOSTNAMES (lexically public, IP-only
672        // classify_host misses them) are blocked, case/dot-insensitively.
673        for h in [
674            "metadata.google.internal",
675            "metadata.google.internal.", // trailing FQDN dot
676            "Metadata.Google.Internal",  // case-insensitive
677            "metadata.goog",
678            "metadata",
679            "instance-data",
680            "instance-data.ec2.internal",
681        ] {
682            assert!(is_blocked_host(h), "metadata host {h:?} must be blocked");
683        }
684        // Non-public IP forms still go through classify_host (incl. the ones
685        // the task lists as examples that must already be covered).
686        for h in ["0.0.0.0", "::1", "169.254.169.254", "127.0.0.1"] {
687            assert!(is_blocked_host(h), "internal IP {h:?} must be blocked");
688        }
689        // Legitimate public hosts are NOT blocked — including a real `.goog`
690        // domain that merely is not the metadata alias.
691        for h in ["example.com", "docs.rs", "abc.goog", "8.8.8.8"] {
692            assert!(!is_blocked_host(h), "public host {h:?} must be allowed");
693        }
694    }
695
696    #[test]
697    fn format_fetch_caps_long_content() {
698        // F46: a huge page body must be truncated with a marker, not dumped whole.
699        let big = "z".repeat(WEB_FETCH_MAX_CHARS * 2);
700        let page = WebFetchResult {
701            title: "T".to_string(),
702            content: big,
703        };
704        let out = format_fetch("https://example.com", &page, None, 2);
705        assert!(
706            out.len() < WEB_FETCH_MAX_CHARS + 256,
707            "content must be capped, got {} bytes",
708            out.len()
709        );
710        assert!(out.contains("truncated"), "expected truncation marker");
711
712        // A short page is emitted intact, with no marker.
713        let small = WebFetchResult {
714            title: "T".to_string(),
715            content: "hello world".to_string(),
716        };
717        let out = format_fetch("https://example.com", &small, None, 2);
718        assert!(out.contains("hello world"));
719        assert!(!out.contains("truncated"));
720    }
721
722    #[test]
723    fn extract_matches_finds_case_insensitive_with_context() {
724        let content = "line one\nline two\nTARGET here\nline four\nline five";
725        let out = extract_matches(content, "target", 1, 20).unwrap();
726        assert!(out.starts_with("1 match for \"target\":"));
727        assert!(out.contains("L2: line two"));
728        assert!(out.contains("L3: TARGET here"));
729        assert!(out.contains("L4: line four"));
730        assert!(!out.contains("L1:"), "context clipped to 1 line: {out}");
731        assert!(!out.contains("L5:"));
732    }
733
734    #[test]
735    fn extract_matches_merges_overlapping_windows() {
736        // Matches on adjacent lines must merge into ONE block (no separator).
737        let content = "a\nhit one\nhit two\nb\nc\nd\ne\nf\ng\nhit three\nz";
738        let out = extract_matches(content, "hit", 1, 20).unwrap();
739        assert!(out.starts_with("3 matches"));
740        assert_eq!(out.matches("---").count(), 1, "two blocks: {out}");
741        // No duplicated lines from the merged windows.
742        assert_eq!(out.matches("hit one").count(), 1);
743    }
744
745    #[test]
746    fn extract_matches_caps_blocks_and_reports_tail() {
747        // 25 matches spaced far apart -> 25 blocks, capped at 20 + tail note.
748        let content = (0..25)
749            .map(|i| format!("match {i}\nx\nx\nx\nx\nx"))
750            .collect::<Vec<_>>()
751            .join("\n");
752        let out = extract_matches(&content, "match", 0, 20).unwrap();
753        assert!(out.starts_with("25 matches"));
754        assert_eq!(out.matches("---").count(), 19, "20 blocks: {out}");
755        assert!(out.contains("(+5 more matches)"), "tail note: {out}");
756    }
757
758    #[test]
759    fn extract_matches_none_and_multibyte() {
760        assert!(extract_matches("nothing here", "absent", 2, 20).is_none());
761        // Multibyte content must not panic and must match case-insensitively.
762        let content = "voil\u{e0} un r\u{e9}sultat\nplain line";
763        let out = extract_matches(content, "R\u{c9}SULTAT", 0, 20).unwrap();
764        assert!(out.contains("L1: voil\u{e0} un r\u{e9}sultat"));
765        // Context 0 keeps only the matching line.
766        assert!(!out.contains("plain line"));
767    }
768
769    #[test]
770    fn format_fetch_pattern_paths() {
771        let page = WebFetchResult {
772            title: "T".to_string(),
773            content: "alpha\nbeta\ngamma".to_string(),
774        };
775        // Match -> report replaces the body.
776        let out = format_fetch("https://example.com", &page, Some("beta"), 1);
777        assert!(out.contains("1 match for \"beta\""));
778        assert!(out.contains("L2: beta"));
779        // No match -> explicit notice + the page head so the model keeps
780        // orientation.
781        let out = format_fetch("https://example.com", &page, Some("nope"), 1);
782        assert!(out.contains("No matches for \"nope\"."));
783        assert!(out.contains("alpha"));
784    }
785
786    #[test]
787    fn find_in_page_runs_before_the_cap() {
788        // A match in the tail of a page longer than the cap must still be
789        // found (matching runs pre-cap; only the report is capped).
790        let mut content = "x\n".repeat(WEB_FETCH_MAX_CHARS / 2);
791        content.push_str("needle in the tail\n");
792        let page = WebFetchResult {
793            title: "T".to_string(),
794            content,
795        };
796        let out = format_fetch("https://example.com", &page, Some("needle"), 1);
797        assert!(out.contains("1 match for \"needle\""), "tail match found");
798        assert!(out.contains("needle in the tail"));
799    }
800
801    #[test]
802    fn parse_queries_single_form() {
803        let args = serde_json::json!({"query": "rust async", "max_results": 3});
804        let q = parse_queries(&args).unwrap();
805        assert_eq!(q.len(), 1);
806        assert_eq!(q[0].0, "rust async");
807        assert_eq!(q[0].1, 3);
808    }
809
810    #[test]
811    fn parse_queries_array_form() {
812        let args = serde_json::json!({"queries": [
813            {"query": "a", "max_results": 2},
814            {"query": "b", "result_count": 5},
815        ]});
816        let q = parse_queries(&args).unwrap();
817        assert_eq!(q.len(), 2);
818        assert_eq!(q[1].1, 5);
819    }
820
821    #[test]
822    fn parse_queries_missing_errors() {
823        let args = serde_json::json!({});
824        assert!(parse_queries(&args).is_err());
825    }
826
827    #[test]
828    fn parse_queries_clamps_count() {
829        let args = serde_json::json!({"query": "q", "max_results": 999});
830        let q = parse_queries(&args).unwrap();
831        assert_eq!(q[0].1, 10);
832        let args = serde_json::json!({"query": "q", "max_results": 0});
833        let q = parse_queries(&args).unwrap();
834        assert_eq!(q[0].1, 1);
835    }
836
837    #[test]
838    fn parse_queries_rejects_excess_fan_out() {
839        // #90: a single call can't request unbounded fan-out.
840        let many: Vec<_> = (0..crate::constants::MAX_BATCH_TOOL_ITEMS + 1)
841            .map(|i| serde_json::json!({"query": format!("q{i}")}))
842            .collect();
843        let args = serde_json::json!({ "queries": many });
844        assert!(parse_queries(&args).is_err());
845
846        // Exactly at the cap is still accepted.
847        let at_cap: Vec<_> = (0..crate::constants::MAX_BATCH_TOOL_ITEMS)
848            .map(|i| serde_json::json!({"query": format!("q{i}")}))
849            .collect();
850        let args = serde_json::json!({ "queries": at_cap });
851        assert_eq!(
852            parse_queries(&args).unwrap().len(),
853            crate::constants::MAX_BATCH_TOOL_ITEMS
854        );
855    }
856
857    #[tokio::test]
858    async fn web_search_batch_survives_empty_and_failed_queries() {
859        use crate::domain::{ToolCallId, ToolStatus, TurnId};
860        use crate::providers::ctx::test_exec_context;
861        use crate::providers::tool::web_client::SearchResult;
862        use async_trait::async_trait;
863        use std::sync::Arc;
864
865        struct Mock;
866        #[async_trait]
867        impl SearchProvider for Mock {
868            async fn search(
869                &self,
870                query: &str,
871                _count: usize,
872            ) -> anyhow::Result<Vec<SearchResult>> {
873                match query {
874                    "boom" => Err(anyhow::anyhow!("backend down")),
875                    "empty" => Ok(Vec::new()),
876                    _ => Ok(vec![SearchResult {
877                        title: "Title".to_string(),
878                        url: "https://example.com".to_string(),
879                        snippet: "snip".to_string(),
880                        full_content: "content".to_string(),
881                    }]),
882                }
883            }
884        }
885
886        let mk = || WebSearchTool {
887            backend: Arc::new(Mock),
888        };
889        let tmp = std::path::PathBuf::from("/tmp");
890
891        // Partial: one good, one empty, one erroring -> success, good kept.
892        let (ctx, _rx) = test_exec_context(TurnId(1), ToolCallId(1), tmp.clone());
893        let out = mk()
894            .execute(
895                serde_json::json!({"queries": [{"query":"good"},{"query":"empty"},{"query":"boom"}]}),
896                ctx,
897            )
898            .await;
899        assert_eq!(
900            out.status,
901            ToolStatus::Success,
902            "a partial batch must not abort"
903        );
904        assert!(
905            out.output().contains("https://example.com"),
906            "keeps the good result"
907        );
908
909        // A single empty query is "no results", not a hard error.
910        let (ctx, _rx) = test_exec_context(TurnId(2), ToolCallId(2), tmp.clone());
911        let out = mk()
912            .execute(serde_json::json!({"query": "empty"}), ctx)
913            .await;
914        assert_eq!(out.status, ToolStatus::Success, "empty is not an error");
915        assert!(out.output().contains("no results"));
916
917        // Every query failing IS a tool error.
918        let (ctx, _rx) = test_exec_context(TurnId(3), ToolCallId(3), tmp);
919        let out = mk()
920            .execute(
921                serde_json::json!({"queries": [{"query":"boom"},{"query":"boom"}]}),
922                ctx,
923            )
924            .await;
925        assert_eq!(out.status, ToolStatus::Error, "total failure is an error");
926    }
927}