mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-08-13 17:16:12 +00:00
fix(read_extract): cap anydoc input size before conversion
The anydoc path from #79781 passed every covered file straight to to_markdown with no pre-check. anydoc loads the whole document through its Rust core and the read_file char budget only applies after conversion, so one large PDF or deck could pin a tool turn and spike RAM. _extract_anydoc now rejects inputs over MAX_ANYDOC_BYTES (50 MB) with ExtractionError before calling the converter, which routes them to the existing read_file fallthrough instead of converting. No timeout is added: the conversion is a synchronous Rust call that cannot be cancelled from Python, so a thread-based deadline would bound the wait but leave the RAM burn running in the background.
This commit is contained in:
@@ -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."""
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user