From 7d760cd10b02dc255fc5556893e60bde0be0703b Mon Sep 17 00:00:00 2001 From: Timur Vafin Date: Fri, 24 Apr 2026 17:12:28 +0300 Subject: [PATCH] feat: accept PDF uploads in agentic mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `.pdf` to the allowed-extensions whitelist so the security middleware lets PDF files through. In `agentic_document`, branch early for PDFs: save the file to `/.uploads/-` and pass Claude an absolute path plus an instruction to use the Read tool. Classic mode is intentionally untouched — PDFs in that path would still hit the UTF-8 decode branch. Co-Authored-By: Claude Opus 4.7 (1M context) --- .gitignore | 1 + src/bot/orchestrator.py | 37 +++++++++++- src/security/validators.py | 2 + tests/unit/test_orchestrator.py | 62 +++++++++++++++++++++ tests/unit/test_security/test_validators.py | 7 +++ 5 files changed, 106 insertions(+), 3 deletions(-) diff --git a/.gitignore b/.gitignore index 6e3390e3..7bd8f44b 100644 --- a/.gitignore +++ b/.gitignore @@ -146,4 +146,5 @@ data/ sessions/ backups/ uploads/ +.uploads/ config/mcp.json \ No newline at end of file diff --git a/src/bot/orchestrator.py b/src/bot/orchestrator.py index 6d9719f0..9a243027 100644 --- a/src/bot/orchestrator.py +++ b/src/bot/orchestrator.py @@ -9,6 +9,7 @@ import re import time from dataclasses import dataclass, field +from datetime import UTC, datetime from pathlib import Path from typing import Any, Callable, Dict, List, Optional @@ -1159,6 +1160,29 @@ async def agentic_text( success=success, ) + async def _save_pdf_and_build_prompt( + self, document: Any, caption: Optional[str] + ) -> str: + """Save PDF to /.uploads/ and build a prompt for Claude. + + Returns prompt with an absolute path so Claude's Read tool works regardless + of cwd. Filename is prefixed with millisecond timestamp to avoid collisions. + """ + uploads_dir = Path(self.settings.approved_directory) / ".uploads" + uploads_dir.mkdir(parents=True, exist_ok=True) + + timestamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S-%f")[:-3] + safe_name = f"{timestamp}-{document.file_name}" + target = uploads_dir / safe_name + + tg_file = await document.get_file() + await tg_file.download_to_drive(str(target)) + + return ( + f"{caption or 'PDF uploaded:'}\n\n" + f"File: `{target}`. Read it via Read tool and answer the user's question." + ) + async def agentic_document( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: @@ -1192,12 +1216,19 @@ async def agentic_document( await chat.send_action("typing") progress_msg = await update.message.reply_text("Working...") + prompt: Optional[str] = None + + # Binary document formats are saved to disk for Claude to read via Read tool. + if document.file_name and document.file_name.lower().endswith(".pdf"): + prompt = await self._save_pdf_and_build_prompt( + document, update.message.caption + ) + # Try enhanced file handler, fall back to basic features = context.bot_data.get("features") file_handler = features.get_file_handler() if features else None - prompt: Optional[str] = None - if file_handler: + if prompt is None and file_handler: try: processed_file = await file_handler.handle_document_upload( document, @@ -1208,7 +1239,7 @@ async def agentic_document( except Exception: file_handler = None - if not file_handler: + if prompt is None and not file_handler: file = await document.get_file() file_bytes = await file.download_as_bytearray() try: diff --git a/src/security/validators.py b/src/security/validators.py index 381ba321..3038accc 100644 --- a/src/security/validators.py +++ b/src/security/validators.py @@ -86,6 +86,8 @@ class SecurityValidator: ".vue", ".svelte", ".lock", + # Document formats (binary, saved to disk for downstream tools) + ".pdf", } # Forbidden filenames and patterns diff --git a/tests/unit/test_orchestrator.py b/tests/unit/test_orchestrator.py index ce5e419e..8b4f2057 100644 --- a/tests/unit/test_orchestrator.py +++ b/tests/unit/test_orchestrator.py @@ -367,6 +367,68 @@ async def test_agentic_document_rejects_large_files(agentic_settings, deps): assert "too large" in call_args.args[0].lower() +async def test_agentic_document_pdf_saves_and_prompts_read(agentic_settings, deps): + """PDF uploads are saved to /.uploads/ and Claude gets + an absolute path + instruction to use Read.""" + orchestrator = MessageOrchestrator(agentic_settings, deps) + + approved_dir = Path(agentic_settings.approved_directory) + pdf_bytes = b"%PDF-1.4\nSMOKE-TOKEN-42\n%%EOF\n" + + async def fake_download(target_path): + Path(target_path).write_bytes(pdf_bytes) + + tg_file = MagicMock() + tg_file.download_to_drive = AsyncMock(side_effect=fake_download) + + update = MagicMock() + update.effective_user.id = 123 + update.message.document.file_name = "ticket.pdf" + update.message.document.file_size = len(pdf_bytes) + update.message.document.get_file = AsyncMock(return_value=tg_file) + update.message.caption = "ticket attached" + update.message.chat.send_action = AsyncMock() + update.message.reply_text = AsyncMock() + + progress_msg = AsyncMock() + progress_msg.edit_text = AsyncMock() + progress_msg.delete = AsyncMock() + update.message.reply_text.return_value = progress_msg + + mock_response = MagicMock() + mock_response.session_id = "pdf-session-1" + mock_response.content = "Read the PDF. Found SMOKE-TOKEN-42." + mock_response.tools_used = [] + + claude_integration = AsyncMock() + claude_integration.run_command = AsyncMock(return_value=mock_response) + + context = MagicMock() + context.user_data = {} + context.bot_data = { + "settings": agentic_settings, + "security_validator": None, + "features": None, + "claude_integration": claude_integration, + } + + await orchestrator.agentic_document(update, context) + + uploads_dir = approved_dir / ".uploads" + saved_files = list(uploads_dir.glob("*-ticket.pdf")) + assert len(saved_files) == 1, f"expected one saved PDF, got {saved_files}" + assert saved_files[0].read_bytes() == pdf_bytes + + claude_integration.run_command.assert_awaited_once() + call_kwargs = claude_integration.run_command.call_args.kwargs + prompt = ( + call_kwargs.get("prompt") or claude_integration.run_command.call_args.args[0] + ) + assert str(saved_files[0]) in prompt + assert "Read" in prompt + assert "ticket attached" in prompt + + async def test_agentic_voice_calls_claude(agentic_settings, deps): """Agentic voice handler transcribes and routes prompt to Claude.""" orchestrator = MessageOrchestrator(agentic_settings, deps) diff --git a/tests/unit/test_security/test_validators.py b/tests/unit/test_security/test_validators.py index a15f5c31..8c6ec4a8 100644 --- a/tests/unit/test_security/test_validators.py +++ b/tests/unit/test_security/test_validators.py @@ -139,6 +139,7 @@ def test_filename_validation_valid(self, validator): "style.css", "data.sql", "build.sh", + "report.pdf", ] for filename in valid_filenames: @@ -146,6 +147,12 @@ def test_filename_validation_valid(self, validator): assert valid is True assert error is None + def test_filename_pdf_exe_suffix_blocked(self, validator): + """Regression: `.pdf.exe` trap name is still blocked by dangerous patterns.""" + valid, error = validator.validate_filename("ticket.pdf.exe") + assert valid is False + assert "not allowed" in error + def test_filename_validation_invalid_extensions(self, validator): """Test rejection of invalid file extensions.""" invalid_filenames = [