mermaid_cli/providers/tool/
web.rs

1//! Web tools: `web_search` and `web_fetch`.
2//!
3//! Both delegate to `web_client::WebSearchClient` — a thin HTTP
4//! client for Ollama Cloud's web API (bearer-token path, via
5//! `OLLAMA_API_KEY`). The wrapper's job is cancellation plumbing +
6//! multi-query fan-out.
7
8use std::sync::Arc;
9
10use async_trait::async_trait;
11
12use crate::domain::{ToolDefinition, ToolMetadata, ToolOutcome, ToolRunMetadata};
13
14use super::super::ctx::{ExecContext, ProgressEvent};
15use super::ToolExecutor;
16use super::web_client::{WebFetchResult, WebSearchClient};
17
18/// `web_search` — query Ollama Cloud's web-search endpoint. Accepts a
19/// single `{query, max_results}` OR a list of `{queries: [{query,
20/// max_results}]}` for parallel fan-out.
21pub struct WebSearchTool {
22    client: Arc<WebSearchClient>,
23}
24
25impl WebSearchTool {
26    pub fn new(api_key: String) -> Self {
27        Self {
28            client: Arc::new(WebSearchClient::new(api_key)),
29        }
30    }
31}
32
33#[async_trait]
34impl ToolExecutor for WebSearchTool {
35    fn name(&self) -> &'static str {
36        "web_search"
37    }
38
39    fn schema(&self) -> ToolDefinition {
40        ToolDefinition {
41            name: "web_search".to_string(),
42            description:
43                "Search the web via Ollama Cloud's search API. Takes either a single `query` + `max_results`, or an array of `queries` for parallel fan-out."
44                    .to_string(),
45            input_schema: serde_json::json!({
46                "type": "object",
47                "properties": {
48                    "query": { "type": "string" },
49                    "max_results": { "type": "integer", "minimum": 1, "maximum": 10, "default": 5 },
50                    "queries": {
51                        "type": "array",
52                        "items": {
53                            "type": "object",
54                            "properties": {
55                                "query": { "type": "string" },
56                                "max_results": { "type": "integer", "minimum": 1, "maximum": 10 }
57                            },
58                            "required": ["query"]
59                        }
60                    }
61                }
62            }),
63        }
64    }
65
66    async fn execute(&self, args: serde_json::Value, ctx: ExecContext) -> ToolOutcome {
67        let queries = match parse_queries(&args) {
68            Ok(q) => q,
69            Err(e) => return ToolOutcome::error(e, 0.0),
70        };
71        if queries.is_empty() {
72            return ToolOutcome::error("web_search requires at least one query", 0.0);
73        }
74        if let Some(blocked) = super::policy_gate::gate_external(
75            &ctx,
76            "web_search",
77            crate::runtime::ToolCategory::Web,
78            format!("web_search ({} queries)", queries.len()),
79            &args,
80        )
81        .await
82        {
83            return blocked;
84        }
85
86        let start = std::time::Instant::now();
87        let mut combined = String::new();
88        let mut result_count = 0usize;
89        let mut sources = Vec::new();
90        for (idx, (query, count)) in queries.iter().enumerate() {
91            let _ = ctx
92                .progress
93                .send(ProgressEvent::Status(format!(
94                    "searching {}/{}: {}",
95                    idx + 1,
96                    queries.len(),
97                    query
98                )))
99                .await;
100
101            let search = self.client.search_query(query, *count);
102            tokio::select! {
103                biased;
104                _ = ctx.token.cancelled() => return ToolOutcome::cancelled(),
105                result = search => {
106                    match result {
107                        Ok(results) => {
108                            result_count += results.len();
109                            sources.extend(results.iter().map(|result| result.url.clone()));
110                            let formatted = self.client.format_results(&results);
111                            if queries.len() > 1 {
112                                combined.push_str(&format!("=== query: {} ===\n{}\n\n", query, formatted));
113                            } else {
114                                combined = formatted;
115                            }
116                        },
117                        Err(e) => {
118                            return ToolOutcome::error(
119                                format!("web_search({}): {}", query, e),
120                                start.elapsed().as_secs_f64(),
121                            );
122                        },
123                    }
124                }
125            }
126        }
127
128        // Cap the aggregate output. Per-result content is already truncated to
129        // WEB_CONTENT_MAX_CHARS, but many results across many queries can still
130        // bloat context (and memory) past what any single result's cap bounds (#28).
131        let combined = crate::utils::truncate_content(
132            &combined,
133            crate::constants::WEB_SEARCH_AGGREGATE_MAX_CHARS,
134        );
135
136        let duration_secs = start.elapsed().as_secs_f64();
137        let requested_count = queries.iter().map(|(_, count)| *count).sum();
138        let query_texts = queries.iter().map(|(query, _)| query.clone()).collect();
139        ToolOutcome::success(
140            combined,
141            format!(
142                "{} {} returned",
143                result_count,
144                if result_count == 1 {
145                    "result"
146                } else {
147                    "results"
148                }
149            ),
150            duration_secs,
151        )
152        .with_metadata(ToolRunMetadata {
153            detail: ToolMetadata::WebSearch {
154                queries: query_texts,
155                requested_count,
156                result_count,
157                sources,
158            },
159            result_count: Some(result_count),
160            ..ToolRunMetadata::default()
161        })
162    }
163}
164
165/// `web_fetch` — retrieve a URL's readable content (Ollama Cloud's
166/// fetch endpoint). Single URL, single response.
167pub struct WebFetchTool {
168    client: Arc<WebSearchClient>,
169}
170
171impl WebFetchTool {
172    pub fn new(api_key: String) -> Self {
173        Self {
174            client: Arc::new(WebSearchClient::new(api_key)),
175        }
176    }
177}
178
179#[async_trait]
180impl ToolExecutor for WebFetchTool {
181    fn name(&self) -> &'static str {
182        "web_fetch"
183    }
184
185    fn schema(&self) -> ToolDefinition {
186        ToolDefinition {
187            name: "web_fetch".to_string(),
188            description: "Retrieve a single URL's main content as text (Ollama Cloud fetch API)."
189                .to_string(),
190            input_schema: serde_json::json!({
191                "type": "object",
192                "properties": { "url": { "type": "string" } },
193                "required": ["url"]
194            }),
195        }
196    }
197
198    async fn execute(&self, args: serde_json::Value, ctx: ExecContext) -> ToolOutcome {
199        let Some(url) = args.get("url").and_then(|v| v.as_str()) else {
200            return ToolOutcome::error("web_fetch requires 'url' (string)", 0.0);
201        };
202        if let Err(reason) = validate_fetch_url(url) {
203            return ToolOutcome::error(format!("web_fetch: {reason}"), 0.0);
204        }
205        if let Some(blocked) = super::policy_gate::gate_external(
206            &ctx,
207            "web_fetch",
208            crate::runtime::ToolCategory::Web,
209            format!("web_fetch {}", url),
210            &args,
211        )
212        .await
213        {
214            return blocked;
215        }
216        let start = std::time::Instant::now();
217        let fetch = self.client.fetch_url(url);
218
219        tokio::select! {
220            biased;
221            _ = ctx.token.cancelled() => ToolOutcome::cancelled(),
222            result = fetch => match result {
223                Ok(page) => {
224                    let output = format_fetch(url, &page);
225                    let duration_secs = start.elapsed().as_secs_f64();
226                    let line_count = output.lines().count();
227                    let byte_count = output.len();
228                    let title = if page.title.is_empty() {
229                        None
230                    } else {
231                        Some(page.title)
232                    };
233                    ToolOutcome::success(
234                        output,
235                        format!("{} {} fetched", line_count, if line_count == 1 { "line" } else { "lines" }),
236                        duration_secs,
237                    )
238                    .with_metadata(ToolRunMetadata {
239                        detail: ToolMetadata::WebFetch {
240                            url: url.to_string(),
241                            title,
242                            line_count,
243                            byte_count,
244                        },
245                        line_count: Some(line_count),
246                        byte_count: Some(byte_count),
247                        ..ToolRunMetadata::default()
248                    })
249                },
250                Err(e) => ToolOutcome::error(
251                    format!("web_fetch({}): {}", url, e),
252                    start.elapsed().as_secs_f64(),
253                ),
254            },
255        }
256    }
257}
258
259/// Cap on a single `web_fetch` body (#F46). The raw fetch is bounded only by the
260/// 16 MB HTTP body limit, so without this one URL could dump megabytes into model
261/// context. A full page warrants more room than a `web_search` snippet
262/// (`WEB_CONTENT_MAX_CHARS`), so this mirrors web_search's per-call aggregate
263/// budget. Applied as a byte cap; truncation is char-boundary safe.
264const WEB_FETCH_MAX_CHARS: usize = crate::constants::WEB_SEARCH_AGGREGATE_MAX_CHARS;
265
266/// Truncate a fetched page body to `WEB_FETCH_MAX_CHARS` bytes, char-boundary
267/// safe, appending a marker — consistent with how `web_search` bounds the
268/// content it returns (#F46). Borrows when no truncation is needed.
269fn cap_fetch_content(content: &str) -> std::borrow::Cow<'_, str> {
270    if content.len() <= WEB_FETCH_MAX_CHARS {
271        return std::borrow::Cow::Borrowed(content);
272    }
273    let cut = content.floor_char_boundary(WEB_FETCH_MAX_CHARS);
274    std::borrow::Cow::Owned(format!("{}\n\n...[content truncated]", &content[..cut]))
275}
276
277fn format_fetch(url: &str, page: &WebFetchResult) -> String {
278    let title = if page.title.is_empty() {
279        "(no title)"
280    } else {
281        page.title.as_str()
282    };
283    let content = cap_fetch_content(&page.content);
284    format!("# {}\n\nURL: {}\n\n{}", title, url, content)
285}
286
287fn parse_queries(args: &serde_json::Value) -> Result<Vec<(String, usize)>, String> {
288    if let Some(arr) = args.get("queries").and_then(|v| v.as_array()) {
289        if arr.len() > crate::constants::MAX_BATCH_TOOL_ITEMS {
290            return Err(format!(
291                "web_search: too many queries ({}); cap is {} per call — split the request",
292                arr.len(),
293                crate::constants::MAX_BATCH_TOOL_ITEMS
294            ));
295        }
296        let mut out = Vec::with_capacity(arr.len());
297        for v in arr {
298            let Some(obj) = v.as_object() else {
299                return Err(
300                    "web_search: 'queries' must be an array of {query, max_results}".to_string(),
301                );
302            };
303            let Some(query) = obj.get("query").and_then(|x| x.as_str()) else {
304                return Err("web_search: each query entry needs 'query' (string)".to_string());
305            };
306            let count = obj
307                .get("max_results")
308                .or_else(|| obj.get("result_count"))
309                .and_then(|x| x.as_u64())
310                .unwrap_or(5)
311                .clamp(1, 10) as usize;
312            out.push((query.to_string(), count));
313        }
314        return Ok(out);
315    }
316    if let Some(query) = args.get("query").and_then(|v| v.as_str()) {
317        let count = args
318            .get("max_results")
319            .or_else(|| args.get("result_count"))
320            .and_then(|v| v.as_u64())
321            .unwrap_or(5)
322            .clamp(1, 10) as usize;
323        return Ok(vec![(query.to_string(), count)]);
324    }
325    Err("web_search requires 'query' (string) or 'queries' (array)".to_string())
326}
327
328/// Reject obviously-unsafe fetch URLs client-side before handing them to the
329/// (server-side) Ollama fetch API: only `http`/`https`, and no loopback /
330/// link-local / private / metadata hosts. Defense-in-depth against SSRF-style
331/// abuse via model-supplied URLs.
332fn validate_fetch_url(url: &str) -> Result<(), String> {
333    let parsed = reqwest::Url::parse(url).map_err(|e| format!("invalid URL: {e}"))?;
334    match parsed.scheme() {
335        "http" | "https" => {},
336        other => {
337            return Err(format!(
338                "unsupported URL scheme '{other}' (only http/https allowed)"
339            ));
340        },
341    }
342    let host = parsed
343        .host_str()
344        .ok_or_else(|| "URL has no host".to_string())?;
345    if is_blocked_host(host) {
346        return Err(format!("refusing to fetch internal/loopback host '{host}'"));
347    }
348    Ok(())
349}
350
351/// Cloud-metadata DNS hostnames that resolve (inside the relevant cloud) to a
352/// link-local metadata IP — `169.254.169.254` and friends — but are LEXICALLY
353/// public, so the IP-only `classify_host` waves them through as
354/// [`crate::utils::HostClass::Public`]. We block them by name as well. Matched
355/// after lowercasing + trimming any surrounding `[]` and trailing FQDN dot.
356const METADATA_HOSTNAMES: &[&str] = &[
357    "metadata.google.internal",   // GCP (canonical)
358    "metadata.goog",              // GCP (alternate)
359    "metadata",                   // GCP/Azure short name (http://metadata/ responds)
360    "instance-data",              // AWS (cloud-init alias)
361    "instance-data.ec2.internal", // AWS
362];
363
364/// F57: client-side SSRF denylist for `web_fetch`. This is DEFENSE-IN-DEPTH
365/// ONLY, NOT the authoritative boundary.
366///
367/// `web_fetch` does NOT retrieve the URL from this process. `WebSearchClient::
368/// fetch_url` POSTs the URL to Ollama Cloud's server-side `/api/web_fetch`
369/// endpoint, and Ollama performs the actual fetch. The in-process reqwest
370/// client only ever connects to `ollama.com`. The AUTHORITATIVE SSRF boundary
371/// is therefore SERVER-SIDE (Ollama): a public DNS name that resolves to an
372/// internal address (the DNS-rebinding hole) cannot be closed here — a no-DNS
373/// lexical check can't see where a name resolves, and resolving it locally
374/// wouldn't bind what the remote server independently resolves later anyway.
375///
376/// What we CAN do cheaply, so a model can't trivially aim the server at an
377/// obvious internal target, is reject:
378/// - every non-public IP form via the shared `classify_host` (loopback,
379///   RFC-1918/ULA, link-local incl. `169.254.169.254`, CGNAT, unspecified
380///   `0.0.0.0`/`::`, plus the IPv4-mapped-IPv6 / ULA / link-local-IPv6 / `[::1]`
381///   forms a hand-rolled IPv4 check would miss); and
382/// - the well-known cloud-metadata HOSTNAMES (`metadata.google.internal`, …)
383///   that are lexically public but front a metadata service.
384fn is_blocked_host(host: &str) -> bool {
385    let normalized = host
386        .trim_start_matches('[')
387        .trim_end_matches(']')
388        .trim_end_matches('.')
389        .to_ascii_lowercase();
390    if METADATA_HOSTNAMES.contains(&normalized.as_str()) {
391        return true;
392    }
393    crate::utils::classify_host(host).is_internal()
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399
400    #[test]
401    fn validate_fetch_url_blocks_unsafe_targets() {
402        // #9: scheme + internal-host guards.
403        for bad in [
404            "file:///etc/passwd",
405            "ftp://example.com/x",
406            "http://localhost/admin",
407            "http://127.0.0.1:8080",
408            "http://169.254.169.254/latest/meta-data/",
409            "http://10.0.0.5/",
410            "http://192.168.1.1/",
411            "http://[::1]/",
412            // #27/#80: IPv6/CGNAT bypasses the old IPv4-centric blocklist missed.
413            "http://[::ffff:169.254.169.254]/latest/meta-data/",
414            "http://[fc00::1]/",
415            "http://[fe80::1]/",
416            "http://100.100.100.200/",
417            // F57: cloud-metadata hostnames are lexically public (IP-only
418            // classify_host waves them through) but front a metadata service.
419            "http://metadata.google.internal/computeMetadata/v1/",
420            "http://metadata.goog/",
421            "http://metadata/",
422            "http://instance-data/latest/meta-data/",
423            "https://METADATA.GOOGLE.INTERNAL./",
424            "not a url",
425        ] {
426            assert!(
427                validate_fetch_url(bad).is_err(),
428                "expected reject for {bad:?}",
429            );
430        }
431        for good in [
432            "https://example.com",
433            "http://example.com/page?x=1",
434            "https://docs.rs/serde",
435        ] {
436            assert!(
437                validate_fetch_url(good).is_ok(),
438                "expected accept for {good:?}",
439            );
440        }
441    }
442
443    #[test]
444    fn is_blocked_host_covers_metadata_names_and_ip_forms() {
445        // F57: cloud-metadata HOSTNAMES (lexically public, IP-only
446        // classify_host misses them) are blocked, case/dot-insensitively.
447        for h in [
448            "metadata.google.internal",
449            "metadata.google.internal.", // trailing FQDN dot
450            "Metadata.Google.Internal",  // case-insensitive
451            "metadata.goog",
452            "metadata",
453            "instance-data",
454            "instance-data.ec2.internal",
455        ] {
456            assert!(is_blocked_host(h), "metadata host {h:?} must be blocked");
457        }
458        // Non-public IP forms still go through classify_host (incl. the ones
459        // the task lists as examples that must already be covered).
460        for h in ["0.0.0.0", "::1", "169.254.169.254", "127.0.0.1"] {
461            assert!(is_blocked_host(h), "internal IP {h:?} must be blocked");
462        }
463        // Legitimate public hosts are NOT blocked — including a real `.goog`
464        // domain that merely is not the metadata alias.
465        for h in ["example.com", "docs.rs", "abc.goog", "8.8.8.8"] {
466            assert!(!is_blocked_host(h), "public host {h:?} must be allowed");
467        }
468    }
469
470    #[test]
471    fn format_fetch_caps_long_content() {
472        // F46: a huge page body must be truncated with a marker, not dumped whole.
473        let big = "z".repeat(WEB_FETCH_MAX_CHARS * 2);
474        let page = WebFetchResult {
475            title: "T".to_string(),
476            content: big,
477        };
478        let out = format_fetch("https://example.com", &page);
479        assert!(
480            out.len() < WEB_FETCH_MAX_CHARS + 256,
481            "content must be capped, got {} bytes",
482            out.len()
483        );
484        assert!(out.contains("truncated"), "expected truncation marker");
485
486        // A short page is emitted intact, with no marker.
487        let small = WebFetchResult {
488            title: "T".to_string(),
489            content: "hello world".to_string(),
490        };
491        let out = format_fetch("https://example.com", &small);
492        assert!(out.contains("hello world"));
493        assert!(!out.contains("truncated"));
494    }
495
496    #[test]
497    fn parse_queries_single_form() {
498        let args = serde_json::json!({"query": "rust async", "max_results": 3});
499        let q = parse_queries(&args).unwrap();
500        assert_eq!(q.len(), 1);
501        assert_eq!(q[0].0, "rust async");
502        assert_eq!(q[0].1, 3);
503    }
504
505    #[test]
506    fn parse_queries_array_form() {
507        let args = serde_json::json!({"queries": [
508            {"query": "a", "max_results": 2},
509            {"query": "b", "result_count": 5},
510        ]});
511        let q = parse_queries(&args).unwrap();
512        assert_eq!(q.len(), 2);
513        assert_eq!(q[1].1, 5);
514    }
515
516    #[test]
517    fn parse_queries_missing_errors() {
518        let args = serde_json::json!({});
519        assert!(parse_queries(&args).is_err());
520    }
521
522    #[test]
523    fn parse_queries_clamps_count() {
524        let args = serde_json::json!({"query": "q", "max_results": 999});
525        let q = parse_queries(&args).unwrap();
526        assert_eq!(q[0].1, 10);
527        let args = serde_json::json!({"query": "q", "max_results": 0});
528        let q = parse_queries(&args).unwrap();
529        assert_eq!(q[0].1, 1);
530    }
531
532    #[test]
533    fn parse_queries_rejects_excess_fan_out() {
534        // #90: a single call can't request unbounded fan-out.
535        let many: Vec<_> = (0..crate::constants::MAX_BATCH_TOOL_ITEMS + 1)
536            .map(|i| serde_json::json!({"query": format!("q{i}")}))
537            .collect();
538        let args = serde_json::json!({ "queries": many });
539        assert!(parse_queries(&args).is_err());
540
541        // Exactly at the cap is still accepted.
542        let at_cap: Vec<_> = (0..crate::constants::MAX_BATCH_TOOL_ITEMS)
543            .map(|i| serde_json::json!({"query": format!("q{i}")}))
544            .collect();
545        let args = serde_json::json!({ "queries": at_cap });
546        assert_eq!(
547            parse_queries(&args).unwrap().len(),
548            crate::constants::MAX_BATCH_TOOL_ITEMS
549        );
550    }
551}
mermaid_cli/providers/tool/web.rs

mermaid_cli/providers/tool/
web.rs
