From 6865cacaf2d34139b74bc4ba8ac9960081b3906c Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Tue, 4 Aug 2026 06:46:22 -0700 Subject: [PATCH] fix: Read safe text attachments consistently Fixes #3130 --- src/document_processor.py | 32 +++---- src/upload_handler.py | 25 ++++-- tests/test_nix_upload_text.py | 152 +++++++++++++++++++++++++++++++++- uv.lock | 3 + 4 files changed, 189 insertions(+), 23 deletions(-) create mode 100644 uv.lock diff --git a/src/document_processor.py b/src/document_processor.py index 8025e22e0..ea07896db 100644 --- a/src/document_processor.py +++ b/src/document_processor.py @@ -9,6 +9,8 @@ import tempfile from typing import List, Dict, Any from src.llm_core import llm_call +from src.markitdown_runtime import is_markitdown_format +from src.upload_handler import is_text_attachment logger = logging.getLogger(__name__) @@ -16,28 +18,25 @@ MAX_INLINE_ATTACHMENT_CHARS = 24000 MIN_INLINE_ATTACHMENT_SLICE = 500 -def _is_text_file(path: str) -> bool: - """Check if file has text extension.""" - return any( - path.lower().endswith(ext) - for ext in (".txt", ".py", ".html", ".htm", ".md", ".json", ".csv", ".log", ".js", ".nix") - ) +def _is_text_file(path: str, content_type: str | None = None) -> bool: + """Check a file using the shared safe text-attachment classification.""" + return is_text_attachment(path, content_type) -def _process_text_file(path: str) -> str: +def _process_text_file(path: str, display_name: str | None = None) -> str: """Process text file with enhanced formatting and metadata.""" language_map = { - ".py": "python", ".js": "javascript", ".html": "html", ".css": "css", + ".py": "python", ".js": "javascript", ".html": "html", ".htm": "html", ".css": "css", ".json": "json", ".md": "markdown", ".txt": "text", ".csv": "csv", ".log": "log", ".sh": "bash", ".bash": "bash", ".nix": "nix", ".yml": "yaml", ".yaml": "yaml", - ".xml": "xml", ".sql": "sql", ".cpp": "cpp", ".c": "c", + ".xml": "xml", ".sql": "sql", ".cpp": "cpp", ".c": "c", ".h": "c", ".java": "java", ".go": "go", ".rs": "rust", ".php": "php", ".rb": "ruby", ".ts": "typescript", ".jsx": "javascript", ".tsx": "typescript", } - filename = os.path.basename(path) - _, ext = os.path.splitext(path.lower()) + filename = os.path.basename(display_name or path) + _, ext = os.path.splitext(filename.lower()) language = language_map.get(ext, "text") max_len = 30000 if ext != ".log" else 10000 @@ -92,9 +91,9 @@ def _process_text_file(path: str) -> str: header += f"[Type: {language}, Lines: {line_count}, Size: {size_str} bytes]" code_extensions = { - ".py", ".js", ".html", ".css", ".json", ".md", ".sh", ".bash", ".nix", + ".py", ".js", ".html", ".htm", ".css", ".json", ".md", ".sh", ".bash", ".nix", ".yml", ".yaml", ".xml", ".sql", ".cpp", ".c", ".java", ".go", ".rs", ".php", ".rb", - ".ts", ".jsx", ".tsx", + ".ts", ".jsx", ".tsx", ".h", } if ext in code_extensions: code_block = f"```{language}\n{content}" @@ -435,6 +434,7 @@ def build_user_content( _, ext = os.path.splitext(path.lower()) mime = upload_info.get("mime") or mimetypes.guess_type(path)[0] or "application/octet-stream" display_name = upload_info.get("name") or upload_info.get("original_name") or path + _, display_ext = os.path.splitext(display_name.lower()) if upload_handler.is_image_file(display_name, mime): try: @@ -472,7 +472,7 @@ def build_user_content( content.insert(0, {"type": "text", "text": "[Audio attached but could not be processed]"}) elif upload_handler.is_document_file(display_name, mime): - if mime == "application/pdf": + if mime.partition(";")[0].strip().lower() == "application/pdf" or display_ext == ".pdf": extracted_text = None if session_id: try: @@ -570,8 +570,8 @@ def build_user_content( logger.warning(f"PDF auto-doc creation failed for {path}: {e}") if extracted_text is None: extracted_text = _process_pdf(path, owner=owner) - elif mime.startswith("text/") or _is_text_file(path): - extracted_text = _process_text_file(path) + elif _is_text_file(display_name, mime) and not is_markitdown_format(display_name): + extracted_text = _process_text_file(path, display_name=display_name) else: extracted_text = _process_office_document( path, diff --git a/src/upload_handler.py b/src/upload_handler.py index ce0b4b129..2754b1405 100644 --- a/src/upload_handler.py +++ b/src/upload_handler.py @@ -64,6 +64,20 @@ ATTACHMENT_REFERENCE_LINE_RE = re.compile( re.IGNORECASE, ) +TEXT_ATTACHMENT_EXTENSIONS = frozenset({ + ".txt", ".py", ".js", ".html", ".htm", ".css", ".json", ".md", + ".csv", ".log", ".xml", ".yml", ".yaml", ".nix", ".sql", ".sh", + ".bash", ".c", ".cpp", ".h", ".java", ".go", ".rs", ".php", ".rb", + ".ts", ".jsx", ".tsx", +}) + + +def is_text_attachment(filename: str, content_type: str | None = None) -> bool: + """Return True for supported text/code extensions or a ``text/*`` MIME.""" + _, ext = os.path.splitext((filename or "").lower()) + mime = (content_type or "").partition(";")[0].strip().lower() + return ext in TEXT_ATTACHMENT_EXTENSIONS or mime.startswith("text/") + def is_valid_upload_id(upload_id: str) -> bool: """Return True when *upload_id* matches the canonical uploads.json id format.""" @@ -309,10 +323,6 @@ class UploadHandler: """Check if a file is a document based on extension or content type.""" document_extensions = { '.pdf', '.docx', '.xlsx', '.pptx', '.xls', '.epub', - '.txt', '.py', '.js', '.html', '.htm', - '.css', '.json', '.md', '.csv', '.log', '.xml', '.yml', - '.yaml', '.nix', '.sql', '.sh', '.bash', '.c', '.cpp', '.h', - '.java', '.go', '.rs', '.php', '.rb', '.ts', '.jsx', '.tsx' } document_mime_types = { 'application/pdf', @@ -321,8 +331,10 @@ class UploadHandler: 'application/vnd.openxmlformats-officedocument.presentationml.presentation', 'application/vnd.ms-excel', 'application/epub+zip', - 'text/plain' } + + if is_text_attachment(filename, content_type): + return True # Check by extension _, ext = os.path.splitext(filename.lower()) @@ -330,7 +342,8 @@ class UploadHandler: return True # Check by content type if provided - if content_type and content_type in document_mime_types: + mime = (content_type or '').partition(';')[0].strip().lower() + if mime in document_mime_types: return True return False diff --git a/tests/test_nix_upload_text.py b/tests/test_nix_upload_text.py index 24de3f9b0..8bef6707d 100644 --- a/tests/test_nix_upload_text.py +++ b/tests/test_nix_upload_text.py @@ -1,5 +1,8 @@ +import pytest + +import src.document_processor as document_processor from src.document_processor import _is_text_file, _process_text_file -from src.upload_handler import UploadHandler +from src.upload_handler import UploadHandler, is_text_attachment def test_nix_files_are_treated_as_readable_documents(tmp_path): @@ -18,3 +21,150 @@ def test_nix_file_processing_includes_content_in_code_block(tmp_path): assert "[Type: nix" in rendered assert "```nix" in rendered assert "services.openssh.enable = true;" in rendered + + +def _build_user_content(tmp_path, *, stored_name, display_name, mime, body): + upload_dir = tmp_path / "uploads" + upload_dir.mkdir(exist_ok=True) + path = upload_dir / stored_name + path.write_bytes(body) + handler = UploadHandler(str(tmp_path), str(upload_dir)) + + return document_processor.build_user_content( + "Read this attachment.", + ["attachment-id"], + str(upload_dir), + handler, + owner="tester", + resolved_uploads={ + "attachment-id": { + "path": str(path), + "name": display_name, + "mime": mime, + } + }, + ) + + +def test_text_mime_without_recognized_extension_reaches_user_content(tmp_path): + handler = UploadHandler(str(tmp_path), str(tmp_path / "uploads")) + mime = "Text/Markdown; charset=utf-8" + + assert is_text_attachment("release-notes.unknown", mime) + assert handler.is_document_file("release-notes.unknown", mime) + + content = _build_user_content( + tmp_path, + stored_name="0123456789abcdef0123456789abcdef", + display_name="release-notes.unknown", + mime=mime, + body=b"Deployment requires a database backup.\n", + ) + + assert "=== File: release-notes.unknown ===" in content + assert "Deployment requires a database backup." in content + + +@pytest.mark.parametrize( + ("filename", "body"), + [ + ("include/widget.h", b"#define WIDGET_LIMIT 8\n"), + ("deploy/config.yaml", b"replicas: 3\n"), + ], +) +def test_text_extension_with_opaque_mime_reaches_user_content(tmp_path, filename, body): + display_name = filename.rsplit("/", 1)[-1] + + assert is_text_attachment(filename, "application/octet-stream") + + content = _build_user_content( + tmp_path, + stored_name=f"0123456789abcdef0123456789abcdef.{display_name.rsplit('.', 1)[-1]}", + display_name=display_name, + mime="application/octet-stream", + body=body, + ) + + assert f"=== File: {display_name} ===" in content + assert body.decode().strip() in content + + +def test_pdf_and_office_files_are_not_classified_as_generic_text(tmp_path): + handler = UploadHandler(str(tmp_path), str(tmp_path / "uploads")) + + assert handler.is_document_file("packet.pdf", "application/pdf") + assert handler.is_document_file( + "report.docx", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ) + assert not is_text_attachment("packet.pdf", "application/pdf") + assert not is_text_attachment( + "report.docx", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ) + + +def test_pdf_and_office_files_keep_specialized_processing_paths(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr( + document_processor, + "_process_pdf", + lambda path, owner=None: calls.append(("pdf", path)) or "\n\n[pdf text]", + ) + monkeypatch.setattr( + document_processor, + "_process_office_document", + lambda path, display_name, **kwargs: calls.append(("office", path)) or "\n\n[office text]", + ) + + pdf_content = _build_user_content( + tmp_path, + stored_name="0123456789abcdef0123456789abcdef.pdf", + display_name="packet.pdf", + mime="application/octet-stream", + body=b"%PDF-1.4 fake", + ) + office_content = _build_user_content( + tmp_path, + stored_name="fedcba9876543210fedcba9876543210.docx", + display_name="report.docx", + mime="application/octet-stream", + body=b"fake office file", + ) + + assert [kind for kind, _path in calls] == ["pdf", "office"] + assert "[pdf text]" in pdf_content + assert "[office text]" in office_content + + +def test_office_extension_takes_precedence_over_text_mime(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr( + document_processor, + "_process_office_document", + lambda path, display_name, **kwargs: calls.append(display_name) or "\n\n[office text]", + ) + + content = _build_user_content( + tmp_path, + stored_name="0123456789abcdef0123456789abcdef.docx", + display_name="report.docx", + mime="text/plain; charset=utf-8", + body=b"fake office file", + ) + + assert calls == ["report.docx"] + assert "[office text]" in content + + +def test_unknown_binary_attachment_is_not_decoded_as_text(tmp_path): + content = _build_user_content( + tmp_path, + stored_name="0123456789abcdef0123456789abcdef.bin", + display_name="payload.bin", + mime="application/octet-stream", + body=b"SECRET_ASCII_INSIDE_BINARY\x00\xff", + ) + + assert "[Attached non-text file]" in content + assert "SECRET_ASCII_INSIDE_BINARY" not in content diff --git a/uv.lock b/uv.lock new file mode 100644 index 000000000..a5bc51476 --- /dev/null +++ b/uv.lock @@ -0,0 +1,3 @@ +version = 1 +revision = 3 +requires-python = ">=3.14"