bamboo_tools/tools/
web_fetch.rs

1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolError, ToolResult};
3use futures::StreamExt;
4use regex::Regex;
5use serde::Deserialize;
6use serde_json::json;
7use std::net::IpAddr;
8use std::time::Duration;
9
10const MAX_RESPONSE_BYTES: usize = 1_000_000;
11
12#[derive(Debug, Deserialize)]
13struct WebFetchArgs {
14    url: String,
15    prompt: String,
16}
17
18pub struct WebFetchTool;
19
20impl WebFetchTool {
21    pub fn new() -> Self {
22        Self
23    }
24
25    fn strip_html(input: &str) -> Result<String, ToolError> {
26        let script_re = Regex::new(r"(?is)<script[^>]*>.*?</script>")
27            .map_err(|e| ToolError::Execution(format!("Failed to compile script regex: {}", e)))?;
28        let style_re = Regex::new(r"(?is)<style[^>]*>.*?</style>")
29            .map_err(|e| ToolError::Execution(format!("Failed to compile style regex: {}", e)))?;
30        let tag_re = Regex::new(r"(?is)<[^>]+>")
31            .map_err(|e| ToolError::Execution(format!("Failed to compile tag regex: {}", e)))?;
32        let whitespace_re = Regex::new(r"[ \t\n\r]+").map_err(|e| {
33            ToolError::Execution(format!("Failed to compile whitespace regex: {}", e))
34        })?;
35
36        let without_scripts = script_re.replace_all(input, " ");
37        let without_styles = style_re.replace_all(&without_scripts, " ");
38        let without_tags = tag_re.replace_all(&without_styles, " ");
39        Ok(whitespace_re
40            .replace_all(&without_tags, " ")
41            .trim()
42            .to_string())
43    }
44
45    fn is_disallowed_ip(ip: IpAddr) -> bool {
46        match ip {
47            IpAddr::V4(ipv4) => {
48                ipv4.is_loopback()
49                    || ipv4.is_private()
50                    || ipv4.is_link_local()
51                    || ipv4.is_multicast()
52                    || ipv4.is_unspecified()
53            }
54            IpAddr::V6(ipv6) => {
55                ipv6.is_loopback()
56                    || ipv6.is_multicast()
57                    || ipv6.is_unspecified()
58                    || ipv6.is_unique_local()
59                    || ipv6.is_unicast_link_local()
60            }
61        }
62    }
63
64    fn is_disallowed_host(host: &str) -> bool {
65        let host = host.trim().to_ascii_lowercase();
66        if host == "localhost" || host.ends_with(".localhost") || host.ends_with(".local") {
67            return true;
68        }
69
70        let Ok(ip) = host.parse::<IpAddr>() else {
71            return false;
72        };
73        Self::is_disallowed_ip(ip)
74    }
75
76    fn resolved_ips_include_disallowed<I>(ips: I) -> bool
77    where
78        I: IntoIterator<Item = IpAddr>,
79    {
80        ips.into_iter().any(Self::is_disallowed_ip)
81    }
82
83    async fn host_resolves_to_disallowed_ip(host: &str, port: u16) -> Result<bool, ToolError> {
84        let addrs = tokio::net::lookup_host((host, port)).await.map_err(|e| {
85            ToolError::Execution(format!("Failed to resolve host '{}': {}", host, e))
86        })?;
87        let ips: Vec<IpAddr> = addrs.map(|addr| addr.ip()).collect();
88        Ok(Self::resolved_ips_include_disallowed(ips))
89    }
90}
91
92impl Default for WebFetchTool {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98#[async_trait]
99impl Tool for WebFetchTool {
100    fn name(&self) -> &str {
101        "WebFetch"
102    }
103
104    fn description(&self) -> &str {
105        "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."
106    }
107
108    fn mutability(&self) -> crate::ToolMutability {
109        crate::ToolMutability::ReadOnly
110    }
111
112    fn concurrency_safe(&self) -> bool {
113        true
114    }
115
116    fn parameters_schema(&self) -> serde_json::Value {
117        json!({
118            "type": "object",
119            "properties": {
120                "url": {
121                    "type": "string",
122                    "format": "uri",
123                    "description": "The URL to fetch"
124                },
125                "prompt": {
126                    "type": "string",
127                    "description": "Caller-supplied extraction intent note; echoed in output for downstream processing"
128                }
129            },
130            "required": ["url", "prompt"],
131            "additionalProperties": false
132        })
133    }
134
135    async fn execute(&self, args: serde_json::Value) -> Result<ToolResult, ToolError> {
136        let parsed: WebFetchArgs = serde_json::from_value(args)
137            .map_err(|e| ToolError::InvalidArguments(format!("Invalid WebFetch args: {}", e)))?;
138        let url = parsed.url.trim();
139        let parsed_url = url::Url::parse(url)
140            .map_err(|e| ToolError::InvalidArguments(format!("Invalid URL: {}", e)))?;
141        let scheme = parsed_url.scheme();
142        if scheme != "http" && scheme != "https" {
143            return Err(ToolError::InvalidArguments(
144                "Only http/https URLs are allowed".to_string(),
145            ));
146        }
147        let Some(host) = parsed_url.host_str() else {
148            return Err(ToolError::InvalidArguments(
149                "URL must include a host".to_string(),
150            ));
151        };
152        if Self::is_disallowed_host(host) {
153            return Err(ToolError::Execution(format!(
154                "Refusing to fetch restricted host: {}",
155                host
156            )));
157        }
158        if host.parse::<IpAddr>().is_err() {
159            let port = parsed_url.port_or_known_default().unwrap_or(80);
160            if Self::host_resolves_to_disallowed_ip(host, port).await? {
161                return Err(ToolError::Execution(format!(
162                    "Refusing to fetch host '{}' because DNS resolved to a restricted IP",
163                    host
164                )));
165            }
166        }
167
168        let client = reqwest::Client::builder()
169            .timeout(Duration::from_secs(30))
170            .build()
171            .map_err(|e| ToolError::Execution(format!("Failed to build HTTP client: {}", e)))?;
172
173        let response = client
174            .get(url)
175            .send()
176            .await
177            .map_err(|e| ToolError::Execution(format!("Failed to fetch URL: {}", e)))?;
178
179        let status = response.status().as_u16();
180        let mut stream = response.bytes_stream();
181        let mut bytes = Vec::with_capacity(64 * 1024);
182        let mut response_truncated = false;
183        while let Some(chunk_result) = stream.next().await {
184            let chunk = chunk_result.map_err(|e| {
185                ToolError::Execution(format!("Failed reading response body: {}", e))
186            })?;
187            if bytes.len() + chunk.len() > MAX_RESPONSE_BYTES {
188                let remaining = MAX_RESPONSE_BYTES.saturating_sub(bytes.len());
189                if remaining > 0 {
190                    bytes.extend_from_slice(&chunk[..remaining]);
191                }
192                response_truncated = true;
193                break;
194            }
195            bytes.extend_from_slice(&chunk);
196        }
197
198        let body = String::from_utf8_lossy(&bytes).to_string();
199
200        let text = Self::strip_html(&body)?;
201        let excerpt: String = text.chars().take(20_000).collect();
202
203        Ok(ToolResult {
204            success: true,
205            result: json!({
206                "url": parsed.url,
207                "status": status,
208                "prompt": parsed.prompt,
209                "content": excerpt,
210                "response_truncated": response_truncated,
211            })
212            .to_string(),
213            display_preference: Some("Collapsible".to_string()),
214        })
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    #[test]
223    fn disallowed_host_rejects_local_and_private_targets() {
224        assert!(WebFetchTool::is_disallowed_host("localhost"));
225        assert!(WebFetchTool::is_disallowed_host("api.localhost"));
226        assert!(WebFetchTool::is_disallowed_host("service.local"));
227        assert!(WebFetchTool::is_disallowed_host("127.0.0.1"));
228        assert!(WebFetchTool::is_disallowed_host("10.0.0.1"));
229        assert!(WebFetchTool::is_disallowed_host("192.168.1.1"));
230        assert!(WebFetchTool::is_disallowed_host("::1"));
231        assert!(!WebFetchTool::is_disallowed_host("example.com"));
232        assert!(!WebFetchTool::is_disallowed_host("8.8.8.8"));
233    }
234
235    #[test]
236    fn resolved_ips_include_disallowed_detects_any_private_or_loopback_ip() {
237        assert!(WebFetchTool::resolved_ips_include_disallowed(vec![
238            "8.8.8.8".parse::<IpAddr>().unwrap(),
239            "10.0.0.8".parse::<IpAddr>().unwrap(),
240        ]));
241        assert!(WebFetchTool::resolved_ips_include_disallowed(vec!["::1"
242            .parse::<IpAddr>()
243            .unwrap(),]));
244        assert!(!WebFetchTool::resolved_ips_include_disallowed(vec![
245            "1.1.1.1".parse::<IpAddr>().unwrap(),
246            "8.8.8.8".parse::<IpAddr>().unwrap(),
247        ]));
248    }
249
250    #[tokio::test]
251    async fn execute_rejects_non_http_schemes() {
252        let tool = WebFetchTool::new();
253        let err = tool
254            .execute(json!({
255                "url": "file:///etc/passwd",
256                "prompt": "read"
257            }))
258            .await
259            .expect_err("non-http scheme should fail");
260
261        assert!(matches!(err, ToolError::InvalidArguments(msg) if msg.contains("http/https")));
262    }
263
264    #[tokio::test]
265    async fn execute_rejects_restricted_hosts_before_network_call() {
266        let tool = WebFetchTool::new();
267        let err = tool
268            .execute(json!({
269                "url": "http://localhost:8080",
270                "prompt": "read"
271            }))
272            .await
273            .expect_err("localhost should be blocked");
274
275        assert!(matches!(err, ToolError::Execution(msg) if msg.contains("restricted host")));
276    }
277}
bamboo_tools/tools/web_fetch.rs

bamboo_tools/tools/
web_fetch.rs
