diff --git a/tests/tools/test_read_extract.py b/tests/tools/test_read_extract.py index 4b656e2de6b7a..038d073bb3e8f 100644 --- a/tests/tools/test_read_extract.py +++ b/tests/tools/test_read_extract.py @@ -134,6 +134,65 @@ class TestAnydocExtraction(unittest.TestCase): self.assertEqual(text, "hello\n") +class TestAnydocSizeCap(unittest.TestCase): + """Oversized inputs must be rejected before anydoc converts them. + Uses a fake binding so it runs regardless of local install state.""" + + def setUp(self): + from tools import read_extract + + self.rex = read_extract + self._saved_module = read_extract._anydoc_module + self._saved_cap = read_extract.MAX_ANYDOC_BYTES + self.tmp = tempfile.mkdtemp(prefix="rex_cap_") + self.calls = [] + + class _FakeAnydoc: + def to_markdown(_self, path): + self.calls.append(path) + return "converted\n" + + read_extract._anydoc_module = _FakeAnydoc() + + def tearDown(self): + import shutil + + self.rex._anydoc_module = self._saved_module + self.rex.MAX_ANYDOC_BYTES = self._saved_cap + shutil.rmtree(self.tmp, ignore_errors=True) + + def _write(self, name, size): + p = os.path.join(self.tmp, name) + with open(p, "wb") as fh: + fh.write(b"x" * size) + return p + + def test_oversized_file_rejected_before_conversion(self): + from tools.read_extract import _extract_anydoc + + self.rex.MAX_ANYDOC_BYTES = 10 + p = self._write("big.pdf", 11) + with self.assertRaises(ExtractionError) as ctx: + _extract_anydoc(p) + self.assertIn("too large", str(ctx.exception)) + self.assertEqual(self.calls, []) + + def test_file_at_limit_converts(self): + from tools.read_extract import _extract_anydoc + + self.rex.MAX_ANYDOC_BYTES = 10 + p = self._write("ok.pdf", 10) + self.assertEqual(_extract_anydoc(p), "converted\n") + self.assertEqual(self.calls, [p]) + + def test_missing_file_raises_extraction_error(self): + from tools.read_extract import _extract_anydoc + + with self.assertRaises(ExtractionError): + _extract_anydoc(os.path.join(self.tmp, "gone.pdf")) + self.assertEqual(self.calls, []) + + class TestAnydocAbsent(unittest.TestCase): """The absent-dep contract, verified regardless of local install state by forcing the cached module handle to None.""" diff --git a/tools/read_extract.py b/tools/read_extract.py index 1230e986adec3..0ec42a17308db 100644 --- a/tools/read_extract.py +++ b/tools/read_extract.py @@ -14,6 +14,7 @@ from __future__ import annotations import importlib import json +import os import posixpath import threading import time @@ -34,6 +35,10 @@ ANYDOC_EXTENSIONS = frozenset({ ".rtf", ".epub", ".pdf", }) MAX_XLSX_BYTES = 50 * 1024 * 1024 +# Refuse to convert huge documents. anydoc loads the whole file through its +# Rust core with no streaming, and the read_file char budget only applies +# after conversion, so an unbounded input can pin a tool turn and spike RAM. +MAX_ANYDOC_BYTES = 50 * 1024 * 1024 _MAX_XLSX_ROWS_PER_SHEET = 5000 _MAX_XLSX_COLS = 256 @@ -120,6 +125,14 @@ def _extract_anydoc(path: str) -> str: mod = _anydoc() if mod is None: raise ExtractionError(f"Unsupported document type: {path!r}") + try: + size = os.path.getsize(path) + except OSError as exc: + raise ExtractionError(str(exc)) from exc + if size > MAX_ANYDOC_BYTES: + raise ExtractionError( + f"Document too large to convert ({size:,} bytes, limit is {MAX_ANYDOC_BYTES:,})" + ) try: text = mod.to_markdown(path) except OSError as exc: