fix: two bugs in document extraction pipeline

Bug 1: _drain_pending did not call extract_documents on follow-up
messages arriving mid-turn. Documents attached to queued messages were
silently dropped because _build_user_content only handles images.
Fix: call extract_documents before _build_user_content in _drain_pending.

Bug 2: extract_documents read the entire file into memory (up to 50 MB)
just to check 16 bytes of magic header for MIME detection.
Fix: read only the first 16 bytes via open()+read(16) instead of
Path.read_bytes().

Added regression tests for both bugs.

Made-with: Cursor
This commit is contained in:
Xubin Ren
2026-04-14 13:15:04 +00:00
parent 92d6fca323
commit c937c07178
4 changed files with 91 additions and 7 deletions
+26
View File
@@ -429,6 +429,32 @@ def test_extract_documents_skips_oversized_files(tmp_path) -> None:
assert image_paths == []
def test_extract_documents_does_not_read_full_file_for_mime(tmp_path) -> None:
"""MIME detection should only read header bytes, not the entire file."""
from pathlib import Path as _Path
big_txt = tmp_path / "big.txt"
big_txt.write_bytes(b"hello world " * 100_000) # ~1.2 MB
original_read_bytes = _Path.read_bytes
read_sizes: list[int] = []
def _tracking_read_bytes(self):
data = original_read_bytes(self)
read_sizes.append(len(data))
return data
import unittest.mock
with unittest.mock.patch.object(_Path, "read_bytes", _tracking_read_bytes):
extract_documents("test", [str(big_txt)])
# If the full file was read for MIME detection, read_sizes would
# contain a >1MB entry. After the fix, only a small header is read.
assert all(size <= 4096 for size in read_sizes), (
f"extract_documents read full file for MIME detection: sizes={read_sizes}"
)
# ---------------------------------------------------------------------------
# DOCX upload test — API saves file, loop layer extracts text
# ---------------------------------------------------------------------------