返回 CodeWhale
extract.rs
根目录 / crates / tui / src / tools / web / extract.rs
1 //! Content-type routing and readable-document extraction for web tools.
2 //!
3 //! Networking deliberately lives elsewhere. This module accepts already
4 //! fetched bytes and turns them into one normalized document so `fetch_url`
5 //! and `web.run` cannot disagree about HTML, Markdown, PDF, or media handling.
6
7 use std::sync::OnceLock;
8
9 use encoding_rs::{Encoding, UTF_8, UTF_16BE, UTF_16LE};
10 use regex::Regex;
11 use tokio_util::sync::CancellationToken;
12
13 use crate::tools::spec::ToolError;
14
15 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
16 pub(crate) enum DocumentKind {
17 Html,
18 Markdown,
19 Text,
20 Pdf,
21 Media,
22 }
23
24 #[derive(Debug, Clone)]
25 pub(crate) struct ExtractedDocument {
26 pub(crate) kind: DocumentKind,
27 pub(crate) title: Option<String>,
28 pub(crate) text: String,
29 pub(crate) markdown: String,
30 /// Readability-cleaned HTML. `web.run` consumes this to retain clickable
31 /// links while avoiding page chrome and consent-banner noise.
32 pub(crate) cleaned_html: Option<String>,
33 pub(crate) pdf_pages: Option<Vec<Vec<String>>>,
34 /// Validated extension for image/audio/video artifacts.
35 pub(crate) media_extension: Option<&'static str>,
36 }
37
38 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
39 struct MediaSignature {
40 extension: &'static str,
41 family: MediaFamily,
42 }
43
44 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
45 enum MediaFamily {
46 Image,
47 Audio,
48 Video,
49 }
50
51 static TITLE_RE: OnceLock<Regex> = OnceLock::new();
52 static FALLBACK_RE: OnceLock<Vec<Regex>> = OnceLock::new();
53 static PAGE_CHROME_RE: OnceLock<Regex> = OnceLock::new();
54 static TAG_RE: OnceLock<Regex> = OnceLock::new();
55 static WHITESPACE_RE: OnceLock<Regex> = OnceLock::new();
56
57 /// HTML's encoding declaration prescan is intentionally small. Keeping the
58 /// bound here prevents a late body string, script, or injected fragment from
59 /// changing how an already-started document is decoded.
60 const HTML_ENCODING_SNIFF_BYTES: usize = 1_024;
61
62 pub(crate) async fn extract_document(
63 url: &str,
64 content_type: Option<&str>,
65 bytes: &[u8],
66 cancel: Option<&CancellationToken>,
67 ) -> Result<ExtractedDocument, ToolError> {
68 extract_document_with_pdf_command(
69 url,
70 content_type,
71 bytes,
72 super::super::pdf::PdfTextCommand::system(cancel),
73 )
74 .await
75 }
76
77 pub(crate) async fn extract_document_with_pdf_command(
78 url: &str,
79 content_type: Option<&str>,
80 bytes: &[u8],
81 pdf_command: super::super::pdf::PdfTextCommand<'_>,
82 ) -> Result<ExtractedDocument, ToolError> {
83 let declared = normalized_content_type(content_type);
84 let declared = declared.as_deref();
85
86 if bytes.is_empty() {
87 return Ok(ExtractedDocument {
88 kind: DocumentKind::Text,
89 title: None,
90 text: String::new(),
91 markdown: String::new(),
92 cleaned_html: None,
93 pdf_pages: None,
94 media_extension: None,
95 });
96 }
97
98 if validate_pdf_response(url, content_type, bytes)? {
99 return extract_pdf(bytes, pdf_command).await;
100 }
101
102 if let Some(signature) = sniff_media(bytes) {
103 if let Some(declared_family) = declared_media_family(declared)
104 && declared_family != signature.family
105 {
106 return Err(ToolError::execution_failed(format!(
107 "Response media type `{}` did not match its bytes",
108 declared.unwrap_or("unknown")
109 )));
110 }
111 return Ok(ExtractedDocument {
112 kind: DocumentKind::Media,
113 title: None,
114 text: String::new(),
115 markdown: String::new(),
116 cleaned_html: None,
117 pdf_pages: None,
118 media_extension: Some(signature.extension),
119 });
120 }
121
122 if declared_media_family(declared).is_some() {
123 return Err(ToolError::execution_failed(format!(
124 "Response claimed media type `{}`, but its bytes did not match a supported media signature",
125 declared.unwrap_or("unknown")
126 )));
127 }
128
129 let sniff_html = should_sniff_html_encoding(declared, url, bytes);
130 let body = decode_response_body(bytes, content_type, sniff_html)?;
131 if sniff_html || is_html(declared, url, &body) {
132 return extract_html(url, &body);
133 }
134 if is_markdown(declared, url) {
135 return Ok(ExtractedDocument {
136 kind: DocumentKind::Markdown,
137 title: markdown_title(&body),
138 text: body.clone(),
139 markdown: body,
140 cleaned_html: None,
141 pdf_pages: None,
142 media_extension: None,
143 });
144 }
145 if is_textual(declared, url) {
146 return Ok(ExtractedDocument {
147 kind: DocumentKind::Text,
148 title: None,
149 text: body.clone(),
150 markdown: body,
151 cleaned_html: None,
152 pdf_pages: None,
153 media_extension: None,
154 });
155 }
156
157 Err(ToolError::execution_failed(format!(
158 "Unsupported binary response type `{}`; use a dedicated download tool",
159 declared.unwrap_or("unknown")
160 )))
161 }
162
163 pub(crate) fn validate_pdf_response(
164 url: &str,
165 content_type: Option<&str>,
166 bytes: &[u8],
167 ) -> Result<bool, ToolError> {
168 let declared = normalized_content_type(content_type);
169 let declared = declared.as_deref();
170 let signed = looks_like_pdf(bytes);
171 if signed && declared_media_family(declared).is_some() {
172 return Err(ToolError::execution_failed(format!(
173 "Response media type `{}` did not match its PDF bytes",
174 declared.unwrap_or("unknown")
175 )));
176 }
177 let claimed = signed || declared == Some("application/pdf") || url_is_pdf(url);
178 if claimed && !signed {
179 return Err(ToolError::execution_failed(
180 "Response claimed to be a PDF, but its bytes did not contain a PDF signature",
181 ));
182 }
183 Ok(claimed)
184 }
185
186 fn extract_html(url: &str, html: &str) -> Result<ExtractedDocument, ToolError> {
187 let parsed_url = reqwest::Url::parse(url)
188 .map_err(|err| ToolError::invalid_input(format!("invalid URL: {err}")))?;
189 let original_title = html_title(html);
190 let mut input = html.as_bytes();
191 let readable = readability::extractor::extract(&mut input, &parsed_url).ok();
192
193 let readable_html = readable
194 .as_ref()
195 .map(|product| product.content.trim())
196 .filter(|content| meaningful_html(content))
197 .map(ToOwned::to_owned);
198 let cleaned_html = readable_html
199 .or_else(|| fallback_main_html(html))
200 .ok_or_else(|| js_required_error(url))?;
201 let markdown = htmd::convert(&cleaned_html).map_err(|err| {
202 ToolError::execution_failed(format!(
203 "Failed to convert readable HTML to Markdown: {err}"
204 ))
205 })?;
206 let text = readable
207 .as_ref()
208 .map(|product| normalize_text(&product.text))
209 .filter(|content| meaningful_text(content))
210 .unwrap_or_else(|| html_to_plain_text(&cleaned_html));
211
212 if !meaningful_text(&text) && !meaningful_text(&markdown) {
213 return Err(js_required_error(url));
214 }
215
216 let title = readable
217 .map(|product| normalize_text(&product.title))
218 .filter(|value| !value.is_empty())
219 .or(original_title);
220
221 Ok(ExtractedDocument {
222 kind: DocumentKind::Html,
223 title,
224 text,
225 markdown,
226 cleaned_html: Some(cleaned_html),
227 pdf_pages: None,
228 media_extension: None,
229 })
230 }
231
232 fn fallback_main_html(html: &str) -> Option<String> {
233 let page_chrome = PAGE_CHROME_RE.get_or_init(|| {
234 Regex::new(concat!(
235 r"(?is)(?:<script(?:\s[^>]*)?>.*?</script\s*>",
236 r"|<style(?:\s[^>]*)?>.*?</style\s*>",
237 r"|<noscript(?:\s[^>]*)?>.*?</noscript\s*>",
238 r"|<nav(?:\s[^>]*)?>.*?</nav\s*>",
239 r"|<header(?:\s[^>]*)?>.*?</header\s*>",
240 r"|<footer(?:\s[^>]*)?>.*?</footer\s*>",
241 r"|<aside(?:\s[^>]*)?>.*?</aside\s*>",
242 r"|<form(?:\s[^>]*)?>.*?</form\s*>)",
243 ))
244 .expect("page chrome regex")
245 });
246 for re in FALLBACK_RE.get_or_init(|| {
247 ["article", "main", "body"]
248 .into_iter()
249 .map(|tag| {
250 Regex::new(&format!(r"(?is)<{tag}(?:\s[^>]*)?>(.*?)</{tag}\s*>"))
251 .expect("fallback element regex")
252 })
253 .collect()
254 }) {
255 let Some(capture) = re.captures(html) else {
256 continue;
257 };
258 let Some(content) = capture.get(1) else {
259 continue;
260 };
261 let without_chrome = page_chrome.replace_all(content.as_str(), "");
262 if meaningful_html(&without_chrome) {
263 return Some(without_chrome.into_owned());
264 }
265 }
266 None
267 }
268
269 fn meaningful_html(html: &str) -> bool {
270 meaningful_text(&html_to_plain_text(html))
271 }
272
273 fn meaningful_text(text: &str) -> bool {
274 text.chars().filter(|ch| !ch.is_whitespace()).count() >= 32
275 && text.split_whitespace().count() >= 5
276 }
277
278 fn html_to_plain_text(html: &str) -> String {
279 let without_tags = TAG_RE
280 .get_or_init(|| Regex::new(r"(?s)<[^>]+>").expect("tag regex"))
281 .replace_all(html, " ");
282 normalize_text(&decode_common_entities(&without_tags))
283 }
284
285 fn normalize_text(text: &str) -> String {
286 WHITESPACE_RE
287 .get_or_init(|| Regex::new(r"\s+").expect("whitespace regex"))
288 .replace_all(text.trim(), " ")
289 .into_owned()
290 }
291
292 fn decode_common_entities(value: &str) -> String {
293 value
294 .replace("&nbsp;", " ")
295 .replace("&amp;", "&")
296 .replace("&lt;", "<")
297 .replace("&gt;", ">")
298 .replace("&quot;", "\"")
299 .replace("&#39;", "'")
300 }
301
302 fn html_title(html: &str) -> Option<String> {
303 let capture = TITLE_RE
304 .get_or_init(|| {
305 Regex::new(r"(?is)<title(?:\s[^>]*)?>(.*?)</title\s*>").expect("title regex")
306 })
307 .captures(html)?;
308 let title = normalize_text(&decode_common_entities(capture.get(1)?.as_str()));
309 (!title.is_empty()).then_some(title)
310 }
311
312 fn markdown_title(body: &str) -> Option<String> {
313 body.lines().find_map(|line| {
314 let title = line.trim().strip_prefix("# ")?.trim();
315 (!title.is_empty()).then(|| title.to_string())
316 })
317 }
318
319 fn js_required_error(url: &str) -> ToolError {
320 ToolError::execution_failed(format!(
321 "No readable page content was found at {url}; the page may require JavaScript. Recovery: use browser automation for this URL."
322 ))
323 }
324
325 /// Decode one response body without guessing from language statistics.
326 ///
327 /// Precedence is receipt-grade and deterministic: BOM, recognized transport
328 /// charset, an HTML-only bounded meta prescan, then UTF-8. Unknown transport
329 /// labels deliberately fall through to a valid HTML declaration. `html_sniff`
330 /// must come from MIME/URL/ASCII markup evidence; JSON and plain text callers
331 /// pass `false`, so a body string cannot impersonate an HTML declaration.
332 pub(crate) fn decode_response_body(
333 bytes: &[u8],
334 content_type: Option<&str>,
335 html_sniff: bool,
336 ) -> Result<String, ToolError> {
337 if let Some((encoding, bom_len)) = Encoding::for_bom(bytes) {
338 let (decoded, _) = encoding.decode_without_bom_handling(&bytes[bom_len..]);
339 reject_binary_nul(bytes, encoding, &decoded)?;
340 return Ok(decoded.into_owned());
341 }
342
343 let transport_encoding = content_type.and_then(content_type_encoding);
344 let encoding = transport_encoding
345 .or_else(|| html_sniff.then(|| html_meta_encoding(bytes)).flatten())
346 .unwrap_or(UTF_8);
347 let (decoded, _) = encoding.decode_without_bom_handling(bytes);
348 reject_binary_nul(bytes, encoding, &decoded)?;
349 Ok(decoded.into_owned())
350 }
351
352 fn reject_binary_nul(
353 bytes: &[u8],
354 encoding: &'static Encoding,
355 decoded: &str,
356 ) -> Result<(), ToolError> {
357 // UTF-16 uses zero bytes structurally for many characters, so inspect its
358 // decoded scalar values. Every other encoding must be NUL-free across the
359 // complete response; neither a BOM nor a late byte may bypass the guard.
360 let contains_nul = if encoding == UTF_16LE || encoding == UTF_16BE {
361 decoded.contains('\0')
362 } else {
363 bytes.contains(&0)
364 };
365 if contains_nul {
366 return Err(ToolError::execution_failed(
367 "Unsupported binary response contained NUL bytes",
368 ));
369 }
370 Ok(())
371 }
372
373 /// Parse only exact semicolon-delimited `charset` parameters. A random
374 /// `charset=` substring inside another parameter is not transport authority.
375 fn content_type_encoding(value: &str) -> Option<&'static Encoding> {
376 value.split(';').skip(1).find_map(|parameter| {
377 let (name, raw_value) = parameter.split_once('=')?;
378 if !name.trim().eq_ignore_ascii_case("charset") {
379 return None;
380 }
381 let value = raw_value.trim();
382 let value = match (value.as_bytes().first(), value.as_bytes().last()) {
383 (Some(b'"'), Some(b'"')) | (Some(b'\''), Some(b'\'')) if value.len() >= 2 => {
384 &value[1..value.len() - 1]
385 }
386 _ if value.contains('"') || value.contains('\'') => return None,
387 _ => value,
388 };
389 let label = value.trim();
390 (!label.is_empty())
391 .then(|| Encoding::for_label(label.as_bytes()))
392 .flatten()
393 })
394 }
395
396 fn should_sniff_html_encoding(content_type: Option<&str>, url: &str, bytes: &[u8]) -> bool {
397 match content_type {
398 Some("text/html" | "application/xhtml+xml") => true,
399 // Explicit non-HTML text and structured formats never consult markup
400 // embedded in their body.
401 Some(value) if value.starts_with("text/") || is_structured_text_type(value) => false,
402 Some("application/octet-stream") | None => {
403 url_path_ends_with(url, &[".html", ".htm"]) || looks_like_html_bytes(bytes)
404 }
405 Some(_) => false,
406 }
407 }
408
409 fn is_structured_text_type(content_type: &str) -> bool {
410 content_type.contains("json")
411 || content_type.contains("xml")
412 || content_type.contains("yaml")
413 || content_type.contains("javascript")
414 }
415
416 fn looks_like_html_bytes(bytes: &[u8]) -> bool {
417 let start = Encoding::for_bom(bytes).map_or(0, |(_, length)| length);
418 let end = bytes
419 .len()
420 .min(start.saturating_add(HTML_ENCODING_SNIFF_BYTES));
421 let ascii = ascii_lowercase_projection(&bytes[start..end]);
422 let Some(prefix) = html_prefix_after_leading_declarations(&ascii) else {
423 return false;
424 };
425 prefix.starts_with("<!doctype html")
426 || prefix.starts_with("<html")
427 || prefix.starts_with("<head")
428 || prefix.starts_with("<meta")
429 }
430
431 fn html_prefix_after_leading_declarations(mut prefix: &str) -> Option<&str> {
432 loop {
433 prefix = prefix.trim_start();
434 if let Some(comment) = prefix.strip_prefix("<!--") {
435 let end = comment.find("-->")?;
436 prefix = &comment[end + 3..];
437 continue;
438 }
439 if let Some(declaration) = prefix.strip_prefix("<?xml") {
440 let end = declaration.find("?>")?;
441 prefix = &declaration[end + 2..];
442 continue;
443 }
444 return Some(prefix);
445 }
446 }
447
448 fn html_meta_encoding(bytes: &[u8]) -> Option<&'static Encoding> {
449 let sniff_len = bytes.len().min(HTML_ENCODING_SNIFF_BYTES);
450 let html = ascii_lowercase_projection(&bytes[..sniff_len]);
451 let mut cursor = 0usize;
452
453 while let Some(relative) = html[cursor..].find('<') {
454 let start = cursor + relative;
455 if html[start..].starts_with("<!--") {
456 cursor = html[start + 4..]
457 .find("-->")
458 .map_or(html.len(), |end| start + 4 + end + 3);
459 continue;
460 }
461 if tag_starts_at(&html, start, "script") || tag_starts_at(&html, start, "style") {
462 let name = if tag_starts_at(&html, start, "script") {
463 "script"
464 } else {
465 "style"
466 };
467 let close = format!("</{name}");
468 cursor = html[start..]
469 .find(&close)
470 .and_then(|close_start| {
471 html[start + close_start..]
472 .find('>')
473 .map(|end| start + close_start + end + 1)
474 })
475 .unwrap_or(html.len());
476 continue;
477 }
478 if !tag_starts_at(&html, start, "meta") {
479 cursor = start + 1;
480 continue;
481 }
482 let relative_end = html[start..].find('>')?;
483 let end = start + relative_end + 1;
484 let tag = &html[start..end];
485 if let Some(label) = html_attribute_value(tag, "charset")
486 && let Some(encoding) = Encoding::for_label(label.as_bytes())
487 {
488 return Some(normalize_meta_encoding(encoding));
489 }
490 let is_content_type = html_attribute_value(tag, "http-equiv")
491 .is_some_and(|value| value.eq_ignore_ascii_case("content-type"));
492 if is_content_type
493 && let Some(content) = html_attribute_value(tag, "content")
494 && let Some(encoding) = content_type_encoding(&content)
495 {
496 return Some(normalize_meta_encoding(encoding));
497 }
498 cursor = end;
499 }
500 None
501 }
502
503 fn normalize_meta_encoding(encoding: &'static Encoding) -> &'static Encoding {
504 if encoding == UTF_16LE || encoding == UTF_16BE {
505 UTF_8
506 } else {
507 encoding
508 }
509 }
510
511 fn ascii_lowercase_projection(bytes: &[u8]) -> String {
512 bytes
513 .iter()
514 .map(|byte| {
515 if byte.is_ascii() {
516 char::from(byte.to_ascii_lowercase())
517 } else {
518 ' '
519 }
520 })
521 .collect()
522 }
523
524 fn tag_starts_at(html: &str, start: usize, name: &str) -> bool {
525 let Some(after_name) = html.get(start + 1 + name.len()..) else {
526 return false;
527 };
528 html[start + 1..].starts_with(name)
529 && after_name
530 .chars()
531 .next()
532 .is_some_and(|ch| ch.is_ascii_whitespace() || matches!(ch, '/' | '>'))
533 }
534
535 fn html_attribute_value(tag: &str, wanted: &str) -> Option<String> {
536 let bytes = tag.as_bytes();
537 let mut cursor = 1usize;
538 while cursor < bytes.len() && !bytes[cursor].is_ascii_whitespace() && bytes[cursor] != b'>' {
539 cursor += 1;
540 }
541 while cursor < bytes.len() {
542 while cursor < bytes.len() && (bytes[cursor].is_ascii_whitespace() || bytes[cursor] == b'/')
543 {
544 cursor += 1;
545 }
546 if cursor >= bytes.len() || bytes[cursor] == b'>' {
547 break;
548 }
549 let name_start = cursor;
550 while cursor < bytes.len()
551 && !bytes[cursor].is_ascii_whitespace()
552 && !matches!(bytes[cursor], b'=' | b'/' | b'>')
553 {
554 cursor += 1;
555 }
556 let name = &tag[name_start..cursor];
557 while cursor < bytes.len() && bytes[cursor].is_ascii_whitespace() {
558 cursor += 1;
559 }
560 if cursor >= bytes.len() || bytes[cursor] != b'=' {
561 continue;
562 }
563 cursor += 1;
564 while cursor < bytes.len() && bytes[cursor].is_ascii_whitespace() {
565 cursor += 1;
566 }
567 if cursor >= bytes.len() {
568 break;
569 }
570 let (value_start, value_end) = if matches!(bytes[cursor], b'"' | b'\'') {
571 let quote = bytes[cursor];
572 cursor += 1;
573 let start = cursor;
574 while cursor < bytes.len() && bytes[cursor] != quote {
575 cursor += 1;
576 }
577 let end = cursor;
578 cursor = cursor.saturating_add(1);
579 (start, end)
580 } else {
581 let start = cursor;
582 while cursor < bytes.len()
583 && !bytes[cursor].is_ascii_whitespace()
584 && bytes[cursor] != b'>'
585 {
586 cursor += 1;
587 }
588 (start, cursor)
589 };
590 if name.eq_ignore_ascii_case(wanted) {
591 return Some(tag[value_start..value_end].trim().to_string());
592 }
593 }
594 None
595 }
596
597 fn normalized_content_type(content_type: Option<&str>) -> Option<String> {
598 content_type
599 .and_then(|value| value.split(';').next())
600 .map(str::trim)
601 .filter(|value| !value.is_empty())
602 .map(str::to_ascii_lowercase)
603 }
604
605 fn is_html(content_type: Option<&str>, url: &str, body: &str) -> bool {
606 matches!(content_type, Some("text/html" | "application/xhtml+xml"))
607 || url_path_ends_with(url, &[".html", ".htm"])
608 || {
609 let prefix = body.trim_start().chars().take(64).collect::<String>();
610 let prefix = prefix.to_ascii_lowercase();
611 prefix.contains("<!doctype html") || prefix.contains("<html")
612 }
613 }
614
615 fn is_markdown(content_type: Option<&str>, url: &str) -> bool {
616 matches!(
617 content_type,
618 Some("text/markdown" | "text/x-markdown" | "application/markdown")
619 ) || url_path_ends_with(url, &[".md", ".markdown"])
620 }
621
622 fn is_textual(content_type: Option<&str>, url: &str) -> bool {
623 content_type.is_some_and(|value| {
624 value.starts_with("text/")
625 || value.contains("json")
626 || value.contains("xml")
627 || value.contains("yaml")
628 || value.contains("javascript")
629 || value == "application/sql"
630 }) || url_path_ends_with(
631 url,
632 &[
633 ".txt", ".json", ".jsonl", ".xml", ".yaml", ".yml", ".csv", ".tsv", ".rs", ".py",
634 ".js", ".ts", ".toml",
635 ],
636 )
637 }
638
639 fn url_is_pdf(url: &str) -> bool {
640 url_path_ends_with(url, &[".pdf"])
641 }
642
643 fn url_path_ends_with(url: &str, extensions: &[&str]) -> bool {
644 reqwest::Url::parse(url)
645 .ok()
646 .map(|parsed| parsed.path().to_ascii_lowercase())
647 .is_some_and(|path| extensions.iter().any(|extension| path.ends_with(extension)))
648 }
649
650 fn looks_like_pdf(bytes: &[u8]) -> bool {
651 bytes.starts_with(b"%PDF-")
652 }
653
654 fn declared_media_family(content_type: Option<&str>) -> Option<MediaFamily> {
655 let content_type = content_type?;
656 if content_type.starts_with("image/") {
657 Some(MediaFamily::Image)
658 } else if content_type.starts_with("audio/") {
659 Some(MediaFamily::Audio)
660 } else if content_type.starts_with("video/") {
661 Some(MediaFamily::Video)
662 } else {
663 None
664 }
665 }
666
667 fn sniff_media(bytes: &[u8]) -> Option<MediaSignature> {
668 let trimmed = bytes
669 .iter()
670 .position(|byte| !byte.is_ascii_whitespace())
671 .map(|start| &bytes[start..])
672 .unwrap_or(bytes);
673 let signature = if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
674 MediaSignature {
675 extension: "png",
676 family: MediaFamily::Image,
677 }
678 } else if bytes.starts_with(b"\xff\xd8\xff") {
679 MediaSignature {
680 extension: "jpg",
681 family: MediaFamily::Image,
682 }
683 } else if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
684 MediaSignature {
685 extension: "gif",
686 family: MediaFamily::Image,
687 }
688 } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
689 MediaSignature {
690 extension: "webp",
691 family: MediaFamily::Image,
692 }
693 } else if bytes.starts_with(b"ID3") || bytes.starts_with(b"\xff\xfb") {
694 MediaSignature {
695 extension: "mp3",
696 family: MediaFamily::Audio,
697 }
698 } else if bytes.starts_with(b"fLaC") {
699 MediaSignature {
700 extension: "flac",
701 family: MediaFamily::Audio,
702 }
703 } else if bytes.starts_with(b"OggS") {
704 MediaSignature {
705 extension: "ogg",
706 family: MediaFamily::Audio,
707 }
708 } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WAVE" {
709 MediaSignature {
710 extension: "wav",
711 family: MediaFamily::Audio,
712 }
713 } else if bytes.len() >= 12 && &bytes[4..8] == b"ftyp" {
714 MediaSignature {
715 extension: "mp4",
716 family: MediaFamily::Video,
717 }
718 } else if bytes.starts_with(b"\x1aE\xdf\xa3") {
719 MediaSignature {
720 extension: "webm",
721 family: MediaFamily::Video,
722 }
723 } else if trimmed.starts_with(b"<svg")
724 || (trimmed.starts_with(b"<?xml")
725 && trimmed
726 .windows(4)
727 .take(1_024)
728 .any(|window| window.eq_ignore_ascii_case(b"<svg")))
729 {
730 MediaSignature {
731 extension: "svg",
732 family: MediaFamily::Image,
733 }
734 } else {
735 return None;
736 };
737 Some(signature)
738 }
739
740 async fn extract_pdf(
741 bytes: &[u8],
742 command: super::super::pdf::PdfTextCommand<'_>,
743 ) -> Result<ExtractedDocument, ToolError> {
744 let text = super::super::pdf::extract_bytes(bytes, command)
745 .await
746 .map_err(super::super::pdf::into_tool_error)?;
747 let pages = split_pdf_pages(&text);
748 let text = pages
749 .iter()
750 .map(|page| page.join("\n"))
751 .collect::<Vec<_>>()
752 .join("\n\n");
753 Ok(ExtractedDocument {
754 kind: DocumentKind::Pdf,
755 title: Some("PDF Document".to_string()),
756 markdown: text.clone(),
757 text,
758 cleaned_html: None,
759 pdf_pages: Some(pages),
760 media_extension: None,
761 })
762 }
763
764 fn split_pdf_pages(text: &str) -> Vec<Vec<String>> {
765 text.split('\x0C')
766 .map(|page| {
767 page.lines()
768 .map(str::trim)
769 .filter(|line| !line.is_empty())
770 .map(ToOwned::to_owned)
771 .collect::<Vec<_>>()
772 })
773 .collect()
774 }
775
776 #[cfg(test)]
777 mod tests {
778 use super::*;
779
780 #[tokio::test]
781 async fn html_becomes_readable_markdown_without_page_chrome() {
782 let html = br#"<!doctype html><html><head><title>Whale &amp; Signal</title></head><body>
783 <nav>Products Pricing Log in Cookies</nav>
784 <article><h1>Fetch once</h1><p>This is the important article body with enough words to be useful.</p>
785 <a href="/proof">Read the proof</a></article>
786 <footer>Privacy Cookies Terms</footer></body></html>"#;
787 let document = extract_document("https://example.com/post", Some("text/html"), html, None)
788 .await
789 .expect("extract html");
790
791 assert_eq!(document.kind, DocumentKind::Html);
792 assert_eq!(document.title.as_deref(), Some("Whale & Signal"));
793 assert!(document.markdown.contains("Fetch once") || document.title.is_some());
794 assert!(
795 document
796 .markdown
797 .contains("[Read the proof](https://example.com/proof)")
798 );
799 assert!(!document.markdown.contains("Products Pricing"));
800 assert!(!document.markdown.contains("Privacy Cookies"));
801 }
802
803 #[tokio::test]
804 async fn sparse_document_uses_article_fallback() {
805 let html = br#"<html><head><title>Fallback</title></head><body><nav>cookie banner</nav>
806 <article><h2>Small source</h2><p>Five useful words survive this compact article fallback path.</p></article>
807 </body></html>"#;
808 let document = extract_document("https://example.com/short", Some("text/html"), html, None)
809 .await
810 .expect("extract fallback");
811
812 assert!(document.markdown.contains("## Small source"));
813 assert!(!document.markdown.contains("cookie banner"));
814 }
815
816 #[tokio::test]
817 async fn javascript_shell_returns_actionable_error() {
818 let error = extract_document(
819 "https://example.com/app",
820 Some("text/html"),
821 b"<html><body><div id='root'></div><script>boot()</script></body></html>",
822 None,
823 )
824 .await
825 .expect_err("empty app shell must fail");
826
827 let message = error.to_string();
828 assert!(message.contains("may require JavaScript"));
829 assert!(message.contains("browser automation"));
830 }
831
832 #[tokio::test]
833 async fn markdown_passes_through_unchanged() {
834 let body = b"# Release note\n\nA complete markdown response remains intact.\n";
835 let document = extract_document(
836 "https://example.com/release.md",
837 Some("text/markdown; charset=utf-8"),
838 body,
839 None,
840 )
841 .await
842 .expect("extract markdown");
843
844 assert_eq!(document.kind, DocumentKind::Markdown);
845 assert_eq!(document.markdown.as_bytes(), body);
846 assert_eq!(document.title.as_deref(), Some("Release note"));
847 }
848
849 #[tokio::test]
850 async fn media_requires_matching_magic_bytes() {
851 let error = extract_document(
852 "https://example.com/not-image.png",
853 Some("image/png"),
854 b"<html>not really an image</html>",
855 None,
856 )
857 .await
858 .expect_err("spoofed media must fail");
859 assert!(error.to_string().contains("did not match"));
860
861 let mut png = b"\x89PNG\r\n\x1a\n".to_vec();
862 png.extend_from_slice(b"fake test payload");
863 let document = extract_document(
864 "https://example.com/image",
865 Some("application/octet-stream"),
866 &png,
867 None,
868 )
869 .await
870 .expect("sniff png");
871 assert_eq!(document.kind, DocumentKind::Media);
872 assert_eq!(document.media_extension, Some("png"));
873 }
874
875 #[tokio::test]
876 async fn arbitrary_binary_is_rejected() {
877 let error = extract_document(
878 "https://example.com/archive.bin",
879 Some("application/octet-stream"),
880 b"PK\x03\x04archive bytes",
881 None,
882 )
883 .await
884 .expect_err("archive must be rejected");
885 assert!(error.to_string().contains("Unsupported binary response"));
886 }
887
888 #[tokio::test]
889 async fn empty_success_body_is_valid_text() {
890 let document = extract_document(
891 "https://example.com/no-content",
892 Some("application/octet-stream"),
893 b"",
894 None,
895 )
896 .await
897 .expect("empty body");
898 assert_eq!(document.kind, DocumentKind::Text);
899 assert!(document.text.is_empty());
900 }
901
902 #[tokio::test]
903 async fn content_type_matching_is_case_insensitive() {
904 let document = extract_document(
905 "https://example.com/document",
906 Some("Application/JSON; Charset=UTF-8"),
907 br#"{"status":"ok"}"#,
908 None,
909 )
910 .await
911 .expect("mixed-case JSON content type");
912
913 assert_eq!(document.kind, DocumentKind::Text);
914 assert_eq!(document.text, r#"{"status":"ok"}"#);
915 }
916
917 #[test]
918 fn bom_wins_over_conflicting_transport_and_is_removed() {
919 let mut utf8 = b"\xef\xbb\xbf".to_vec();
920 utf8.extend_from_slice("café".as_bytes());
921 assert_eq!(
922 decode_response_body(&utf8, Some("text/html; charset=windows-1252"), true)
923 .expect("UTF-8 BOM"),
924 "café"
925 );
926
927 let mut utf16 = vec![0xff, 0xfe];
928 for unit in "BOM 日本語".encode_utf16() {
929 utf16.extend_from_slice(&unit.to_le_bytes());
930 }
931 assert_eq!(
932 decode_response_body(&utf16, Some("text/plain; charset=windows-1252"), false)
933 .expect("UTF-16 BOM"),
934 "BOM 日本語"
935 );
936 }
937
938 #[test]
939 fn content_type_charset_is_exact_recognized_and_order_independent() {
940 let (bytes, _, _) = encoding_rs::WINDOWS_1252.encode("café");
941 for content_type in [
942 "text/plain; charset=windows-1252",
943 "TEXT/PLAIN; boundary=x; CHARSET = \"windows-1252\"; q=1",
944 "text/plain; q=1; charset='windows-1252'",
945 ] {
946 assert_eq!(
947 decode_response_body(&bytes, Some(content_type), false).expect("declared charset"),
948 "café",
949 "{content_type}"
950 );
951 }
952
953 for malformed in [
954 "text/plain; note=charset=windows-1252",
955 "text/plain; charset=\"windows-1252",
956 "text/plain; charset=definitely-not-an-encoding",
957 ] {
958 let decoded =
959 decode_response_body(&bytes, Some(malformed), false).expect("UTF-8 fallback");
960 assert!(decoded.contains('\u{fffd}'), "{malformed}: {decoded}");
961 }
962 }
963
964 #[test]
965 fn invalid_header_falls_through_to_direct_and_legacy_html_meta() {
966 let direct = r#"<html><head><meta charset="gbk"></head><body>中文</body></html>"#;
967 let (direct_bytes, _, _) = encoding_rs::GBK.encode(direct);
968 assert!(
969 decode_response_body(&direct_bytes, Some("text/html; charset=not-real"), true,)
970 .expect("direct meta")
971 .contains("中文")
972 );
973
974 let legacy = r#"<html><head><meta content="text/html; charset=windows-1252" http-equiv="Content-Type"></head><body>café</body></html>"#;
975 let (legacy_bytes, _, _) = encoding_rs::WINDOWS_1252.encode(legacy);
976 assert!(
977 decode_response_body(&legacy_bytes, Some("text/html"), true)
978 .expect("legacy meta")
979 .contains("café")
980 );
981 }
982
983 #[test]
984 fn recognized_transport_charset_beats_conflicting_meta() {
985 let html = r#"<html><head><meta charset="shift_jis"></head><body>中文</body></html>"#;
986 let (bytes, _, _) = encoding_rs::GBK.encode(html);
987 let decoded = decode_response_body(&bytes, Some("text/html; charset=gbk"), true)
988 .expect("transport charset");
989 assert!(decoded.contains("中文"), "{decoded}");
990 }
991
992 #[test]
993 fn html_prescan_ignores_comments_scripts_and_late_meta() {
994 let cases = [
995 "<!-- <meta charset=windows-1252> --><html><body>café</body></html>".to_string(),
996 "<script>\"<meta charset=windows-1252>\"</script><html><body>café</body></html>"
997 .to_string(),
998 format!(
999 "<html><head>{}<meta charset=windows-1252></head><body>café</body></html>",
1000 " ".repeat(HTML_ENCODING_SNIFF_BYTES)
1001 ),
1002 ];
1003 for html in cases {
1004 let (bytes, _, _) = encoding_rs::WINDOWS_1252.encode(&html);
1005 let decoded = decode_response_body(&bytes, Some("text/html"), true)
1006 .expect("bounded HTML fallback");
1007 assert!(
1008 decoded.contains('\u{fffd}'),
1009 "late/ignored meta changed decoding: {decoded}"
1010 );
1011 }
1012 }
1013
1014 #[test]
1015 fn non_html_bodies_never_sniff_meta_markup() {
1016 let plain = "literal <meta charset=windows-1252> café";
1017 let (bytes, _, _) = encoding_rs::WINDOWS_1252.encode(plain);
1018 for content_type in ["text/plain", "application/json"] {
1019 let decoded = decode_response_body(&bytes, Some(content_type), false)
1020 .expect("non-HTML UTF-8 fallback");
1021 assert!(decoded.contains('\u{fffd}'), "{content_type}: {decoded}");
1022 }
1023 }
1024
1025 #[test]
1026 fn declared_gbk_shift_jis_and_windows_1252_decode_deterministically() {
1027 let cases = [
1028 (encoding_rs::GBK, "中文", "gbk"),
1029 (encoding_rs::SHIFT_JIS, "日本語", "shift_jis"),
1030 (encoding_rs::WINDOWS_1252, "café", "windows-1252"),
1031 ];
1032 for (encoding, text, label) in cases {
1033 let (bytes, _, had_errors) = encoding.encode(text);
1034 assert!(!had_errors, "fixture must be representable in {label}");
1035 assert_eq!(
1036 decode_response_body(&bytes, Some(&format!("text/plain; charset={label}")), false,)
1037 .expect("decode declared encoding"),
1038 text
1039 );
1040 }
1041 }
1042
1043 #[test]
1044 fn nul_binary_is_rejected_but_utf16_bom_text_is_not() {
1045 let error = decode_response_body(b"PK\0\x03\x04archive", Some("text/plain"), false)
1046 .expect_err("NUL binary must fail");
1047 assert!(error.to_string().contains("NUL bytes"));
1048
1049 let bom_binary = b"\xef\xbb\xbfapparently text\0binary";
1050 let error = decode_response_body(bom_binary, Some("text/plain"), false)
1051 .expect_err("a BOM must not bypass the NUL guard");
1052 assert!(error.to_string().contains("NUL bytes"));
1053
1054 let mut late_binary = vec![b'x'; 8_193];
1055 late_binary.push(0);
1056 let error = decode_response_body(&late_binary, Some("text/plain"), false)
1057 .expect_err("a late NUL must not bypass the full-body guard");
1058 assert!(error.to_string().contains("NUL bytes"));
1059
1060 let utf16 = [0xff, 0xfe, b'O', 0, b'K', 0];
1061 assert_eq!(
1062 decode_response_body(&utf16, Some("application/octet-stream"), false)
1063 .expect("BOM proves UTF-16 text"),
1064 "OK"
1065 );
1066
1067 let utf16_nul = [0xff, 0xfe, b'O', 0, 0, 0, b'K', 0];
1068 let error = decode_response_body(&utf16_nul, Some("text/plain"), false)
1069 .expect_err("decoded UTF-16 NUL must remain binary");
1070 assert!(error.to_string().contains("NUL bytes"));
1071 }
1072
1073 #[tokio::test]
1074 async fn extensionless_html_sniff_skips_leading_comments_and_xml_declarations() {
1075 let cases = [
1076 r#"<!-- deployment marker --><html><head><meta charset="windows-1252"><title>Café release notes</title></head><body><article><h1>Café release notes</h1><p>This extensionless page contains enough meaningful text for deterministic extraction.</p></article></body></html>"#,
1077 r#"<?xml version="1.0"?><!-- marker --><head><meta charset="windows-1252"><title>Café release notes</title></head><body><article><h1>Café release notes</h1><p>This extensionless page contains enough meaningful text for deterministic extraction.</p></article></body>"#,
1078 ];
1079 for html in cases {
1080 let (bytes, _, _) = encoding_rs::WINDOWS_1252.encode(html);
1081 let document =
1082 extract_document("https://example.com/extensionless", None, &bytes, None)
1083 .await
1084 .expect("leading declarations preserve extensionless HTML sniffing");
1085 assert_eq!(document.kind, DocumentKind::Html);
1086 assert_eq!(document.title.as_deref(), Some("Café release notes"));
1087 }
1088 }
1089
1090 #[tokio::test]
1091 async fn svg_requires_and_accepts_svg_markup_signature() {
1092 let svg = br#"<?xml version="1.0"?><svg xmlns="http://www.w3.org/2000/svg"></svg>"#;
1093 let document = extract_document(
1094 "https://example.com/diagram",
1095 Some("image/svg+xml"),
1096 svg,
1097 None,
1098 )
1099 .await
1100 .expect("sniff svg");
1101 assert_eq!(document.kind, DocumentKind::Media);
1102 assert_eq!(document.media_extension, Some("svg"));
1103 }
1104 }
1105
1105 lines RUST