Skip to main content

brainos_backends/
fetch.rs

1//! URL-fetch backend — bounded HTTP GET + HTML→text reduction.
2//!
3//! Used by `cortex::actions::ActionDispatcher` to inline the body of
4//! user-pasted URLs alongside web search results. The fetcher imposes
5//! its own timeout and body cap so a runaway page can't blow the
6//! LLM context window or stall the answer.
7
8use 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
16/// Hard cap on bytes read from any single page. Anything bigger is
17/// truncated *after* HTML stripping so most real pages still produce
18/// useful text.
19const MAX_BODY_BYTES: usize = 256 * 1024;
20/// Hard cap on cleaned-text bytes returned to the dispatcher.
21const MAX_TEXT_BYTES: usize = 4 * 1024;
22/// Connect/read timeout for a single fetch. Tight on purpose — the
23/// answering LLM is waiting and a slow site shouldn't hold up the
24/// whole response.
25const DEFAULT_FETCH_TIMEOUT: Duration = Duration::from_secs(8);
26
27/// Default HTTP fetcher: GET, follow redirects, strip HTML, return
28/// truncated text. No JavaScript execution, no auth, no cookies.
29pub 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
119/// Pull the `<title>` content from raw HTML. Case-insensitive on the
120/// tag name so `<TITLE>` and `<Title>` both match.
121fn extract_title(html: &str) -> Option<String> {
122    let lower = html.to_ascii_lowercase();
123    let start = lower.find("<title")?;
124    // Skip past the closing `>` of the opening tag (which may include attrs).
125    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
136/// Strip HTML tags + scripts/styles + decode common entities. Not a
137/// full parser — good enough to feed an LLM a readable excerpt.
138fn 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
159/// Remove every `<tag …>…</tag>` block (case-insensitive on the tag).
160fn 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
193/// Decode the handful of HTML entities common enough to matter for
194/// readable LLM context. Anything else is passed through verbatim.
195fn decode_entities(s: &str) -> String {
196    s.replace("&nbsp;", " ")
197        .replace("&amp;", "&")
198        .replace("&lt;", "<")
199        .replace("&gt;", ">")
200        .replace("&quot;", "\"")
201        .replace("&#39;", "'")
202        .replace("&apos;", "'")
203}
204
205/// Collapse runs of whitespace (including newlines) into single spaces,
206/// preserving paragraph breaks (a sequence with two or more newlines
207/// becomes a single `\n`).
208fn 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 &amp; jerry &lt;3 &nbsp; 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        // Multi-newline sequences collapse to a single \n; runs of spaces collapse to one.
275        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 &amp; 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}