bamboo_tools/tools/
web_fetch.rs

1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolClass, ToolCtx, ToolError, ToolOutcome, ToolResult};
3use futures::StreamExt;
4use regex::Regex;
5use serde::Deserialize;
6use serde_json::json;
7use std::net::IpAddr;
8use std::sync::OnceLock;
9use std::time::Duration;
10
11const MAX_RESPONSE_BYTES: usize = 1_000_000;
12
13// Static, compile-time-constant patterns: compile each exactly once and reuse.
14// `expect` is safe here because the patterns are hardcoded and verified valid.
15static SCRIPT_RE: OnceLock<Regex> = OnceLock::new();
16static STYLE_RE: OnceLock<Regex> = OnceLock::new();
17static TAG_RE: OnceLock<Regex> = OnceLock::new();
18static WHITESPACE_RE: OnceLock<Regex> = OnceLock::new();
19
20#[derive(Debug, Deserialize)]
21struct WebFetchArgs {
22    url: String,
23    prompt: String,
24}
25
26pub struct WebFetchTool;
27
28impl WebFetchTool {
29    pub fn new() -> Self {
30        Self
31    }
32
33    fn strip_html(input: &str) -> Result<String, ToolError> {
34        let script_re = SCRIPT_RE.get_or_init(|| {
35            Regex::new(r"(?is)<script[^>]*>.*?</script>").expect("valid static regex")
36        });
37        let style_re = STYLE_RE.get_or_init(|| {
38            Regex::new(r"(?is)<style[^>]*>.*?</style>").expect("valid static regex")
39        });
40        let tag_re =
41            TAG_RE.get_or_init(|| Regex::new(r"(?is)<[^>]+>").expect("valid static regex"));
42        let whitespace_re =
43            WHITESPACE_RE.get_or_init(|| Regex::new(r"[ \t\n\r]+").expect("valid static regex"));
44
45        let without_scripts = script_re.replace_all(input, " ");
46        let without_styles = style_re.replace_all(&without_scripts, " ");
47        let without_tags = tag_re.replace_all(&without_styles, " ");
48        Ok(whitespace_re
49            .replace_all(&without_tags, " ")
50            .trim()
51            .to_string())
52    }
53
54    fn is_disallowed_ip(ip: IpAddr) -> bool {
55        match ip {
56            IpAddr::V4(ipv4) => {
57                ipv4.is_loopback()
58                    || ipv4.is_private()
59                    || ipv4.is_link_local()
60                    || ipv4.is_multicast()
61                    || ipv4.is_unspecified()
62            }
63            IpAddr::V6(ipv6) => {
64                let segments = ipv6.segments();
65                let first = segments[0];
66                let is_unique_local = (first & 0xfe00) == 0xfc00;
67                let is_unicast_link_local = (first & 0xffc0) == 0xfe80;
68                ipv6.is_loopback()
69                    || ipv6.is_multicast()
70                    || ipv6.is_unspecified()
71                    || is_unique_local
72                    || is_unicast_link_local
73            }
74        }
75    }
76
77    fn is_disallowed_host(host: &str) -> bool {
78        let host = host.trim().to_ascii_lowercase();
79        if host == "localhost" || host.ends_with(".localhost") || host.ends_with(".local") {
80            return true;
81        }
82
83        let Ok(ip) = host.parse::<IpAddr>() else {
84            return false;
85        };
86        Self::is_disallowed_ip(ip)
87    }
88
89    fn resolved_ips_include_disallowed<I>(ips: I) -> bool
90    where
91        I: IntoIterator<Item = IpAddr>,
92    {
93        ips.into_iter().any(Self::is_disallowed_ip)
94    }
95
96    async fn host_resolves_to_disallowed_ip(host: &str, port: u16) -> Result<bool, ToolError> {
97        let addrs = tokio::net::lookup_host((host, port)).await.map_err(|e| {
98            ToolError::Execution(format!("Failed to resolve host '{}': {}", host, e))
99        })?;
100        let ips: Vec<IpAddr> = addrs.map(|addr| addr.ip()).collect();
101        Ok(Self::resolved_ips_include_disallowed(ips))
102    }
103}
104
105impl Default for WebFetchTool {
106    fn default() -> Self {
107        Self::new()
108    }
109}
110
111#[async_trait]
112impl Tool for WebFetchTool {
113    fn name(&self) -> &str {
114        "WebFetch"
115    }
116
117    fn description(&self) -> &str {
118        "Fetch an HTTP(S) URL and return a cleaned text excerpt plus metadata. The `prompt` field is caller context only; this tool does not run an extra model."
119    }
120
121    fn classify(&self, _args: &serde_json::Value) -> ToolClass {
122        ToolClass::READONLY_PARALLEL.promotable()
123    }
124
125    fn parameters_schema(&self) -> serde_json::Value {
126        json!({
127            "type": "object",
128            "properties": {
129                "url": {
130                    "type": "string",
131                    "format": "uri",
132                    "description": "The URL to fetch"
133                },
134                "prompt": {
135                    "type": "string",
136                    "description": "Caller-supplied extraction intent note; echoed in output for downstream processing"
137                }
138            },
139            "required": ["url", "prompt"],
140            "additionalProperties": false
141        })
142    }
143
144    async fn invoke(
145        &self,
146        args: serde_json::Value,
147        _ctx: ToolCtx,
148    ) -> Result<ToolOutcome, ToolError> {
149        let parsed: WebFetchArgs = serde_json::from_value(args)
150            .map_err(|e| ToolError::InvalidArguments(format!("Invalid WebFetch args: {}", e)))?;
151        let url = parsed.url.trim();
152        let parsed_url = url::Url::parse(url)
153            .map_err(|e| ToolError::InvalidArguments(format!("Invalid URL: {}", e)))?;
154        let scheme = parsed_url.scheme();
155        if scheme != "http" && scheme != "https" {
156            return Err(ToolError::InvalidArguments(
157                "Only http/https URLs are allowed".to_string(),
158            ));
159        }
160        let Some(host) = parsed_url.host_str() else {
161            return Err(ToolError::InvalidArguments(
162                "URL must include a host".to_string(),
163            ));
164        };
165        if Self::is_disallowed_host(host) {
166            return Err(ToolError::Execution(format!(
167                "Refusing to fetch restricted host: {}",
168                host
169            )));
170        }
171        if host.parse::<IpAddr>().is_err() {
172            let port = parsed_url.port_or_known_default().unwrap_or(80);
173            if Self::host_resolves_to_disallowed_ip(host, port).await? {
174                return Err(ToolError::Execution(format!(
175                    "Refusing to fetch host '{}' because DNS resolved to a restricted IP",
176                    host
177                )));
178            }
179        }
180
181        let client = reqwest::Client::builder()
182            .timeout(Duration::from_secs(30))
183            .build()
184            .map_err(|e| ToolError::Execution(format!("Failed to build HTTP client: {}", e)))?;
185
186        let response = client
187            .get(url)
188            .send()
189            .await
190            .map_err(|e| ToolError::Execution(format!("Failed to fetch URL: {}", e)))?;
191
192        let status = response.status().as_u16();
193        let mut stream = response.bytes_stream();
194        let mut bytes = Vec::with_capacity(64 * 1024);
195        let mut response_truncated = false;
196        while let Some(chunk_result) = stream.next().await {
197            let chunk = chunk_result.map_err(|e| {
198                ToolError::Execution(format!("Failed reading response body: {}", e))
199            })?;
200            if bytes.len() + chunk.len() > MAX_RESPONSE_BYTES {
201                let remaining = MAX_RESPONSE_BYTES.saturating_sub(bytes.len());
202                if remaining > 0 {
203                    bytes.extend_from_slice(&chunk[..remaining]);
204                }
205                response_truncated = true;
206                break;
207            }
208            bytes.extend_from_slice(&chunk);
209        }
210
211        let body = String::from_utf8_lossy(&bytes).to_string();
212
213        let text = Self::strip_html(&body)?;
214        let excerpt: String = text.chars().take(20_000).collect();
215
216        Ok(ToolOutcome::Completed(ToolResult {
217            success: true,
218            result: json!({
219                "url": parsed.url,
220                "status": status,
221                "prompt": parsed.prompt,
222                "content": excerpt,
223                "response_truncated": response_truncated,
224            })
225            .to_string(),
226            display_preference: Some("Collapsible".to_string()),
227            images: Vec::new(),
228        }))
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[test]
237    fn strip_html_strips_scripts_styles_tags_and_collapses_whitespace() {
238        // Scripts and styles are dropped, tags are stripped, and runs of
239        // whitespace are collapsed to a single space, then trimmed.
240        let html = "<html><head><style>body{color:red}</style></head><body>\
241<script>alert(1)</script><h1>Title</h1><p>Hello   world</p></body></html>";
242        assert_eq!(WebFetchTool::strip_html(html).unwrap(), "Title Hello world");
243
244        // A second call exercises the already-initialized (cached) static regexes
245        // and must produce identical semantics.
246        let html2 = "<div>  <b>A</b>  <i>B</i>  </div>";
247        assert_eq!(WebFetchTool::strip_html(html2).unwrap(), "A B");
248    }
249
250    #[test]
251    fn disallowed_host_rejects_local_and_private_targets() {
252        assert!(WebFetchTool::is_disallowed_host("localhost"));
253        assert!(WebFetchTool::is_disallowed_host("api.localhost"));
254        assert!(WebFetchTool::is_disallowed_host("service.local"));
255        assert!(WebFetchTool::is_disallowed_host("127.0.0.1"));
256        assert!(WebFetchTool::is_disallowed_host("10.0.0.1"));
257        assert!(WebFetchTool::is_disallowed_host("192.168.1.1"));
258        assert!(WebFetchTool::is_disallowed_host("::1"));
259        assert!(!WebFetchTool::is_disallowed_host("example.com"));
260        assert!(!WebFetchTool::is_disallowed_host("8.8.8.8"));
261    }
262
263    #[test]
264    fn resolved_ips_include_disallowed_detects_any_private_or_loopback_ip() {
265        assert!(WebFetchTool::resolved_ips_include_disallowed(vec![
266            "8.8.8.8".parse::<IpAddr>().unwrap(),
267            "10.0.0.8".parse::<IpAddr>().unwrap(),
268        ]));
269        assert!(WebFetchTool::resolved_ips_include_disallowed(vec!["::1"
270            .parse::<IpAddr>()
271            .unwrap(),]));
272        assert!(!WebFetchTool::resolved_ips_include_disallowed(vec![
273            "1.1.1.1".parse::<IpAddr>().unwrap(),
274            "8.8.8.8".parse::<IpAddr>().unwrap(),
275        ]));
276    }
277
278    #[tokio::test]
279    async fn execute_rejects_non_http_schemes() {
280        let tool = WebFetchTool::new();
281        let err = tool
282            .invoke(
283                json!({
284                    "url": "file:///etc/passwd",
285                    "prompt": "read"
286                }),
287                ToolCtx::none("t"),
288            )
289            .await
290            .expect_err("non-http scheme should fail");
291
292        assert!(matches!(err, ToolError::InvalidArguments(msg) if msg.contains("http/https")));
293    }
294
295    #[tokio::test]
296    async fn execute_rejects_restricted_hosts_before_network_call() {
297        let tool = WebFetchTool::new();
298        let err = tool
299            .invoke(
300                json!({
301                    "url": "http://localhost:8080",
302                    "prompt": "read"
303                }),
304                ToolCtx::none("t"),
305            )
306            .await
307            .expect_err("localhost should be blocked");
308
309        assert!(matches!(err, ToolError::Execution(msg) if msg.contains("restricted host")));
310    }
311}
bamboo_tools/tools/web_fetch.rs

bamboo_tools/tools/
web_fetch.rs
