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(&self, args: serde_json::Value, _ctx: ToolCtx) -> Result<ToolOutcome, ToolError> {
145        let parsed: WebFetchArgs = serde_json::from_value(args)
146            .map_err(|e| ToolError::InvalidArguments(format!("Invalid WebFetch args: {}", e)))?;
147        let url = parsed.url.trim();
148        let parsed_url = url::Url::parse(url)
149            .map_err(|e| ToolError::InvalidArguments(format!("Invalid URL: {}", e)))?;
150        let scheme = parsed_url.scheme();
151        if scheme != "http" && scheme != "https" {
152            return Err(ToolError::InvalidArguments(
153                "Only http/https URLs are allowed".to_string(),
154            ));
155        }
156        let Some(host) = parsed_url.host_str() else {
157            return Err(ToolError::InvalidArguments(
158                "URL must include a host".to_string(),
159            ));
160        };
161        if Self::is_disallowed_host(host) {
162            return Err(ToolError::Execution(format!(
163                "Refusing to fetch restricted host: {}",
164                host
165            )));
166        }
167        if host.parse::<IpAddr>().is_err() {
168            let port = parsed_url.port_or_known_default().unwrap_or(80);
169            if Self::host_resolves_to_disallowed_ip(host, port).await? {
170                return Err(ToolError::Execution(format!(
171                    "Refusing to fetch host '{}' because DNS resolved to a restricted IP",
172                    host
173                )));
174            }
175        }
176
177        let client = reqwest::Client::builder()
178            .timeout(Duration::from_secs(30))
179            .build()
180            .map_err(|e| ToolError::Execution(format!("Failed to build HTTP client: {}", e)))?;
181
182        let response = client
183            .get(url)
184            .send()
185            .await
186            .map_err(|e| ToolError::Execution(format!("Failed to fetch URL: {}", e)))?;
187
188        let status = response.status().as_u16();
189        let mut stream = response.bytes_stream();
190        let mut bytes = Vec::with_capacity(64 * 1024);
191        let mut response_truncated = false;
192        while let Some(chunk_result) = stream.next().await {
193            let chunk = chunk_result.map_err(|e| {
194                ToolError::Execution(format!("Failed reading response body: {}", e))
195            })?;
196            if bytes.len() + chunk.len() > MAX_RESPONSE_BYTES {
197                let remaining = MAX_RESPONSE_BYTES.saturating_sub(bytes.len());
198                if remaining > 0 {
199                    bytes.extend_from_slice(&chunk[..remaining]);
200                }
201                response_truncated = true;
202                break;
203            }
204            bytes.extend_from_slice(&chunk);
205        }
206
207        let body = String::from_utf8_lossy(&bytes).to_string();
208
209        let text = Self::strip_html(&body)?;
210        let excerpt: String = text.chars().take(20_000).collect();
211
212        Ok(ToolOutcome::Completed(ToolResult {
213            success: true,
214            result: json!({
215                "url": parsed.url,
216                "status": status,
217                "prompt": parsed.prompt,
218                "content": excerpt,
219                "response_truncated": response_truncated,
220            })
221            .to_string(),
222            display_preference: Some("Collapsible".to_string()),
223            images: Vec::new(),
224        }))
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[test]
233    fn strip_html_strips_scripts_styles_tags_and_collapses_whitespace() {
234        // Scripts and styles are dropped, tags are stripped, and runs of
235        // whitespace are collapsed to a single space, then trimmed.
236        let html = "<html><head><style>body{color:red}</style></head><body>\
237<script>alert(1)</script><h1>Title</h1><p>Hello   world</p></body></html>";
238        assert_eq!(WebFetchTool::strip_html(html).unwrap(), "Title Hello world");
239
240        // A second call exercises the already-initialized (cached) static regexes
241        // and must produce identical semantics.
242        let html2 = "<div>  <b>A</b>  <i>B</i>  </div>";
243        assert_eq!(WebFetchTool::strip_html(html2).unwrap(), "A B");
244    }
245
246    #[test]
247    fn disallowed_host_rejects_local_and_private_targets() {
248        assert!(WebFetchTool::is_disallowed_host("localhost"));
249        assert!(WebFetchTool::is_disallowed_host("api.localhost"));
250        assert!(WebFetchTool::is_disallowed_host("service.local"));
251        assert!(WebFetchTool::is_disallowed_host("127.0.0.1"));
252        assert!(WebFetchTool::is_disallowed_host("10.0.0.1"));
253        assert!(WebFetchTool::is_disallowed_host("192.168.1.1"));
254        assert!(WebFetchTool::is_disallowed_host("::1"));
255        assert!(!WebFetchTool::is_disallowed_host("example.com"));
256        assert!(!WebFetchTool::is_disallowed_host("8.8.8.8"));
257    }
258
259    #[test]
260    fn resolved_ips_include_disallowed_detects_any_private_or_loopback_ip() {
261        assert!(WebFetchTool::resolved_ips_include_disallowed(vec![
262            "8.8.8.8".parse::<IpAddr>().unwrap(),
263            "10.0.0.8".parse::<IpAddr>().unwrap(),
264        ]));
265        assert!(WebFetchTool::resolved_ips_include_disallowed(vec!["::1"
266            .parse::<IpAddr>()
267            .unwrap(),]));
268        assert!(!WebFetchTool::resolved_ips_include_disallowed(vec![
269            "1.1.1.1".parse::<IpAddr>().unwrap(),
270            "8.8.8.8".parse::<IpAddr>().unwrap(),
271        ]));
272    }
273
274    #[tokio::test]
275    async fn execute_rejects_non_http_schemes() {
276        let tool = WebFetchTool::new();
277        let err = tool
278            .invoke(
279                json!({
280                    "url": "file:///etc/passwd",
281                    "prompt": "read"
282                }),
283                ToolCtx::none("t"),
284            )
285            .await
286            .expect_err("non-http scheme should fail");
287
288        assert!(matches!(err, ToolError::InvalidArguments(msg) if msg.contains("http/https")));
289    }
290
291    #[tokio::test]
292    async fn execute_rejects_restricted_hosts_before_network_call() {
293        let tool = WebFetchTool::new();
294        let err = tool
295            .invoke(
296                json!({
297                    "url": "http://localhost:8080",
298                    "prompt": "read"
299                }),
300                ToolCtx::none("t"),
301            )
302            .await
303            .expect_err("localhost should be blocked");
304
305        assert!(matches!(err, ToolError::Execution(msg) if msg.contains("restricted host")));
306    }
307}
bamboo_tools/tools/web_fetch.rs

bamboo_tools/tools/
web_fetch.rs
