brainos_backends/
fetch.rs1use std::sync::Arc;
9use std::time::Duration;
10
11use brain_core::metrics::SubsystemMetrics;
12use cortex::actions::{ActionError, FetchedPage, UrlFetchBackend};
13
14use crate::resilience::{resilient_send, CircuitBreaker};
15
16const MAX_BODY_BYTES: usize = 256 * 1024;
20const MAX_TEXT_BYTES: usize = 4 * 1024;
22const DEFAULT_FETCH_TIMEOUT: Duration = Duration::from_secs(8);
26
27pub struct BasicUrlFetcher {
30 client: reqwest::Client,
31 circuit_breaker: Arc<CircuitBreaker>,
32 max_retries: u32,
33 retry_base_ms: u64,
34}
35
36impl BasicUrlFetcher {
37 pub fn new(resilience: &brain_core::config::ResilienceConfig) -> anyhow::Result<Self> {
38 Self::new_with_metrics(resilience, None)
39 }
40
41 pub fn new_with_metrics(
42 resilience: &brain_core::config::ResilienceConfig,
43 metrics: Option<Arc<SubsystemMetrics>>,
44 ) -> anyhow::Result<Self> {
45 let client = reqwest::Client::builder()
46 .timeout(DEFAULT_FETCH_TIMEOUT)
47 .user_agent("brainos/url-fetch (+https://github.com/keshavashiya/brain)")
48 .redirect(reqwest::redirect::Policy::limited(5))
49 .build()
50 .map_err(|e| anyhow::anyhow!("URL fetch client init failed: {e}"))?;
51 let mut cb = CircuitBreaker::new(
52 "url-fetch",
53 resilience.circuit_breaker_threshold,
54 resilience.circuit_breaker_cooldown_secs,
55 );
56 if let Some(m) = metrics {
57 cb = cb.with_metrics(m);
58 }
59 Ok(Self {
60 client,
61 circuit_breaker: Arc::new(cb),
62 max_retries: resilience.max_retries,
63 retry_base_ms: resilience.retry_base_ms,
64 })
65 }
66}
67
68#[async_trait::async_trait]
69impl UrlFetchBackend for BasicUrlFetcher {
70 async fn fetch(&self, url: &str) -> Result<FetchedPage, ActionError> {
71 if !(url.starts_with("http://") || url.starts_with("https://")) {
72 return Err(ActionError::InvalidArguments(format!(
73 "url-fetch requires http(s):// scheme, got {url}"
74 )));
75 }
76 let client = self.client.clone();
77 let url_owned = url.to_string();
78 let response = resilient_send(
79 || client.get(&url_owned),
80 &self.circuit_breaker,
81 self.max_retries,
82 self.retry_base_ms,
83 )
84 .await?;
85
86 if !response.status().is_success() {
87 return Err(ActionError::ExecutionFailed(format!(
88 "{url} returned HTTP {}",
89 response.status()
90 )));
91 }
92
93 let final_url = response.url().to_string();
94 let bytes = response
95 .bytes()
96 .await
97 .map_err(|e| ActionError::ExecutionFailed(format!("body read failed: {e}")))?;
98
99 let head = &bytes[..bytes.len().min(MAX_BODY_BYTES)];
100 let raw = String::from_utf8_lossy(head);
101 let title = extract_title(&raw).unwrap_or_else(|| final_url.clone());
102 let text = html_to_text(&raw);
103 let text = if text.len() > MAX_TEXT_BYTES {
104 let mut t: String = text.chars().take(MAX_TEXT_BYTES).collect();
105 t.push_str("\n…[truncated]");
106 t
107 } else {
108 text
109 };
110
111 Ok(FetchedPage {
112 url: final_url,
113 title,
114 text,
115 })
116 }
117}
118
119fn extract_title(html: &str) -> Option<String> {
122 let lower = html.to_ascii_lowercase();
123 let start = lower.find("<title")?;
124 let after_open = lower[start..].find('>').map(|i| start + i + 1)?;
126 let end_rel = lower[after_open..].find("</title>")?;
127 let raw = &html[after_open..after_open + end_rel];
128 let cleaned = decode_entities(raw).trim().to_string();
129 if cleaned.is_empty() {
130 None
131 } else {
132 Some(cleaned)
133 }
134}
135
136fn html_to_text(html: &str) -> String {
139 let mut s = String::with_capacity(html.len());
140 let stripped = strip_block(html, "script");
141 let stripped = strip_block(&stripped, "style");
142 let stripped = strip_block(&stripped, "noscript");
143 let mut in_tag = false;
144 for ch in stripped.chars() {
145 match ch {
146 '<' => in_tag = true,
147 '>' => {
148 in_tag = false;
149 s.push(' ');
150 }
151 _ if in_tag => {}
152 _ => s.push(ch),
153 }
154 }
155 let decoded = decode_entities(&s);
156 collapse_whitespace(&decoded)
157}
158
159fn strip_block(input: &str, tag: &str) -> String {
161 let lower = input.to_ascii_lowercase();
162 let open_pat = format!("<{tag}");
163 let close_pat = format!("</{tag}>");
164 let mut out = String::with_capacity(input.len());
165 let mut cursor = 0usize;
166 while cursor < input.len() {
167 match lower[cursor..].find(&open_pat) {
168 None => {
169 out.push_str(&input[cursor..]);
170 break;
171 }
172 Some(rel) => {
173 let abs_open = cursor + rel;
174 out.push_str(&input[cursor..abs_open]);
175 let after_open = match lower[abs_open..].find('>') {
176 Some(i) => abs_open + i + 1,
177 None => break,
178 };
179 match lower[after_open..].find(&close_pat) {
180 Some(end_rel) => {
181 cursor = after_open + end_rel + close_pat.len();
182 }
183 None => {
184 cursor = after_open;
185 }
186 }
187 }
188 }
189 }
190 out
191}
192
193fn decode_entities(s: &str) -> String {
196 s.replace(" ", " ")
197 .replace("&", "&")
198 .replace("<", "<")
199 .replace(">", ">")
200 .replace(""", "\"")
201 .replace("'", "'")
202 .replace("'", "'")
203}
204
205fn collapse_whitespace(s: &str) -> String {
209 let mut out = String::with_capacity(s.len());
210 let mut last_was_space = false;
211 let mut newline_run = 0usize;
212 for ch in s.chars() {
213 if ch == '\n' {
214 newline_run += 1;
215 continue;
216 }
217 if ch.is_whitespace() {
218 if !last_was_space && newline_run == 0 {
219 out.push(' ');
220 last_was_space = true;
221 }
222 continue;
223 }
224 if newline_run >= 2 {
225 out.push('\n');
226 } else if newline_run == 1 && !last_was_space {
227 out.push(' ');
228 }
229 newline_run = 0;
230 last_was_space = false;
231 out.push(ch);
232 }
233 out.trim().to_string()
234}
235
236#[cfg(test)]
237mod tests {
238 use super::*;
239
240 #[test]
241 fn extract_title_handles_attrs_and_case() {
242 let html = "<html><head><TITLE id=\"x\">Hello World</TITLE></head><body/></html>";
243 assert_eq!(extract_title(html).as_deref(), Some("Hello World"));
244 }
245
246 #[test]
247 fn extract_title_returns_none_when_missing() {
248 assert!(extract_title("<html><body>no title</body></html>").is_none());
249 }
250
251 #[test]
252 fn html_to_text_strips_scripts_and_styles() {
253 let html = "<html><head><style>body{color:red}</style></head>\
254 <body><script>alert(1)</script><p>Hello <b>world</b>!</p></body></html>";
255 let txt = html_to_text(html);
256 assert!(!txt.contains("alert"));
257 assert!(!txt.contains("color:red"));
258 assert!(txt.contains("Hello"));
259 assert!(txt.contains("world"));
260 }
261
262 #[test]
263 fn html_to_text_decodes_common_entities() {
264 let html = "<p>tom & jerry <3 coffee</p>";
265 let txt = html_to_text(html);
266 assert!(txt.contains("tom & jerry"));
267 assert!(txt.contains("<3"));
268 }
269
270 #[test]
271 fn collapse_whitespace_keeps_paragraph_breaks() {
272 let s = "alpha\n\nbeta gamma\n\n\ndelta";
273 let out = collapse_whitespace(s);
274 assert_eq!(out, "alpha\nbeta gamma\ndelta");
276 }
277
278 fn fast_resilience() -> brain_core::config::ResilienceConfig {
279 brain_core::config::ResilienceConfig {
280 max_retries: 0,
281 retry_base_ms: 1,
282 circuit_breaker_threshold: 5,
283 circuit_breaker_cooldown_secs: 60,
284 }
285 }
286
287 #[tokio::test]
288 async fn fetch_returns_text_and_title() {
289 let mut server = mockito::Server::new_async().await;
290 let _mock = server
291 .mock("GET", "/page")
292 .with_status(200)
293 .with_header("content-type", "text/html; charset=utf-8")
294 .with_body(
295 "<html><head><title>Demo Page</title></head>\
296 <body><script>x=1</script><h1>Hi</h1><p>Some & text.</p></body></html>",
297 )
298 .create_async()
299 .await;
300
301 let fetcher = BasicUrlFetcher::new(&fast_resilience()).unwrap();
302 let url = format!("{}/page", server.url());
303 let page = fetcher.fetch(&url).await.unwrap();
304 assert_eq!(page.title, "Demo Page");
305 assert!(page.text.contains("Hi"));
306 assert!(page.text.contains("Some & text"));
307 assert!(!page.text.contains("x=1"));
308 }
309
310 #[tokio::test]
311 async fn fetch_rejects_non_http_schemes() {
312 let fetcher = BasicUrlFetcher::new(&fast_resilience()).unwrap();
313 let err = fetcher.fetch("file:///etc/passwd").await.unwrap_err();
314 assert!(matches!(err, ActionError::InvalidArguments(_)));
315 }
316
317 #[tokio::test]
318 async fn fetch_surfaces_non_2xx_as_error() {
319 let mut server = mockito::Server::new_async().await;
320 let _mock = server
321 .mock("GET", "/missing")
322 .with_status(404)
323 .create_async()
324 .await;
325 let fetcher = BasicUrlFetcher::new(&fast_resilience()).unwrap();
326 let url = format!("{}/missing", server.url());
327 assert!(fetcher.fetch(&url).await.is_err());
328 }
329}