diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 19c1c8aefa..d655c89397 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -1381,17 +1381,12 @@ _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion) # sections stripped during conversion; 512 KB reaches article content even # where alone is ~200 KB. _MAX_FETCH_BYTES = 512 * 1024 -# Chars that don't occur in real text: undecodable bytes (U+FFFD) plus control -# chars (C0 minus tab/newline/CR and ESC, DEL, C1). ESC is excluded so a text -# page of ANSI-colored terminal output isn't mistaken for binary. A decoded page -# is treated as binary when more than 1/_BINARY_CHAR_DIVISOR (12.5%) of its chars -# are these, tolerating up to _MIN_BINARY_CHARS so minor glitches don't drop it. +# Undecodable bytes and controls, excluding text whitespace and ESC for ANSI logs. +# More than 12.5% is binary after allowing 16 minor encoding glitches. _BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]") _MIN_BINARY_CHARS = 16 _BINARY_CHAR_DIVISOR = 8 -# Leading bytes of common binary formats that servers sometimes mislabel as -# text/*; matched by signature because their replacement-char density alone can -# be low (e.g. a PDF whose first chunk is mostly ASCII object/xref syntax). +# Common binary signatures that can otherwise look text-heavy when mislabeled. _BINARY_MAGIC = ( b"%PDF-", # PDF b"PK\x03\x04", # zip / docx / xlsx / pptx / epub / jar @@ -1405,17 +1400,13 @@ _BINARY_MAGIC = ( b"\x28\xb5\x2f\xfd", # zstd ) -# Undeclared legacy Western text should have substantial ASCII structure even -# when its non-ASCII bytes are cp1252. Requiring at least 75% ASCII text bytes -# prevents arbitrary high-byte binary from becoming printable-looking gibberish -# merely because cp1252 defines a character for nearly every byte. +# A cp1252 retry needs 75% ASCII structure so it cannot rescue high-byte binary. _MIN_SINGLE_BYTE_ASCII_RATIO = 3 / 4 _ASCII_TEXT_BYTES = frozenset((*range(0x20, 0x7F), 0x09, 0x0A, 0x0D, 0x1B)) def _looks_binary(text: str) -> bool: - """True when more than 1/_BINARY_CHAR_DIVISOR of ``text`` is binary chars - (control/undecodable, per _BINARY_CHAR_RE), past the _MIN_BINARY_CHARS floor.""" + """Whether control or undecodable characters exceed the binary threshold.""" return len(_BINARY_CHAR_RE.findall(text)) > max( _MIN_BINARY_CHARS, len(text) // _BINARY_CHAR_DIVISOR ) @@ -1536,27 +1527,16 @@ _TEXT_APPLICATION_SUBTYPES = frozenset( ) -def _is_texty_content_type(content_type: str | None) -> bool: - """True for MIME types safe to decode or inspect as possible text. - - This includes HTML, plain text, JSON/XML/YAML feeds, and generic - ``application/octet-stream`` downloads whose bytes must be sniffed. Known - binary types (PDF, images, archives, Office/ZIP) return False so the fetcher - never decodes them into replacement chars that poison the model context - (unslothai/unsloth#7084). - - A missing Content-Type coerces to ``text/plain`` upstream, so unlabeled - bodies pass here and are caught instead by the binary-char fallback. - """ +def _is_text_candidate_content_type(content_type: str | None) -> bool: + """Whether a MIME type is textual or ambiguous enough for byte sniffing.""" ct = (content_type or "").partition(";")[0].strip().lower() if not ct: - return True # unlabeled: let the binary-char fallback decide + return True if ct.startswith("text/"): return True if ct.startswith("application/"): - # Match the exact subtype, not a loose substring, so a .docx labeled - # application/vnd.openxmlformats-... isn't read as xml. RFC 6839 - # +json/+xml suffixes (ld+json, xhtml+xml, ...) pass too. + # Exact matches avoid treating "openxmlformats" as XML. Structured + # +json/+xml suffixes remain valid text candidates. subtype = ct[len("application/") :].removeprefix("x-") return ( subtype == "octet-stream" @@ -1646,33 +1626,24 @@ def _fetch_page_text( else: return "Failed to fetch URL: too many redirects." - # Reject binary bodies (PDF, image, archive): decoding them as text - # floods the model context with U+FFFD replacement chars (#7084). + # Reject MIME types known to be binary before decoding. content_type = resp.headers.get_content_type() - if not _is_texty_content_type(content_type): - # Trim to a clean MIME token: get_content_type() can echo control - # chars from an obs-folded header, and this string is returned to - # the model. + if not _is_text_candidate_content_type(content_type): + # Only echo a clean MIME token back to the model. m = re.match(r"[\w.+-]+/[\w.+-]+", content_type or "") safe_type = m.group(0) if m else "unknown type" return f"(non-text content: {safe_type}, {len(raw_bytes)} bytes; not readable as text)" - # A PDF/image/archive mislabeled as text/*: catch by signature, since a - # text-heavy first chunk can slip past the replacement-char density check. + # Catch text-labeled binary whose header and first chunk look textual. if raw_bytes.startswith(_BINARY_MAGIC): return f"(binary content, {len(raw_bytes)} bytes; not readable as text)" declared = resp.headers.get_content_charset() raw_html = raw_bytes.decode(declared or "utf-8", errors = "replace") - # Fallback for binary mislabeled as text/* or sent with no Content-Type, - # including valid-UTF-8 binary (NUL/control-heavy payloads) that decodes - # without replacement chars: a real text page has few binary chars. + # Catch mislabeled or unlabeled binary, including valid UTF-8 controls. if _looks_binary(raw_html): - # An undeclared non-UTF-8 page (latin-1/cp1252) also decodes to many - # U+FFFD. Retry as cp1252 only when the bytes have substantial ASCII - # text structure; otherwise high-byte binary would become printable - # looking gibberish because cp1252 maps nearly every byte. + # Rescue undeclared cp1252 only when the bytes have text structure. alt = ( raw_bytes.decode("cp1252", "replace") if declared is None and _has_single_byte_text_evidence(raw_bytes) diff --git a/studio/backend/tests/test_web_fetch_binary_guard.py b/studio/backend/tests/test_web_fetch_binary_guard.py index d69526e898..1feaf066b2 100644 --- a/studio/backend/tests/test_web_fetch_binary_guard.py +++ b/studio/backend/tests/test_web_fetch_binary_guard.py @@ -1,11 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Regression for unslothai/unsloth#7084: the web_search fetcher must not decode -a binary body into control characters or a flood of U+FFFD replacement chars -that poison the model context. It rejects known binary Content-Types up front -and sniffs generic, mislabeled, or unlabeled bodies. Real text is unaffected. -""" +"""Regression tests for binary bodies poisoning web_search model context (#7084).""" from __future__ import annotations @@ -58,9 +54,6 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str: return tools._fetch_page_text("https://example.com/thing", timeout = 5) -# ── content-type classifier ── - - @pytest.mark.parametrize( "content_type,expected", [ @@ -71,27 +64,22 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str: ("application/xml", True), ("application/xhtml+xml", True), ("application/ld+json", True), - ("application/yaml", True), # raw configs / API specs + ("application/yaml", True), ("application/x-yaml", True), - ("application/x-ndjson", True), # newline-delimited JSON is text + ("application/x-ndjson", True), ("application/ndjson", True), ("application/pdf", False), ("image/png", False), - ("image/svg+xml", False), # SVG source isn't extracted downstream; reject - # Generic downloads need byte sniffing because they can contain text. + ("image/svg+xml", False), ("application/octet-stream", True), ("application/zip", False), - # A .docx ZIP must not pass just because "xml" is inside "openxmlformats". ("application/vnd.openxmlformats-officedocument.wordprocessingml.document", False), - ("", True), # unlabeled: defer to the binary-char fallback + ("", True), (None, True), ], ) -def test_is_texty_content_type(content_type, expected): - assert tools._is_texty_content_type(content_type) is expected - - -# ── fetcher end-to-end (mocked network) ── +def test_is_text_candidate_content_type(content_type, expected): + assert tools._is_text_candidate_content_type(content_type) is expected def test_pdf_rejected_by_content_type(monkeypatch): @@ -100,12 +88,6 @@ def test_pdf_rejected_by_content_type(monkeypatch): assert "non-text content" in out and "application/pdf" in out -def test_image_rejected_by_content_type(monkeypatch): - out = _fetch_with(monkeypatch, b"\x89PNG\r\n\x1a\n" + bytes(range(256)) * 4, "image/png") - assert "�" not in out - assert "non-text content" in out and "image/png" in out - - def test_text_octet_stream_kept_after_sniffing(monkeypatch): body = b"level=info\nmessage=plain text artifact\n" * 100 out = _fetch_with(monkeypatch, body, "application/octet-stream") @@ -113,54 +95,31 @@ def test_text_octet_stream_kept_after_sniffing(monkeypatch): assert "non-text content" not in out and "binary content" not in out -def test_binary_octet_stream_rejected_after_sniffing(monkeypatch): - out = _fetch_with(monkeypatch, bytes(range(256)) * 20, "application/octet-stream") - assert "binary content" in out - - -def test_binary_mislabeled_as_text_caught_by_fallback(monkeypatch): - # A server sends binary but labels it text/plain -> the type check passes, - # so the binary-char fallback must catch it. - out = _fetch_with(monkeypatch, bytes(range(256)) * 20, "text/plain") +@pytest.mark.parametrize("content_type", ["application/octet-stream", "text/plain", None]) +def test_binary_candidates_rejected_after_sniffing(monkeypatch, content_type): + out = _fetch_with(monkeypatch, bytes(range(256)) * 20, content_type) assert "�" not in out assert "binary content" in out def test_valid_utf8_binary_caught_by_control_chars(monkeypatch): - # NUL/control-heavy binary is valid UTF-8, so it decodes with no U+FFFD; - # the fallback must still catch it via control-char density (codex #2). - body = bytes([0, 1, 2, 3, 4, 5, 6, 7]) * 400 # all valid UTF-8, 0 replacement chars + # These controls are valid UTF-8 and therefore produce no replacement chars. + body = bytes([0, 1, 2, 3, 4, 5, 6, 7]) * 400 out = _fetch_with(monkeypatch, body, "text/plain") assert "binary content" in out -def test_pdf_mislabeled_as_text_caught_by_magic(monkeypatch): - # A PDF served as text/plain: the content-type gate passes, so the magic-byte - # sniff must catch it (its ASCII-heavy head can stay under the ratio). - pdf = b"%PDF-1.7\n" + b"1 0 obj<>endobj\n" * 300 - out = _fetch_with(monkeypatch, pdf, "text/plain") - assert "binary content" in out - - -def test_zip_mislabeled_as_text_caught_by_magic(monkeypatch): - out = _fetch_with(monkeypatch, b"PK\x03\x04" + b"filename.txt content " * 200, "text/plain") - assert "binary content" in out - - @pytest.mark.parametrize( "magic", - [b"\x1f\x8b", b"BZh", b"\xfd7zXZ\x00", b"\x28\xb5\x2f\xfd"], # gzip, bzip2, xz, zstd + [b"%PDF-", b"PK\x03\x04", b"\x1f\x8b", b"BZh", b"\xfd7zXZ\x00", b"\x28\xb5\x2f\xfd"], ) -def test_compression_mislabeled_as_text_caught_by_magic(monkeypatch, magic): - # A small compressed body mislabeled as text can slip past the density check; - # the signature must reject it. - out = _fetch_with(monkeypatch, magic + b"some short compressed-looking body", "text/plain") +def test_text_labeled_binary_caught_by_magic(monkeypatch, magic): + out = _fetch_with(monkeypatch, magic + b" printable text-heavy body" * 100, "text/plain") assert "binary content" in out def test_latin1_text_without_charset_kept(monkeypatch): - # Accent-dense cp1252 text with no declared charset decodes to many U+FFFD as - # UTF-8; the cp1252 retry must keep it instead of dropping it as binary. + # The cp1252 retry should rescue accent-heavy text with ASCII structure. body = ( "Muller lauft uber die Strasse: schoene, groesse. MARKERWORD ".replace("ue", "ü") + "äöüß éèà " @@ -170,36 +129,20 @@ def test_latin1_text_without_charset_kept(monkeypatch): assert "MARKERWORD" in out -def test_control_heavy_binary_survives_cp1252_retry(monkeypatch): - # Genuine control-heavy binary (no magic) must stay rejected even after the - # cp1252 retry, which cannot rescue real binary. - out = _fetch_with(monkeypatch, bytes([0, 1, 2, 3, 4, 5, 6, 7]) * 400, "text/plain") - assert "binary content" in out - - def test_high_byte_binary_not_rescued_as_cp1252(monkeypatch): - # cp1252 maps almost every high byte to a printable character. Without the - # ASCII-structure requirement, this binary payload would be returned as text. + # cp1252 maps these bytes to printable characters, but they lack ASCII structure. body = bytes(range(0xA0, 0x100)) * 40 out = _fetch_with(monkeypatch, body, "text/plain") assert "binary content" in out def test_ansi_colored_text_log_kept(monkeypatch): - # A text log with per-token ANSI color codes is text; ESC is excluded from - # the binary-char set so it isn't dropped as binary. + # ESC is excluded from the binary set so ANSI logs remain readable. line = "".join(f"\x1b[32m+{i}\x1b[0m\n" for i in range(300)).encode() out = _fetch_with(monkeypatch, line, "text/plain") assert "binary content" not in out -def test_binary_unlabeled_caught_by_fallback(monkeypatch): - # No Content-Type coerces to text/plain upstream; the fallback still catches it. - out = _fetch_with(monkeypatch, bytes(range(256)) * 20, None) - assert "�" not in out - assert "binary content" in out - - def test_html_page_unaffected(monkeypatch): html = b"

Hello

Real text content here.

" out = _fetch_with(monkeypatch, html, "text/html; charset=utf-8") @@ -208,8 +151,7 @@ def test_html_page_unaffected(monkeypatch): def test_content_type_sanitized_in_message(monkeypatch): - # An obs-folded Content-Type can smuggle control chars into get_content_type(); - # the returned message must be trimmed to a clean MIME token. + # Do not echo obs-folded header content into the model response. out = _fetch_with(monkeypatch, b"\x00\x01\x02" * 500, "application/pdf\r\n data: injected") assert "\n" not in out and "\r" not in out assert "injected" not in out @@ -219,23 +161,18 @@ def test_content_type_sanitized_in_message(monkeypatch): @pytest.mark.parametrize( "n_bad,n_total,expect_binary", [ - # Straddle the 12.5% ratio (well above the 16-char floor): just under vs - # just over len//8. Locks the divisor so it can't silently drift. - (120, 1000, False), # 120 <= 1000//8 (125) -> kept - (130, 1000, True), # 130 > 1000//8 (125) -> binary + (120, 1000, False), + (130, 1000, True), ], ) def test_binary_char_ratio_boundary(monkeypatch, n_bad, n_total, expect_binary): - # Body of n_total chars: n_bad NUL control bytes + ASCII filler. NUL stays - # binary through the cp1252 retry, so only the ratio threshold decides here. body = b"\x00" * n_bad + b"a" * (n_total - n_bad) out = _fetch_with(monkeypatch, body, "text/plain") assert ("binary content" in out) is expect_binary def test_text_with_a_few_stray_replacement_chars_kept(monkeypatch): - # A mostly-clean page with a handful of bad bytes stays (below the floor), - # so we don't drop legitimate pages over minor encoding glitches. + # Minor encoding glitches below the floor should not drop a real page. body = ("Real article text. " * 200).encode() + b"\xff\xfe\xff" out = _fetch_with(monkeypatch, body, "text/html") assert "Real article text." in out